Add HBG prepared epoch pipeline

docs/task-flow.md

@@ -64,8 +64,9 @@ to C++:
 | `w3.submit_sub(handle, …)` dispatched to a SUB child | `LOCAL_PYTHON` | child resolves digest to a Python callable and calls `fn(args)` |
 
 All three paths share one mailbox wire format: `MAILBOX_OFF_CALLABLE` is
-reserved, and the 32-byte digest prefixes the args blob. The receiving child
-does the digest-to-slot resolve in its own address space.
+reserved, the run's generation-safe pipeline lease follows `CallConfig`, and
+the 32-byte digest prefixes the args blob. The receiving child does the
+digest-to-slot resolve in its own address space.
 
 The proposed remote L3 path keeps the same callable identity contract, but
 sends it in a versioned TASK frame. The remote endpoint resolves the digest
@@ -292,12 +293,71 @@ own AICore stream and retires it on every exit path, and no record of which
 image a stream last ran is load-bearing. The AICPU stream carries no such state
 and stays with its slot.
 
-This is dormant capacity at this layer. The ordinary synchronous entry point
-continues to use slot 0, and the chip child's mailbox loop passes no lease, so
-every production run is unleased. This contract does not enable a second
-mailbox frame, a second device execution, or cross-run publication overlap.
-Carrying a lease across the mailbox, and deciding when slot 1 may be leased at
-all, belong to whole-run admission.
+#### Whole-run FIFO admission
+
+L3 graph callbacks remain synchronous and serialized, but native admission
+allows two live run reservations:
+
+```text
+run N:     EXECUTING
+run N+1:   PREPARED
+run N+2:   blocked in begin_run before its graph callback
+```
+
+`begin_run` acquires a generation-safe lease before invoking the callback.
+The FIFO head may enter `EXECUTING` while its callback is still building, which
+preserves orchestration callbacks that submit device work and wait for L2
+communication before returning. A non-head run remains `BUILDING` or becomes
+`PREPARED` when graph construction closes; it cannot execute until every prior
+run is terminal. The scheduler observes only the ready-queue partition belonging
+to that single active FIFO head. A run's device effects therefore cannot
+interleave with another run even when both graphs contain ready tasks. TensorMap
+keys remain `(run_id, tensor_key)`, so adjacent runs may reuse the same tensor
+address without creating cross-run dependencies.
+
+The terminal transition releases the reservation and lease exactly once,
+wakes a blocked third submission, and activates the next prepared run. Empty
+runs take the same transition immediately. If graph construction fails, every
+unstarted slot is poisoned and consumed, its ready-queue partition is erased,
+and the lease is returned without dispatching device work.
+
+Each direct chip child publishes its runtime contract's `pipeline_depth` and
+endpoint capability bits in the startup mailbox before `INIT_READY`. The parent
+configures admission to the minimum published depth and enables prepare/activate
+per endpoint only when the loaded backend reports the complete prepared-run ABI.
+Backends without a depth-two contract therefore keep depth-one serial behavior
+instead of receiving an invalid slot-1 lease.
+
+#### Prepare/activate endpoint lane
+
+A local two-frame endpoint may separately advertise
+`supports_prepare_activate`. The scheduler then stages at most one ready
+NEXT_LEVEL task slot from the prepared FIFO successor. A group is staged as one
+task slot across all of its target workers; an endpoint set that cannot stage
+the whole group leaves it on the normal ready queue.
+
+The prepared frame uses a distinct protocol path:
+
+```text
+PREPARE_READY -> BACKEND_READY -> ACTIVATE -> TASK_ACTIVE -> TASK_DONE | TASK_FAILED
+```
+
+`BACKEND_READY` reports only that unpublished backend state exists. It does not
+satisfy the launch-acceptance fence and does not permit device work. The
+WorkerThread waits on an activation condition variable while retaining the
+frame and its generation-safe identity. Whole-run FIFO promotion latches
+`ACTIVATE`; the latch may arrive before backend preparation finishes, but the
+endpoint cannot publish it to the child until `BACKEND_READY` is observed for
+the same run, slot, generation, and dispatch id.
+
+Only the first eligible, non-diagnostic dispatch in a run uses this lane. Other
+tasks from the same run retain the existing sequential compatibility path.
+Diagnostic execution takes an exclusive backend gate so its runner-global
+collectors cannot overlap Host preparation. Remote, SUB, A5, TMR, sim,
+single-frame, and endpoints without the capability remain unchanged. Onboard
+a2a3 HBG is the first backend to advertise support: it builds into the
+lease-selected arena bank, reports `BACKEND_READY` without publishing device
+work, and launches only after the matching `ACTIVATE`.
 
 Simulation implements the same depth, so the contract means the same thing on
 both platforms: its runner owns one arena bank and one retained temporary
@@ -382,6 +442,9 @@ Where the data goes after submit:
    Tags are consumed during the same submit call for dep inference and
    **never carried further**.
 3. `CallConfig` — copied into `slot.config` (parent heap, POD)
+4. `PipelineSlotLease` — copied from the owning run into
+   `slot.pipeline_lease`; local chip mailboxes forward `{slot_id, generation}`
+   to `ChipWorker::run_with_lease`.
 
 For the full submit mechanics (ring alloc, TensorMap lookup/insert, scope ref,
 fanout wiring), see [orchestrator.md](orchestrator.md).
@@ -390,7 +453,8 @@ fanout wiring), see [orchestrator.md](orchestrator.md).
 
 For local endpoints, after the Scheduler resolves the submitted NEXT_LEVEL
 target (or chooses an idle SUB worker), `LocalMailboxEndpoint` encodes
-`(callable digest, CallConfig, TaskArgs)` into the per-worker shm mailbox and
+`(callable digest, CallConfig, PipelineSlotLease, TaskArgs)` into the
+per-worker shm mailbox and
 the forked child decodes it. Remote NEXT_LEVEL dispatch through
 `RemoteL3Endpoint` serializes the same logical payload into a framed TASK
 request instead.
@@ -409,7 +473,8 @@ Local mailbox path:
 ```text
 slot.callable.digest ─┐
 slot.config          ─┼─► memcpy into shm mailbox ─► child resolves digest
-slot.task_args       ─┘    (dispatch_process)         and runs local slot
+slot.pipeline_lease  ─┤    (dispatch_process)         and runs local slot
+slot.task_args       ─┘
 ```
 
 For SUB children the same mailbox layout is reused; the Python child

python/bindings/task_interface.cpp

@@ -1456,6 +1456,7 @@ NB_MODULE(_task_interface, m) {
             [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity,
                const CallConfig &config, uint64_t accepted_state_addr, int32_t accepted_value, uint32_t pipeline_slot,
                uint64_t pipeline_generation) {
+                nb::gil_scoped_release release;
                 // The mailbox region is the on-wire format `write_blob` produced;
                 // `read_blob` is the matching reader that returns a zero-copy
                 // TaskArgsView into the caller-owned bytes. Forwards to the
@@ -1483,6 +1484,46 @@ NB_MODULE(_task_interface, m) {
             "of the per-task tensor/scalar layout). The blob must be in the format "
             "produced by `write_blob`; read_blob enforces capacity bounds against shm corruption."
         )
+        .def(
+            "prepare_from_blob",
+            [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity,
+               const CallConfig &config, uint32_t pipeline_slot, uint64_t pipeline_generation, uint64_t run_id,
+               uint64_t dispatch_id) {
+                nb::gil_scoped_release release;
+                TaskArgsView view = read_blob(reinterpret_cast<const uint8_t *>(args_blob_ptr), blob_capacity);
+                ChipStorageTaskArgs storage = view_to_chip_storage(view);
+                self.prepare_with_lease(
+                    callable_id, &storage, config, PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id,
+                    dispatch_id
+                );
+            },
+            nb::arg("callable_id"), nb::arg("args_blob_ptr"), nb::arg("blob_capacity"), nb::arg("config"),
+            nb::arg("pipeline_slot"), nb::arg("pipeline_generation"), nb::arg("run_id"), nb::arg("dispatch_id"),
+            "Build one unpublished prepared run in a generation-safe pipeline slot."
+        )
+        .def(
+            "execute_prepared",
+            [](ChipWorker &self, int32_t callable_id, const CallConfig &config, uint32_t pipeline_slot,
+               uint64_t pipeline_generation, uint64_t run_id, uint64_t dispatch_id) {
+                nb::gil_scoped_release release;
+                self.execute_prepared(
+                    callable_id, config, PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id, dispatch_id
+                );
+            },
+            nb::arg("callable_id"), nb::arg("config"), nb::arg("pipeline_slot"), nb::arg("pipeline_generation"),
+            nb::arg("run_id"), nb::arg("dispatch_id"),
+            "Publish and complete a matching prepared run on the sequential executor."
+        )
+        .def(
+            "abort_prepared",
+            [](ChipWorker &self, uint32_t pipeline_slot, uint64_t pipeline_generation, uint64_t run_id,
+               uint64_t dispatch_id) {
+                nb::gil_scoped_release release;
+                self.abort_prepared(PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id, dispatch_id);
+            },
+            nb::arg("pipeline_slot"), nb::arg("pipeline_generation"), nb::arg("run_id"), nb::arg("dispatch_id"),
+            "Reclaim an unpublished prepared run after a local protocol failure."
+        )
         .def(
             "unregister_callable",
             [](ChipWorker &self, int32_t callable_id) {
@@ -1497,6 +1538,7 @@ NB_MODULE(_task_interface, m) {
         .def_prop_ro("initialized", &ChipWorker::initialized)
         .def_prop_ro("pipeline_depth", &ChipWorker::pipeline_depth)
         .def_prop_ro("runtime_slot_count", &ChipWorker::runtime_slot_count)
+        .def_prop_ro("supports_prepared_runs", &ChipWorker::supports_prepared_runs)
         .def_prop_ro(
             "runtime_buffer_addrs", &ChipWorker::runtime_buffer_addrs,
             "Host Runtime staging buffer address of every copy the runtime's "

python/bindings/worker_bind.h

@@ -350,7 +350,7 @@ inline void bind_worker(nb::module_ &m) {
         .def("scope_end", &Orchestrator::scope_end, "Close the innermost scope. Non-blocking.")
         .def("_scope_begin", &Orchestrator::scope_begin)
         .def("_scope_end", &Orchestrator::scope_end)
-        .def("_begin_run", &Orchestrator::begin_run)
+        .def("_begin_run", &Orchestrator::begin_run, nb::call_guard<nb::gil_scoped_release>())
         .def("_close_run_submission", &Orchestrator::close_run_submission, nb::arg("run_id"))
         .def(
             "_fail_run_submission",
@@ -393,22 +393,35 @@ inline void bind_worker(nb::module_ &m) {
 
         .def(
             "add_next_level_worker",
-            [](Worker &self, uint64_t mailbox_ptr, int child_pid) {
-                self.add_worker(WorkerType::NEXT_LEVEL, reinterpret_cast<void *>(mailbox_ptr), child_pid);
+            [](Worker &self, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count,
+               bool supports_prepare_activate) {
+                self.add_worker(
+                    WorkerType::NEXT_LEVEL, reinterpret_cast<void *>(mailbox_ptr), child_pid, task_frame_count,
+                    supports_prepare_activate
+                );
             },
-            nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1,
+            nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, nb::arg("task_frame_count") = 1,
+            nb::arg("supports_prepare_activate") = false,
             "Add a NEXT_LEVEL sub-worker. `mailbox_ptr` is the address of a "
             "MAILBOX_SIZE-byte MAP_SHARED region; the child process loop is "
             "Python-managed (fork + _chip_process_loop). `child_pid` is that "
             "forked child, used to detect an exit before mailbox completion."
         )
         .def(
             "add_next_level_worker_at",
-            [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, int child_pid) {
-                self.add_next_level_worker(worker_id, reinterpret_cast<void *>(mailbox_ptr), child_pid);
+            [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count,
+               bool supports_prepare_activate) {
+                self.add_next_level_worker(
+                    worker_id, reinterpret_cast<void *>(mailbox_ptr), child_pid, task_frame_count,
+                    supports_prepare_activate
+                );
             },
-            nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1,
-            "Add a NEXT_LEVEL sub-worker with an explicit worker id."
+            nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, nb::arg("task_frame_count") = 1,
+            nb::arg("supports_prepare_activate") = false, "Add a NEXT_LEVEL sub-worker with an explicit worker id."
+        )
+        .def(
+            "configure_pipeline_depth", &Worker::configure_pipeline_depth, nb::arg("depth"),
+            "Set run admission depth from the minimum direct-chip runtime capability before init."
         )
         .def(
             "add_sub_worker",
@@ -764,8 +777,10 @@ inline void bind_worker(nb::module_ &m) {
 
     m.attr("DEFAULT_HEAP_RING_SIZE") = static_cast<uint64_t>(DEFAULT_HEAP_RING_SIZE);
     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_ERROR_MSG_SIZE") = static_cast<int>(MAILBOX_ERROR_MSG_SIZE);
+    m.attr("PTO_PIPELINE_MAX_DEPTH") = static_cast<uint32_t>(PTO_PIPELINE_MAX_DEPTH);
     m.attr("MAX_RING_DEPTH") = static_cast<int32_t>(MAX_RING_DEPTH);
     m.attr("MAX_SCOPE_DEPTH") = static_cast<int32_t>(MAX_SCOPE_DEPTH);
 

python/simpler/task_interface.py

@@ -41,6 +41,7 @@
 import _task_interface as _ti_module  # pyright: ignore[reportMissingImports]
 from _task_interface import (  # pyright: ignore[reportMissingImports]
     MAILBOX_ERROR_MSG_SIZE,
+    MAILBOX_FRAME_SIZE,
     MAILBOX_OFF_ERROR_MSG,
     MAILBOX_SIZE,
     MAX_REGISTERED_CALLABLE_IDS,
@@ -156,6 +157,7 @@ def _assert_bindings_match_source_tree() -> None:
     "TaskState",
     "_Worker",
     "MAILBOX_SIZE",
+    "MAILBOX_FRAME_SIZE",
     "MAILBOX_OFF_ERROR_MSG",
     "MAILBOX_ERROR_MSG_SIZE",
     "read_args_from_blob",