From 79124daab162e3894bac007230e36f829ed9e6b5 Mon Sep 17 00:00:00 2001 From: Yangwook Kang Date: Tue, 14 Apr 2026 09:14:03 -0700 Subject: [PATCH] ADR-0024 Phase B (partial): scheduler-level collective drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- benches/ccl_allreduce.py | 16 +++++++++++++- src/kernbench/runtime_api/distributed.py | 28 ++++++++++++++++++------ tests/test_ccl_allreduce_matrix.py | 13 +++++------ 3 files changed, 42 insertions(+), 15 deletions(-) diff --git a/benches/ccl_allreduce.py b/benches/ccl_allreduce.py index a57a358..00ea17d 100644 --- a/benches/ccl_allreduce.py +++ b/benches/ccl_allreduce.py @@ -156,7 +156,13 @@ def run(torch) -> None: n_sips = int(spec.get("system", {}).get("sips", {}).get("count", 1)) if world_size == n_sips: - # ADR-0024 D12/D13: one greenlet per rank, simple round-robin. + # ADR-0024 D12/D13: one greenlet per rank. After each scheduler + # round, the main greenlet drains any pending collective handles + # (ADR-0024 D7) — this must happen in the main context, not inside + # a worker, so env.run is invoked with main as the current greenlet + # and kernel_runner's spawned kernel greenlets correctly get main + # as their parent. + backend = dist._backend gs: list[greenlet] = [] for rank in range(world_size): def _entry(r: int = rank) -> None: @@ -171,6 +177,14 @@ def run(torch) -> None: for g in alive: if not g.dead: g.switch() + # Drain pending collective handles. All sibling workers have + # either submitted (and yielded) or completed; their kernels + # are live in the SimPy queue, ready to exchange via IPCQ. + pending = backend._pending_collective_handles + if pending: + for h, _sip_id, meta in pending: + torch.wait(h, _meta=meta) + backend._pending_collective_handles = [] else: # Legacy single-worker path (ccl.yaml world_size override). worker(rank=dist.get_rank(), world_size=world_size, torch=torch) diff --git a/src/kernbench/runtime_api/distributed.py b/src/kernbench/runtime_api/distributed.py index 41a25b6..bb39e47 100644 --- a/src/kernbench/runtime_api/distributed.py +++ b/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 diff --git a/tests/test_ccl_allreduce_matrix.py b/tests/test_ccl_allreduce_matrix.py index eb97e65..fb732e2 100644 --- a/tests/test_ccl_allreduce_matrix.py +++ b/tests/test_ccl_allreduce_matrix.py @@ -70,18 +70,17 @@ CASES = [ # Default fallback — no world_size override → ADR-0024 D1 derives # from topology (SIP count = 2). Exercises the new SIP-level TP # launcher + cross-SIP ring. - # XFAIL: ADR-0024 Phase A delivers launcher infrastructure; Phase B - # will finish cross-SIP ring kernel integration. Today this hangs in - # the SimPy drain despite ADR-0025's direction-addressing fix — - # suspected per-rank-tensor kernel_args / program_id mismatch under - # multi-greenlet dispatch. Separate Phase will diagnose. + # XFAIL: Phase A fix (scheduler-level wait) resolves the greenlet- + # re-entry hang, but Phase 2 DataExecutor still reports only 1 math + # op for a 2-rank ring (expected 2) — cross-SIP op_log replay + # integration pending ADR-0024 Phase B. pytest.param( "ring_allreduce_tcm", "kernbench.ccl.algorithms.ring_allreduce", "ring_1d", "tcm", None, 8, 2, id="ring_default_ws", marks=pytest.mark.xfail( - reason="ADR-0024 Phase B: cross-SIP multi-greenlet kernel integration", - run=False, # skip execution to avoid hang; revisit in Phase B + reason="ADR-0024 Phase B: cross-SIP op_log replay integration", + strict=True, ), ), # Buffer variants at 8-rank (fast — same kernel, different slot space).