Fix: align full-rank chip run activation

docs/chip-level-arch.md

@@ -109,7 +109,7 @@ runner.set_executors(aicpu_binary, aicore_binary);   // once, at init time
 std::unique_ptr<DeviceRunnerBase::PreparedExecution> prepared;
 runner.prepare_execution(runtime, config, pipeline_slot, identity, &prepared);
 auto launched = runner.launch_execution(std::move(prepared), std::move(permit));
-runner.drain_execution(*launched.active);            // child progress path owns progress
+runner.drain_execution(*launched.active);            // resident lane lifecycle owner waits/finalizes
 runner.finalize();
 ```
 

docs/task-flow.md

@@ -379,17 +379,20 @@ prepared-but-not-launched run.
 
 A direct A2/A3 chip endpoint with a negotiated depth of at least two uses two
 task frames and advertises `supports_frame_staging`. One `WorkerThread` owns
-both frames and drives them through a non-blocking progress interface; the
-child process likewise has one loop that services control traffic, both task
-frames, and the bounded active/prepared native lifecycles. There is no thread
-per frame.
+both frames and drives them through a non-blocking progress interface. In the
+child process, the mailbox loop services control traffic and both task frames,
+while one resident C++ lifecycle thread per chip lane blocks on the launched
+native token's device completion. There is no thread per frame or per run.
 
 The active and successor paths are:
 
 ```text
 IDLE -> TASK_READY    -> FRAME_STAGED -> TASK_LAUNCHED -> TASK_DONE | TASK_FAILED
 IDLE -> PREPARE_READY -> FRAME_STAGED -> ACTIVATE -> TASK_LAUNCHED
                                                -> TASK_DONE | TASK_FAILED
+IDLE -> NATIVE_PREPARE_READY -> FRAME_STAGED -> ACTIVATE -> TASK_LAUNCHED
+                                                      -> TASK_DONE | TASK_FAILED
+                                     -> ABANDON -> TASK_FAILED
 ```
 
 `FRAME_STAGED` means that the child owns an immutable frame snapshot; it does
@@ -411,16 +414,29 @@ between these meanings.
 
 An HBG successor's prepared token remains unlaunched and unaccepted until
 `ACTIVATE`, and activation still cannot launch it until the predecessor has
-polled complete and finalized. The sticky acceptance word therefore remains
-zero throughout preparation. Shutdown, stale activation, and pre-launch
-failure finalize the token exactly once before the frame becomes terminal.
+completed and finalized. The resident lifecycle thread owns the launched
+native token's blocking wait and finalizes outside the lane mutex; it takes the
+mutex only to publish terminal state and launch the successor. The mailbox
+loop uses a short, bounded wait for launched handles so the binding releases
+the GIL and does not starve that lifecycle thread, while unlaunched handles
+remain non-blocking status probes. The sticky acceptance word therefore
+remains zero throughout preparation. Shutdown, stale activation, and
+pre-launch failure finalize the token exactly once before the frame becomes
+terminal.
 
 The scheduler stages only the first eligible single NEXT_LEVEL task from the
-prepared FIFO successor. Tasks from the active run use only the active lane, so
-the second frame cannot create same-device execution overlap. Prepared groups
-remain on their normal queue and dispatch synchronously after FIFO promotion.
-Remote, SUB, A5, simulation, nested-worker, and single-frame endpoints retain
-the blocking compatibility path.
+prepared FIFO successor. A prepared group never crosses the whole-run FIFO;
+the second frame therefore cannot create same-device execution overlap. A
+non-diagnostic NEXT_LEVEL group whose members all target local two-frame endpoints uses a
+group barrier: the scheduler publishes every member as
+`NATIVE_PREPARE_READY`, keeps all target workers reserved, and publishes no
+`ACTIVATE` until every member has completed native bind/prepare and then
+reported `FRAME_STAGED` with disposition `NATIVE_PREPARED`. It activates every
+member in the same scheduler progress pass. If preparation of one member fails, staged peers receive
+`ABANDON`; the child releases each unlaunched native run and acknowledges
+`TASK_FAILED`, so a partial group can never launch. Diagnostic groups and
+remote, SUB, A5, simulation, nested-worker, and single-frame endpoints retain
+the compatibility path.
 
 The child validates newly visible metadata from both frames before selecting
 the next active `dispatch_id`. It prepares that active token first; preparation
@@ -750,9 +766,9 @@ Step-by-step (one chip worker):
 | 2 | `Worker::run` | `scope_begin` → call `my_orch(&orch_, args.view(), cfg)` |
 | 3 | `Orchestrator::submit_next_level` | `slot = ring.alloc()`; move `chip_args` into `slot.task_args`; walk tags → `tensormap.lookup(a.data)`, `tensormap.lookup(b.data)`, `tensormap.insert(c.data, slot)`; push ready |
 | 4 | Scheduler thread | pop `slot` from worker 0's FIFO; resolve stable worker ID 0 to WT_chip_0; dispatch |
-| 5 | WT_chip_0 parent side | encode one leased task frame: write `config`, digest prefix, and the args blob; publish `TASK_READY` for the active lane or `PREPARE_READY` for a staged successor |
+| 5 | WT_chip_0 parent side | encode one leased task frame: write `config`, digest prefix, and the args blob; publish `TASK_READY` for the active lane, `PREPARE_READY` for a staged successor, or `NATIVE_PREPARE_READY` for a local group barrier member |
 | 6 | chip_0 child process | validate the frame and resolve its digest; ordinary HBG with an active predecessor also prepares the leased inactive arena bank before publishing `FRAME_STAGED`, while a frame with no active predecessor, diagnostic HBG, and TMR publish after validation and defer native prepare |
-| 7 | chip_0 native-run path | after activation and the predecessor's finalization fence, launch an already-prepared HBG run or finish deferred native preparation and then launch; poll it to completion and finalize it before another staged frame may launch. Compatibility endpoints perform the equivalent operation through blocking `ChipWorker::run` |
+| 7 | chip_0 native-run path | after activation and the predecessor's finalization fence, launch an already-prepared HBG run or finish deferred native preparation and then launch; the resident lifecycle owner blocks on device completion, finalizes outside the lane mutex, publishes terminal state, and launches the staged successor. Compatibility endpoints perform the equivalent operation through blocking `ChipWorker::run` |
 | 8 | runtime.so | translate host ptrs → device ptrs; dispatch AICPU / AICore; write output into `c`'s shm |
 | 9 | chip_0 child | native finalization returns; write `TASK_DONE` |
 | 10 | WT_chip_0 parent | observe `TASK_DONE`; push success completion |

python/bindings/task_interface.cpp

@@ -2759,6 +2759,7 @@ NB_MODULE(_task_interface, m) {
 
     nb::class_<ChipRun>(m, "_ChipRun")
         .def("done", &ChipRun::done)
+        .def("prepare", &ChipRun::prepare, nb::call_guard<nb::gil_scoped_release>())
         .def("activate", &ChipRun::activate)
         .def("abandon", &ChipRun::abandon, nb::call_guard<nb::gil_scoped_release>())
         .def_prop_ro("launched", &ChipRun::launched)

python/bindings/worker_bind.h

@@ -234,6 +234,12 @@ inline void mailbox_store_i32(uint64_t addr, int32_t v) {
 #endif
 }
 
+inline void mailbox_notify_i32(uint64_t addr) {
+    auto *ptr = reinterpret_cast<int32_t *>(addr);
+    (void)__atomic_add_fetch(ptr, 1, __ATOMIC_RELEASE);
+    mpi_group_mailbox::wake_word(ptr);
+}
+
 inline void bind_worker(nb::module_ &m) {
     // --- WorkerType ---
     nb::enum_<WorkerType>(m, "WorkerType").value("NEXT_LEVEL", WorkerType::NEXT_LEVEL).value("SUB", WorkerType::SUB);
@@ -859,6 +865,7 @@ inline void bind_worker(nb::module_ &m) {
     m.attr("MAILBOX_SIZE") = static_cast<int>(MAILBOX_SIZE);
     m.attr("MAILBOX_FRAME_SIZE") = static_cast<int>(MAILBOX_FRAME_SIZE);
     m.attr("MAILBOX_OFF_ERROR_MSG") = static_cast<int>(MAILBOX_OFF_ERROR_MSG);
+    m.attr("MAILBOX_OFF_NOTIFICATION") = static_cast<int>(MAILBOX_OFF_NOTIFICATION);
     m.attr("MAILBOX_ERROR_MSG_SIZE") = static_cast<int>(MAILBOX_ERROR_MSG_SIZE);
     // The MailboxState values as the C++ side defines them, keyed by
     // enumerator name. They are a cross-process wire contract: the word at
@@ -881,6 +888,8 @@ inline void bind_worker(nb::module_ &m) {
     mailbox_states["TASK_FAILED"] = static_cast<int32_t>(MailboxState::TASK_FAILED);
     mailbox_states["ACTIVATE"] = static_cast<int32_t>(MailboxState::ACTIVATE);
     mailbox_states["PREPARE_READY"] = static_cast<int32_t>(MailboxState::PREPARE_READY);
+    mailbox_states["ABANDON"] = static_cast<int32_t>(MailboxState::ABANDON);
+    mailbox_states["NATIVE_PREPARE_READY"] = static_cast<int32_t>(MailboxState::NATIVE_PREPARE_READY);
     m.attr("MAILBOX_STATE_VALUES") = mailbox_states;
     nb::dict mailbox_dispositions;
     mailbox_dispositions["NONE"] = static_cast<int32_t>(MailboxPreparationDisposition::NONE);
@@ -908,6 +917,13 @@ inline void bind_worker(nb::module_ &m) {
         },
         nb::arg("addr"), nb::arg("value"), "Release-store a 32-bit mailbox word at `addr`."
     );
+    m.def(
+        "_mailbox_notify_i32",
+        [](uint64_t addr) {
+            mailbox_notify_i32(addr);
+        },
+        nb::arg("addr"), "Increment and wake a 32-bit mailbox notification generation at `addr`."
+    );
     m.def(
         "_mailbox_wait_i32",
         [](uint64_t addr, int32_t expected, double timeout_s) {

python/simpler/worker.py

@@ -88,13 +88,16 @@ def my_l4_orch(orch, args, config):
 import cloudpickle
 from _task_interface import (  # pyright: ignore[reportMissingImports]
     HOST_STRACE_ENABLED,
+    MAILBOX_OFF_NOTIFICATION,
     MAX_REGISTERED_CALLABLE_IDS,
     PTO_PIPELINE_MAX_DEPTH,
     RUNTIME_ENV_RING_COUNT,
     WorkerType,
     _emit_host_span,
     _mailbox_load_i32,
+    _mailbox_notify_i32,
     _mailbox_store_i32,
+    _mailbox_wait_i32,
     _read_control_copy_request,
     _set_host_span_level_prefix,
     _worker_host_mapped_region_ack_cleanup_error,
@@ -335,7 +338,7 @@ def _host_spans_active() -> bool:
 _OFF_FRAME_SLOT_ID = _OFF_ACCEPTED - 24
 _OFF_FRAME_GENERATION = _OFF_ACCEPTED - 16
 _OFF_FRAME_DISPATCH_ID = _OFF_ACCEPTED - 8
-_TASK_PROTOCOL_VERSION = 3
+_TASK_PROTOCOL_VERSION = 5
 # Mirrors MAILBOX_OFF_SHUTDOWN / MAILBOX_SHUTDOWN_REQUESTED: termination is a
 # sticky one-way word on the control frame, not a MailboxState. _OFF_STATE has
 # three writers (parent CONTROL_REQUEST, child CONTROL_DONE, C++
@@ -369,6 +372,8 @@ def _host_spans_active() -> bool:
 _TASK_FAILED = 10
 _ACTIVATE = 11
 _PREPARE_READY = 12
+_ABANDON = 13
+_NATIVE_PREPARE_READY = 14
 _TASK_FRAME_COUNT = 2
 
 
@@ -397,6 +402,8 @@ def _assert_mailbox_wire_constants() -> None:
         "TASK_FAILED": _TASK_FAILED,
         "ACTIVATE": _ACTIVATE,
         "PREPARE_READY": _PREPARE_READY,
+        "ABANDON": _ABANDON,
+        "NATIVE_PREPARE_READY": _NATIVE_PREPARE_READY,
     }
     dispositions = {
         "NONE": _DISPOSITION_NONE,
@@ -2048,6 +2055,7 @@ def _request_child_shutdown(buf) -> None:
     """
     _mailbox_store_i32(_buffer_field_addr(buf, _OFF_SHUTDOWN), _SHUTDOWN_REQUESTED)
     _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN)
+    _mailbox_notify_i32(_buffer_field_addr(buf, MAILBOX_OFF_NOTIFICATION))
 
 
 def _write_error(buf, code: int, msg: str = "") -> None:
@@ -2112,6 +2120,7 @@ def _reexport_args_from_mailbox(buf, worker: Worker) -> TaskArgs:
 # orphan is reaped before it is noticeable, cheap enough to be lost in the
 # noise of the poll itself.
 _PARENT_LIVENESS_POLL_INTERVAL = 1000
+_CHIP_RUN_PROGRESS_WAIT_S = 0.001
 
 
 def _run_mailbox_loop(
@@ -2970,6 +2979,7 @@ class _StagedFrame:
             cid: int
             config: CallConfig
             activated: bool
+            require_native_prepare: bool
             chip_run: Any = None
             launched_published: bool = False
 
@@ -2985,7 +2995,15 @@ def read_identity(frame_buf: memoryview) -> tuple[int, int, int, int, int]:
             )
 
         def task_frame_references_digest(digest: bytes) -> bool:
-            live_states = (_TASK_READY, _PREPARE_READY, _ACTIVATE, _FRAME_STAGED, _TASK_LAUNCHED)
+            live_states = (
+                _TASK_READY,
+                _PREPARE_READY,
+                _NATIVE_PREPARE_READY,
+                _ACTIVATE,
+                _FRAME_STAGED,
+                _TASK_LAUNCHED,
+                _ABANDON,
+            )
             for index, frame_buf in enumerate(frame_bufs):
                 if _mailbox_load_i32(frame_addrs[index] + _OFF_STATE) not in live_states:
                     continue
@@ -3037,6 +3055,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None:
                     cid=int(cid),
                     config=_read_config_from_mailbox(frame_buf),
                     activated=initial_state in (_TASK_READY, _ACTIVATE),
+                    require_native_prepare=initial_state == _NATIVE_PREPARE_READY,
                 )
             except Exception as e:  # noqa: BLE001
                 _write_error(frame_buf, 1, _format_exc(f"chip_process dev={device_id} frame={index}", e))
@@ -3066,6 +3085,8 @@ def submit_frame(frame: _StagedFrame) -> None:
                 _TASK_ACCEPTED,
                 False,
             )
+            if frame.require_native_prepare:
+                frame.chip_run.prepare()
             raw_disposition = frame.chip_run.preparation_disposition
             disposition = int(getattr(raw_disposition, "value", raw_disposition))
             if disposition not in (_VALIDATED_ONLY, _NATIVE_PREPARED):
@@ -3080,8 +3101,10 @@ def submit_frame(frame: _StagedFrame) -> None:
         liveness_countdown = _PARENT_LIVENESS_POLL_INTERVAL
         shutdown_message = f"chip_process dev={device_id}: task loop shut down"
         shutdown_addr = _buffer_field_addr(buf, _OFF_SHUTDOWN)
+        notification_addr = _buffer_field_addr(buf, MAILBOX_OFF_NOTIFICATION)
         try:
             while True:
+                notification_snapshot = _mailbox_load_i32(notification_addr)
                 control_state = _mailbox_load_i32(state_addr)
                 if control_state == _SHUTDOWN or _mailbox_load_i32(shutdown_addr) == _SHUTDOWN_REQUESTED:
                     break
@@ -3101,7 +3124,12 @@ def submit_frame(frame: _StagedFrame) -> None:
                     frame_state = _mailbox_load_i32(frame_addrs[index] + _OFF_STATE)
                     staged = staged_frames.get(index)
                     if staged is None:
-                        if frame_state in (_TASK_READY, _PREPARE_READY, _ACTIVATE):
+                        if frame_state in (
+                            _TASK_READY,
+                            _PREPARE_READY,
+                            _NATIVE_PREPARE_READY,
+                            _ACTIVATE,
+                        ):
                             staged = stage_frame(index, frame_state)
                             if staged is not None:
                                 new_frames.append(staged)
@@ -3124,6 +3152,15 @@ def submit_frame(frame: _StagedFrame) -> None:
                         except Exception as e:  # noqa: BLE001
                             shutdown_message = _format_exc(f"chip_process dev={device_id}: native activation", e)
                             break
+                    if frame_state == _ABANDON and not staged.activated:
+                        try:
+                            staged.chip_run.abandon()
+                        except Exception as e:  # noqa: BLE001
+                            shutdown_message = _format_exc(f"chip_process dev={device_id}: native abandonment", e)
+                            break
+                        _mailbox_store_i32(staged.frame_addr + _OFF_STATE, _TASK_FAILED)
+                        staged_frames.pop(index, None)
+                        continue
                 else:
                     for staged in sorted(new_frames, key=lambda frame: frame.identity[4]):
                         try:
@@ -3140,7 +3177,15 @@ def submit_frame(frame: _StagedFrame) -> None:
                                 _mailbox_store_i32(staged.frame_addr + _OFF_STATE, _TASK_LAUNCHED)
                                 staged.launched_published = True
                                 continue
-                            run_complete = bool(staged.chip_run.done())
+                            # A launched run has a resident C++ completion
+                            # owner. Its bounded wait releases the GIL while
+                            # preserving mailbox responsiveness. Unlaunched
+                            # staged work remains a nonblocking status probe.
+                            run_complete = bool(
+                                staged.chip_run.wait(_CHIP_RUN_PROGRESS_WAIT_S)
+                                if staged.chip_run.launched
+                                else staged.chip_run.done()
+                            )
                             if not run_complete and staged.chip_run.launched and not staged.launched_published:
                                 _mailbox_store_i32(staged.frame_addr + _OFF_STATE, _TASK_LAUNCHED)
                                 staged.launched_published = True
@@ -3181,6 +3226,14 @@ def submit_frame(frame: _StagedFrame) -> None:
                             if os.getppid() != parent_pid:
                                 shutdown_message = f"chip_process dev={device_id}: parent exited"
                                 break
+                        # With no launched run, every interesting transition is
+                        # a parent mailbox publication.  Park on the shared
+                        # generation instead of burning a core in Python.  A
+                        # publication racing this scan changes the expected
+                        # value, so FUTEX_WAIT returns immediately rather than
+                        # losing the wake.
+                        if not any(frame.chip_run.launched for frame in staged_frames.values()):
+                            _mailbox_wait_i32(notification_addr, notification_snapshot, 0.01)
                         continue
                 break
         finally:
@@ -3195,8 +3248,10 @@ def submit_frame(frame: _StagedFrame) -> None:
                 if _mailbox_load_i32(frame_state_addr) in (
                     _TASK_READY,
                     _PREPARE_READY,
+                    _NATIVE_PREPARE_READY,
                     _ACTIVATE,
                     _FRAME_STAGED,
+                    _ABANDON,
                     _TASK_LAUNCHED,
                 ):
                     _write_error(frame_buf, 1, shutdown_message)