Fix: reserve NEXT_LEVEL workers for ready groups

docs/directed-next-level-scheduling.md

@@ -44,14 +44,15 @@ routing operation.
 
 Each Scheduler iteration:
 
-- launches the group FIFO head only when every one of its target workers is
-  idle — all members or none, with no partial reservation;
-- dispatches the head of each idle worker's single FIFO;
+- checks the group FIFO head and either launches every member together or
+  reserves all of its target workers until the complete set is idle;
+- dispatches the head of each idle, unreserved worker's single FIFO;
 - dispatches READY SUB work through the existing free-selection path.
 
-Only the group FIFO head is examined, and a launchable group is tried before
-conflicting singles. The runtime adds no fairness, aging, priority, or
-reservation policy; callers choose worker sets that make acceptable progress.
+Only the group FIFO head reserves workers. A reservation blocks new singles on
+the group's targets while their running work drains; singles on other workers
+continue normally. The reservation is released when the complete group
+launches. Later groups remain behind the FIFO head and do not reserve workers.
 
 ## Invariants
 
@@ -62,13 +63,15 @@ reservation policy; callers choose worker sets that make acceptable progress.
   terminal state.
 - The NEXT_LEVEL worker set is fixed after `init()`; a worker lookup that fails
   after submit-time validation is an invariant violation, not a fallback.
+- A NEXT_LEVEL single does not wait for a not-yet-dispatched peer at the same
+  scheduler level. Work that requires concurrent placement uses the group API.
 - No queue or scheduler mutex is held while an endpoint executes.
 
 ## Non-goals
 
 - No SUB worker selection.
 - No priorities, work stealing, rebinding, queue scanning, aging, quotas,
-  preemption, starvation prevention, or partial group reservation.
+  preemption, or fairness beyond the group-head reservation.
 - No compatibility flags, environment variables, or macros.
 
 ## Related documents

docs/orchestrator.md

@@ -239,6 +239,10 @@ A group task is a single DAG node that executes in parallel on N workers.
 Each worker gets its own `TaskArgs`; the node only reaches COMPLETED when all
 N finish.
 
+Callers submit tasks that wait for same-level peers as one complete group.
+Submitting those members as independent singles can start one member before
+its peers are READY, leaving the running member unable to finish.
+
 ```cpp
 SubmitResult Orchestrator::submit_next_level_group(
     const CallableIdentity &callable, const std::vector<TaskArgs> &args_list,
@@ -259,10 +263,11 @@ routing.
 
 At dispatch time the Scheduler checks the group FIFO head and resolves every
 entry in `workers` to that exact stable worker ID. It dispatches only if the
-entire target set is idle; a blocked group reserves no partial worker set and
-does not cause a scan past the FIFO head. Each WorkerThread runs `worker->run`
-with its own `task_args_list[i]`. Completion remains aggregated at the group
-slot, so downstream consumers are released once after every member is terminal.
+entire target set is idle. A blocked group reserves all of its targets against
+new singles but does not cause a scan past the FIFO head. Each WorkerThread
+runs `worker->run` with its own `task_args_list[i]`. Completion remains
+aggregated at the group slot, so downstream consumers are released once after
+every member is terminal.
 
 ---
 

docs/scheduler.md

@@ -54,8 +54,8 @@ while (true) {
         on_task_complete(item);
     }
 
-    dispatch_next_level_group();
-    dispatch_next_level_singles();
+    reserved_workers = dispatch_next_level_group();
+    dispatch_next_level_singles(reserved_workers);
     dispatch_sub_ready();
 
     if (stop_requested && all workers are idle) {
@@ -86,8 +86,8 @@ if target worker is idle and its FIFO is non-empty:
 ```
 
 There is no idle-worker search, rebinding, work stealing, or scan into another
-worker's queue. FIFO is independent per worker, so a busy worker A does not
-block READY work for worker B.
+worker's queue. Outside a group-head reservation, each worker FIFO progresses
+independently, so a busy worker A does not block READY work for worker B.
 
 ### Group tasks
 
@@ -103,13 +103,15 @@ if every target is idle:
     dispatch member i to target_worker_ids[i]
 else:
     leave the group at the FIFO head
+    reserve every target against single-task dispatch
 ```
 
-The check is all-or-nothing: a blocked group reserves no partial worker set.
-The Scheduler does not scan later groups. It continues to single-task queues,
-so the runtime adds no fairness, aging, priority, or reservation policy beyond
-trying a launchable group before singles in each iteration. Users are
-responsible for choosing worker sets that make acceptable progress.
+Dispatch remains all-or-nothing: no group member launches until the complete
+target set is idle. A blocked FIFO head reserves every target against new
+single-task dispatch while existing work drains. Singles on other workers
+continue normally, and later groups do not reserve workers because the
+Scheduler does not scan past the FIFO head. The reservation is released when
+the group launches.
 
 ## 4. SUB dispatch
 
@@ -184,6 +186,9 @@ The scheduling invariants are:
 5. Only the Scheduler calls `WorkerThread::dispatch`.
 6. Only one successful PENDING-to-READY transition enqueues a consumer.
 7. A group produces one aggregate DAG completion regardless of member count.
+8. A NEXT_LEVEL single does not wait for a not-yet-dispatched peer at the same
+   scheduler level; work that requires concurrent placement is submitted as a
+   group.
 
 ## 8. Related documents
 

examples/workers/l3/allreduce/README.md

@@ -9,8 +9,8 @@ Walks through the full flow:
 3. **`worker.init()`** — fork chip children; lazy base-communication init
 4. **`orch.allocate_domain(...)`** — allocate a communication domain with a
    `CommBufferSpec` scratch window
-5. **`orch.submit_next_level(chip_handle, chip_args, cfg, worker=i)`** —
-   submit the allreduce task for each rank
+5. **`orch.submit_next_level_group(chip_handle, args_list, cfg, workers=...)`** —
+   submit every mutually waiting allreduce rank as one group
 6. **`worker.run(orch_fn, ...)`** — execute the DAG and golden-check against
    the known expected sum
 

examples/workers/l3/allreduce/main.py

@@ -151,6 +151,7 @@ def orch_fn(orch, _args, cfg):
                 window_size=window_size,
                 buffers=[CommBufferSpec(name="scratch", dtype="float32", count=float_elems, nbytes=scratch_nbytes)],
             ) as handle:
+                args_list = []
                 for i in range(nranks):
                     domain = handle[i]
                     chip_args = TaskArgs()
@@ -167,7 +168,8 @@ def orch_fn(orch, _args, cfg):
                     )
                     chip_args.add_scalar(domain.domain_size)
                     chip_args.add_scalar(domain.device_ctx)
-                    orch.submit_next_level(chip_handle, chip_args, cfg, worker=i)
+                    args_list.append(chip_args)
+                orch.submit_next_level_group(chip_handle, args_list, cfg, workers=list(range(nranks)))
 
         print(f"[allreduce] running {nranks}-chip allreduce DAG...")
         worker.run(orch_fn, args=None, config=CallConfig())

examples/workers/l3/domain_rank_map/README.md

@@ -23,7 +23,9 @@ tail = workers [1, 2]        # chip 2 is in both
 After the inspection pass, each domain runs its own small allreduce — in its
 **own `worker.run()`**, so chip 2 never juggles two collectives at once. That
 separation is deliberate; see `dual_domain_overlap` for the case where both
-domains are live across the same DAG.
+domains are live across the same DAG. The ranks within each allreduce are one
+`submit_next_level_group`, so every peer that participates in the device
+barrier is dispatched as a complete set.
 
 ## Run
 

examples/workers/l3/domain_rank_map/main.py

@@ -210,13 +210,15 @@ def _orch_fn(orch, _args, cfg):
                 window_size=WINDOW_SIZE,
                 buffers=_scratch_buffers(),
             ) as handle:
+                args_list = []
                 for worker_idx in DOMAINS[domain_name]:
                     domain = handle[worker_idx]
                     args = TaskArgs()
                     args.add_tensor(make_tensor_arg(host_inputs[worker_idx]), TensorArgType.INPUT)
                     args.add_tensor(make_tensor_arg(outputs[domain_name][worker_idx]), TensorArgType.OUTPUT_EXISTING)
                     _add_domain_scratch(args, domain)
-                    orch.submit_next_level(allreduce_handle, args, cfg, worker=worker_idx)
+                    args_list.append(args)
+                orch.submit_next_level_group(allreduce_handle, args_list, cfg, workers=DOMAINS[domain_name])
 
         return _orch_fn
 

examples/workers/l3/dual_domain_overlap/README.md

@@ -20,7 +20,7 @@ through both and then computes on the results.
 | ------- | --- |
 | **Per-domain identity** | Chip 1 is rank 1 in `left` and rank 0 in `right`. The `workers` list order defines the dense rank, so the same chip legitimately holds two different ranks at once. |
 | **Domains allocated inside the orch function** | `with orch.allocate_domain(name=..., workers=..., window_size=..., buffers=[CommBufferSpec(...)])` — created and released within one orchestration, not configured on the `Worker`. |
-| **`submit_next_level_group`** | The affine stage submits one `TaskArgs` per member in a single call, with `workers=worker_indices`, instead of a loop of `submit_next_level`. |
+| **`submit_next_level_group`** | Both the peer-waiting allreduce and the affine stage submit one `TaskArgs` per member in a single call with `workers=worker_indices`. |
 | **Compute that depends only on its own domain's result** | Each affine task reads `reduce_out[domain][chip]`. The dependency is implicit — same `buffer.addr` as the reduce output — so `left`'s affine work can never consume `right`'s reduction. |
 
 ## Run

examples/workers/l3/dual_domain_overlap/main.py

@@ -208,6 +208,7 @@ def _orch_fn(orch, _args, cfg):
                     window_size=WINDOW_SIZE,
                     buffers=_scratch_buffers(),
                 ) as handle:
+                    args_list = []
                     for worker_idx in worker_indices:
                         domain = handle[worker_idx]
                         print(
@@ -221,7 +222,8 @@ def _orch_fn(orch, _args, cfg):
                             make_tensor_arg(reduce_out[domain_name][worker_idx]), TensorArgType.OUTPUT_EXISTING
                         )
                         _add_domain_scratch(args, domain)
-                        orch.submit_next_level(allreduce_handle, args, cfg, worker=worker_idx)
+                        args_list.append(args)
+                    orch.submit_next_level_group(allreduce_handle, args_list, cfg, workers=worker_indices)
 
             return _orch_fn
 

examples/workers/l3/ep_dispatch_combine/README.md

@@ -56,6 +56,7 @@ scatters into `routed_y_buf` without clearing it first.
 | Concept | How |
 | ------- | --- |
 | **Three children under one `ChipCallable`** | `children=[(0, dispatch), (1, local_expert), (2, combine)]` — the integers are `func_id`s, matching the `rt_submit_aiv_task(0/1/2, …)` calls in the orchestration. Each child declares only the args it consumes; the orchestration signature is the union. |
+| **Collective group dispatch** | The complete per-rank callable is submitted as one NEXT_LEVEL group because its dispatch and combine phases wait on peer ranks. |
 | **Ordering without dependencies** | The three tasks run back-to-back because `rt_submit_aiv_task` dispatches in submission order, not because any tensor edge forces it. |
 | **Chaining through host-backed tensors** | `recv_x_out` / `recv_w_out` / `recv_count_out` are `OUTPUT_EXISTING` for dispatch and inputs to `local_expert`; `recv_y` likewise feeds `combine`. |
 | **A hand-laid-out window** | `SCRATCH_NBYTES` sums every region — counts table, two signal areas, three receive windows, the combine push destination, a third signal — and must match the `kOff*` offsets in the kernels. |

examples/workers/l3/ep_dispatch_combine/main.py

@@ -548,6 +548,7 @@ def orch_fn(orch, _args, cfg):
                     )
                 ],
             ) as handle:
+                args_list = []
                 for i in range(nranks):
                     domain = handle[i]
                     print(
@@ -577,7 +578,8 @@ def orch_fn(orch, _args, cfg):
                     )
                     chip_args.add_scalar(domain.domain_size)
                     chip_args.add_scalar(domain.device_ctx)
-                    orch.submit_next_level(chip_handle, chip_args, cfg, worker=i)
+                    args_list.append(chip_args)
+                orch.submit_next_level_group(chip_handle, args_list, cfg, workers=list(range(nranks)))
 
         print("[ep_dispatch] running 2-chip dispatch DAG...")
         worker.run(orch_fn, args=None, config=CallConfig())

examples/workers/l3/ffn_tp_parallel/README.md

@@ -17,6 +17,7 @@ Per rank:
 | Concept | How |
 | ------- | --- |
 | **Implicit producer/consumer edge** | `host_partial[i]` is `OUTPUT_EXISTING` on the stage-1 submit and `INPUT` on the stage-2 submit. Both carry the same `buffer.addr`, so TensorMap links the two tasks itself — there is no barrier, no event, and no ordering call in the orch function. |
+| **Collective group dispatch** | Stage 2 is one `submit_next_level_group`, so it becomes READY after every rank's stage-1 output exists and dispatches all mutually waiting ranks together. |
 | **Mixed core types in one DAG** | Stage 1 compiles with `core_type="aic"` and its orchestration calls `rt_submit_aic_task`; stage 2 uses `core_type="aiv"` and `rt_submit_aiv_task`. One `Worker`, one `run()`. |
 | **`func_id`, not core id** | The integer in `children=[(0, core_callable)]` is the `func_id` the orchestration passes to `rt_submit_*_task(func_id, params)` — here `0` for the matmul and `1` for the reduce. It selects *which child kernel*, not which core type. |
 | **Cross-rank exchange through a domain buffer** | The stage-2 kernel reduces over a `scratch` buffer in the communication window: a mailbox of `nranks * M * N` floats followed by a signal tail of `nranks` int32 slots. |

examples/workers/l3/ffn_tp_parallel/main.py

@@ -201,6 +201,7 @@ def orch_fn(orch, _args, cfg):
                 window_size=window_size,
                 buffers=[CommBufferSpec(name="scratch", dtype="float32", count=scratch_count, nbytes=scratch_nbytes)],
             ) as handle:
+                allreduce_args = []
                 for i in range(nranks):
                     domain = handle[i]
                     print(
@@ -233,7 +234,8 @@ def orch_fn(orch, _args, cfg):
                     )
                     a2.add_scalar(domain.domain_size)
                     a2.add_scalar(domain.device_ctx)
-                    orch.submit_next_level(allreduce_handle, a2, cfg, worker=i)
+                    allreduce_args.append(a2)
+                orch.submit_next_level_group(allreduce_handle, allreduce_args, cfg, workers=list(range(nranks)))
 
         print("[ffn_tp_parallel] running 2-chip 2-stage DAG...")
         worker.run(orch_fn, args=None, config=CallConfig())

src/common/hierarchical/scheduler.cpp

@@ -11,7 +11,6 @@
 
 #include "scheduler.h"
 
-#include <algorithm>
 #include <stdexcept>
 #include <utility>
 
@@ -350,8 +349,8 @@ void Scheduler::try_consume(TaskSlot slot) {
 // sched_thread_ with no surrounding handler, any throw is fatal to the whole
 // worker tree (std::terminate), not a per-task failure.
 void Scheduler::dispatch_ready() {
-    dispatch_next_level_group();
-    dispatch_next_level_singles();
+    const std::unordered_set<int32_t> reserved_worker_ids = dispatch_next_level_group();
+    dispatch_next_level_singles(reserved_worker_ids);
     dispatch_sub_ready();
 }
 
@@ -392,7 +391,7 @@ void Scheduler::dispatch_sub_ready() {
     }
 }
 
-void Scheduler::dispatch_next_level_group() {
+std::unordered_set<int32_t> Scheduler::dispatch_next_level_group() {
     TaskSlot slot;
     while (cfg_.ready_next_level_queues->try_front_group(slot)) {
         TaskSlotState &s = *cfg_.ring->slot_state(slot);
@@ -408,18 +407,22 @@ void Scheduler::dispatch_next_level_group() {
         const int32_t group_size = s.group_size();
         std::vector<WorkerThread *> workers;
         workers.reserve(static_cast<size_t>(group_size));
+        std::unordered_set<int32_t> target_worker_ids;
+        target_worker_ids.reserve(static_cast<size_t>(group_size));
+        bool all_workers_idle = true;
         for (int32_t i = 0; i < group_size; ++i) {
             const int32_t worker_id = s.target_worker_id(i);
             WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id);
             if (worker == nullptr) {
                 throw std::runtime_error("Scheduler::dispatch_next_level_group: invalid target worker");
             }
-            if (std::find(workers.begin(), workers.end(), worker) != workers.end()) {
+            if (!target_worker_ids.insert(worker_id).second) {
                 throw std::runtime_error("Scheduler::dispatch_next_level_group: duplicate target worker");
             }
-            if (!worker->idle()) return;
+            if (!worker->idle()) all_workers_idle = false;
             workers.push_back(worker);
         }
+        if (!all_workers_idle) return target_worker_ids;
 
         TaskSlot popped;
         if (!cfg_.ready_next_level_queues->try_pop_group(popped) || popped != slot) {
@@ -432,10 +435,13 @@ void Scheduler::dispatch_next_level_group() {
             workers[static_cast<size_t>(i)]->dispatch(WorkerDispatch{slot, i});
         }
     }
+    return {};
 }
 
-void Scheduler::dispatch_next_level_singles() {
+void Scheduler::dispatch_next_level_singles(const std::unordered_set<int32_t> &reserved_worker_ids) {
     for (int32_t worker_id : cfg_.ready_next_level_queues->worker_ids()) {
+        if (reserved_worker_ids.find(worker_id) != reserved_worker_ids.end()) continue;
+
         WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id);
         if (worker == nullptr) {
             throw std::runtime_error(

src/common/hierarchical/scheduler.h

@@ -40,6 +40,7 @@
 #include <queue>
 #include <string>
 #include <thread>
+#include <unordered_set>
 
 #include "types.h"
 
@@ -103,7 +104,7 @@ class Scheduler {
     void poison_task(TaskSlot slot, const std::string &root_message);
     void try_consume(TaskSlot slot);
     void dispatch_ready();
-    void dispatch_next_level_group();
-    void dispatch_next_level_singles();
+    std::unordered_set<int32_t> dispatch_next_level_group();
+    void dispatch_next_level_singles(const std::unordered_set<int32_t> &reserved_worker_ids);
     void dispatch_sub_ready();
 };