diff --git a/collectivex/README.md b/collectivex/README.md index 6d5aae0705..4b1f866714 100644 --- a/collectivex/README.md +++ b/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), diff --git a/collectivex/bench/ep_nccl.py b/collectivex/bench/ep_nccl.py index 69d9749a96..5d062971fe 100644 --- a/collectivex/bench/ep_nccl.py +++ b/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,9 +129,10 @@ 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 @@ -136,7 +140,7 @@ class NCCLEPBackend(EPBackend): # 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. diff --git a/collectivex/docs/methodology-summary.md b/collectivex/docs/methodology-summary.md index f319994b1c..17dc448fb0 100644 --- a/collectivex/docs/methodology-summary.md +++ b/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. diff --git a/collectivex/docs/methodology.md b/collectivex/docs/methodology.md index dc8b1d08ec..1cc6f3298d 100644 --- a/collectivex/docs/methodology.md +++ b/collectivex/docs/methodology.md @@ -80,11 +80,9 @@ for a precision it does not support. Every FP8 codec here is per-128-block e4m3 (vLLM main a9fe30e9f6 `fp8_utils.py:663-664`): byte-identical and timing-neutral, but not bit-identical to vLLM's quantize. - NCCL EP's FP8 is its internal DS_FP8E3M4 recipe, for low-latency decode only, staged (never - zero-copy) with a BF16 combine. It is the one realizable (backend, precision) pair off every - deployed path: TensorRT-LLM's `NcclEP`, the only engine integration, sends DeepSeek-FP8 activations - in BF16 (`uses_internal_fp8_dispatch = False`). So `sweep_matrix.py`'s `OFF_PATH_PRECISIONS` keeps - it out of the default matrix, and naming the precision (`--precisions fp8`) opts it back in: the - one place a precision filter adds rows rather than only removing them. NCCL EP's normal mode and + zero-copy) with a BF16 combine. It runs in the default matrix as a library row: TensorRT-LLM's + `NcclEP`, the only engine integration, sends DeepSeek-FP8 activations in BF16 + (`uses_internal_fp8_dispatch = False`). NCCL EP's normal mode and prefill stay BF16-only here, although the library also offers HT a caller-quantized FP8 path with forwarded scales (`QUANT_FWD`). @@ -129,8 +127,11 @@ roundtrip inverted the fp8-vs-bf16 verdict in 39 of 51 comparisons. ## Backends Each backend runs the configuration its reference engine selects (vLLM, except TensorRT-LLM for -NCCL EP low-latency), so a cross-backend comparison compares deployed contracts. `implementation.maturity` is `production` where an engine can select the configuration -today and `candidate` where the row describes the library (see [Result Document](#result-document)). +NCCL EP low-latency), except where a library runs better than its engine runs it: NCCL EP's graphed +low-latency rows zero-copy inside the scale-up domain, which TensorRT-LLM does not, so they are +library rows (see [NCCL EP](nccl-ep.md)). `implementation.maturity` is `production` where an +engine can select the configuration today and `candidate` where the row describes the library (see +[Result Document](#result-document)); compare `production` rows for deployed contracts. ### DeepEP V2 @@ -344,10 +345,10 @@ immediately and the kernel reads the cached routing inside the timed dispatch. `low-latency` runs TensorRT-LLM's `NcclEP` configuration: rank-major receive with caller-side weighted expert pre-reduction and an unweighted rank-sum combine, one group sized to the ladder's -top (TensorRT-LLM sizes it to `min(max_num_tokens, moe_max_num_tokens)`), and the engine's zero-copy -policy: graphed rows dispatch into a plain tensor that the library stages, as TensorRT-LLM does -under CUDA graph capture, eager rows write a registered window directly within one LSA/MNNVL domain, -and scale-out stages in both regimes. v0.1's combine recv pipeline, a port of DeepEP's pre-#642 code +top (TensorRT-LLM sizes it to `min(max_num_tokens, moe_max_num_tokens)`). Dispatch writes a +registered window directly within one LSA/MNNVL domain, graphed or eager, and scale-out stages in +both regimes. TensorRT-LLM stages under CUDA graph capture, so graphed zero-copy rows are library +rows. v0.1's combine recv pipeline, a port of DeepEP's pre-#642 code missing the `fence.proxy.async.shared::cta`, raced on every rung (1-in-5 bimodal corruption at T=256 on gb300 EP8, error 0.47 vs 0.0039; a T<=128 ladder clamp reduced exposure but was never a safety boundary). The v0.2 wheel ships that fence (`fence_view_async_shared` before the `emptyBarriers` @@ -466,9 +467,8 @@ changing its contract. Prefill stays eager on every backend, as the engines run `COLLX_CUDA_GRAPH=0` restores the eager pipeline everywhere. - **deepep-v2** `normal` and `low-latency` decode (see [DeepEP V2](#deepep-v2)). -- **nccl-ep** low-latency and HT decode. Graphed low-latency dispatch writes a plain tensor that the - library stages, as TensorRT-LLM's captured decode does, so it carries no `-zc`; eager rows within - one domain take the window path (`-zc`). HT's per-step routing `ncclAllGather` is the one host +- **nccl-ep** low-latency and HT decode. Graphed low-latency dispatch within one domain takes the + window path (`-zc`), as eager does; scale-out stages. HT's per-step routing `ncclAllGather` is the one host NCCL collective in any capture, so nccl-ep's communicator runs NCCL graph usage mode 1 (one graph at a time): the default mixing mode wraps every captured collective in an external event wait and record, which left ~14 µs idle before each all-gather (h100 EP8 trace) and put the graphed HT @@ -1004,7 +1004,7 @@ generation: | deepep-v2 | `v2-elastic-buffer-vllm-nosync`, plus `-relaxed-ordering` on h200 EP16; `legacy-buffer-ll-hook`, plus `-ue8m0` for Blackwell FP8 | | uccl-ep | `uccl-legacy-buffer-vllm-ht`; `uccl-legacy-buffer-ll-hook`, plus `-ue8m0` for B200 FP8 | | mori | `intranode-vllm`, `inter-node-v1-vllm` and `inter-node-v1ll-vllm`, each plus `-sdma` when `MORI_ENABLE_SDMA` was set on | -| nccl-ep | `nccl-ep-v02-ht-routed-zc-static` (plus `-gum1` when graphed), `nccl-ep-v02-ll-rm` (plus `-zc` on the eager window path) or `nccl-ep-v02-ll-em`, then `-fp8-ds-e3m4` on FP8 rows and `-cap-` ([nccl-ep.md](nccl-ep.md)) | +| nccl-ep | `nccl-ep-v02-ht-routed-zc-static` (plus `-gum1` when graphed), `nccl-ep-v02-ll-rm` (plus `-zc` on the window path) or `nccl-ep-v02-ll-em`, then `-fp8-ds-e3m4` on FP8 rows and `-cap-` ([nccl-ep.md](nccl-ep.md)) | | flashinfer-ep | `flashinfer-mnnvl-one-sided-staged-` | Backend source pins live in `runtime/common.sh` and are enforced by exact fetched-commit comparison. diff --git a/collectivex/docs/nccl-ep.md b/collectivex/docs/nccl-ep.md index b32aebef67..7f037da198 100644 --- a/collectivex/docs/nccl-ep.md +++ b/collectivex/docs/nccl-ep.md @@ -11,14 +11,24 @@ before the 2026-10-05 nightly. They have not been rechecked against the pinned 2 | Row | Configuration | `maturity` | |---|---|---| -| `low-latency`, rank-major, BF16 | TensorRT-LLM's `NcclEP` | `production` | -| `low-latency`, rank-major, DS-FP8 (`--precisions fp8`) | library | `candidate` | +| `low-latency`, rank-major, BF16, eager inside one scale-up domain (zero-copy) | TensorRT-LLM's `NcclEP` | `production` | +| `low-latency`, rank-major, BF16, scale-out (staged, graphed or eager) | TensorRT-LLM's `NcclEP` | `production` | +| `low-latency`, rank-major, BF16, graphed inside one scale-up domain (zero-copy, the default) | library | `candidate` | +| `low-latency`, rank-major, BF16, graphed inside one scale-up domain, `COLLX_NCCL_LL_ZERO_COPY=0` (staged) | TensorRT-LLM's `NcclEP` | `production` | +| `low-latency`, rank-major, BF16, eager inside one scale-up domain, `COLLX_NCCL_LL_ZERO_COPY=0` (staged) | library | `candidate` | +| `low-latency`, rank-major, DS-FP8 | library | `candidate` | | `low-latency`, expert-major (`COLLX_NCCL_LL_LAYOUT=expert-major`) | library | `candidate` | | `normal` (HT) | library | `candidate` | `maturity` is set per row, so `implementation.maturity` separates engine rows from library rows. The registry's per-backend `backend_maturity` still says `candidate`. +Scale-out means the EP group is larger than the case's scale-up domain (`world_size > +scale_up_domain`, `_internode` in `ep_nccl.py`; the case's `scope` is `scale-out`), so dispatch +crosses the GIN network path. Node count does not decide it: GB200/GB300 EP8 and EP16 span two and +four nodes inside one 72-GPU MNNVL domain and zero-copy. H100, H200, B200 and B300 EP16 are +scale-out. + ## Low-latency: TensorRT-LLM's configuration - **Group.** `LOW_LATENCY`, the routed expert count, `max_dispatch_tokens_per_rank` = capacity, @@ -30,16 +40,24 @@ The registry's per-backend `backend_maturity` still says `candidate`. admits every rung, so low-latency runs the 512 rung that the other backends' fixed 256-token receive drops. Capacity changes what LL measures: each rank-major dispatch resets every routing slot of the receive plane, `EP x capacity x top-k` entries (`ll_dispatch_lsa.cuh:652-659`). -- **Zero-copy follows the regime.** TensorRT-LLM dispatches into a registered NCCL window only - when it is not capturing. Under CUDA graph capture it dispatches into a plain tensor, and the - library copies every received token out of its staging buffer (`nccl_ep.py:305-327`, - `ll_dispatch_lsa.cuh:484-490`). So graphed rows, the default, run the staged dispatch, and eager - rows (`COLLX_CUDA_GRAPH=0`) write the window directly. The window path exists only within one - NVLink domain, for rank-major dispatch without internal quantization (`nccl_ep.cc:4589-4600`). - Scale-out, expert-major and DS-FP8 rows stage in both regimes. A zero-copy row creates its group - with `ZeroCopyMode.ON`, which fails any dispatch that lacks the window, so the `-zc` label cannot - name a path that did not run. A staged row uses `OFF`, which equals TensorRT-LLM's default `AUTO` - at this pin. Combine always stages in low-latency mode. +- **Zero-copy wherever the library supports it.** Low-latency dispatch writes a registered NCCL + window directly, graphed or eager. The window path exists only within one NVLink/MNNVL scale-up + domain, for rank-major dispatch without internal quantization (`nccl_ep.cc:4589-4600`), so + scale-out, expert-major and DS-FP8 rows stage in both regimes. TensorRT-LLM differs under CUDA + graph capture: it dispatches into a plain tensor and the library copies every received token out + of its staging buffer (`nccl_ep.py:305-327`, `ll_dispatch_lsa.cuh:484-490`). Graphed zero-copy + rows are therefore library rows; eager rows and scale-out rows match TensorRT-LLM. A zero-copy + row creates its group with `ZeroCopyMode.ON`, which fails any dispatch that lacks the window, so + the `-zc` label cannot name a path that did not run. A staged row uses `OFF`, which equals + TensorRT-LLM's default `AUTO` at this pin. Combine always stages in low-latency mode. +- **`COLLX_NCCL_LL_ZERO_COPY=0`.** Turns low-latency zero-copy off (default `1`; any other value is + rejected), so the graphed in-domain row runs TensorRT-LLM's graphed staged dispatch and carries + `engine_contract` `tensorrt-llm`. The same setting makes the eager in-domain row a library row, + since TensorRT-LLM zero-copies eager. HT ignores it. The row records it twice: the family loses + `-zc`, and `implementation.library_config.ll_zero_copy_opt_out` is `true`. Like + `COLLX_NCCL_LL_LAYOUT`, it is read from the launching shell's environment, which every rank + inherits (`srun --export=ALL`), and the `case_id` does not change, so a run with it replaces the + case's committed default result. - **Handle.** One handle, rebound to each ladder point's routing in untimed code. TensorRT-LLM also creates one handle and updates it every step. A low-latency update does no device work (`nccl_ep.cc:4014-4020`). @@ -64,19 +82,19 @@ scales (`QUANT_FWD`), which this suite does not measure. FP8 uses NCCL EP's internal `DS_FP8E3M4` recipe, for low-latency decode only. TensorRT-LLM never selects it: `NcclEP` sets `uses_internal_fp8_dispatch = False` and transports DeepSeek-FP8 -activations in BF16 (`nccl_ep.py:78-79`). The default matrix therefore omits nccl-ep FP8 -(`sweep_matrix.OFF_PATH_PRECISIONS`, which lists no other precision for any backend). -`--precisions fp8` adds it back. +activations in BF16 (`nccl_ep.py:78-79`), so FP8 rows are library rows. They run in the default +matrix, like every other backend's FP8. ## Row identity `kernel_generation` is a family, `nccl-ep-v02-ht-routed-zc-static` (HT, plus `-gum1` when graphed), -`nccl-ep-v02-ll-rm` (plus `-zc` on the eager window path) or `nccl-ep-v02-ll-em`, then +`nccl-ep-v02-ll-rm` (plus `-zc` on the window path) or `nccl-ep-v02-ll-em`, then `-fp8-ds-e3m4` on FP8 rows, `-cap-`, and the harness's regime suffix. HT's `-zc` is part of its family: HT always takes the window. Examples: -- `nccl-ep-v02-ll-rm-cap512-0.2.0.dev20261013-cudagraph-spin-ar`: graphed low-latency, staged. -- `nccl-ep-v02-ll-rm-zc-cap512-0.2.0.dev20261013-eager-spin-ar`: eager low-latency, zero-copy. +- `nccl-ep-v02-ll-rm-zc-cap512-0.2.0.dev20261013-cudagraph-spin-ar`: graphed low-latency, zero-copy. +- `nccl-ep-v02-ll-rm-cap512-0.2.0.dev20261013-cudagraph-spin-ar`: graphed low-latency, staged + (scale-out, or `COLLX_NCCL_LL_ZERO_COPY=0`; `library_config.ll_zero_copy_opt_out` tells them apart). - `nccl-ep-v02-ht-routed-zc-static-gum1-cap512-0.2.0.dev20261013-cudagraph-spin-ar`: graphed HT decode. - `nccl-ep-v02-ht-routed-zc-static-cap8192-0.2.0.dev20261013-eager-spin-ar`: HT prefill. @@ -87,7 +105,8 @@ but the mode only matters to that capture. `EPBackend.implementation_details()` returns the realized configuration, which the result document publishes as `implementation.library_config`: `capacity_tokens_per_rank`, `zero_copy`, `num_sms` -(`dispatch` and `combine`), and `engine_contract` (`tensorrt-llm`, or null for a library row). The +(`dispatch` and `combine`), `ll_zero_copy_opt_out` (`COLLX_NCCL_LL_ZERO_COPY=0` on a low-latency +row), and `engine_contract` (`tensorrt-llm`, or null for a library row). The wheel version stays in `implementation.library_version`. `num_sms` repeats the library's resolution: 16 for HT and every SM for low-latency, then `NCCL_EP_COMM_SMS`, then `NCCL_EP_DISPATCH_SMS` and `NCCL_EP_COMBINE_SMS`, each ignored outside 1 to the device's SM count diff --git a/collectivex/results/cases.json b/collectivex/results/cases.json index b8efe5a46c..3f3b010c0b 100644 --- a/collectivex/results/cases.json +++ b/collectivex/results/cases.json @@ -380,6 +380,44 @@ "reason": null, "sku": "b200-nscale" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "b200-nscale-nccl-ep-deepseek-v4-pro-low-latency-decode-ep8-uniform-fp8", + "ep": 8, + "experts": 384, + "gpus_per_node": 8, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 1, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "b200-nvlink-island", + "transport": "nvlink", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "b200-nscale" + }, { "case": { "backend": "deepep-v2", @@ -874,6 +912,44 @@ "reason": null, "sku": "b200-nscale" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "b200-nscale-nccl-ep-deepseek-v4-pro-low-latency-decode-ep16-uniform-fp8", + "ep": 16, + "experts": 384, + "gpus_per_node": 8, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 2, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": "rdma", + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "scope": "scale-out", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "b200-nvlink-rdma", + "transport": "nvlink-rdma", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "b200-nscale" + }, { "case": { "backend": "deepep-v2", @@ -1216,6 +1292,44 @@ "reason": null, "sku": "b300" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "b300-nccl-ep-deepseek-v4-pro-low-latency-decode-ep8-uniform-fp8", + "ep": 8, + "experts": 384, + "gpus_per_node": 8, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 1, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "b300-nvlink-island", + "transport": "nvlink", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "b300" + }, { "case": { "backend": "deepep-v2", @@ -1786,6 +1900,44 @@ "reason": null, "sku": "gb200" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "gb200-nccl-ep-deepseek-v4-pro-low-latency-decode-ep8-uniform-fp8", + "ep": 8, + "experts": 384, + "gpus_per_node": 4, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 2, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "gb200-nvl72-mnnvl", + "transport": "mnnvl", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "gb200" + }, { "case": { "backend": "flashinfer-ep", @@ -2280,6 +2432,44 @@ "reason": null, "sku": "gb200" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "gb200-nccl-ep-deepseek-v4-pro-low-latency-decode-ep16-uniform-fp8", + "ep": 16, + "experts": 384, + "gpus_per_node": 4, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 4, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "gb200-nvl72-mnnvl", + "transport": "mnnvl", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "gb200" + }, { "case": { "backend": "flashinfer-ep", @@ -2774,6 +2964,44 @@ "reason": null, "sku": "gb300" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "gb300-nccl-ep-deepseek-v4-pro-low-latency-decode-ep8-uniform-fp8", + "ep": 8, + "experts": 384, + "gpus_per_node": 4, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 2, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "gb300-nvl72-mnnvl", + "transport": "mnnvl", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "gb300" + }, { "case": { "backend": "flashinfer-ep", @@ -3268,6 +3496,44 @@ "reason": null, "sku": "gb300" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "gb300-nccl-ep-deepseek-v4-pro-low-latency-decode-ep16-uniform-fp8", + "ep": 16, + "experts": 384, + "gpus_per_node": 4, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 4, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 72, + "scale_up_transport": "mnnvl", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "gb300-nvl72-mnnvl", + "transport": "mnnvl", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "gb300" + }, { "case": { "backend": "flashinfer-ep", @@ -3838,6 +4104,44 @@ "reason": null, "sku": "h100-dgxc" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "h100-dgxc-nccl-ep-deepseek-v4-pro-low-latency-decode-ep8-uniform-fp8", + "ep": 8, + "experts": 384, + "gpus_per_node": 8, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 1, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "h100-nvlink-island", + "transport": "nvlink", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "h100-dgxc" + }, { "case": { "backend": "deepep-v2", @@ -4712,6 +5016,44 @@ "reason": null, "sku": "h200-dgxc" }, + { + "case": { + "backend": "nccl-ep", + "case_id": "h200-dgxc-nccl-ep-deepseek-v4-pro-low-latency-decode-ep8-uniform-fp8", + "ep": 8, + "experts": 384, + "gpus_per_node": 8, + "hidden": 7168, + "ladder": "1 2 4 8 16 32 64 128 256 512", + "mode": "low-latency", + "nodes": 1, + "phase": "decode", + "precision": "fp8", + "routing": "uniform", + "scale_out_transport": null, + "scale_up_domain": 8, + "scale_up_transport": "nvlink", + "scope": "scale-up", + "seed": 67, + "suite": "ep-core", + "timing": { + "chain_drop": 16, + "chain_iters_per_trial": 128, + "chain_trials_per_point": 4, + "iters_per_trial": 8, + "trials_per_point": 256, + "warmup_iters_per_trial": 32 + }, + "topk": 6, + "topology_class": "h200-nvlink-island", + "transport": "nvlink", + "workload": "deepseek-v4-pro" + }, + "detail": null, + "disposition": "runnable", + "reason": null, + "sku": "h200-dgxc" + }, { "case": { "backend": "deepep-v2", diff --git a/collectivex/sweep_matrix.py b/collectivex/sweep_matrix.py index 7ebbe17f05..ce3f4cd809 100644 --- a/collectivex/sweep_matrix.py +++ b/collectivex/sweep_matrix.py @@ -44,22 +44,17 @@ def _load_config(name: str) -> dict[str, Any]: "deepep-v2": ("bf16", "fp8"), "mori": ("bf16", "fp8"), "uccl-ep": ("bf16", "fp8"), - # Native DS_FP8E3M4 dispatch, low-latency decode only (the mode/phase gate below), and off - # every deployed path -- see OFF_PATH_PRECISIONS. + # Native DS_FP8E3M4 dispatch, low-latency decode only (the mode/phase gate below). "nccl-ep": ("bf16", "fp8"), # FlashInfer FP8 is dispatch-side only (scales as a fourth payload, combine stays BF16), # and uses the same per-128-block e4m3 recipe as deepep-v2/uccl-ep so the axis is # comparable. vLLM selects it on this transport for DeepSeek blockwise FP8 from v0.29.0. "flashinfer-ep": ("bf16", "fp8"), } -# Precisions a backend REALIZES but that no serving engine can select on that transport. Kept -# out of the default matrix so a production sweep measures deployable configurations, and still -# reachable by naming the precision explicitly (`--precisions fp8`) for transport comparison. -# TensorRT-LLM, the one engine with an NCCL EP selector, never asks NCCL EP to quantize: its NcclEP -# sends DeepSeek-FP8 activations in BF16 (`uses_internal_fp8_dispatch = False`), so nccl-ep's -# DS_FP8E3M4 row is the library's. Declared here rather than on the adapter because this -# generator must resolve the matrix with no vendor imports. -OFF_PATH_PRECISIONS = {"nccl-ep": ("fp8",)} +# Precisions a backend realizes but that no serving engine can select on that transport, kept out +# of the default matrix and still reachable by naming them (`--precisions fp8`). Declared here +# rather than on the adapter because this generator must resolve the matrix with no vendor imports. +OFF_PATH_PRECISIONS: dict[str, tuple[str, ...]] = {} # Short shard-ID slug per non-normal mode. Normal-mode shard IDs carry no mode # segment so existing references stay valid; a low-latency shard adds "-ll". _MODE_SLUG = {"low-latency": "ll"} diff --git a/collectivex/tests/test_nccl_ep_contract.py b/collectivex/tests/test_nccl_ep_contract.py index ccd917773c..6c363b5b45 100644 --- a/collectivex/tests/test_nccl_ep_contract.py +++ b/collectivex/tests/test_nccl_ep_contract.py @@ -51,12 +51,15 @@ def _module(): return module -def _backend(mode="low-latency", world_size=8, phase="decode", precision="bf16", capacity=512, env=None): +def _backend(mode="low-latency", world_size=8, phase="decode", precision="bf16", capacity=512, env=None, + scale_up_domain=8): """(module, backend, handle, details) for one case built and sized as the harness drives it, under - `env` alone (graphs on unless it sets COLLX_CUDA_GRAPH=0).""" + `env` alone (graphs on unless it sets COLLX_CUDA_GRAPH=0). `scale_up_domain=72` is a GB200/GB300 + MNNVL rack, where EP16 spans four 4-GPU nodes inside one scale-up domain.""" module = _module() options = types.SimpleNamespace( - mode=mode, precision=precision, phase=phase, experts=384, hidden=7168, topk=6, scale_up_domain=8, + mode=mode, precision=precision, phase=phase, experts=384, hidden=7168, topk=6, + scale_up_domain=scale_up_domain, ) handle = FakeHandle() group = types.SimpleNamespace(create_handle=lambda *a, **k: handle, destroy=lambda: None) @@ -79,37 +82,41 @@ def _dispatch(backend, steps=1): backend.dispatch(problem) -class LowLatencyFollowsTensorRTLLM(unittest.TestCase): - """TensorRT-LLM dispatches into a registered window only when it is not capturing - (nccl_ep.py:305-327): graphed LL stages, eager LL zero-copies, and `-zc` names what ran.""" +class LowLatencyZeroCopies(unittest.TestCase): + """Inside one NVLink domain LL rank-major dispatches into the registered window, graphed or + eager, and `-zc` names what ran.""" - def test_graphed_rank_major_dispatches_into_a_plain_tensor(self): - module, backend, handle, details = _backend() - _dispatch(backend) - - self.assertFalse(backend.zero_copy) - self.assertEqual(module.recorded["group"][0]["zero_copy"], module.ZeroCopyMode.OFF) - backend._comm.register_window.assert_not_called() - self.assertIsNone(handle.outputs.tokens.window) - self.assertFalse(details["zero_copy"]) - self.assertEqual(backend.kernel_generation, f"nccl-ep-v02-ll-rm-cap512-{VERSION}") + def test_rank_major_dispatches_into_the_zero_copy_window_in_both_regimes(self): + for graph in ("1", "0"): + with self.subTest(graph=graph): + module, backend, handle, details = _backend(env={"COLLX_CUDA_GRAPH": graph}) + _dispatch(backend) - def test_eager_rank_major_keeps_the_zero_copy_window(self): - module, backend, handle, details = _backend(env={"COLLX_CUDA_GRAPH": "0"}) - _dispatch(backend) + self.assertTrue(backend.zero_copy) + self.assertEqual(module.recorded["group"][0]["zero_copy"], module.ZeroCopyMode.ON) + backend._comm.register_window.assert_called_once() + self.assertIs(handle.outputs.tokens.window, backend._recv_window) + self.assertTrue(details["zero_copy"]) + self.assertEqual(backend.kernel_generation, f"nccl-ep-v02-ll-rm-zc-cap512-{VERSION}") - self.assertTrue(backend.zero_copy) - self.assertEqual(module.recorded["group"][0]["zero_copy"], module.ZeroCopyMode.ON) - backend._comm.register_window.assert_called_once() - self.assertIs(handle.outputs.tokens.window, backend._recv_window) - self.assertTrue(details["zero_copy"]) - self.assertEqual(backend.kernel_generation, f"nccl-ep-v02-ll-rm-zc-cap512-{VERSION}") + def test_ep16_inside_an_mnnvl_domain_is_not_scale_out(self): + # GB200/GB300 EP16 spans four nodes but one 72-GPU MNNVL domain: the window path applies. + for graph in ("1", "0"): + with self.subTest(graph=graph): + module, backend, _, details = _backend( + world_size=16, scale_up_domain=72, env={"COLLX_CUDA_GRAPH": graph}, + ) + self.assertFalse(backend._internode) + self.assertTrue(backend.zero_copy) + self.assertEqual(module.recorded["group"][0]["zero_copy"], module.ZeroCopyMode.ON) + self.assertEqual(backend.kernel_generation, f"nccl-ep-v02-ll-rm-zc-cap512-{VERSION}") def test_paths_the_library_cannot_zero_copy_stage_in_both_regimes(self): cases = { "scale-out": {"world_size": 16}, "expert-major": {"env": {"COLLX_NCCL_LL_LAYOUT": "expert-major"}}, "ds-fp8": {"precision": "fp8"}, + "opt-out": {"env": {"COLLX_NCCL_LL_ZERO_COPY": "0"}}, } for name, case in cases.items(): for graph in ("1", "0"): @@ -119,6 +126,7 @@ def test_paths_the_library_cannot_zero_copy_stage_in_both_regimes(self): world_size=case.get("world_size", 8), precision=case.get("precision", "bf16"), env=env, ) self.assertFalse(backend.zero_copy) + self.assertEqual(details["ll_zero_copy_opt_out"], name == "opt-out") self.assertFalse(details["zero_copy"]) self.assertEqual(module.recorded["group"][0]["zero_copy"], module.ZeroCopyMode.OFF) self.assertNotIn("-zc", backend.kernel_generation) @@ -184,13 +192,25 @@ def test_low_latency_dispatch_carries_no_update(self): class RowLabels(unittest.TestCase): def test_only_the_trtllm_configuration_is_production(self): + eager = {"COLLX_CUDA_GRAPH": "0"} + staged = {"COLLX_NCCL_LL_ZERO_COPY": "0"} cases = { - "ll-rank-major": ({}, "production", "tensorrt-llm"), - "ll-rank-major-eager": ({"env": {"COLLX_CUDA_GRAPH": "0"}}, "production", "tensorrt-llm"), - "ll-scale-out": ({"world_size": 16}, "production", "tensorrt-llm"), + # TensorRT-LLM zero-copies eager and stages under capture: a row is its configuration + # only when its zero-copy choice matches the regime. + "ll-rank-major-graphed-zc": ({}, "candidate", None), + "ll-rank-major-eager-zc": ({"env": eager}, "production", "tensorrt-llm"), + "ll-mnnvl-ep16-graphed-zc": ({"world_size": 16, "scale_up_domain": 72}, "candidate", None), + "ll-mnnvl-ep16-eager-zc": ( + {"world_size": 16, "scale_up_domain": 72, "env": eager}, "production", "tensorrt-llm"), + "ll-scale-out-graphed": ({"world_size": 16}, "production", "tensorrt-llm"), + "ll-scale-out-eager": ({"world_size": 16, "env": eager}, "production", "tensorrt-llm"), + "ll-opt-out-graphed": ({"env": staged}, "production", "tensorrt-llm"), + "ll-opt-out-eager": ({"env": {**staged, **eager}}, "candidate", None), + "ll-opt-out-scale-out": ({"world_size": 16, "env": staged}, "production", "tensorrt-llm"), "ll-expert-major": ({"env": {"COLLX_NCCL_LL_LAYOUT": "expert-major"}}, "candidate", None), "ll-ds-fp8": ({"precision": "fp8"}, "candidate", None), "ht": ({"mode": "normal"}, "candidate", None), + "ht-opt-out-ignored": ({"mode": "normal", "env": staged}, "candidate", None), } for name, (case, maturity, engine) in cases.items(): with self.subTest(case=name): @@ -198,6 +218,19 @@ def test_only_the_trtllm_configuration_is_production(self): self.assertEqual(backend.maturity, maturity) self.assertEqual(details["engine_contract"], engine) + def test_the_opt_out_runs_trtllm_graphed_dispatch_and_says_so(self): + module, backend, handle, details = _backend(env={"COLLX_NCCL_LL_ZERO_COPY": "0"}) + _dispatch(backend) + self.assertEqual(module.recorded["group"][0]["zero_copy"], module.ZeroCopyMode.OFF) + backend._comm.register_window.assert_not_called() + self.assertEqual(backend.kernel_generation, f"nccl-ep-v02-ll-rm-cap512-{VERSION}") + self.assertTrue(details["ll_zero_copy_opt_out"]) + self.assertEqual(details["engine_contract"], "tensorrt-llm") + + def test_the_opt_out_rejects_other_values(self): + with self.assertRaisesRegex(ValueError, "COLLX_NCCL_LL_ZERO_COPY"): + _backend(env={"COLLX_NCCL_LL_ZERO_COPY": "off"}) + def test_generation_and_library_version_name_the_installed_wheel(self): _, backend, _, _ = _backend(precision="fp8") self.assertEqual(backend.library_version, VERSION) @@ -205,11 +238,12 @@ def test_generation_and_library_version_name_the_installed_wheel(self): def test_library_config_is_one_json_object(self): # Published as `implementation.library_config`; the wheel stays in `library_version`. - _, _, _, details = _backend() + _, _, _, details = _backend(env={"COLLX_CUDA_GRAPH": "0"}) self.assertEqual(json.loads(json.dumps(details)), { "capacity_tokens_per_rank": 512, - "zero_copy": False, + "zero_copy": True, "num_sms": {"dispatch": DEVICE_SMS, "combine": DEVICE_SMS}, + "ll_zero_copy_opt_out": False, "engine_contract": "tensorrt-llm", })