From 8ae76c6d50af603ee6d2af4f099c8ca15d4e59ec Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:25:57 +0800 Subject: [PATCH 1/2] CollectiveX: bind one nccl-ep handle per group, not one per token count The nccl-ep adapter created a handle per problem shape and cached it on the problem, so a token ladder left up to nine handles live on one EP group and the randomised trial order interleaved them. That is not safe: `buffer_idx`, the LL double-buffer parity selector, is per-HANDLE state, but the buffers it selects are offsets into the per-GROUP rdma_buffer. Two handles built from the same group config resolve to the SAME parity-0/parity-1 count and flag slots (which themselves share one offset) and advance their parity independently, so one handle's "next buffer, safe to clean" is the other's "current buffer, in flight". Cleans land on live signals and polls observe a sibling handle's leftovers. That corrupts the signalling whether or not it hangs outright, which makes every latency drawn from a multi-handle run suspect -- not just the runs that wedged. On a stock nccl4py 0.3.1 wheel, GB200, 4 ranks on one tray, with the handle count as the only variable: decode ladder [1] (one handle) completes with zero receive timeouts and correctness passing, while ladder [1, 2] (two handles) takes 64 dispatch and 6 combine receive timeouts and dies with cudaErrorLaunchFailure. Filed upstream as NVIDIA/nccl#2303. Rebinding is the intended lifecycle: `ncclEpInitHandle` takes no token count, and `ncclEpUpdateHandle` is documented as a "per-step collective: prepare the handle for the given top-k routing decisions". The wheel exposes it as `Handle.update`, which rebinds top-k indices without reallocating buffers. This also brings nccl-ep in line with the other three backends, which each allocate one object sized to the ladder maximum and vary the token count per call. Keep the per-problem namespace -- it holds only host-side tensor descriptors and, for HT, that problem's counters -- and move the handle to the group, rebinding when the bound shape changes. Both create and update are collective, so they must run in the same order on every rank; they only run on a shape change, and because every timed component is preceded by an untimed warm() on its own problem, the collective always lands in warm or in the oracle passes rather than inside a timed window. Re-entering the already-bound problem returns without a collective or a sync, so the steady state of a timed loop is unchanged. HT re-reads its received-token count on rebind because the metadata exchange recomputes it for that routing; the value is deterministic per problem, so returning to an earlier shape reproduces it. Teardown is now a single explicit `Handle.destroy` instead of waiting on per-problem GC. The new tests stub torch and nccl so they run without the benchmark image, and lock the contract: one create_handle across nine rungs, update on a shape switch, no collective when the bound shape is re-entered, no layout_info on an LL rebind, HT rebinds carrying that problem's own counters, and idempotent teardown. They withdraw their stubs after importing the adapter so the genuinely torch-dependent modules still skip rather than error. --- experimental/CollectiveX/bench/ep_nccl.py | 90 +++++++-- .../CollectiveX/tests/test_ep_nccl_handle.py | 171 ++++++++++++++++++ 2 files changed, 242 insertions(+), 19 deletions(-) create mode 100644 experimental/CollectiveX/tests/test_ep_nccl_handle.py diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index f1ce89c056..14b4c2de9e 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -114,6 +114,9 @@ def __init__(self, args, rank, world_size, local_rank, device): self._combine_cfg = CombineConfig(send_only=0) self._comm = None self._ep_group = None + # Exactly ONE handle per group, rebound per problem shape — see _ensure_handle. + self._handle = None + self._bound = None def buffer_cap(self, args): if self._ll: @@ -238,15 +241,34 @@ def create_buffer(self, spec): self._ht_disp_counts_t = self._t(self._ht_disp_counts) def _ensure_handle(self, p): - """Create (once, cached on the problem) the reusable per-step handle for p's routing. - - create_handle is collective and, in HT, performs the metadata exchange that fixes the - received-token count — so it must run in the same order on every rank. It first runs - inside Pass 1's untimed warm (ladder order, identical across ranks); the timed passes - then reuse the cached handle and never enter a collective here. + """Bind the group's single handle to p's routing, creating it on first use. + + ONE handle per group, rebound per shape — never one handle per shape. `buffer_idx`, the + LL double-buffer parity selector, is per-HANDLE state, but the buffers it selects are + offsets into the per-GROUP rdma_buffer: two handles built from the same group config + resolve to the SAME parity-0/parity-1 count+flag slots and advance their parity + independently, so one handle's "next buffer, safe to clean" is the other's "current + buffer, in flight". Interleaving handles therefore corrupts the signalling even when it + does not hang outright, which makes every latency drawn from such a run suspect. Filed + upstream as NVIDIA/nccl#2303; reproduced on a stock wheel by ladder [1] (one handle, + clean) vs ladder [1, 2] (two handles, 64 dispatch + 6 combine receive timeouts -> 719). + + `ncclEpInitHandle` takes no token count and `ncclEpUpdateHandle` is documented as a + "per-step collective: prepare the handle for the given top-k routing decisions", so + rebinding IS the intended lifecycle. This also matches the other three backends, which + each allocate one object sized to the ladder maximum and vary the token count per call. + + Both create_handle and update are collective, and HT additionally performs the metadata + exchange that fixes the received-token count, so they must run in the same order on + every rank. They only ever run on a shape CHANGE, and every timed component is preceded + by an untimed warm() on its own problem — so the collective always lands in warm (or in + the oracle passes), never inside a timed window. Re-entering with the already-bound + problem returns immediately without a collective or a sync. """ cached = getattr(p, "_nccl", None) if cached is not None: + if self._bound is not cached: + self._rebind(cached) return cached stream = self._stream() topk_idx_t = self._t(p.topk_idx) @@ -271,18 +293,45 @@ def _ensure_handle(self, p): expert_counters=self._t(h.recv_experts), recv_total_counter=self._t(h.recv_total), ) - h.handle = self._ep_group.create_handle( - self._layout, - topk_idx_t, - layout_info=ht_layout_info, - config=HandleConfig(), - stream=stream, + # LL takes layout_info only on dispatch (the API forbids it on create/update); HT needs + # it here so this problem's counters receive its own metadata-exchange results. + h.layout_info = ht_layout_info + if self._handle is None: + self._handle = self._ep_group.create_handle( + self._layout, + topk_idx_t, + layout_info=ht_layout_info, + config=HandleConfig(), + stream=stream, + ) + h.handle = self._handle + torch.cuda.synchronize() + if not self._ll: + h.count = int(h.recv_total.item()) + self._bound = h + else: + h.handle = self._handle + self._rebind(h) + p._nccl = h + return h + + def _rebind(self, h): + """Point the single handle at h's routing (collective; untimed callers only). + + Rebinding does not reallocate: `Handle.update` only swaps in the new top-k indices. + HT re-reads its received-token count because the metadata exchange recomputes it for + this routing; the value is deterministic per problem, so a later rebind to the same + problem reproduces it. + """ + self._handle.update( + h.topk_idx_t, + layout_info=None if self._ll else h.layout_info, + stream=self._stream(), ) torch.cuda.synchronize() if not self._ll: h.count = int(h.recv_total.item()) - p._nccl = h - return h + self._bound = h # ---- transport contract ------------------------------------------------------------------ @@ -464,8 +513,11 @@ def finalize(self, rc): return rc def _destroy_handles(self): - # Per-problem handles are cached on the problem namespaces; the group keeps no registry, - # so there is nothing to walk here. Handles are released when their problems are GC'd - # (Handle.destroy runs in the binding's __del__); the group/comm destroy below reclaims - # the device buffers. Kept as a seam in case bring-up needs explicit handle teardown. - return + # One handle for the whole group, so teardown is a single explicit destroy rather than + # waiting on per-problem GC. The problem namespaces still hold a reference to it for + # their dispatch/combine calls; dropping _bound first keeps a late rebind from touching + # a destroyed handle. The group/comm destroy below reclaims the device buffers. + self._bound = None + if self._handle is not None: + self._handle.destroy() + self._handle = None diff --git a/experimental/CollectiveX/tests/test_ep_nccl_handle.py b/experimental/CollectiveX/tests/test_ep_nccl_handle.py new file mode 100644 index 0000000000..943625ee7d --- /dev/null +++ b/experimental/CollectiveX/tests/test_ep_nccl_handle.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +"""Torch-free tests for the nccl-ep single-handle contract. + +The invariant under test: ONE handle per group, rebound per problem shape. Two handles built +from the same group config resolve to the same LL parity signal slots while advancing their +parity independently, which corrupts signalling (NVIDIA/nccl#2303) -- so a regression that +reintroduces per-shape handles must fail loudly here rather than on a cluster. + +torch and nccl are stubbed so this runs without the benchmark image. +""" +from __future__ import annotations + +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[1] + + +def _stub_modules(): + """Fake torch / nccl modules so `import ep_nccl` succeeds without the benchmark image.""" + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.int32 = "int32" + torch.empty = lambda *a, **k: types.SimpleNamespace(shape=a[0] if a else ()) + torch.zeros = lambda *a, **k: types.SimpleNamespace(item=lambda: 7) + torch.cuda = types.SimpleNamespace(synchronize=lambda: None) + dist = types.ModuleType("torch.distributed") + torch.distributed = dist + + ep = types.ModuleType("nccl.ep") + for name in ( + "Algorithm", "CombineConfig", "CombineInputs", "CombineOutputs", "DispatchConfig", + "DispatchInputs", "DispatchOutputs", "GroupConfig", "HandleConfig", "Layout", + "LayoutInfo", "Tensor", + ): + setattr(ep, name, type(name, (), {"__init__": lambda self, *a, **k: None})) + ep.Algorithm = types.SimpleNamespace(LOW_LATENCY="LL", HIGH_THROUGHPUT="HT") + ep.Layout = types.SimpleNamespace(EXPERT_MAJOR="EM", FLAT="FLAT") + core = types.ModuleType("nccl.core") + pkg = types.ModuleType("nccl") + pkg.ep, pkg.core = ep, core + return { + "torch": torch, "torch.distributed": dist, + "nccl": pkg, "nccl.ep": ep, "nccl.core": core, + } + + +sys.path[:0] = [str(ROOT), str(ROOT / "bench")] + +# Import ep_nccl against the stubs, then withdraw them: leaving a fake torch in sys.modules +# makes the genuinely torch-dependent modules in this process (test_runtime, test_ll_oracle) +# error instead of skipping. Dropping ep_nccl too keeps the stub-built module private to us. +with mock.patch.dict(sys.modules, _stub_modules()): + import ep_nccl # noqa: E402 + + sys.modules.pop("ep_nccl", None) + + +class FakeHandle: + """Records every rebind so the tests can assert on the collective call pattern.""" + + def __init__(self): + self.updates = [] + self.destroyed = False + + def update(self, topk_idx, *, layout_info=None, stream=None): + self.updates.append((topk_idx, layout_info)) + + def destroy(self): + self.destroyed = True + + +class FakeGroup: + def __init__(self): + self.created = 0 + self.handle = FakeHandle() + + def create_handle(self, layout, topk_idx, *, layout_info=None, config=None, stream=None): + self.created += 1 + return self.handle + + +def backend(ll=True): + """An NCCLEPBackend with just the fields _ensure_handle touches (no __init__, no GPU).""" + b = object.__new__(ep_nccl.NCCLEPBackend) + b._ll = ll + b._layout = "EM" if ll else "FLAT" + b._handle = None + b._bound = None + b._ep_group = FakeGroup() + b.device = "cuda:0" + b.num_local_experts = 4 + b.args = types.SimpleNamespace(hidden=16) + b._t = lambda x: x + b._stream = lambda: 0 + return b + + +def problem(T): + return types.SimpleNamespace( + T=T, dispatch_x=f"x{T}", topk_idx=f"idx{T}", topk_weights=f"w{T}" + ) + + +class TestSingleHandle(unittest.TestCase): + def test_one_handle_across_many_shapes(self): + """Nine ladder rungs must still produce exactly one create_handle.""" + b = backend() + for T in (1, 2, 4, 8, 16, 32, 64, 128, 256): + b._ensure_handle(problem(T)) + self.assertEqual(b._ep_group.created, 1) + + def test_shape_change_rebinds_and_repeat_does_not(self): + """update() on a shape switch; no collective when the bound shape is re-entered.""" + b = backend() + pa, pb = problem(1), problem(2) + b._ensure_handle(pa) + self.assertEqual(len(b._ep_group.handle.updates), 0) # first bind is the create + + b._ensure_handle(pb) + self.assertEqual(len(b._ep_group.handle.updates), 1) + + # Re-entering the bound problem repeatedly -- the timed loop's steady state -- must not + # enter a collective, otherwise every iteration gains a rank-synchronising step. + for _ in range(8): + b._ensure_handle(pb) + self.assertEqual(len(b._ep_group.handle.updates), 1) + + # Returning to an earlier shape rebinds again (its cached namespace is reused). + b._ensure_handle(pa) + self.assertEqual(len(b._ep_group.handle.updates), 2) + + def test_every_problem_shares_the_one_handle(self): + b = backend() + handles = {id(b._ensure_handle(problem(T)).handle) for T in (1, 2, 4)} + self.assertEqual(len(handles), 1) + self.assertIs(b._ensure_handle(problem(1)).handle, b._handle) + + def test_ll_never_passes_layout_info_on_rebind(self): + """The API forbids layout_info on create/update in LL mode.""" + b = backend(ll=True) + b._ensure_handle(problem(1)) + b._ensure_handle(problem(2)) + self.assertEqual([info for _, info in b._ep_group.handle.updates], [None]) + + def test_ht_rebind_carries_that_problems_counters(self): + """HT re-runs the metadata exchange into the rebound problem's own counter tensors.""" + b = backend(ll=False) + ha = b._ensure_handle(problem(1)) + hb = b._ensure_handle(problem(2)) + self.assertEqual(len(b._ep_group.handle.updates), 1) + self.assertIs(b._ep_group.handle.updates[0][1], hb.layout_info) + self.assertIsNot(ha.layout_info, hb.layout_info) + self.assertEqual(hb.count, 7) # re-read after the exchange + + def test_destroy_releases_the_handle_once(self): + b = backend() + b._ensure_handle(problem(1)) + handle = b._handle + b._destroy_handles() + self.assertTrue(handle.destroyed) + self.assertIsNone(b._handle) + self.assertIsNone(b._bound) + b._destroy_handles() # idempotent + + +if __name__ == "__main__": + unittest.main() From 7d3b01d3f247c1a433142a537156d6007cbe0ab2 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:37:40 +0800 Subject: [PATCH 2/2] CollectiveX: restore the nccl-ep low-latency rows on all six NVIDIA SKUs The rows were dropped while every low-latency leg wedged, on the reading that the wheel's signal protocol was at fault and only a fixed wheel could bring them back. That reading was wrong about the cause. The wedge needs TWO low-latency handles live on one EP group. `buffer_idx`, the LL double-buffer parity selector, is per-handle state, but the buffers it selects are offsets into the per-group rdma_buffer, so two handles built from the same group config alias one another's parity-0/parity-1 count and flag slots while advancing their parity independently. The adapter used to create a handle per token count, which left up to nine of them interleaved on one group by the randomised trial order. It now binds one handle and rebinds it per shape (preceding commit), which removes the aliasing without touching the wheel. Isolated on a stock nccl4py 0.3.1 wheel with the handle count as the only variable: decode ladder [1] (one handle) completes with zero receive timeouts, ladder [1, 2] (two handles) takes 64 dispatch and 6 combine receive timeouts and dies with cudaErrorLaunchFailure. NVIDIA's own ep_bench never reproduced it across eight configurations -- barrier removed, injected rank stall, dispatch-only and combine-only loops, up to 20 000 iterations -- because it uses a single handle. The upstream report has been corrected accordingly (NVIDIA/nccl#2303). Low-latency is decode-only and EP8-only, so the GB SKUs carry an EP8 row here even though normal mode runs EP16 on them, and b300/gb200/gb300 gain their first ll_backends entry. Restoring the rows is what makes the claim testable: high-throughput CI already passes on the single-handle adapter (h100 EP8, decode and prefill, correctness green, latency within noise of the multi-handle baseline), but nothing exercises the low-latency path while the matrix emits no low-latency shard. If a leg still wedges, the rows must come back out and the aliasing was not the whole story. --- .../CollectiveX/configs/platform_config.json | 9 +++++--- experimental/CollectiveX/tests/test_matrix.py | 23 +++++++++++-------- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index bc17013fb4..3babb66536 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -10,7 +10,7 @@ "scale_up_transport": "nvlink", "launcher": "single-slurm", "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, - "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, + "ll_backends": {"deepep-v2": [8], "uccl-ep": [8], "nccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 2x200GbE", "switch": "Arista 7060DX5-64S (Tomahawk4, 25.6T)"}, "operator": { "partition": "hpc-gpu-1", @@ -34,7 +34,7 @@ "scale_up_transport": "nvlink", "launcher": "single-slurm", "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, - "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, + "ll_backends": {"deepep-v2": [8], "uccl-ep": [8], "nccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400G", "switch": "NVIDIA Quantum-2 QM9790 (25.6T, InfiniBand)"}, "operator": { "partition": "main", @@ -54,7 +54,7 @@ "scale_up_transport": "nvlink", "launcher": "single-slurm", "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, - "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, + "ll_backends": {"deepep-v2": [8], "uccl-ep": [8], "nccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400GbE", "switch": "Whitebox Tomahawk3 leaf + Tomahawk4 (RoCE)"}, "operator": { "partition": "gpu-2", @@ -77,6 +77,7 @@ "scale_up_transport": "nvlink", "launcher": "single-slurm", "backends": {"deepep-v2": [8, 16], "nccl-ep": [8]}, + "ll_backends": {"nccl-ep": [8]}, "fabric": {"nic": "ConnectX-8 2x400GbE", "switch": "NVIDIA Spectrum-X SN5600 (51.2T)"}, "operator": { "partition": "batch_1", @@ -101,6 +102,7 @@ "scale_up_transport": "mnnvl", "launcher": "gb-nv", "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, + "ll_backends": {"nccl-ep": [8]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { "partition": "batch", @@ -120,6 +122,7 @@ "scale_up_transport": "mnnvl", "launcher": "gb-nv", "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, + "ll_backends": {"nccl-ep": [8]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { "partition": "batch_1", diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index 07e3749f66..83d3c2a4b7 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -213,14 +213,15 @@ def test_nccl_ep_rollout_shape(self): # both stay inside the 72-GPU scale-up domain (world <= scale_up_domain => LSA, no GIN), # so EP16 works over MNNVL where the RDMA-GIN path walls. # * AMD SKUs: no rows (NCCL EP is NVIDIA-only; AMD runs mori). - # * LOW-LATENCY: no rows on ANY SKU. The wheel's LL count/flag protocol consumes stale - # double-buffer signals (values carry no generation, so a signal from two calls earlier is - # bit-identical at a repeating workload); a rank that slips one parity cycle gets lapped and - # the pipeline wedges on dispatch/combine receive timeouts, ending in cudaErrorLaunchFailure. - # Reported as NVIDIA/nccl#2303, fixed by NVIDIA/nccl#2306, unfixable here (we install the - # published wheel). Observed first on GB, then reproduced on every x86 SKU — 5/5 over SSH and - # 4/4 in sweep 30155842613 — so ll_backends carries no nccl-ep row anywhere until a fixed - # wheel ships. Restore per-SKU rows only alongside a spec bump that contains the fix. + # * LOW-LATENCY: EP8 on all six NVIDIA SKUs. These rows were dropped while every LL leg + # wedged, on the reading that only a fixed wheel could restore them. That reading was + # wrong about the cause: the wedge needed TWO LL handles live on one group. `buffer_idx` + # is per-handle but its buffers are offsets into the per-group rdma_buffer, so same-config + # handles alias one another's parity count/flag slots. The adapter now binds a single + # handle and rebinds it per shape, which removes the aliasing without a wheel bump — + # proven on a stock wheel by ladder [1] (one handle, clean) vs [1, 2] (two handles, + # 64 dispatch + 6 combine receive timeouts). LL is decode-only and EP8-only: GB stays at + # EP8 here even though normal mode runs EP16, because LL adds no EP16 row on any SKU. # BF16 only — no FP8 case (NCCL EP FP8 unsupported this release). document = matrix(backend="all") runnable = { @@ -255,14 +256,16 @@ def test_nccl_ep_rollout_shape(self): }, {"bf16"}, ) - # Low-latency: no nccl-ep case on any SKU while the wheel carries the stale-signal wedge. + # Low-latency: EP8 on every NVIDIA SKU, and EP8 only — a stray EP16 LL row would dispatch a + # shape the mode does not define. ll = { (item["sku"], item["case"]["ep"]) for item in document["requested_cases"] if item["case"]["backend"] == "nccl-ep" and item["case"]["mode"] == "low-latency" + and item["disposition"] == "runnable" } - self.assertEqual(ll, set()) + self.assertEqual(ll, {(sku, 8) for sku in rdma_skus | gb_skus}) def test_invalid_filters_fail_closed(self): for options in (