CollectiveX: one nccl-ep handle per group, and restore the low-latency rows - #2407
Conversation
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.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
LGTM — this is a well-scoped, thoroughly-tested fix to the nccl-ep adapter's handle lifecycle.
What was reviewed: the group-level single-handle bind/rebind path (_ensure_handle/_rebind) for both HT and LL algorithms, the teardown ordering in _destroy_handles, and the new test suite's coverage of the create-once/rebind-on-shape-change/no-collective-on-reentry contract. Traced the per-problem _nccl cache vs. the new group-level _handle/_bound state and confirmed every code path (first-ever bind, first bind of a new shape, re-entry to an already-bound shape) updates _bound and calls create/update exactly as the docstring and tests describe. No correctness issues found beyond what the bug-hunting pass already reported (none).
Extended reasoning...
Overview
This PR changes the NCCL-EP benchmark adapter (experimental/CollectiveX/bench/ep_nccl.py) so that exactly one nccl-ep Handle is created per EP group and rebound (via Handle.update) across problem shapes, instead of creating and caching a distinct handle per token-count shape. It adds a new torch/nccl-stubbed unit test file that locks the create-once / rebind-on-shape-change / no-collective-on-reentry / idempotent-teardown contract. This is confined to experimental benchmarking code and does not touch any production inference path.
Security risks
None. This is a benchmark harness for measuring MoE all-to-all communication performance; there is no user input, auth, or externally reachable surface involved. The change is purely about internal handle/buffer lifecycle management against a vendored NCCL EP library.
Level of scrutiny
Moderate-to-high scrutiny is warranted given the subtlety of the underlying bug (per-handle double-buffer parity state aliasing over per-group buffer offsets, causing an NVIDIA library-level race that only manifests with multiple live handles), but the blast radius is limited: this is experimental benchmarking code, not a production critical path, and the PR description documents concrete on-metal evidence (GB200, reproducible receive-timeout failures with 2 handles vs. clean with 1) plus an upstream-filed bug (NVIDIA/nccl#2303). I traced the handle bind/rebind control flow by hand across all three cases (first-ever handle creation, first bind of a not-yet-seen shape once a handle already exists, and re-entry into an already-bound shape) and confirmed the _bound tracking and collective (create/update) calls are consistent with the documented invariants, matching what the new tests assert.
Other factors
The new test file (test_ep_nccl_handle.py) stubs torch/nccl and directly exercises the exact contract this PR is fixing — one create_handle across nine ladder rungs, update only on shape change, no collective on re-entry, no layout_info on LL rebind, HT counter carry-through on rebind, and idempotent teardown — which gives good confidence the logic matches intent. The PR explicitly notes no on-metal re-run has happened yet and that nccl-ep LL remains walled off in platform_config.json, so this change does not affect any currently-published benchmark numbers, further lowering risk. The bug-hunting system found no issues, and my own review did not surface anything beyond that.
CI validation (HT, h100 EP8 single-node)Both runs green; artifacts checked rather than just the status.
Round-trip latency,
Decode is flat within noise both ways, so the per-shape The prefill p99 improvement is consistent across all four rungs while p50 is flat, which is the signature of removed tail interference rather than a faster path: with one handle there are no sibling handles racing on the shared parity count/flag slots. That is suggestive, not established — two runs cannot separate it from pool variance, and I would not quote it as a result without repeats. Not covered: LL. |
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.
Oseltamivir
changed the title
…ling diff Review caught the same false claim I corrected in the README surviving in docs/methodology.md's Matrix section: NCCL EP "carries no ll_backends row on any SKU ... the cells stay out of the matrix until a fixed wheel ships". #2407 enabled those rows on all six NVIDIA SKUs. Corrected, with the B300/GB200/GB300 decode gap stated there too. Also tightened the diff after re-reading every added line: * The production/candidate definition was written out three times (ep_backend, ep_harness, sweep_matrix). It now lives once on EPBackend.maturity; the other two point at it. * The B300 low-latency explanation in the backend table ran to 945 characters inside one cell. Kept the diagnostic facts — IBGDA self-enabled on a single-node run, the ibgda.cpp:2234 address-handle failure, rc255, and that NVSHMEM_DISABLE_IB=1 was tested and does not help — dropped the four-step error chain, the narrative around the attempted fix, and a speculative sentence about what a real fix would need. * Dropped a parenthetical in the ll_backends note that repeated the backend table's own candidate cell. 90 -> 89 added lines, and the ones that remain each say something the others do not. 中文：修正 methodology 中过时的 Matrix 说明，并精简标注改动的 diff。 Review 发现我已在 README 中更正的同一处错误描述仍存在于 docs/methodology.md 的 Matrix 章节： 称 NCCL EP「在任何 SKU 上都没有 ll_backends 行……相关单元格在修复版 wheel 发布前不进入矩阵」。 而 #2407 已在全部六个 NVIDIA SKU 上启用这些行。现已更正，并同时写明 B300/GB200/GB300 的解码缺口。 在逐行复查新增内容后也做了精简： * production/candidate 的定义此前被写了三遍（ep_backend、ep_harness、sweep_matrix）。现在只 保留在 EPBackend.maturity 一处，另外两处指向它。 * 后端表中 B300 低延迟的说明在单个单元格里达到 945 字符。保留了诊断事实 —— 单节点运行下 IBGDA 自行启用、ibgda.cpp:2234 的 address handle 失败、rc255，以及 NVSHMEM_DISABLE_IB=1 经测试无效 —— 删去了四步错误链、围绕尝试修复的叙述，以及关于「真正的修复需要什么」的推测性句子。 * 删去 ll_backends 说明中一处与后端表 candidate 单元格重复的插入语。 新增行数从 90 降到 89，且保留下来的每一行都在陈述其他行没有说的内容。
1 participant
Why
The nccl-ep adapter created one handle per problem shape and cached it on the problem, so walking a token ladder left up to nine handles live on a single EP group, interleaved by the randomised trial order.
That is unsafe.
buffer_idx— the LL double-buffer parity selector — is per-handle state, but the buffers it selects are offsets into the per-grouprdma_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.This is the real cause of the low-latency wedge we had attributed to the wheel. Isolated on a stock
nccl4py[cu13]==0.3.1wheel, GB200, 4 ranks on one tray, with the handle count as the only variable:[1][1, 2]cudaErrorLaunchFailureNVIDIA's own
ep_benchnever reproduces it across eight configurations (per-iterationMPI_Barrierremoved, an injected 500 ms rank stall, dispatch-only and combine-only loops, up to 20 000 iterations) purely because it uses a single handle. The upstream report has been corrected accordingly — NVIDIA/nccl#2303.It also corrupts signalling without hanging, so latencies from any multi-handle run were suspect, not only the runs that wedged.
What changed
Rebinding is the intended lifecycle, not a workaround:
ncclEpInitHandletakes no token count, andncclEpUpdateHandleis documented as a "per-step collective: prepare the handle for the given top-k routing decisions". The wheel exposes it asHandle.update, which rebinds top-k indices without reallocating buffers. This also brings nccl-ep in line with the other three backends, each of which allocates one object sized to the ladder maximum and varies the token count per call.One handle per group (
bench/ep_nccl.py). The per-problem namespace stays — it holds only host-side tensor descriptors and, for HT, that problem's counters — while the handle moves to the group and is rebound 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 since every timed component is preceded by an untimedwarm()on its own problem, the collective always lands inwarmor the oracle passes, never inside a timed window. Re-entering the already-bound problem returns with no collective and no sync, so a timed loop's steady state 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. Teardown is a single explicitHandle.destroy.Low-latency rows restored on all six NVIDIA SKUs (
configs/platform_config.json). These were dropped in d58d565 while every LL leg wedged, on the reading that only a fixed wheel could restore them — that reading was wrong about the cause, and the single-handle adapter removes the aliasing on the stock wheel.COLLX_NCCL4PY_SPECis untouched. LL is decode-only and EP8-only, so the GB SKUs carry an EP8 row even though normal mode runs EP16 there, and b300/gb200/gb300 gain their firstll_backendsentry.test_nccl_ep_rollout_shapepreviously pinned the absence of these rows and now pins their presence, plus the absence of any EP16 LL row.CI validation
High throughput (h100 EP8, decode + prefill) against a
mainbaseline on the same pool — PR run / baseline. Bothsuccess, every rung correct, 2048 samples per point, max relative error unchanged. Round-trip p50 within ±1.3% and decode p99 within ±2.1% — so the per-shapeupdatecollective does not perturb the steady state. Prefill p99 came in 5.7–9.9% lower on all four rungs with p50 flat, which is the signature of removed tail interference; suggestive only, two runs cannot separate it from pool variance.Low latency — run 30440418355, all six SKUs green, 9 rungs each (decode T=1…256), every rung
correct=True, max relative error 0.00389 throughout:The two GB rows are EP8 across two trays over MNNVL — the configuration from NVIDIA/nccl#2303, which had never completed on any wheel — and the four x86 SKUs previously wedged 5/5 over SSH and 4/4 in sweep 30155842613.
Unit tests:
python3 -m unittest discover -s tests— 64 passed, 6 skipped (torch-dependent).Notes
tests/test_ep_nccl_handle.pystubs torch and nccl so it runs without the benchmark image, and locks the contract: onecreate_handleacross nine rungs,updateon a shape switch, no collective when the bound shape is re-entered, nolayout_infoon an LL rebind, HT rebinds carrying that problem's own counters, and idempotent teardown. The stubs are withdrawn after the import so the genuinely torch-dependent modules still skip rather than error.The upstream patch (NVIDIA/nccl#2306) is now a defence-in-depth measure rather than the thing that unblocks us: it makes a foreign or stale signal inert, which is worth having, but this tree no longer depends on it.