CollectiveX: NCCL EP v0.2 via nccl-extensions — lift the LL clamp, enable LL rows and B200 EP16

experimental/CollectiveX/README.md

@@ -39,8 +39,11 @@ in one of two modes:
   `kNumMaxWarpGroups` is 16, since upstream raised `kNumMaxTopK` 9 -> 16 (uccl#1016, 2026-07-13) and
   our pin is six days later. The product is 16 for every CU count, so this is a dated regression
   rather than a hardware limit, and the AMD SKUs keep UCCL-EP normal mode without LL),
-  and NCCL EP EP8 on all six NVIDIA SKUs, restored once the single-handle fix removed the
-  [NVIDIA/nccl#2303](https://github.com/NVIDIA/nccl/issues/2303) signal aliasing that had wedged them.
+  and NCCL EP at EP8 on H100/H200/B300 and at EP8 *and EP16* on B200/GB200/GB300, enabled with the
+  v0.2 upgrade (the `nccl-extensions` wheel ships the combine-recv fence whose absence had held
+  these rows — see the backend table below; the earlier
+  [NVIDIA/nccl#2303](https://github.com/NVIDIA/nccl/issues/2303) signal aliasing was already fixed
+  by the single-handle change).
   B300 carries the `candidate` NCCL EP as its *only* low-latency row, so it has no production
   decode coverage. DeepEP V2 emits no LL row at all on B300 (the IBGDA address-handle wall in the
   backend table below), and `_ll_runnable` adds only runnable cells, so that wall is prose here
@@ -131,7 +134,7 @@ scale-up domain.
 | DeepEP V2 | `production`, with vLLM `--all2all-backend deepep_v2`, SGLang `--moe-a2a-backend deepep` | `normal` mode is PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out. FP8 dispatch via `use_fp8_dispatch` (blockwise e4m3fn) alongside BF16. `low-latency` mode is the legacy `deep_ep.Buffer` IBGDA decode kernels (per-expert padded layout, weighted combine, `use_fp8` e4m3fn), decode only, with EP8 wherever enabled, plus EP16 on GB200/GB300 (inside the MNNVL domain) and on B200's nscale bare-metal pool (IBGDA over native IB rails with `/dev/gdrdrv`, although the prior virtualized b200 pool could never run it). B300 is an unsupported coverage row in `low-latency`: the legacy Buffer self-enables NVSHMEM IBGDA even for a single-node EP8 run, and on B300 address-handle creation fails (`ibgda.cpp:2234 Unable to create ah`), rc255 on all eight ranks. `NVSHMEM_DISABLE_IB=1` does not help. The Buffer re-enables IBGDA regardless, and the run fails identically with it set and unset (measured on b300-002 and b300-011) |
 | MoRI | `production`, with vLLM `--all2all-backend mori_*`, SGLang `--moe-a2a-backend mori` | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU. EP16 remains an unsupported coverage row on all three CDNA SKUs. Part of the old ROCm/mori#475 corruption was this harness passing dispatch's returned recv-slot indices to `combine()` instead of the rank's own routing (root-caused upstream, guarded by ROCm/mori#546, kernels unchanged) — with the corrected call, single-shot InterNodeV1 is clean through T=512 on mi355x — but a residual stochastic corruption remains from T~128 up under repeated execution and is near-certain at prefill sizes (run 33045314017/33050026476; unaffected by per-pair drains, so not a buffer-reuse race; upstream cannot reproduce on ionic driver 26.03 vs our 25.11). The tw pairs additionally have no cross-node GPU fabric, and mi355x EP16 has no publishable transport today: uccl-ep's CPU-proxy RDMA is functional on Pollara but ~13x under its documented bandwidth (~6 GB/s vs 82; unchanged by registration mode or traffic class — ionic-driver suspect, same 25.11-vs-26.03 delta as the mori residual). `low-latency` mode selects the `AsyncLL` decode kernel driven split-phase (dispatch send+recv, combine send+recv — the kernel SGLang deploys; same compact layout and unweighted combine as `IntraNode`; earlier rows measured `IntraNodeLL`, discriminated by `kernel_generation`), decode/EP8 only. FP8 dispatch is caller-prequantized (per-SKU e4m3fnuz on gfx942, e4m3fn on gfx950). Combine stays BF16 (`quant_type=none`) alongside BF16 dispatch |
 | UCCL-EP | `candidate` (no engine exposes a UCCL-EP selector) | [UCCL](https://github.com/uccl-project/uccl) EP: a drop-in, API-identical DeepEP replacement whose CPU proxies issue GPUDirect RDMA over plain `libibverbs` (no NVSHMEM/IBGDA), with software message ordering, atomics, and flow control. Scale-up is single-node `cudaIpc` over NVLink/XGMI (never MNNVL). `normal` mode is the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum). `low-latency` reuses the legacy `low_latency_dispatch`/`low_latency_combine` decode kernels (weighted combine), decode/EP8 only. FP8 dispatch is caller-prequantized in `normal` mode (blockwise e4m3fn, per-SKU e4m3fnuz on gfx942). In `low-latency` mode the caller sends BF16 and the decode kernel quantizes to e4m3 internally (`use_fp8`). Combine is BF16. Runs on NVIDIA and AMD (H100/H200/B200 + MI300X/MI325X/MI355X), EP8 scale-up. Cross-node EP16 is functional (the internode RDMA path connects and the light case passes correctness) but its CPU-proxy throughput overruns the standardized per-case wall-clock budget on heavy token counts, so EP16 is an unsupported coverage row for now |
-| NCCL EP | `candidate` (NVIDIA's own library, but no engine exposes an NCCL-EP selector) | [NCCL EP](https://github.com/NVIDIA/nccl/tree/master/contrib/nccl_ep): NVIDIA's native MoE dispatch/combine on the NCCL Device API, using LSA (NVLink load/store) intra-node and GIN (GPU-Initiated Networking) inter-node, driven through the `nccl4py` bindings. `normal` mode selects the `HIGH_THROUGHPUT` algorithm (FLAT `[N, hidden]` receive, unweighted rank-sum combine). The `LOW_LATENCY` algorithm carries an EP8 `ll_backends` row on all six NVIDIA SKUs, restored once the single-handle fix removed the NVIDIA/nccl#2303 signal aliasing. That LL decode ladder is clamped to T<=128, below its 256-slot receive: `nccl_ep`'s combine recv pipeline is a port of DeepEP's pre-#642 kernel and is missing the same shared-memory fence before `mbarrier_arrive`, which corrupted T=256 on GB300 in 1 of 5 executions. It was bimodal, with healthy rows at 0.0039 relative error against 0.4704 on the failure. The fence is absent at NVIDIA/nccl master, so it is unfixed upstream. The clamp lowers exposure and is **not** a safety boundary: the fence is missing on every combine recv and T=256 is merely the rung with the most pipeline iterations, so lower rungs are less likely to hit the race rather than immune. Restore when a fixed wheel ships BF16 only: `contrib/nccl_ep/RELEASE.md` says "No FP8 support", so no FP8 case is emitted. That note is worth re-testing rather than trusting, because the C library at our pinned commit does read `inputs->scales` and switch on e4m3/e5m2, the two documented FP8 exclusions are expert-major layouts we do not use, and `NVIDIA/nccl` has not moved since 2026-06-11 while `NVIDIA/nccl-extensions` has replaced that row outright. NVIDIA-only and CUDA 13 only. EP8 scale-up on H100/H200/B200/B300 plus EP8 and EP16 on GB200/GB300, where EP16 stays inside the MNNVL scale-up domain. x86 EP16 scale-out is an unsupported coverage row: the cross-node GIN path faults inside `nccl_ep.cc` identically on RoCE and IB across four SKUs, a GDAKI limit rather than a fabric-selection one |
+| NCCL EP | `candidate` (NVIDIA's own library, but no engine exposes an NCCL-EP selector) | [NCCL EP v0.2](https://github.com/NVIDIA/nccl-extensions): NVIDIA's native MoE dispatch/combine on the NCCL Device API, using LSA (NVLink load/store) intra-node and GIN (GPU-Initiated Networking) inter-node. Since v0.2 it ships as the `nccl-extensions` PyPI wheel (which owns `nccl.ep`; `nccl4py` — pinned alongside — provides only `nccl.core`), replacing the `contrib/nccl_ep` tree in `NVIDIA/nccl`, which has not moved since 2026-06-11. `normal` mode selects the `HIGH_THROUGHPUT` algorithm (FLAT `[N, hidden]` receive, unweighted rank-sum combine). The `LOW_LATENCY` algorithm carries `ll_backends` rows at EP8 on H100/H200/B300 and EP8+EP16 on B200/GB200/GB300. Its decode ladder runs unclamped to the full 256-slot receive: v0.1's combine recv pipeline was a port of DeepEP's pre-#642 kernel, missing the shared-memory fence before `mbarrier_arrive` (bimodal T=256 corruption on GB300, 1 of 5 executions at 0.4704 relative error against healthy 0.0039), and the v0.2 wheel ships that fence (`fence_view_async_shared` before the `emptyBarriers` arrive in the shipped `ll_ep.cuh`), the documented restore condition for the T<=128 clamp it forced. BF16 only for now: v0.2 exposes FP8 dispatch (`DS_FP8E3M4`) and an experimental NVFP4 combine recipe, but wiring those into the fp8_consume model is separate bring-up, so no FP8 case is emitted yet. NVIDIA-only and CUDA 13 only. EP8 scale-up on H100/H200/B300, EP8 and EP16 on B200 (v0.2 fixes the RDMA EP16 GIN fault this pool hit under v0.1) and on GB200/GB300, where EP16 stays inside the MNNVL scale-up domain. x86 EP16 scale-out elsewhere (H100/H200/B300) remains an unsupported coverage row pending revalidation: under v0.1 the cross-node GIN path faulted inside `nccl_ep.cc` identically on RoCE and IB, and only the B200 configuration is claimed fixed |
 | FlashInfer EP | `production`, with vLLM `--all2all-backend flashinfer_nvlink_one_sided` | [FlashInfer](https://github.com/flashinfer-ai/flashinfer) `MoeAlltoAll`: TensorRT-LLM's one-sided MNNVL all-to-all, where each rank writes tokens straight into its peers' workspace windows and combine reads them back, with no send/recv pairing and no NVSHMEM. `normal` mode only (there is one kernel family and no separate decode path), and GB200/GB300 only, since the transport is MNNVL. FP8 dispatch is caller-prequantized blockwise e4m3fn, carried as a fourth dispatch payload alongside its per-128-block FP32 scales, with the combine plane forced to BF16. The C++ `toNvDataType` accepts only fp16/bf16/fp32 for combine, so an FP8 combine buffer would raise rather than corrupt. EP8 and EP16, both inside the scale-up domain. Unlike every other backend here, its combine accumulates in the PAYLOAD dtype rather than FP32: wheels before 0.6.16 reduce the top-k contributions with a pairwise BF16 tree that rounds at every level, so the oracle models that reduction directly (`combine_reduction = "topk-slot-tree"`) instead of widening the tolerance. 0.6.16 moved the accumulator to FP32, and the adapter switches models on the installed version |
 
 DeepEP V2 means the `ElasticBuffer` implementation introduced by

experimental/CollectiveX/bench/ep_nccl.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """NCCL EP adapter: NVIDIA's native MoE dispatch/combine on the NCCL Device API.
 
-NCCL EP (github.com/NVIDIA/nccl/contrib/nccl_ep, arXiv 2603.13606) is a ground-up MoE
+NCCL EP (github.com/NVIDIA/nccl-extensions, arXiv 2603.13606) is a ground-up MoE
 communication library built on NCCL's Device API — LSA (NVLink load/store) intra-node and
 GIN (GPU-Initiated Networking) inter-node — with two algorithms selected per case:
   normal      -> HIGH_THROUGHPUT (HT), the Hybrid-EP-derived prefill/train path. FLAT recv
@@ -13,19 +13,21 @@
 Because both map exactly onto the two combine contracts CollectiveX already models, the
 scale_up_domain two-level combine oracle in ep_harness applies unchanged.
 
-BF16 only: `contrib/nccl_ep/RELEASE.md` says "No FP8 support", so this adapter does not
-override the FP8 encode hooks (SUPPORTED_PRECISIONS=("bf16",)). Re-test before trusting that
-note — the C library at our pinned commit reads `inputs->scales` and switches on e4m3/e5m2,
-and the two documented FP8 exclusions are expert-major layouts we do not use.
+BF16 only for now: v0.2 grows a real quantization surface (DispatchQuantizationRecipe.FWD /
+DS_FP8E3M4 for FP8 dispatch, experimental CombineQuantizationRecipe.NVFP4), but wiring it into
+the fp8_consume model is its own bring-up, so this adapter still does not override the FP8
+encode hooks (SUPPORTED_PRECISIONS=("bf16",)).
 
 Communicator bootstrap: NCCL EP forms its OWN NCCL communicator (separate from PyTorch's
 process group) via ``Communicator.init(nranks, rank, unique_id)``. Upstream broadcasts the
 unique id with MPI; CollectiveX has no MPI, so rank 0 generates the id and we broadcast its
 bytes over the already-initialized torch process group (see ``_bootstrap_comm``).
 
-Python bindings are ``nccl4py`` (``import nccl.core`` + ``nccl.ep``); the API surface used
-here is verified against upstream ``bindings/nccl4py/nccl/ep`` and driven exactly as
-``contrib/nccl_ep/ep_test.py`` drives it.
+Python bindings are split across two wheels since v0.2: ``nccl-extensions`` owns ``nccl.ep``
+(the libnccl_ep.so JIT runtime + Cython bindings) and ``nccl4py`` provides ``nccl.core``
+(Communicator/UniqueId). The API surface used here is verified against the published
+nccl-extensions wheel and driven exactly as upstream's ep_test.py drives it — every class and
+signature this adapter touches is unchanged from v0.1.
 """
 from __future__ import annotations
 
@@ -69,23 +71,16 @@
 # token counts are measured. Separating them lets the ladder be clamped around a kernel defect
 # without moving the footprint and silently re-basing the rungs that remain.
 #
-# The ladder sits below the buffer because nccl_ep's low_latency.cu is a port of DeepEP's
-# PRE-FIX low-latency combine: in the combine recv pipeline the reduction warps read shared
-# memory and then mbarrier_arrive(emptyBarriers[stageIdx]) with no fence.proxy.async.shared::cta
-# between, so the producer's next TMA load can overwrite a stage while consumer reads are still
-# in flight. DeepEP closed exactly this with a one-line fence in PR #642; the fence is absent
-# both at our pin and at NVIDIA/nccl master, so it is unfixed upstream.
-#
-# Observed on gb300 EP8 BF16 at T=256: 1 failure in 5 executions, bimodal -- healthy rows give
-# max relative error 0.0039, the failure gave 0.4704, with nothing between, which is a discrete
-# corrupted write rather than tolerance noise.
-#
-# THIS CLAMP IS NOT A SAFETY BOUNDARY. The fence is missing on every combine recv; T=256 is only
-# the rung with the most pipeline iterations, and the receive plane is not even full there. Lower
-# rungs are LESS LIKELY to hit the race, not immune. Restore _LL_LADDER_CAP to _LL_BUFFER_CAP
-# once a fixed wheel ships.
+# The ladder is RESTORED to the full buffer under nccl-ep v0.2. The v0.1 port of DeepEP's
+# low-latency combine lacked the PR #642 fence — reduction warps read shared memory, then
+# mbarrier_arrive(emptyBarriers[stageIdx]) with no fence between, letting the producer's next
+# TMA load overwrite a stage mid-read (observed on gb300 EP8 BF16 at T=256: bimodal 1-in-5
+# discrete corruption, max rel err 0.4704 vs healthy 0.0039). The v0.2 wheel ships the fence:
+# fence_view_async_shared() before the elect_one_sync mbarrier_arrive in the shipped headers'
+# ll_ep.cuh combine recv pipeline — exactly the documented restore condition for this clamp.
+# The T=256 rung is back on the ladder; the correctness oracle re-verifies it on every run.
 _LL_BUFFER_CAP = 256
-_LL_LADDER_CAP = 128
+_LL_LADDER_CAP = _LL_BUFFER_CAP
 
 
 class NCCLEPBackend(EPBackend):
@@ -98,8 +93,10 @@ class NCCLEPBackend(EPBackend):
     # "-routed" marks the generation whose timed HT dispatch charges the per-step
     # ncclEpUpdateHandle (see dispatch()); earlier "nccl-ep-ht" rows excluded it and the
     # docs record that those cannot be separated by any other field — this suffix is the
-    # per-row discriminator that change lacked.
-    kernel_generation = "nccl-ep-ht-routed"
+    # per-row discriminator that change lacked. "v02" marks the nccl-extensions v0.2 mover
+    # (new kernels: LL combine fence, B200 EP16 fix, HT gains) so pre-upgrade rows never
+    # pool with post-upgrade rows.
+    kernel_generation = "nccl-ep-v02-ht-routed"
     SUPPORTED_MODES = ("normal", "low-latency")
     SUPPORTED_PRECISIONS = ("bf16",)
     stage_device_work = False
@@ -125,7 +122,7 @@ def __init__(self, args, rank, world_size, local_rank, device):
             # LL decode kernels apply the top-k gate at the source (weighted), not an
             # unweighted rank sum — the benchmark stages the UNWEIGHTED per-expert transform
             # and the kernel multiplies by the gate. Same contract as deepep-v2 low-latency.
-            self.kernel_generation = "nccl-ep-ll"
+            self.kernel_generation = "nccl-ep-v02-ll"
             self.receive_layout = "token-expert"
             self.combine_weight_semantics = "weighted-kernel-sum"
         # NCCL EP's handle is explicitly reusable across dispatch/combine cycles (ep_test.py
@@ -149,9 +146,9 @@ def __init__(self, args, rank, world_size, local_rank, device):
 
     def buffer_cap(self, args):
         if self._ll:
-            # Bounds which token counts are MEASURED. Below _LL_BUFFER_CAP today because the
-            # combine recv pipeline races (see the constants above); the harness reports every
-            # dropped rung rather than silently truncating.
+            # Bounds which token counts are MEASURED. Equal to _LL_BUFFER_CAP under v0.2 (the
+            # combine recv fence shipped — see the constants above); the two names stay separate
+            # so a future defect can clamp the ladder without moving the footprint.
             return _LL_LADDER_CAP
         return None
 

experimental/CollectiveX/configs/platform_config.json

@@ -11,7 +11,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",
@@ -35,7 +35,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,8 +54,8 @@
       "scale_up_domain": 8,
       "scale_up_transport": "nvlink",
       "launcher": "single-slurm",
-      "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16], "nccl-ep": [8]},
-      "ll_backends": {"deepep-v2": [8, 16], "uccl-ep": [8]},
+      "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16], "nccl-ep": [8, 16]},
+      "ll_backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8, 16]},
       "fabric": {"nic": "8x ConnectX-7-class 400G InfiniBand rails (bare metal, gdrdrv)", "switch": "InfiniBand (model not surveyed)"},
       "operator": {
         "partition": "batch_1",
@@ -77,7 +77,7 @@
       "scale_up_transport": "nvlink",
       "launcher": "single-slurm",
       "backends": {"deepep-v2": [8, 16], "nccl-ep": [8]},
-      "ll_backends": {"deepep-v2": [8]},
+      "ll_backends": {"deepep-v2": [8], "nccl-ep": [8]},
       "fabric": {"nic": "ConnectX-8 2x400GbE", "switch": "NVIDIA Spectrum-X SN5600 (51.2T)"},
       "operator": {
         "partition": "batch_1",
@@ -104,7 +104,7 @@
       "scale_up_transport": "mnnvl",
       "launcher": "gb-nv",
       "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]},
-      "ll_backends": {"deepep-v2": [8, 16]},
+      "ll_backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]},
       "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"},
       "operator": {
         "partition": "batch",
@@ -124,7 +124,7 @@
       "scale_up_transport": "mnnvl",
       "launcher": "gb-nv",
       "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]},
-      "ll_backends": {"deepep-v2": [8, 16]},
+      "ll_backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]},
       "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"},
       "operator": {
         "partition": "batch_1",