ADR-0024 Phase B (partial): scheduler-level collective drain

Root cause (hang diagnosis):
`kernel_runner.run()` captures `greenlet.getcurrent()` at spawn time as
the kernel greenlet's `_parent`. When a worker greenlet (say g0) calls
`dist.all_reduce` → `ctx.wait(h)` → `env.run(until=h0)`, the SimPy
scheduler steps pe_cpu processes, which in turn spawn kernel greenlets.
Those kernels' `_parent` becomes g0 (current greenlet at spawn). When a
kernel yields via switch_to_simpy, control jumps back up to g0's LAST
switch point — which is the main scheduler's `g.switch()` call — rather
than the kernel_runner's generator frame. Main then re-enters its
`for g in alive: g.switch()` loop mid-wait, producing nested greenlet
re-entry. Scheduler spins: g0 never completes, g1 appears to complete
out of order, infinite loop at 100% CPU.

Fix:
- AhbmCCLBackend.all_reduce: in multi-greenlet mode, submit via
  launch(_defer_wait=True), extend backend._pending_collective_handles,
  and yield to the parent greenlet. Worker does NOT call wait.
- benches/ccl_allreduce.py run(): after each scheduler round, the MAIN
  greenlet drains backend._pending_collective_handles. This keeps
  env.run invocation in the main context, so kernel_runner's spawned
  kernel greenlets have main as their _parent — no nested re-entry.
- Legacy single-driver path (no bench scheduler): all_reduce falls back
  to inline wait when g.parent is None.

Result:
- Multi-greenlet cross-SIP ring no longer hangs (was 100% CPU infinite
  loop in kernel_runner._switch_kernel).
- ring_default_ws still xfail(strict=True): now fails as a data
  correctness issue — DataExecutor reports only 1 math op for a 2-rank
  ring (expected 2). Cross-SIP op_log replay integration is the
  remaining Phase B task.

514 passed, 1 xfailed (strict).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/kernbench/runtime_api/distributed.py

@@ -40,6 +40,12 @@ class AhbmCCLBackend:
         self._merged = resolve_algorithm_config(self._cfg_all)
         self._algo_module = importlib.import_module(self._merged["module"])
         self._world_size = self._resolve_world_size()
+        # ADR-0024 D7: handles pending drain by the main scheduler.
+        # Worker greenlets extend this list after submitting their collective
+        # kernel, then yield. The bench `run()` loop drains the list after
+        # all workers yielded (so all sibling kernels are live in SimPy
+        # before any rank waits, avoiding cross-rank deadlock).
+        self._pending_collective_handles: list = []
 
         # Eager IPCQ install — ``init_process_group`` time. Mirrors NCCL
         # communicator creation: done once, reused across every subsequent
@@ -103,11 +109,13 @@ class AhbmCCLBackend:
         n_elem = shards[0].nbytes // tensor.itemsize
         kernel_fn = self._algo_module.kernel
         kernel_args = self._algo_module.kernel_args(self._world_size, n_elem)
-        # ADR-0024 D7: submit + yield + wait. All sibling ranks must submit
-        # their CCL kernels before any of them starts waiting, otherwise the
-        # first rank's wait drains SimPy while peer kernels are missing →
-        # IpcqDeadlock. The yield hands control back to the bench scheduler
-        # so other worker greenlets can submit too.
+        # ADR-0024 D7: submit + yield. When running under the multi-greenlet
+        # bench launcher, the scheduler (not the worker) drains the pending
+        # handles. This is required because env.run must be invoked from the
+        # MAIN greenlet — otherwise kernel_runner's spawned kernel-greenlet
+        # captures the worker-greenlet as its `_parent`, and kernel
+        # switch_to_simpy() returns control to the main scheduler loop
+        # mid-wait, causing nested re-entry and the scheduler to spin.
         pending = self.ctx.launch(
             self._merged["algorithm"], kernel_fn, tensor, *kernel_args,
             _defer_wait=True,
@@ -115,9 +123,15 @@ class AhbmCCLBackend:
         from greenlet import getcurrent
         g = getcurrent()
         if g.parent is not None and not g.parent.dead:
+            # Multi-greenlet mode: hand pending to the backend-level queue so
+            # the main scheduler drains. Worker just yields.
+            self._pending_collective_handles.extend(pending)
             g.parent.switch()
-        for h, _sip_id, meta in pending:
-            self.ctx.wait(h, _meta=meta)
+            # On resume, all pending handles have been drained by main.
+        else:
+            # Single-driver (no bench scheduler): drain inline.
+            for h, _sip_id, meta in pending:
+                self.ctx.wait(h, _meta=meta)
 
     def barrier(self) -> None:
         # Single-driver model → no cross-process sync needed. Keeping the