CollectiveX: restore nccl-ep zero-copy under CUDA graphs and FP8 in the default matrix / CollectiveX：恢复 nccl-ep 在 CUDA graph 下的零拷贝，并让 FP8 重新进入默认矩阵

collectivex/README.md

@@ -28,18 +28,17 @@ runs the decode and prefill ladders; `low-latency` is decode-only.
 | Backend | Engine selector (`maturity`) | `normal` | `low-latency` |
 |---|---|---|---|
 | `deepep-v2` | vLLM `--all2all-backend deepep_v2`, SGLang `--moe-a2a-backend deepep` (`production`) | EP8, EP16: H100, H200, B200, B300, GB200, GB300 | EP8: H200. EP8, EP16: B200, GB200, GB300 |
-| `mori` | vLLM `--all2all-backend mori_*`, SGLang `--moe-a2a-backend mori` (`production`) | EP8: MI300X, MI325X, MI355X. EP16: MI355X | EP16: MI355X |
-| `uccl-ep` | none; engines load it as a drop-in `deep_ep` under vLLM's DeepEP backends (`candidate`) | EP8: H100, H200, B200, MI300X, MI325X, MI355X. EP16: B200 | EP8: H100, H200, B200 |
-| `nccl-ep` | TensorRT-LLM `NcclEP` (`production` for low-latency rank-major BF16, `candidate` otherwise) | EP8, EP16: H100, H200, B200, B300, GB200, GB300 | EP8: H100, H200, B300. EP8, EP16: B200, GB200, GB300 |
+| `mori` | vLLM `--all2all-backend mori_*`, SGLang `--moe-a2a-backend mori` (`production`) | EP8, EP16: MI300X, MI325X, MI355X | EP16: MI300X, MI325X, MI355X |
+| `uccl-ep` | none; engines load it as a drop-in `deep_ep` under vLLM's DeepEP backends (`candidate`) | EP8: H100, H200, B200, MI300X, MI325X, MI355X. EP16: B200, MI300X, MI325X | EP8: H100, H200, B200. EP16: MI300X, MI325X |
+| `nccl-ep` | TensorRT-LLM `NcclEP` (`production` for low-latency rank-major BF16 where it matches, `candidate` otherwise) | EP8, EP16: H100, H200, B200, B300, GB200, GB300 | EP8: H100, H200, B300. EP8, EP16: B200, GB200, GB300 |
 | `flashinfer-ep` | vLLM `--all2all-backend flashinfer_nvlink_one_sided` (`production`) | EP8, EP16: GB200, GB300 | none |
 
-Every backend runs a BF16 dispatch control and an FP8 dispatch, except that nccl-ep's FP8 (its
-internal DS_FP8E3M4 recipe, low-latency decode only) is off every engine's path and runs only when
-`--precisions fp8` names it. Combine is always BF16. Coverage is uniform routing only.
+Every backend runs a BF16 dispatch control and an FP8 dispatch; nccl-ep's FP8 is its internal
+DS_FP8E3M4 recipe, low-latency decode only. Combine is always BF16. Coverage is uniform routing only.
 
-`results/cases.json` is the committed copy of the full matrix: 195 requested cases, of which 163 are
-runnable EP cases, 18 are swap-blocks cases (two per pool) and 14 are unsupported EP cells (UCCL-EP
-EP16 on H100, H200, MI300X, MI325X and MI355X, and MoRI EP16 on MI300X and MI325X, once per phase).
+`results/cases.json` is the committed copy of the full matrix: 220 requested cases, of which 196 are
+runnable EP cases, 18 are swap-blocks cases (two per pool) and 6 are unsupported EP cells (UCCL-EP
+EP16 on H100, H200 and MI355X, once per phase).
 A cell is listed as unsupported only when its backend runs on that SKU at another EP size, and
 low-latency adds only runnable cells, so a missing low-latency cell (DeepEP V2 on B300, for example)
 appears nowhere; [docs/methodology.md](docs/methodology.md#backends) gives each reason. The NCCL EP
@@ -69,14 +68,16 @@ pools also need an operator document that names a stage directory (see
 [Runner Configuration](#runner-configuration)).
 
 These variables change what a row measures. The result records each one (in `kernel_generation`,
-`implementation.fp8_consume` or `implementation.rank_alignment`), but the `case_id` does not, so a
-run with one replaces the case's committed default-configuration result:
+`implementation.fp8_consume`, `implementation.rank_alignment` or `implementation.library_config`),
+but the `case_id` does not, so a run with one replaces the case's committed default-configuration
+result:
 
 | Variable | Default | Effect |
 |---|---|---|
 | `COLLX_CUDA_GRAPH=0` | `1` | times graph-capable decode cells eager (`-eager-spin-ar`) |
 | `CX_FP8_CONSUME=dequant` | `native` | stages every FP8 pair inside the chain, a verification hatch |
 | `COLLX_NCCL_LL_LAYOUT=expert-major` | `rank-major` | runs nccl-ep low-latency on its expert-major layout |
+| `COLLX_NCCL_LL_ZERO_COPY=0` | `1` | stages nccl-ep low-latency dispatch, TensorRT-LLM's graphed path (drops `-zc`) |
 | `COLLX_ALIGN_SPIN_US` | `2000` | the per-sample rank-alignment spin |
 
 `COLLX_TIME` (allocation minutes: 300 per EP shard and 45 per swap-blocks shard unless set),

collectivex/bench/ep_nccl.py

@@ -15,13 +15,16 @@
 communication strategy (tensorrt_llm/_torch/moe/fused_moe/communication/nccl_ep.py at main
 a6142c920), auto-selected after its NVLink strategies or forced with
 TRTLLM_FORCE_COMM_METHOD=NCCL_EP, runs LL rank-major only and never NCCL EP's internal
-DS_FP8E3M4 quantization: DeepSeek-FP8 activations cross it in BF16. LL rank-major BF16 rows run
-that configuration: its group capacity, its staged dispatch under CUDA graph capture and its
-window-backed zero-copy dispatch eager. HT, expert-major LL and DS-FP8 are modes no engine runs;
-they stay library (`candidate`) rows.
-
-FP8 uses NCCL EP's native DS_FP8E3M4 recipe, for low-latency decode only. No engine selects it,
-so the matrix runs it only when named (sweep_matrix.OFF_PATH_PRECISIONS). Received E4M3 values
+DS_FP8E3M4 quantization: DeepSeek-FP8 activations cross it in BF16. LL rank-major BF16 rows
+zero-copy wherever the library can (one NVLink/MNNVL scale-up domain), graphed or eager, as NCCL EP
+is built to run. TensorRT-LLM stages its dispatch under CUDA graph capture instead, so the rows that
+match its configuration (eager, or scale-out where every path stages) carry its contract; graphed
+zero-copy rows, HT, expert-major LL and DS-FP8 are library (`candidate`) rows.
+COLLX_NCCL_LL_ZERO_COPY=0 stages every LL dispatch, which makes graphed in-domain LL rank-major
+BF16 TensorRT-LLM's graphed configuration (and eager in-domain a library row).
+
+FP8 uses NCCL EP's native DS_FP8E3M4 recipe, for low-latency decode only, and runs in the
+default matrix like every other backend's FP8. Received E4M3 values
 are dequantized into a BF16 staging plane, so the BF16 combine and correctness oracle remain the
 common comparison contract. Normal mode remains BF16-only.
 
@@ -126,17 +129,18 @@ def _blockwise_cast_back_into(
 
 class NCCLEPBackend(EPBackend):
     name = "nccl-ep"
-    # Set per row in __init__: "production" for the one configuration an engine selects
-    # (TensorRT-LLM's NcclEP, LL rank-major with BF16 transport; communication_factory.py:254-278
-    # and 438-459), "candidate" for the library modes no engine runs (HT, expert-major LL, DS-FP8).
+    # Set per row in __init__: "production" where the row runs the configuration an engine selects
+    # (TensorRT-LLM's NcclEP, LL rank-major with BF16 transport, communication_factory.py:254-278
+    # and 438-459, with TensorRT-LLM's zero-copy choice for the regime), "candidate" for every other
+    # row (graphed zero-copy LL, HT, expert-major LL, DS-FP8).
     maturity = "candidate"
     # One library, two algorithms selected by args.mode. kernel_generation and the combine
     # semantics are switched to their LL values in __init__ (mirrors ep_deepep_v2), and
     # create_buffer appends what changes a row without changing its code path: "-cap<N>", the
     # group's token capacity (HT's per-step update and LL's rank-major routing reset are sized by
     # it, not by T), then the installed nccl-extensions version with its nightly date.
     #   normal      -> HT / FLAT layout / unweighted-rank-sum combine.
-    #   low-latency -> LL / RANK_MAJOR layout / TensorRT-LLM's NcclEP contract.
+    #   low-latency -> LL / RANK_MAJOR layout (TensorRT-LLM's NcclEP's) / unweighted-rank-sum combine.
     # "-routed" marks the generation whose timed HT dispatch charges the per-step
     # ncclEpUpdateHandle (see dispatch()); "v02" marks the nccl-extensions v0.2 kernels.
     # "-static" marks the combine input bound to the full static receive plane (see
@@ -149,6 +153,7 @@ class NCCLEPBackend(EPBackend):
     CUDA_GRAPH_MODES = ("normal", "low-latency")
     zero_copy = True
     _ll_expert_major = False
+    _ll_zero_copy_opt_out = False
 
     @property
     def library_version(self) -> str:
@@ -157,14 +162,16 @@ def library_version(self) -> str:
 
     def implementation_details(self):
         """The group's token capacity, whether dispatch writes a registered window (HT's combine
-        reads one too; LL's combine always stages), the SM budgets the group resolved, and the
-        engine whose configuration the row runs (None for a library row)."""
+        reads one too; LL's combine always stages), the SM budgets the group resolved, whether
+        COLLX_NCCL_LL_ZERO_COPY=0 turned LL zero-copy off, and the engine whose configuration the
+        row runs (None for a library row)."""
         device_sms = torch.cuda.get_device_properties(self.device).multi_processor_count
         default = device_sms if self._ll else _HT_DEFAULT_NUM_SMS
         return {
             "capacity_tokens_per_rank": self.max_dispatch,
             "zero_copy": bool(self.zero_copy),
             "num_sms": _resolve_num_sms(default, device_sms, self._sm_env),
+            "ll_zero_copy_opt_out": bool(self._ll_zero_copy_opt_out),
             "engine_contract": self.engine_contract,
         }
 
@@ -200,21 +207,32 @@ def __init__(self, args, rank, world_size, local_rank, device):
         if layout not in ("rank-major", "expert-major"):
             raise ValueError(f"COLLX_NCCL_LL_LAYOUT must be rank-major or expert-major, got {layout!r}")
         self._ll_expert_major = self._ll and layout == "expert-major"
-        # The one configuration an engine runs (module docstring); every other row is the library's.
-        self.engine_contract = (
-            "tensorrt-llm" if self._ll and not self._ll_expert_major and not self._fp8 else None
-        )
-        self.maturity = "production" if self.engine_contract else "candidate"
-        # HT needs its windows in both regimes (`_bind_ht_recv_count`). LL rank-major follows
-        # TensorRT-LLM: eager, dispatch writes into a registered NCCL window, but under CUDA graph
-        # capture it writes into a plain context-owned tensor, which the library stages
-        # (nccl_ep.py:305-327). So graphed rows run the staged dispatch and eager rows
-        # (COLLX_CUDA_GRAPH=0) the window. Either way the library zero-copies only NVLink-only,
-        # rank-major NONE/FWD dispatch (nccl_ep.cc:4589-4600): scale-out, expert-major and
+        # COLLX_NCCL_LL_ZERO_COPY=0 stages every LL dispatch, so the graphed in-domain row can run
+        # TensorRT-LLM's graphed (staged) dispatch. HT ignores it.
+        ll_zero_copy = os.environ.get("COLLX_NCCL_LL_ZERO_COPY", "1")
+        if ll_zero_copy not in ("0", "1"):
+            raise ValueError(f"COLLX_NCCL_LL_ZERO_COPY must be '0' or '1', got {ll_zero_copy!r}")
+        self._ll_zero_copy_opt_out = self._ll and ll_zero_copy == "0"
+        # HT needs its windows in both regimes (`_bind_ht_recv_count`). LL zero-copies wherever the
+        # library can, graphed or eager: it zero-copies only NVLink-only, rank-major NONE/FWD
+        # dispatch (nccl_ep.cc:4589-4600), so scale-out (EP beyond the scale-up domain, whose GIN
+        # path has no window; an MNNVL domain spanning nodes is not scale-out), expert-major and
         # DS_FP8E3M4, which generates its scales internally, always stage.
         self.zero_copy = not self._ll or not (
-            self._internode or self._ll_expert_major or self._fp8 or self.cuda_graph_enabled
+            self._internode or self._ll_expert_major or self._fp8 or self._ll_zero_copy_opt_out
         )
+        # The one configuration an engine runs (module docstring): TensorRT-LLM's LL rank-major
+        # BF16, which dispatches into a registered window eager (where the library can) but into a
+        # plain context-owned tensor under CUDA graph capture (nccl_ep.py:305-327). A row carries
+        # its contract only when its zero-copy choice is TensorRT-LLM's for the regime.
+        trtllm_zero_copy = not self.cuda_graph_enabled and not self._internode
+        self.engine_contract = (
+            "tensorrt-llm"
+            if self._ll and not self._ll_expert_major and not self._fp8
+            and self.zero_copy == trtllm_zero_copy
+            else None
+        )
+        self.maturity = "production" if self.engine_contract else "candidate"
         if self._ll_expert_major:
             # Weighted source-side combine over a per-expert padded receive: deepep-v2 LL's
             # contract, so this is the like-for-like row against the DeepEP-API backends.

collectivex/docs/methodology-summary.md

@@ -21,8 +21,8 @@ not rank or recommend.
 - **`T` is tokens per rank**: decode 1-512 and prefill 1024-8192, in powers of two. A rung a backend
   cannot reach goes in `ladder_dropped` rather than vanishing.
 - **Precision**: dispatch is BF16, or FP8 where the backend supports it upstream. Combine is always
-  BF16. Every FP8 codec is per-128-block e4m3 with FP32 scales on the wire. nccl-ep's FP8 is off
-  every engine's path, so only `--precisions fp8` runs it.
+  BF16. Every FP8 codec is per-128-block e4m3 with FP32 scales on the wire. nccl-ep's FP8 is its
+  internal DS_FP8E3M4 recipe, which no engine selects, so its rows are library rows.
 - **Scale**: EP8 is one 8-GPU node (2x4 MNNVL on GB200/GB300). EP16 is 2x8 over RDMA on the x86 and
   AMD pools and 4x4 MNNVL on GB200/GB300, so GB EP16 stays scale-up.
 - **Coverage**: `configs/platform_config.json` says which backend runs where (the README's coverage
@@ -31,20 +31,23 @@ not rank or recommend.
 ## Backends run as engines run them
 
 Each backend runs the configuration its reference engine selects (vLLM, and TensorRT-LLM for
-nccl-ep low-latency), not a common denominator, so a cross-backend comparison
-compares deployed contracts.
+nccl-ep low-latency), not a common denominator, so a cross-backend comparison of `production` rows
+compares deployed contracts. The exception is nccl-ep's graphed low-latency default, which
+zero-copies inside the scale-up domain where TensorRT-LLM stages, and is therefore `candidate`.
 
 | Backend | Maturity | Launch | Contract |
 |---|---|---|---|
 | deepep-v2 | production | graph; normal prefill eager | vLLM's `deepep_v2` in both phases: overlap off (64 SMs), a single reduction, no CPU sync, receive sized to the next power of two of `T`. LL as vLLM's `deepep_low_latency`: send kernel plus receive hook, MNNVL on GB only, UE8M0 FP8 scales on Blackwell. The LL receive is fixed at 256, so LL stops at T=256. |
-| nccl-ep | production for LL rank-major BF16 (TensorRT-LLM's `NcclEP`); candidate otherwise | graph; HT prefill eager | LL runs TensorRT-LLM's configuration: rank-major, pre-reduced, one group sized to the ladder's top. Graphed dispatch stages, because TensorRT-LLM dispatches into a plain tensor under capture; eager dispatch is zero-copy inside one NVLink/MNNVL domain. HT is a library row: dispatch includes the per-step routing update, whose memsets, all-gather and scan are sized by the capacity (`-cap512` decode, `-cap8192` prefill), not by T. The communicator uses graph usage mode 1. FP8 (DS_FP8E3M4, LL decode) is off by default because no engine selects it. |
+| nccl-ep | production for LL rank-major BF16 where it matches TensorRT-LLM's `NcclEP` (eager in-domain, scale-out, or graphed with `COLLX_NCCL_LL_ZERO_COPY=0`); candidate otherwise | graph; HT prefill eager | LL runs TensorRT-LLM's configuration: rank-major, pre-reduced, one group sized to the ladder's top. Dispatch is zero-copy inside one NVLink/MNNVL scale-up domain, graphed or eager (`COLLX_NCCL_LL_ZERO_COPY=0` stages it); TensorRT-LLM stages under capture, so graphed zero-copy rows are library rows. HT is a library row: dispatch includes the per-step routing update, whose memsets, all-gather and scan are sized by the capacity (`-cap512` decode, `-cap8192` prefill), not by T. The communicator uses graph usage mode 1. FP8 (DS_FP8E3M4, LL decode) runs in the default matrix as a library row. |
 | uccl-ep | candidate | normal eager (host-synced); LL graph | DeepEP's legacy API over CPU-proxy RDMA (proxies run only at b200 EP16), in vLLM's DeepEP configuration on UCCL's drop-in `deep_ep`: `deepep_ht` launch configs at 20 SMs and 1 GiB buffers; LL called as deepep-v2's, with a 256-token receive. |
 | mori | production | graph decode; eager prefill | vLLM's MoRI configuration: IntraNode at 80 blocks, 16 warps (EP8), InterNodeV1 at 64/32/8 on STATIC_HEAP (EP16), external input buffer, one QP per PE. LL is vLLM's `mori_low_latency`: InterNodeV1LL at 64/32/8 on STATIC_HEAP, EP16 only (one node runs the normal rows' IntraNode). `MORI_ENABLE_SDMA` unset, as vLLM leaves it. |
 | flashinfer-ep | production | graph decode; eager prefill | GB200/GB300 only (one-sided MNNVL). vLLM's integration: combine stages a caller-owned payload (`payload_in_workspace=False`) into a caller-owned output. FP8 is DeepSeek blockwise, which vLLM selects from v0.29.0. `kernel_generation` carries the wheel version. |
 
 `production` means an engine can select the configuration today: vLLM or SGLang, or, for a library
-vLLM does not expose, the engine that does (TensorRT-LLM's `NcclEP` for nccl-ep's low-latency
-rank-major BF16 rows). `candidate` rows describe the library, not a deployable configuration.
+vLLM does not expose, the engine that does (TensorRT-LLM's `NcclEP` for the nccl-ep low-latency
+rank-major BF16 rows whose zero-copy choice matches it; graphed zero-copy rows, the default inside
+one scale-up domain, are `candidate`). `candidate` rows describe the library, not a deployable
+configuration.
 nccl-ep sets maturity per row. uccl-ep runs vLLM's DeepEP configuration but stays `candidate`,
 because no engine ships its selector.