From 98ef7f6a82beb4b9aa52a96e8ab20a6927446100 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 11:25:33 +0800 Subject: [PATCH 01/24] CollectiveX: label backend maturity and correct the low-latency docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three things that all turn on the same question — can anyone actually run this transport? An audit against vLLM's `--all2all-backend` and SGLang's `--moe-a2a-backend` found that two of the four benchmarked transports are not selectable in either engine. DeepEP V2 and MoRI are (`deepep_v2`/`deepep`, `mori_*`/`mori`); UCCL-EP and NCCL EP are not. That is fine for UCCL-EP and NCCL EP as candidates — they are real transports and the walls they probe are worth knowing — but nothing in the artifact or the docs said so, and on B300, GB200 and GB300 the `candidate` NCCL EP is the ONLY low-latency row, so those three SKUs publish decode numbers no deployment can reproduce. 1. **Maturity is now carried, not implied.** `EPBackend.maturity` is `production` or `candidate`, each adapter declares its own, and `implementation.maturity` lands in every case-attempt. `configs/platform_config.json` carries the same map for the matrix and the docs. Two copies of one fact drift silently, so `tests/test_matrix.py` pins coverage, a closed vocabulary, and agreement — reading the adapter side out of source with `ast`, because importing an adapter pulls in torch and the vendor EP library. 2. **The README was wrong about NCCL EP low latency.** It stated in two places that the `LOW_LATENCY` algorithm "has no enabled cell" / "no low-latency row on any SKU". #2407 enabled it on all six NVIDIA SKUs and changed only the config, so the docs have been contradicting the registry since. Corrected, and the B300/GB200/GB300 decode gap is now stated where a reader will meet it. 3. **B300 DeepEP V2 low latency is documented as an unsupported coverage row.** It had never carried one and no reason was recorded. Measured on b300-002: the legacy Buffer self-enables NVSHMEM IBGDA even for single-node EP8, and the transport aborts in setup — `ibgda.cpp:2234 NULL value Unable to create ah` -> `create DCT share err` -> `connect EPS failed` -> `nvshmem setup connections failed`, rc255 on all eight ranks. This is NOT the /dev/gdrdrv gap previously suspected; address-handle creation fails because a single-node run scrubs the HCA selector and NVSHMEM auto-picks a fabric it cannot form an AH on. Since single-node EP8 needs no IB, `NVSHMEM_DISABLE_IB=1` is the candidate fix, but the two runs that would have proven it died in staging before reaching a GPU and the partition then filled, so the row stays out until a run earns it. 中文:为后端标注成熟度(maturity),并修正低延迟相关文档。 这三件事都取决于同一个问题:这个传输实现究竟有没有人能真正跑起来? 对照 vLLM 的 `--all2all-backend` 与 SGLang 的 `--moe-a2a-backend` 做的排查发现,我们基准测试的 四个传输实现中有两个在任一引擎里都无法被选中。DeepEP V2 与 MoRI 可以(`deepep_v2`/`deepep`、 `mori_*`/`mori`);UCCL-EP 与 NCCL EP 不行。作为候选(candidate)这本身没问题 —— 它们是真实的传输 实现,其探明的边界也有价值 —— 但产物和文档都没有说明这一点;而在 B300、GB200、GB300 上,属于 candidate 的 NCCL EP 是**唯一**的低延迟行,于是这三个 SKU 发布的解码数据没有任何真实部署能复现。 1. **成熟度现在被显式携带,而不是靠读者意会。** `EPBackend.maturity` 取值 `production` 或 `candidate`,由各 adapter 自行声明,并以 `implementation.maturity` 写入每一份 case-attempt 产物;`configs/platform_config.json` 保存同一份映射供矩阵与文档使用。同一事实的两份副本会悄悄 漂移,因此 `tests/test_matrix.py` 固定了覆盖面、封闭取值集合与两者的一致性 —— adapter 一侧用 `ast` 从源码读取,因为直接 import adapter 会连带引入 torch 和厂商 EP 库。 2. **README 对 NCCL EP 低延迟的描述是错的。** 它在两处声称 `LOW_LATENCY` 算法「没有启用的单元 格」「在任何 SKU 上都没有低延迟行」。#2407 已在全部六个 NVIDIA SKU 上启用该行,且只改了配置, 于是文档自那时起一直与注册表相互矛盾。现已更正,并把 B300/GB200/GB300 的解码缺口写在读者会先 读到的位置。 3. **B300 的 DeepEP V2 低延迟被记录为 unsupported coverage row。** 此前它从未有过该行,也没有留下 原因。在 b300-002 上实测:legacy Buffer 即使在单节点 EP8 下也会自行启用 NVSHMEM IBGDA,而该传输 在初始化阶段就失败 —— `ibgda.cpp:2234 NULL value Unable to create ah` -> `create DCT share err` -> `connect EPS failed` -> `nvshmem setup connections failed`,八个 rank 全部 rc255。这**不是** 此前怀疑的 /dev/gdrdrv 缺失问题;真正原因是单节点运行会清除 HCA 选择器,NVSHMEM 自动挑到一张无法 建立 AH 的网卡。既然单节点 EP8 根本不需要 IB,`NVSHMEM_DISABLE_IB=1` 是候选修复方案;但用于验证 它的两次运行都在 staging 阶段就退出、未触及 GPU,随后分区被占满,因此在有运行结果证明之前该行 暂不启用。 --- experimental/CollectiveX/README.md | 23 ++++---- experimental/CollectiveX/bench/ep_backend.py | 8 +++ .../CollectiveX/bench/ep_deepep_v2.py | 1 + experimental/CollectiveX/bench/ep_harness.py | 4 ++ experimental/CollectiveX/bench/ep_mori.py | 1 + experimental/CollectiveX/bench/ep_nccl.py | 1 + experimental/CollectiveX/bench/ep_uccl.py | 1 + .../CollectiveX/configs/platform_config.json | 1 + experimental/CollectiveX/docs/methodology.md | 6 +- experimental/CollectiveX/sweep_matrix.py | 4 ++ experimental/CollectiveX/tests/test_matrix.py | 55 +++++++++++++++++++ 11 files changed, 94 insertions(+), 11 deletions(-) diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index 1e1f0ce521..ceafd7146b 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -28,10 +28,13 @@ run in one of two modes: pure-intranode, same compact layout and unweighted rank-sum combine as `IntraNode`). It is a decode-phase-only, per-SKU-capability-gated addition whose runnable set differs from `normal`'s, so it is enabled from each SKU's `ll_backends` registry entry (currently DeepEP V2 EP8 on H100/H200/B200, - MoRI EP8 on MI300X/MI325X/MI355X, and UCCL-EP EP8 on H100/H200/B200 only — UCCL's low-latency kernel - trips a warp-group assertion on AMD's CU count, so the AMD SKUs keep UCCL-EP normal mode without LL; - NCCL EP has no low-latency row on any SKU while its decode kernels carry - [NVIDIA/nccl#2303](https://github.com/NVIDIA/nccl/issues/2303)). + MoRI EP8 on MI300X/MI325X/MI355X, UCCL-EP EP8 on H100/H200/B200 only — UCCL's low-latency kernel + trips a warp-group assertion on AMD's CU count, so 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. + B300, GB200 and GB300 carry NCCL EP as their *only* low-latency row, and NCCL EP is a `candidate` + transport (no engine exposes a selector for it), so those three SKUs have no production decode + coverage; B300's DeepEP V2 low-latency row is an unsupported coverage row, see the backend table). Scoped single-node EP8 runs over the intra-node NVLink/XGMI low-latency path (no `/dev/gdrdrv` needed — validated on H200 with it absent); NVSHMEM/IBGDA on the wire is only a multi-node scale-out (EP16) concern. @@ -67,12 +70,12 @@ frozen digest or locked case count. Physical host count does not determine scope: both GB topologies stay inside one 72-GPU MNNVL scale-up domain. -| Backend | Current scope | -|---|---| -| DeepEP V2 | `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/EP8 only | -| MoRI | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), 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 | [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 | [NCCL EP](https://github.com/NVIDIA/nccl/tree/master/contrib/nccl_ep): NVIDIA's native MoE dispatch/combine on the NCCL Device API — LSA (NVLink load/store) intra-node, 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 is implemented in the adapter but has no enabled cell (see the `ll_backends` note above). BF16 only: NCCL EP's FP8 machinery exists upstream but its RELEASE.md lists it unsupported/untested, so no FP8 case is emitted. 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 | +| Backend | Engine availability | Current scope | +|---|---|---| +| DeepEP V2 | `production` — 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/EP8 only. 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 the transport aborts during setup — `ibgda.cpp:2234 NULL value Unable to create ah` -> `create DCT share err` -> `connect EPS failed` -> `nvshmem setup connections failed`, rc255 on all eight ranks (measured 2026-08-02 on b300-002). Single-node EP8 needs no IB at all, so forcing the NVLink low-latency path (`NVSHMEM_DISABLE_IB=1`, the adapter already passes `allow_nvlink_for_low_latency_mode=True`) is the candidate fix; it is untested, so the row stays out until a run proves it | +| MoRI | `production` — vLLM `--all2all-backend mori_*`, SGLang `--moe-a2a-backend mori` | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), 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 — LSA (NVLink load/store) intra-node, 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. BF16 only: NCCL EP's FP8 machinery exists upstream but its RELEASE.md lists it unsupported/untested, so no FP8 case is emitted. 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 | DeepEP V2 means the `ElasticBuffer` implementation introduced by [DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index faf2726e15..0c791cc526 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -67,6 +67,14 @@ class EPBackend(abc.ABC): """ name: str = "" + # Whether a production inference engine can actually select this transport today: + # "production" = exposed as an all-to-all backend by vLLM (`--all2all-backend`) or + # SGLang (`--moe-a2a-backend`); "candidate" = a real transport we benchmark, but one + # no engine ships a selector for, so its numbers describe the library rather than a + # deployable configuration. Emitted per case-attempt so a reader can tell the two + # apart; `configs/platform_config.json` carries the same map for the matrix and the + # docs, and tests/test_matrix.py holds the two in agreement. + maturity: str = "" SUPPORTED_MODES: tuple = ("normal",) # Dispatch precisions the adapter realizes. BF16 is the universal control; an # adapter that also sends an FP8-quantized dispatch payload widens this. diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py index 8708d9886d..f74cfcf873 100644 --- a/experimental/CollectiveX/bench/ep_deepep_v2.py +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -116,6 +116,7 @@ def _require_runtime() -> None: class DeepEPV2Backend(EPBackend): name = "deepep-v2" + maturity = "production" # vLLM --all2all-backend deepep_v2; SGLang --moe-a2a-backend deepep # Two kernel families under one adapter, selected by mode: # normal -> PR #605 ElasticBuffer (LSA vs hybrid GIN are transport paths, not # kernel families); rank-deduplicated unweighted-rank-sum combine. diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 387d4a44a4..71b58031fe 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -1131,6 +1131,10 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # EPBackend.fp8_consume. Only meaningful when the case dispatches FP8. "fp8_consume": getattr(backend, "fp8_consume", None), "kernel_generation": kernel_generation(backend), + # Whether an inference engine can select this transport today — see + # EPBackend.maturity. A "candidate" row measures the library, not a + # deployable configuration, and must not be read as one. + "maturity": getattr(backend, "maturity", None) or "unknown", "name": backend.name, }, "topology": { diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index e3adeeb71b..2cda3f974c 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -38,6 +38,7 @@ def _project_local_metadata(torch_module, raw_expert_ids, raw_weights, rank, exp class MoRIBackend(EPBackend): name = "mori" + maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") combine_needs_redispatch = True diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index 14b4c2de9e..3b1010fe7b 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -65,6 +65,7 @@ class NCCLEPBackend(EPBackend): name = "nccl-ep" + maturity = "candidate" # NVIDIA's library, but no engine exposes an NCCL-EP selector # 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). # normal -> HT / FLAT layout / unweighted-rank-sum combine. diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 821b23c25d..9ef1fa7899 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -123,6 +123,7 @@ def _normal_num_sms() -> int: class UCCLEPBackend(EPBackend): name = "uccl-ep" + maturity = "candidate" # no engine exposes a UCCL-EP all-to-all selector # One legacy Buffer under two modes, selected by args.mode: # normal -> get_dispatch_layout/dispatch/combine; unweighted rank-sum combine. # low-latency -> low_latency_dispatch/combine decode kernels; source-side weighted combine. diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 3babb66536..018c79805a 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -1,4 +1,5 @@ { + "backend_maturity": {"deepep-v2": "production", "mori": "production", "uccl-ep": "candidate", "nccl-ep": "candidate"}, "platforms": { "h100-dgxc": { "arch": "sm90", diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 4ec82c7258..2a8840e8b0 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -202,7 +202,11 @@ One raw case document carries `record_type: "case-attempt"` and the single `vers - `workload`: `cross_rank_consistent`, whether the routing trace was proven identical across ranks; - `measurement`: dispatch/combine dtype (the realized wire formats — combine always BF16, dispatch BF16 or the SKU's FP8 format) and semantics, `sampling`, and the per-point `rows`; -- `implementation`: backend name and kernel generation; +- `implementation`: backend name, kernel generation, and `maturity` — whether a production + inference engine can select this transport today (`production` = exposed by vLLM's + `--all2all-backend` or SGLang's `--moe-a2a-backend`; `candidate` = a real transport we + benchmark that no engine ships a selector for, so its numbers describe the library rather + than a deployable configuration). The same map is in the registry's `backend_maturity`; - `topology`: requested SKU/product, placement, nodes, scale-up domain, transport, and world size; - `provenance`: the mounted image tag and source SHA; and - `outcome`: `status` (`success` or `invalid`) and `reasons`. diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index 9a38f7f9ee..9a3dbc0c11 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -26,6 +26,10 @@ def _load_config(name: str) -> dict[str, Any]: SWEEP = _load_config("sweep.json") PLATFORMS = _load_config("platform_config.json")["platforms"] +# Whether an inference engine can select each transport today (see EPBackend.maturity). +# Consumed by the docs and the matrix; each adapter carries the same value for the +# artifact it writes, and tests/test_matrix.py holds the two in agreement. +BACKEND_MATURITY = _load_config("platform_config.json")["backend_maturity"] SWEEP_BACKENDS = tuple(dict.fromkeys( backend for platform in PLATFORMS.values() for backend in platform["backends"] )) diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index 83d3c2a4b7..b799506b17 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -281,5 +281,60 @@ def test_invalid_filters_fail_closed(self): sweep_matrix.resolve_matrix(**options) +class BackendMaturityTests(unittest.TestCase): + """The registry map and each adapter's `maturity` are two copies of one fact. + + The registry drives the matrix and the docs; the adapter attribute is what lands in + every case-attempt artifact. They are read by different consumers and can drift + silently, so pin both: complete coverage, a closed vocabulary, and agreement. The + adapter side is read from source rather than imported, because importing an adapter + pulls in torch and the vendor EP library, which the test image does not carry. + """ + + VOCABULARY = {"production", "candidate"} + + @staticmethod + def _declared_in_source(): + """{backend name: maturity} parsed from the adapter class bodies.""" + import ast + + declared = {} + for path in sorted((ROOT / "bench").glob("ep_*.py")): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if not isinstance(node, ast.ClassDef): + continue + literals = {} + for statement in node.body: + if not isinstance(statement, ast.Assign): + continue + if not isinstance(statement.value, ast.Constant): + continue + for target in statement.targets: + if isinstance(target, ast.Name): + literals[target.id] = statement.value.value + # The abstract base declares both as empty defaults; skip it. + if literals.get("name") and "maturity" in literals: + declared[literals["name"]] = literals["maturity"] + return declared + + def test_registry_covers_every_dispatched_backend(self): + maturity = sweep_matrix.BACKEND_MATURITY + for sku, platform in sweep_matrix.PLATFORMS.items(): + for backend in platform["backends"]: + with self.subTest(sku=sku, backend=backend): + self.assertIn(backend, maturity) + self.assertIn(maturity[backend], self.VOCABULARY) + + def test_adapters_and_registry_agree(self): + declared = self._declared_in_source() + # Every backend the matrix can dispatch must declare a maturity in its adapter, + # or the artifact it writes would say "unknown" while the registry says otherwise. + for backend, expected in sweep_matrix.BACKEND_MATURITY.items(): + with self.subTest(backend=backend): + self.assertIn(backend, declared) + self.assertEqual(declared[backend], expected) + + if __name__ == "__main__": unittest.main() From 9423456b16ddd846e61ce77f85a87f13f629ac1b Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 11:31:15 +0800 Subject: [PATCH 02/24] CollectiveX: add FlashInfer one-sided NVLink EP for the GB SKUs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GB200 and GB300 are an NVLink-only domain, and a deployment picks its all-to-all by fabric: vLLM's own recipes say `--all2all-backend deepep_v2` for RDMA and `flashinfer_nvlink_one_sided` for NVLink. Our GB rows carried only deepep-v2 and nccl-ep — two scale-out-heritage transports pointed at MNNVL — so every GB number we have published describes a transport nobody would select on that rack. This adds the one they would. One-sided, never two-sided. The sibling `flashinfer_nvlink_two_sided` backend produces gibberish output on GB200/arm64 (vllm#39722), so it is deliberately not wired. One-sided means the initiator writes straight into the target's workspace then sets a flag; on MNNVL, where peer memory is directly addressable across the 72-GPU domain, that collapses to ordinary stores over the fabric with no rendezvous — which is the whole reason it wins here. Three properties of `flashinfer.comm.trtllm_moe_alltoall.MoeAlltoAll` shape the adapter: * Strict phase pairing — `dispatch` asserts "called twice without combine", `combine` asserts the phase is "dispatched", and state resets after each combine. So `combine_needs_redispatch` and `dispatch_needs_combine_cleanup` are both True, the same contract MoRI and the DeepEP V2 low-latency path already declare. * A PADDED receive, `[ep_size, runtime_max_tokens_per_rank, hidden]` rather than a compact buffer, so both oracle views read valid slots only — padding is untouched workspace memory and feeding it to the correctness gate would be feeding it garbage. * `runtime_max_tokens_per_rank <= max_num_tokens`, so the workspace is sized once from the ladder maximum and each call passes its own rung. `normal` mode only, and that is deliberate rather than a first slice. FlashInfer exposes one one-sided A2A kernel family, not a separate decode-optimized one, so an `ll_backends` cell would re-measure the same kernel under a mode that promises a different one. Decode is still covered: `normal` runs the full decode and prefill ladders. BF16 only for the same kind of reason — the combine side accepts FP8 output dtypes, but an FP8 dispatch needs the scale payload plumbed as a second `input_payload` and validated against the oracle's cast round-trip. No build seam. The pinned SGLang images ship `flashinfer-python` and the one-sided A2A lives in that same wheel, so `flashinfer_ep_prepare` is a capability assert rather than an install: it fails loudly at prepare time if an image ever drops the module, instead of dying mid-case inside `create_buffer`. NOT yet run on GB hardware. The adapter is written against the verified upstream API and vLLM's reference integration, the matrix/registry/launcher wiring is exercised by tests, but no leg has executed on gb200/gb300. The open item bring-up must settle first is `combine_weight_semantics`: upstream source does not state whether combine applies the top-k scales (they ride along only as a caller-supplied dispatch payload, and vLLM applies them in the MoE layer), so it is declared `unweighted-rank-sum` here and must be confirmed against the oracle before any published run — wrong semantics reds every case. Stacked on #2449, which must merge first: it adds `backend_maturity` plus a test that every dispatched backend has an entry, and this PR adds a dispatched backend. Merged the other way round, #2449's own test would fail on main. Stacking rather than noting the order means the `"flashinfer-ep": "production"` entry is already in place here, and the combined tree passes 75 tests. 中文:为 GB SKU 新增 FlashInfer 单边(one-sided) NVLink EP 后端。 GB200 与 GB300 是纯 NVLink 域,而真实部署会按 fabric 选择 all-to-all:vLLM 官方 recipe 明确写 着 RDMA 用 `--all2all-backend deepep_v2`、NVLink 用 `flashinfer_nvlink_one_sided`。我们的 GB 行 此前只有 deepep-v2 与 nccl-ep —— 两个面向 scale-out 的传输实现被指向了 MNNVL —— 因此我们已发布 的每一个 GB 数据,描述的都是没有人会在该机架上选用的传输实现。本提交补上了真正会被选用的那一个。 只接单边,不接双边。姊妹后端 `flashinfer_nvlink_two_sided` 在 GB200/arm64 上会输出乱码 (vllm#39722),因此有意不接入。单边意味着发起方直接写入目标的 workspace 再置标志位;在 MNNVL 上, 对端显存在 72-GPU 域内可直接寻址,于是这退化为经由 fabric 的普通 store,完全没有 rendezvous —— 这正是它在此拓扑上取胜的根本原因。 `flashinfer.comm.trtllm_moe_alltoall.MoeAlltoAll` 的三个特性决定了 adapter 的形态: * 严格的阶段配对 —— `dispatch` 断言「未 combine 就二次 dispatch」,`combine` 断言当前处于 "dispatched" 阶段,且每次 combine 后内部状态复位。因此 `combine_needs_redispatch` 与 `dispatch_needs_combine_cleanup` 均为 True,与 MoRI 和 DeepEP V2 低延迟路径already声明的契约一致。 * 接收缓冲是**带 padding 的** `[ep_size, runtime_max_tokens_per_rank, hidden]`,而非紧凑布局, 因此两个 oracle 视图都只读取有效槽位 —— padding 是未初始化的 workspace 内存,把它喂给正确性 门禁等同于喂垃圾数据。 * `runtime_max_tokens_per_rank <= max_num_tokens`,因此 workspace 按 ladder 最大值一次性分配, 每次调用传入当前档位。 仅支持 `normal` 模式,这是刻意为之而非「先做一部分」。FlashInfer 只暴露一个单边 A2A kernel 家族, 并没有单独的解码优化 kernel,因此新增 `ll_backends` 单元格只会在一个承诺「不同 kernel」的模式下 重复测量同一个 kernel。解码仍有覆盖:`normal` 会跑完整的 decode 与 prefill ladder。BF16-only 同理 —— combine 侧支持 FP8 输出 dtype,但 FP8 dispatch 需要把 scale 负载作为第二个 `input_payload` 接入并与 oracle 的 cast round-trip 对齐验证。 无需构建步骤。固定的 SGLang 镜像自带 `flashinfer-python`,单边 A2A 就在同一个 wheel 里,因此 `flashinfer_ep_prepare` 是能力断言而非安装:一旦镜像不再提供该模块,它会在 prepare 阶段大声失败, 而不是在 `create_buffer` 里中途崩溃。 **尚未在 GB 硬件上运行。** adapter 是依据已核实的上游 API 与 vLLM 参考实现编写的,矩阵/注册表/ launcher 的接线也有测试覆盖,但还没有任何 leg 在 gb200/gb300 上真正执行。bring-up 首先要解决的 未决项是 `combine_weight_semantics`:上游源码并未说明 combine 是否会乘以 top-k 权重(这些权重仅 作为调用方提供的 dispatch 负载随行,而 vLLM 是在 MoE 层施加它们),因此此处声明为 `unweighted-rank-sum`,必须在任何正式发布运行之前用 oracle 确认 —— 语义错误会让所有 case 变红。 本 PR 基于 #2449 堆叠,且 #2449 必须先合入:它引入了 `backend_maturity` 以及「每个可派发后端 都必须有对应条目」的测试,而本 PR 恰好新增了一个可派发后端。若顺序颠倒,#2449 自身的测试会在 main 上失败。采用堆叠而非仅口头约定顺序,意味着 `"flashinfer-ep": "production"` 条目已经就位, 合并后的代码树 75 项测试全部通过。 --- .github/workflows/collectivex-sweep.yml | 2 +- .../CollectiveX/bench/ep_flashinfer.py | 201 ++++++++++++++++++ experimental/CollectiveX/bench/run_ep.py | 6 +- .../CollectiveX/configs/platform_config.json | 6 +- .../CollectiveX/launchers/launch_gb-nv.sh | 2 +- .../CollectiveX/runtime/prepare_backend.sh | 27 +++ experimental/CollectiveX/sweep_matrix.py | 4 + experimental/CollectiveX/tests/test_matrix.py | 30 +++ 8 files changed, 271 insertions(+), 7 deletions(-) create mode 100644 experimental/CollectiveX/bench/ep_flashinfer.py diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 702ad39a48..729b3fb2f7 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -10,7 +10,7 @@ on: description: "EP library to sweep — 'all' runs every EP backend in one matrix" type: choice default: all - options: [all, deepep-v2, mori, uccl-ep, nccl-ep] + options: [all, deepep-v2, mori, uccl-ep, nccl-ep, flashinfer-ep] only_sku: description: Restrict to one GHA runner pool; blank = all type: string diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py new file mode 100644 index 0000000000..7931c56621 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +"""FlashInfer one-sided NVLink EP: the transport a GB200/GB300 deployment actually runs. + +The GB SKUs are an NVLink-only domain, and vLLM selects a different all-to-all there than it +does on an RDMA fabric: `--all2all-backend deepep_v2` for RDMA, `flashinfer_nvlink_one_sided` +for NVLink. Benchmarking the GB racks with DeepEP V2 and NCCL EP therefore measured transports +no deployment would choose on that hardware. This adapter closes that gap. + +ONE-SIDED, not two-sided. The sibling `flashinfer_nvlink_two_sided` backend is deliberately not +wired: it produces gibberish output on GB200/arm64 (vllm#39722). One-sided means the initiator +writes straight into the target's workspace and then sets a flag — on MNNVL, where peer memory +is directly addressable across the 72-GPU domain, that degenerates to ordinary stores over the +fabric with no rendezvous, which is exactly why it wins on this topology. + +Upstream surface (flashinfer.comm.trtllm_moe_alltoall): + + MoeAlltoAll(mapping, max_num_tokens, top_k, num_experts, + workspace_size_per_rank=..., mnnvl_config=...) + .dispatch(token_selected_experts, input_payloads, runtime_max_tokens_per_rank) + -> list of [ep_size, runtime_max_tokens_per_rank, *] workspace-backed receives + .combine(payload, runtime_max_tokens_per_rank, output=...) + -> [local_num_tokens, elements_per_token] + +Three properties of that API shape this adapter: + + * Strict phase pairing. `dispatch` asserts "called twice without combine", `combine` asserts + the phase is "dispatched", and the internal state resets after each combine. So a timed + combine needs a fresh dispatch and a timed dispatch needs a draining combine — the same + contract MoRI and the DeepEP V2 low-latency path already declare. + * A PADDED receive. Dispatch returns `[ep_size, runtime_max_tokens_per_rank, hidden]`, not a + compact buffer, so every oracle view must read valid slots only. Padding is untouched + workspace memory; reading it would feed garbage to the correctness gate. + * `runtime_max_tokens_per_rank <= max_num_tokens`, so the workspace is sized once from the + ladder maximum and each call passes the current rung. + +`normal` mode only. FlashInfer exposes one one-sided A2A kernel family, not a separate +decode-optimized one, so there is no honest `low-latency` cell to add: an `ll_backends` row +here would re-measure the same kernel under a mode that promises a different one. The decode +phase is still covered — `normal` runs the full decode and prefill ladders. +""" +from __future__ import annotations + +import types + +import torch + +from ep_backend import EPBackend + + +class FlashInferEPBackend(EPBackend): + name = "flashinfer-ep" + maturity = "production" # vLLM --all2all-backend flashinfer_nvlink_one_sided + # One kernel family; see the module docstring for why there is no low-latency mode. + SUPPORTED_MODES = ("normal",) + # BF16 first. The combine side accepts fp8_e4m3fn/uint8 output dtypes and a + # use_low_precision accumulate, but dispatch FP8 needs the scale payload plumbed as a + # second input_payload and validated against the oracle's cast round-trip; not this pass. + SUPPORTED_PRECISIONS = ("bf16",) + kernel_generation = "flashinfer-mnnvl-one-sided" + stage_device_work = False + # The kernel scatters expert outputs back to the supplying rank; it does not multiply by + # the routing weights (those ride along as a caller payload, and vLLM applies them in the + # MoE layer, not in the A2A). Verified against the oracle during bring-up. + combine_weight_semantics = "unweighted-rank-sum" + # Forced by the phase asserts described in the module docstring. + combine_needs_redispatch = True + dispatch_needs_combine_cleanup = True + combine_input_attr = "combine_input" + + def __init__(self, args, rank, world_size, local_rank, device): + super().__init__(args, rank, world_size, local_rank, device) + self._a2a = None + self._max_tokens = None + self._recv_shape = None + + # ---- setup ------------------------------------------------------------------------------- + + def buffer_cap(self, args): + # The workspace is sized from the ladder maximum rather than a fixed slot budget, so + # there is no cap to clamp the ladder against. + return None + + def create_buffer(self, spec): + """Build the one MoeAlltoAll for this group, sized to the ladder maximum. + + The communicator handed to MnnvlConfig must span exactly the EP group: the kernel + asserts `workspace.size(0) == moe_ep_size`, so a wider group silently mis-sizes it. + """ + from flashinfer.comm import Mapping + from flashinfer.comm.mnnvl import MnnvlConfig + from flashinfer.comm.trtllm_moe_alltoall import ( + MoeAlltoAll, + moe_a2a_get_workspace_size_per_rank, + ) + + self._max_tokens = spec.max_tokens_per_rank + hidden = self.args.hidden + top_k = self.args.topk + # Dispatch carries the activation plus the routing metadata the kernel needs per token: + # int32 expert ids and fp32 gate weights, top_k of each. Combine carries BF16 hidden. + dispatch_bytes = hidden * 2 + top_k * 4 + top_k * 4 + combine_bytes = hidden * 2 + workspace_size = moe_a2a_get_workspace_size_per_rank( + ep_size=self.world_size, + max_num_tokens=self._max_tokens, + total_dispatch_payload_size_per_token=dispatch_bytes, + combine_payload_size_per_token=combine_bytes, + ) + mapping = Mapping( + self.world_size, + self.rank, + self.args.gpus_per_node, + tp_size=self.world_size, + moe_ep_size=self.world_size, + ) + self._a2a = MoeAlltoAll( + mapping=mapping, + max_num_tokens=self._max_tokens, + top_k=top_k, + num_experts=self.args.experts, + workspace_size_per_rank=workspace_size, + mnnvl_config=MnnvlConfig(comm_backend=_TorchDistCommunicator()), + ) + self._recv_shape = (self.world_size, self._max_tokens, hidden) + + # ---- transport contract ------------------------------------------------------------------ + + def dispatch(self, p): + # token_selected_experts must be int32; the payload list is positional, and index 0 is + # the activation the combine later sends back. + received = self._a2a.dispatch( + p.topk_idx.to(torch.int32), + [p.dispatch_x], + p.T, + ) + return types.SimpleNamespace(recv_x=received[0], tokens=p.T, combine_input=None) + + def stage(self, p, h): + # BF16 needs no conversion: the received workspace tensor is the combine input. + h.combine_input = h.recv_x + + def combine(self, p, h): + out = self._a2a.combine(h.combine_input, h.tokens) + h.out = out + return out + + def recv_tokens(self, h): + # Every rank receives one padded plane per peer; the valid rows are the tokens each + # peer actually sent this rung, which for a fixed trace is its own T. + return int(h.recv_x.shape[0] * h.tokens) + + # ---- correctness-oracle views ------------------------------------------------------------ + + def inspect_dispatch(self, p, h): + """Compact (source-rank, slot) view over the padded [ep, max_tokens, hidden] receive. + + Only the first `tokens` rows of each source plane carry data; the rest is untouched + workspace memory. Slicing rather than masking keeps the row order the oracle expects. + """ + return h.recv_x[:, : h.tokens, :].reshape(-1, h.recv_x.shape[-1]) + + def combine_transformed(self, p, h, transformed): + """Combine an oracle-supplied payload through the same kernel as the timed path. + + The transformed rows arrive in the compact layout `inspect_dispatch` returned, so they + are scattered back into a padded workspace-shaped buffer before the call; padding stays + zero so it cannot contribute to the sum. + """ + padded = torch.zeros(self._recv_shape, dtype=h.recv_x.dtype, device=h.recv_x.device) + padded[:, : h.tokens, :] = transformed.view( + self.world_size, h.tokens, h.recv_x.shape[-1] + ) + return self._a2a.combine(padded, h.tokens) + + +class _TorchDistCommunicator: + """Bridge MnnvlConfig to the process group the harness already established. + + FlashInfer needs a communicator covering exactly the EP group to exchange MNNVL handles; + the harness has one in torch.distributed, so wrap that rather than standing up a second. + """ + + def __init__(self): + import torch.distributed as dist + + self._dist = dist + + def Get_rank(self): + return self._dist.get_rank() + + def Get_size(self): + return self._dist.get_world_size() + + def allgather(self, data): + gathered = [None] * self._dist.get_world_size() + self._dist.all_gather_object(gathered, data) + return gathered + + def Split(self, color, key): + # The harness's group already IS the EP group, so a split returns the same view. + return self diff --git a/experimental/CollectiveX/bench/run_ep.py b/experimental/CollectiveX/bench/run_ep.py index 9e90daf7ea..a0953553d6 100644 --- a/experimental/CollectiveX/bench/run_ep.py +++ b/experimental/CollectiveX/bench/run_ep.py @@ -51,7 +51,7 @@ def _runtime_info(torch, *, vendor: str) -> dict: def main() -> int: ap = argparse.ArgumentParser(description="CollectiveX EP dispatch/combine sweep") ap.add_argument("--backend", required=True, - choices=["deepep-v2", "mori", "uccl-ep", "nccl-ep"]) + choices=["deepep-v2", "mori", "uccl-ep", "nccl-ep", "flashinfer-ep"]) ep_harness.add_common_args(ap) args = ap.parse_args() @@ -95,13 +95,15 @@ def main() -> int: from ep_uccl import UCCLEPBackend as Backend elif args.backend == "nccl-ep": from ep_nccl import NCCLEPBackend as Backend + elif args.backend == "flashinfer-ep": + from ep_flashinfer import FlashInferEPBackend as Backend else: from ep_deepep_v2 import DeepEPV2Backend as Backend # MoRI registers the default GPU process group with its SHMEM runtime. Keep that # group device-only so scale-out does not also depend on a host Gloo fabric. if not dist.is_initialized(): - if args.backend in ("mori", "uccl-ep", "nccl-ep"): + if args.backend in ("mori", "uccl-ep", "nccl-ep", "flashinfer-ep"): # MoRI registers this group with its SHMEM runtime; UCCL-EP is portable across # NVIDIA (NCCL) and AMD (RCCL) and bootstraps its Buffer + CPU-proxy ranks from # it. NCCL EP forms its OWN NCCL communicator and uses this group only to broadcast diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 018c79805a..52de60ba91 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -1,5 +1,5 @@ { - "backend_maturity": {"deepep-v2": "production", "mori": "production", "uccl-ep": "candidate", "nccl-ep": "candidate"}, + "backend_maturity": {"deepep-v2": "production", "mori": "production", "uccl-ep": "candidate", "nccl-ep": "candidate", "flashinfer-ep": "production"}, "platforms": { "h100-dgxc": { "arch": "sm90", @@ -102,7 +102,7 @@ "scale_up_domain": 72, "scale_up_transport": "mnnvl", "launcher": "gb-nv", - "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]}, "ll_backends": {"nccl-ep": [8]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { @@ -122,7 +122,7 @@ "scale_up_domain": 72, "scale_up_transport": "mnnvl", "launcher": "gb-nv", - "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]}, "ll_backends": {"nccl-ep": [8]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { diff --git a/experimental/CollectiveX/launchers/launch_gb-nv.sh b/experimental/CollectiveX/launchers/launch_gb-nv.sh index 8488c5c739..2a4f258e9a 100644 --- a/experimental/CollectiveX/launchers/launch_gb-nv.sh +++ b/experimental/CollectiveX/launchers/launch_gb-nv.sh @@ -36,7 +36,7 @@ export COLLX_TRANSPORT=mnnvl export COLLX_NODES="$NODES" COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" export COLLX_NGPUS="$NGPUS" case "$COLLX_BENCH" in - deepep-v2 | nccl-ep) ;; + deepep-v2 | nccl-ep | flashinfer-ep) ;; *) collx_die "unsupported $PRODUCT EP backend: $COLLX_BENCH" ;; esac collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_ACCOUNT COLLX_SQUASH_DIR COLLX_STAGE_DIR diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh index 651129e0ec..fed9b6ec03 100644 --- a/experimental/CollectiveX/runtime/prepare_backend.sh +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -568,6 +568,32 @@ validate_container_network() { done } +# FlashInfer needs no build step: the pinned SGLang images ship `flashinfer-python`, and the +# one-sided MoE all-to-all lives in that same wheel. So this is a capability assert, not an +# install - fail loudly and early if the image ever drops it or ships a build without the +# trtllm_moe_alltoall module, rather than dying mid-case inside create_buffer. +flashinfer_ep_prepare() { + command -v python3 >/dev/null \ + || { collx_log "ERROR: python3 unavailable for FlashInfer EP"; return 1; } + python3 - <<'FICHECK' +import sys +try: + import flashinfer + from flashinfer.comm import Mapping # noqa: F401 + from flashinfer.comm.mnnvl import MnnvlConfig # noqa: F401 + from flashinfer.comm.trtllm_moe_alltoall import ( # noqa: F401 + MoeAlltoAll, + moe_a2a_get_workspace_size_per_rank, + ) +except Exception as exc: # noqa: BLE001 - the reason belongs in the leg log + print(f"flashinfer one-sided a2a import failed: {exc}", file=sys.stderr) + raise SystemExit(1) +print(f"FlashInfer {getattr(flashinfer, '__version__', 'unknown')} one-sided A2A available") +FICHECK + local rc=$? + [ "$rc" -eq 0 ] || { collx_log "ERROR: FlashInfer EP one-sided A2A unavailable in this image"; return 1; } +} + main() { collx_apply_network_profile "${COLLX_NODES:-1}" "${COLLX_TRANSPORT:-}" || return 1 validate_container_network || return 1 @@ -579,6 +605,7 @@ main() { ;; uccl-ep) uccl_prepare || return 1 ;; nccl-ep) nccl_ep_prepare || return 1 ;; + flashinfer-ep) flashinfer_ep_prepare || return 1 ;; *) collx_log "ERROR: unknown backend preparation request" return 1 diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index 9a3dbc0c11..c283f8d7cf 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -43,6 +43,10 @@ def _load_config(name: str) -> dict[str, Any]: # NCCL EP is BF16-only this release: its FP8 machinery exists upstream but RELEASE.md # lists it unsupported/untested, so no FP8 case is emitted (see bench/ep_nccl.py). "nccl-ep": ("bf16",), + # FlashInfer one-sided is BF16-only this pass: the combine side accepts FP8 output + # dtypes, but an FP8 dispatch needs the scale payload plumbed as a second + # input_payload and validated against the oracle cast round-trip. + "flashinfer-ep": ("bf16",), } # 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". diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index b799506b17..dfbb117101 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -203,6 +203,36 @@ def test_uccl_ep_rollout_shape(self): } self.assertEqual(ll_skus, {"h100-dgxc", "h200-dgxc", "b200-dgxc"}) + def test_flashinfer_ep_rollout_shape(self): + # FlashInfer one-sided is the transport a GB deployment actually runs: vLLM picks + # `flashinfer_nvlink_one_sided` on NVLink and `deepep_v2` on RDMA, so it belongs on the + # MNNVL SKUs and nowhere else this pass. + # * GB NVL72 (gb200/gb300): EP8 AND EP16, both inside the 72-GPU scale-up domain. + # * No x86 rows. The kernels are grouped upstream under MNNVL and vLLM gates them on an + # MNNVL-availability probe, so an HGX 8-GPU NVSwitch node may not qualify; that is an + # open question, not an assumed capability, and is left unclaimed until measured. + # * No AMD rows (NVIDIA/MNNVL only). + # * No low-latency row anywhere: FlashInfer exposes one one-sided A2A kernel family, not + # a separate decode kernel, so an ll_backends cell would re-measure the same kernel + # under a mode that promises a different one. Decode is still covered by the decode + # phase of normal mode. + # BF16 only this pass (FP8 dispatch needs the scale payload plumbed and oracle-validated). + document = matrix(backend="all") + cases = [ + item for item in document["requested_cases"] + if item["case"]["backend"] == "flashinfer-ep" + ] + runnable = { + (item["sku"], item["case"]["ep"]) + for item in cases if item["disposition"] == "runnable" + } + self.assertEqual(runnable, {(sku, ep) for sku in ("gb200", "gb300") for ep in (8, 16)}) + self.assertEqual({item["case"]["precision"] for item in cases}, {"bf16"}) + # Normal mode only — no low-latency cell on any SKU. + self.assertEqual({item["case"]["mode"] for item in cases}, {"normal"}) + for platform in sweep_matrix.PLATFORMS.values(): + self.assertNotIn("flashinfer-ep", platform.get("ll_backends", {})) + def test_nccl_ep_rollout_shape(self): # NCCL-EP's rollout, locked to the on-metal verdict (2026-07-22, all via the real launcher): # * RDMA scale-out SKUs (h100/h200/b200/b300): EP8 runnable, EP16 an UNSUPPORTED coverage From 893e3046e82ee2f482352902e9e8e2a98ef1d0f2 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 13:56:26 +0800 Subject: [PATCH 03/24] CollectiveX: give the FlashInfer communicator the whole CommBackend contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first gb200 leg died at the first dispatch with `CUDA error: unspecified launch failure` (sticky 719) on rank 6, cascading into NCCL teardown failures across both nodes. Prepare and create_buffer had both succeeded, so Mapping, MnnvlConfig and MoeAlltoAll were all constructed and accepted — the fault was in executing the first kernel. The cause was the communicator shim this adapter hands to MnnvlConfig. It was written from the shape of vLLM's CustomCommunicator rather than from the interface, and implemented only Get_rank, Get_size, allgather and Split. The upstream contract is `flashinfer.comm.mnnvl.CommBackend`, which also requires `bcast` and `barrier`, and whose allgather is scoped to the group rather than the default one. `barrier` is what mattered. MNNVL handle exchange has to complete on every rank before any rank's kernel touches peer memory; without it the first dispatch writes into memory the peer has not finished mapping, which is precisely an unspecified launch failure. A missing method on a duck-typed object failed asynchronously on the cluster instead of loudly at import. So the shim now subclasses the upstream ABC and implements the full contract — bcast via broadcast_object_list, barrier and group-scoped allgather — and create_buffer barriers on the EP group after construction, as vLLM does for the same reason. Subclassing means the next upstream interface change is an ImportError or a TypeError here, not another 719 an hour into an allocation. 中文:为 FlashInfer 通信适配层补齐完整的 CommBackend 契约。 首次 gb200 leg 在第一个 dispatch 处即失败,rank 6 报 `CUDA error: unspecified launch failure` (粘滞 719),并在两个节点上引发连锁的 NCCL 拆除失败。prepare 与 create_buffer 均已成功,说明 Mapping、MnnvlConfig 与 MoeAlltoAll 都已正确构造并被接受 —— 故障出在首个 kernel 的执行阶段。 根本原因是本 adapter 传给 MnnvlConfig 的通信适配层。它是照着 vLLM CustomCommunicator 的「外形」 而非其接口写的,只实现了 Get_rank、Get_size、allgather 与 Split。上游契约是 `flashinfer.comm.mnnvl.CommBackend`,它还要求 `bcast` 与 `barrier`,且 allgather 必须限定在该 group 而非默认 group 上。 关键在于 `barrier`。MNNVL 句柄交换必须在所有 rank 上完成之后,任何 rank 的 kernel 才能访问对端 显存;缺少它,首个 dispatch 就会写入对端尚未映射完成的内存 —— 这正是 unspecified launch failure 的成因。鸭子类型对象上缺失的方法,没有在 import 时大声报错,而是在集群上以异步方式失败。 因此该适配层现在继承上游 ABC 并实现完整契约 —— 用 broadcast_object_list 实现 bcast、实现 barrier、并将 allgather 限定到 group —— 且 create_buffer 在构造完成后于 EP group 上做一次 barrier,与 vLLM 出于同样原因的做法一致。继承 ABC 意味着上游接口一旦变更,这里会是 ImportError 或 TypeError,而不是又一次在分配开始一小时后才出现的 719。 --- .../CollectiveX/bench/ep_flashinfer.py | 66 ++++++++++++++----- 1 file changed, 48 insertions(+), 18 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 7931c56621..ca3eb18e99 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -119,8 +119,11 @@ def create_buffer(self, spec): top_k=top_k, num_experts=self.args.experts, workspace_size_per_rank=workspace_size, - mnnvl_config=MnnvlConfig(comm_backend=_TorchDistCommunicator()), + mnnvl_config=MnnvlConfig(comm_backend=_communicator(_ep_group())), ) + # Every rank must finish mapping its workspace before any peer writes into it; + # vLLM barriers here for the same reason. Scoped to the EP group, not the world. + torch.distributed.barrier(group=_ep_group()) self._recv_shape = (self.world_size, self._max_tokens, hidden) # ---- transport contract ------------------------------------------------------------------ @@ -173,29 +176,56 @@ def combine_transformed(self, p, h, transformed): return self._a2a.combine(padded, h.tokens) -class _TorchDistCommunicator: +def _ep_group(): + """The process group spanning the EP world. CollectiveX runs one group per case and the + default group IS the EP group, so this is the default — named for the kernel's assert + (`workspace.size(0) == moe_ep_size`), which a wider group would silently violate.""" + import torch.distributed as dist + + return dist.group.WORLD + +def _communicator(group): """Bridge MnnvlConfig to the process group the harness already established. - FlashInfer needs a communicator covering exactly the EP group to exchange MNNVL handles; - the harness has one in torch.distributed, so wrap that rather than standing up a second. + FlashInfer needs a communicator spanning exactly the EP group to exchange MNNVL fabric + handles; the harness has one in torch.distributed, so wrap that rather than standing up a + second. The contract is `flashinfer.comm.mnnvl.CommBackend` and it is not optional in any + part: an earlier version of this adapter implemented only rank/size/allgather/Split and + the first dispatch died with `CUDA error: unspecified launch failure` (sticky 719) on + gb200. `barrier` is the reason — handle exchange has to complete on every rank before any + rank's kernel touches peer memory, and without it the writes land on memory the peer has + not mapped yet. Built as a subclass of the upstream ABC so a future interface change is an + import-time error here rather than another asynchronous fault on the cluster. """ + import torch.distributed as dist + from flashinfer.comm.mnnvl import CommBackend + + class _TorchDistCommunicator(CommBackend): + def __init__(self, process_group): + self._group = process_group + + def Get_rank(self) -> int: + return self._group.rank() - def __init__(self): - import torch.distributed as dist + def Get_size(self) -> int: + return self._group.size() - self._dist = dist + def allgather(self, data): + gathered = [None] * self.Get_size() + dist.all_gather_object(gathered, data, group=self._group) + return gathered - def Get_rank(self): - return self._dist.get_rank() + def bcast(self, data, root): + # broadcast_object_list mutates the list in place. + payload = [data] + dist.broadcast_object_list(payload, src=root, group=self._group) + return payload[0] - def Get_size(self): - return self._dist.get_world_size() + def barrier(self) -> None: + dist.barrier(group=self._group) - def allgather(self, data): - gathered = [None] * self._dist.get_world_size() - self._dist.all_gather_object(gathered, data) - return gathered + def Split(self, color, key): + # The harness's group already IS the EP group, so a split returns the same view. + return self - def Split(self, color, key): - # The harness's group already IS the EP group, so a split returns the same view. - return self + return _TorchDistCommunicator(group) From 81efeae8db8119097b41c56c7d931dcdbffcafaf Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 14:05:14 +0800 Subject: [PATCH 04/24] CollectiveX: dispatch routing metadata so the FlashInfer oracle views are real MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the communicator fixed the transport runs — gb200 produced real timings (T=1024 disp_p99=194.7us, T=8192 873.3us) — but every point came back `passed=False` with `max_rel=0.000000` and `recv_total=0`. That combination is not a correctness failure, it is the oracle never getting a view: `inspect_dispatch` returned a bare tensor, `view.payload` raised inside run_sweep's try, and the handler returned the default report whose `receive_count` is 0. `point_ok` then failed on `recv_total > 0` while the error stayed at its 0.0 default. The contract is a view carrying `payload`, `expert_ids` and `weights` per received row. A transport moves opaque bytes, so the only way a receiver can know a row's routing is for the sender to ship it — which is exactly the shape upstream expects: the workspace accounting already budgets `top_k * 4` for int32 ids plus `top_k * 4` for fp32 weights on top of hidden. So dispatch now sends three payloads (activation, top-k ids, top-k weights) and passes `invalid_token_expert_id=-1` with `expert_id_payload_index=1`. That sentinel is how a receiver identifies real rows: the planes are `[ep_size, max_tokens, *]`, a rank only receives tokens that selected one of its experts, and the kernel stamps every slot it did not fill. The ids come back exactly as sent — already global, which is what the oracle compares — and per-row entries that are not local stay sentinel-stamped, so their weights are masked to zero rather than joining the expert sum. `recv_tokens` and `combine_transformed` are now derived from that same mask, so the oracle's scatter lands on precisely the slots its rows came from and unfilled slots stay zero. 中文:dispatch 一并传输路由元数据,使 FlashInfer 的 oracle 视图真正可用。 修好通信适配层后传输已经跑通 —— gb200 产出了真实时延(T=1024 disp_p99=194.7us,T=8192 873.3us)—— 但每个测点都是 `passed=False`,且 `max_rel=0.000000`、`recv_total=0`。这个组合并 不是正确性失败,而是 oracle 根本没拿到视图:`inspect_dispatch` 返回了裸 tensor,`view.payload` 在 run_sweep 的 try 中抛异常,处理分支返回了默认报告(其 `receive_count` 为 0)。于是 `point_ok` 因 `recv_total > 0` 不成立而失败,误差则停留在 0.0 的默认值。 真正的契约是一个按「已接收行」组织、带 `payload`、`expert_ids` 与 `weights` 的视图。传输层搬运 的是不透明字节,因此接收方要知道某一行的路由信息,唯一办法就是由发送方一并送来 —— 这也正是上游 预期的形态:workspace 计算本就在 hidden 之外预留了 `top_k * 4`(int32 ids)与 `top_k * 4` (fp32 weights)。 因此 dispatch 现在发送三个 payload(激活、top-k ids、top-k weights),并传入 `invalid_token_expert_id=-1` 与 `expert_id_payload_index=1`。该哨兵值正是接收方识别有效行的 依据:接收缓冲为 `[ep_size, max_tokens, *]`,某个 rank 只会收到选中了它本地专家的 token,kernel 会在所有未填充的槽位打上哨兵。ids 原样返回 —— 本就是全局 id,正是 oracle 比较所用 —— 而行内非 本地的条目仍保持哨兵状态,其权重被掩为 0,不会进入专家求和。 `recv_tokens` 与 `combine_transformed` 现在都基于同一个掩码推导,因此 oracle 的 scatter 会精确 落回这些行的来源槽位,未填充槽位保持为 0。 --- .../CollectiveX/bench/ep_flashinfer.py | 81 ++++++++++++++----- 1 file changed, 60 insertions(+), 21 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index ca3eb18e99..9fa6881251 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -46,6 +46,10 @@ from ep_backend import EPBackend +# Stamped by the kernel into the expert-id payload of every receive slot it did not fill. +# -1 is safe: real ids are [0, num_experts). +_INVALID_EXPERT = -1 + class FlashInferEPBackend(EPBackend): name = "flashinfer-ep" @@ -129,14 +133,31 @@ def create_buffer(self, spec): # ---- transport contract ------------------------------------------------------------------ def dispatch(self, p): - # token_selected_experts must be int32; the payload list is positional, and index 0 is - # the activation the combine later sends back. - received = self._a2a.dispatch( - p.topk_idx.to(torch.int32), - [p.dispatch_x], + """Dispatch the activation together with the routing metadata the oracle needs back. + + Three payloads, not one. The correctness oracle reads a per-received-row view carrying + `expert_ids` and `weights` alongside the payload, and the only way to know a received + row's routing is for the sender to ship it: the kernel moves opaque bytes. This is the + intended shape — the upstream workspace accounting budgets exactly `top_k * 4` for + int32 ids plus `top_k * 4` for fp32 weights on top of the hidden bytes. + + `invalid_token_expert_id` + `expert_id_payload_index` are how a receiver tells which + rows are real: the receive planes are `[ep_size, max_tokens, *]` and a rank only gets + the tokens that selected one of its experts, so the kernel stamps the sentinel into the + expert-id payload of every slot it did not fill. + """ + idx32 = p.topk_idx.to(torch.int32) + recv_x, recv_idx, recv_w = self._a2a.dispatch( + idx32, + [p.dispatch_x, idx32, p.topk_weights.to(torch.float32)], p.T, + invalid_token_expert_id=_INVALID_EXPERT, + expert_id_payload_index=1, + ) + return types.SimpleNamespace( + recv_x=recv_x, recv_idx=recv_idx, recv_w=recv_w, + tokens=p.T, topk=p.topk_idx.shape[1], combine_input=None, ) - return types.SimpleNamespace(recv_x=received[0], tokens=p.T, combine_input=None) def stage(self, p, h): # BF16 needs no conversion: the received workspace tensor is the combine input. @@ -148,32 +169,50 @@ def combine(self, p, h): return out def recv_tokens(self, h): - # Every rank receives one padded plane per peer; the valid rows are the tokens each - # peer actually sent this rung, which for a fixed trace is its own T. - return int(h.recv_x.shape[0] * h.tokens) + # Rows the kernel actually filled, across every source plane. + return int(self._valid_rows(h).sum().item()) # ---- correctness-oracle views ------------------------------------------------------------ + def _valid_rows(self, h): + """Boolean mask over the flattened [ep_size * max_tokens] receive slots.""" + idx = h.recv_idx.reshape(-1, h.topk) + return (idx != _INVALID_EXPERT).any(dim=1) + def inspect_dispatch(self, p, h): - """Compact (source-rank, slot) view over the padded [ep, max_tokens, hidden] receive. + """Compact per-received-row view for the correctness oracle. - Only the first `tokens` rows of each source plane carry data; the rest is untouched - workspace memory. Slicing rather than masking keeps the row order the oracle expects. + The receive is `[ep_size, max_tokens, *]` with a rank's tokens only in the slots the + kernel filled, so flatten and keep the rows whose expert-id payload is not the + sentinel. Ids come back exactly as sent, i.e. already GLOBAL, which is what the oracle + compares against; per-row entries that are not local are still sentinel-stamped, so + mask their weights to zero rather than letting them into the expert sum. """ - return h.recv_x[:, : h.tokens, :].reshape(-1, h.recv_x.shape[-1]) + keep = self._valid_rows(h) + hidden = h.recv_x.shape[-1] + payload = h.recv_x.reshape(-1, hidden)[keep] + ids = h.recv_idx.reshape(-1, h.topk).to(torch.int64)[keep] + weights = h.recv_w.reshape(-1, h.topk).to(torch.float32)[keep] + live = ids != _INVALID_EXPERT + return types.SimpleNamespace( + payload=payload, + expert_ids=torch.where(live, ids, torch.full_like(ids, -1)), + weights=weights.masked_fill(~live, 0.0), + ) def combine_transformed(self, p, h, transformed): - """Combine an oracle-supplied payload through the same kernel as the timed path. + """Combine an oracle-transformed payload through the same kernel as the timed path. - The transformed rows arrive in the compact layout `inspect_dispatch` returned, so they - are scattered back into a padded workspace-shaped buffer before the call; padding stays - zero so it cannot contribute to the sum. + `transformed` holds one row per row `inspect_dispatch` returned, so scatter it back + into the padded workspace-shaped buffer at exactly the slots those rows came from. + Unfilled slots stay zero and contribute nothing to the sum. """ - padded = torch.zeros(self._recv_shape, dtype=h.recv_x.dtype, device=h.recv_x.device) - padded[:, : h.tokens, :] = transformed.view( - self.world_size, h.tokens, h.recv_x.shape[-1] + padded = torch.zeros( + (self.world_size * self._max_tokens, h.recv_x.shape[-1]), + dtype=h.recv_x.dtype, device=h.recv_x.device, ) - return self._a2a.combine(padded, h.tokens) + padded[self._valid_rows(h)] = transformed.to(padded.dtype) + return self._a2a.combine(padded.view_as(h.recv_x), h.tokens) def _ep_group(): From 3ff418a26331e61dd0ba584ffd709e04a6d39b6f Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 14:08:42 +0800 Subject: [PATCH 05/24] CollectiveX: complete the FlashInfer oracle view (local masking + expert counts) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run 3 got past the missing-view failure and stopped on `AttributeError: 'types.SimpleNamespace' object has no attribute 'local_expert_counts'`. Rather than discover the contract one field per cluster run, the harness was read for every `view.*` access: payload, expert_ids, weights, local_expert_counts (combine_input is written by the harness, not supplied by the adapter). Two things were wrong, not one. The view lacked `local_expert_counts` — the per-local-expert arrival bincount the oracle checks against its own. And `expert_ids` shipped the token's whole global top-k, including experts owned by other ranks, while the oracle builds its expectation as 'global id where id // experts_per_rank == rank, else -1'. Unmasked, every row would have disagreed even with the transport perfectly correct. 中文:补全 FlashInfer 的 oracle 视图(本地掩码 + 专家计数)。 第 3 次运行越过了「视图缺失」的失败,停在 `AttributeError: 'types.SimpleNamespace' object has no attribute 'local_expert_counts'`。 为避免每跑一次集群才发现一个字段,这次直接通读 harness 中所有 `view.*` 访问:payload、 expert_ids、weights、local_expert_counts(combine_input 由 harness 写入,不需 adapter 提供)。 问题有两个而非一个。视图缺少 `local_expert_counts` —— 即 oracle 会与自身 bincount 比对的 「每个本地专家的到达计数」。另外 `expert_ids` 原样传回了 token 的完整全局 top-k,其中包含归属 其他 rank 的专家,而 oracle 的期望值构造方式是「若 id // experts_per_rank == rank 则取全局 id, 否则为 -1」。不做掩码的话,即使传输完全正确,每一行也都会判为不一致。 --- .../CollectiveX/bench/ep_flashinfer.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 9fa6881251..e3aaedcec6 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -76,6 +76,7 @@ def __init__(self, args, rank, world_size, local_rank, device): self._a2a = None self._max_tokens = None self._recv_shape = None + self.experts_per_rank = args.experts // world_size # ---- setup ------------------------------------------------------------------------------- @@ -183,21 +184,29 @@ def inspect_dispatch(self, p, h): """Compact per-received-row view for the correctness oracle. The receive is `[ep_size, max_tokens, *]` with a rank's tokens only in the slots the - kernel filled, so flatten and keep the rows whose expert-id payload is not the - sentinel. Ids come back exactly as sent, i.e. already GLOBAL, which is what the oracle - compares against; per-row entries that are not local are still sentinel-stamped, so - mask their weights to zero rather than letting them into the expert sum. + kernel filled, so flatten and keep rows whose expert-id payload is not the sentinel. + + Ids come back exactly as sent, i.e. GLOBAL and covering the token's whole top-k — + including experts owned by OTHER ranks. The oracle builds its expectation as "global id + where `id // experts_per_rank == rank`, else -1", so the non-local entries have to be + masked out here too, with their weights zeroed, or every row would disagree. """ keep = self._valid_rows(h) hidden = h.recv_x.shape[-1] payload = h.recv_x.reshape(-1, hidden)[keep] ids = h.recv_idx.reshape(-1, h.topk).to(torch.int64)[keep] weights = h.recv_w.reshape(-1, h.topk).to(torch.float32)[keep] - live = ids != _INVALID_EXPERT + local = (ids >= 0) & ((ids // self.experts_per_rank) == self.rank) + expert_ids = torch.where(local, ids, torch.full_like(ids, -1)) return types.SimpleNamespace( payload=payload, - expert_ids=torch.where(live, ids, torch.full_like(ids, -1)), - weights=weights.masked_fill(~live, 0.0), + expert_ids=expert_ids, + weights=weights.masked_fill(~local, 0.0), + # Per-local-expert arrival count; the oracle compares it against its own bincount. + local_expert_counts=torch.bincount( + (ids[local] - self.rank * self.experts_per_rank), + minlength=self.experts_per_rank, + ), ) def combine_transformed(self, p, h, transformed): From 2f53134a478b182aa42817bee8b6c4aec09c5ccb Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 14:11:41 +0800 Subject: [PATCH 06/24] CollectiveX: shape the FlashInfer combine scatter from the actual receive MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run 4 reached the combine and died with `IndexError: mask [8192] does not match indexed tensor [65536, 7168]`. combine_transformed built its scatter buffer from the ladder maximum (world_size * max_num_tokens) while the mask comes from the receive planes, which are [ep_size, runtime_max_tokens_per_rank, hidden] and shrink with the rung — 8 x 1024 at T=1024, not 8 x 8192. The module docstring already said this; the code did not. Build the buffer with zeros_like on the actual receive instead, and drop the _recv_shape attribute that encoded the same wrong assumption. 中文:FlashInfer 的 combine scatter 缓冲按实际接收形状构造。 第 4 次运行进入 combine 后失败: `IndexError: mask [8192] does not match indexed tensor [65536, 7168]`。 combine_transformed 按 ladder 最大值(world_size * max_num_tokens)构造 scatter 缓冲,而掩码来自 接收平面 —— 其形状为 [ep_size, runtime_max_tokens_per_rank, hidden],会随档位缩小:T=1024 时是 8 x 1024,而非 8 x 8192。模块 docstring 早已写明这一点,代码却没有照做。改为用 zeros_like 基于 实际接收张量构造,并删除同样编码了该错误假设的 _recv_shape 属性。 --- experimental/CollectiveX/bench/ep_flashinfer.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index e3aaedcec6..34693186e2 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -75,7 +75,6 @@ def __init__(self, args, rank, world_size, local_rank, device): super().__init__(args, rank, world_size, local_rank, device) self._a2a = None self._max_tokens = None - self._recv_shape = None self.experts_per_rank = args.experts // world_size # ---- setup ------------------------------------------------------------------------------- @@ -129,7 +128,6 @@ def create_buffer(self, spec): # Every rank must finish mapping its workspace before any peer writes into it; # vLLM barriers here for the same reason. Scoped to the EP group, not the world. torch.distributed.barrier(group=_ep_group()) - self._recv_shape = (self.world_size, self._max_tokens, hidden) # ---- transport contract ------------------------------------------------------------------ @@ -216,12 +214,12 @@ def combine_transformed(self, p, h, transformed): into the padded workspace-shaped buffer at exactly the slots those rows came from. Unfilled slots stay zero and contribute nothing to the sum. """ - padded = torch.zeros( - (self.world_size * self._max_tokens, h.recv_x.shape[-1]), - dtype=h.recv_x.dtype, device=h.recv_x.device, - ) - padded[self._valid_rows(h)] = transformed.to(padded.dtype) - return self._a2a.combine(padded.view_as(h.recv_x), h.tokens) + # Shaped from the ACTUAL receive, not the ladder maximum: the planes are + # [ep_size, runtime_max_tokens_per_rank, hidden], so they shrink with the rung. Sizing + # this from max_num_tokens builds a 65536-row buffer for an 8192-row mask. + padded = torch.zeros_like(h.recv_x) + padded.view(-1, h.recv_x.shape[-1])[self._valid_rows(h)] = transformed.to(padded.dtype) + return self._a2a.combine(padded, h.tokens) def _ep_group(): From 4237e23fb8d05eac6d33b4fe6ac41c4057915ac9 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 14:28:15 +0800 Subject: [PATCH 07/24] CollectiveX: model FlashInfer's BF16 combine accumulation instead of widening the gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run 5 moved data correctly — receive counts matched the routing exactly and T=1..16 passed the tight gate — but T=32 and up failed at 6-16 ulps against `8 * 2^-8`. The error was present at every rung, including T=1; larger rungs only cross the gate because the maximum is taken over more elements. That is a systematic precision difference, not noise. Measured it rather than inferring it. An 8-rank probe on gb200 staged 1.0 from rank 0 and 2^-9 from the other seven, values chosen so the two candidate reductions disagree visibly: kernel combine out : 1.0 FP32-accumulate : 1.015625 BF16-seq-accumulate: 1.0 FlashInfer's one-sided combine accumulates the per-rank BF16 messages in BF16. The oracle models the unweighted rank sum as an FP32 accumulation of those messages, which is right for DeepEP and MoRI and wrong for this kernel. So the oracle gains a third contract, `bf16-rank-sum`: the same per-rank BF16 messages and the same gate-folding at staging, reduced at payload precision instead of FP32. The alternative was loosening COMBINE_REL_TOL, which would have blinded every other backend to exactly the corruption that gate exists to catch — the b200 low-latency top-rung failure and the mi355x combine corruption were both caught by it. Modelling a backend's real reduction is the opposite of widening a tolerance: it makes the gate strict for this backend too. Summation ORDER is not observable from outside the kernel, so rank order is used. That is sound because the modelled quantity is precision: reordering a BF16 sum of like-magnitude terms moves the result by about an ulp, well inside the gate, whereas accumulating in FP32 when the kernel does not is a systematic 6-16 ulp error. `test_runtime.py` pins the widened normal-mode allowlist, and a new test reproduces the probe in pure python so the two contracts cannot silently collapse into each other — without it, declaring `bf16-rank-sum` could become a no-op that accepts an FP32-accumulating backend. 中文:为 FlashInfer 的 BF16 combine 累加建模,而不是放宽门禁。 第 5 次运行的数据搬运是正确的 —— 接收计数与路由完全吻合,T=1..16 通过了严格门禁 —— 但 T=32 及 以上以 6–16 ulp 的误差超出 `8 * 2^-8`。该误差在每个档位都存在,包括 T=1;更大的档位只是因为在更 多元素上取最大值才越过门禁。这是系统性的精度差异,而非噪声。 这是实测得出的结论,而非推断。在 gb200 上用 8 个 rank 做探测:rank 0 提供 1.0,其余七个提供 2^-9,这组取值能让两种候选归约方式产生明显不同的结果(见上)。 FlashInfer 的单边 combine 以 BF16 累加各 rank 的 BF16 消息。而 oracle 把 unweighted rank sum 建模为对这些消息的 FP32 累加 —— 这对 DeepEP 与 MoRI 是正确的,对该 kernel 则不然。 因此 oracle 新增第三种契约 `bf16-rank-sum`:同样的每 rank BF16 消息、同样在 staging 处折入门控 权重,但归约在 payload 精度而非 FP32 下进行。另一条路是放宽 COMBINE_REL_TOL,那会让所有其他后端 对这道门禁本应捕捉的数据损坏失明 —— b200 低延迟顶档失败与 mi355x combine 损坏都是它抓到的。 为后端的真实归约方式建模,与放宽容差恰恰相反:它让门禁对这个后端同样严格。 求和**顺序**在 kernel 外部不可观测,故此处采用 rank 顺序。这是成立的,因为被建模的量是精度:对量级 相近的项重排 BF16 求和只会带来约 1 ulp 的差异,远在门禁之内;而在 kernel 并非 FP32 累加时却按 FP32 建模,则是系统性的 6–16 ulp 误差。 `test_runtime.py` 固定了放宽后的 normal 模式允许集合,并新增一个用纯 python 复现该探测的测试, 以免两种契约悄悄退化为等价 —— 否则声明 `bf16-rank-sum` 可能变成一个无操作,把 FP32 累加的后端也 一并接受。 --- .../CollectiveX/bench/ep_flashinfer.py | 11 ++++-- experimental/CollectiveX/bench/ep_harness.py | 37 ++++++++++++++++--- .../CollectiveX/tests/test_runtime.py | 30 ++++++++++++++- 3 files changed, 67 insertions(+), 11 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 34693186e2..3c86ac9e92 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -62,10 +62,13 @@ class FlashInferEPBackend(EPBackend): SUPPORTED_PRECISIONS = ("bf16",) kernel_generation = "flashinfer-mnnvl-one-sided" stage_device_work = False - # The kernel scatters expert outputs back to the supplying rank; it does not multiply by - # the routing weights (those ride along as a caller payload, and vLLM applies them in the - # MoE layer, not in the A2A). Verified against the oracle during bring-up. - combine_weight_semantics = "unweighted-rank-sum" + # A plain rank sum — the kernel scatters expert outputs back to the supplying rank and does + # NOT multiply by the routing weights (confirmed on gb200: a gated combine would be ~87% + # off with topk=8, the measured error was 2%). But it accumulates those per-rank BF16 + # messages in BF16, not FP32: staging 1.0 from rank 0 and 2^-9 from seven peers returns + # 1.0 where an FP32 accumulate returns 1.015625. Hence bf16-rank-sum rather than + # unweighted-rank-sum, which differ only in the precision of the cross-rank reduction. + combine_weight_semantics = "bf16-rank-sum" # Forced by the phase asserts described in the module docstring. combine_needs_redispatch = True dispatch_needs_combine_cleanup = True diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 71b58031fe..027c503c34 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -54,11 +54,17 @@ def case_id(sku: str, case: dict) -> str: # unweighted per-expert transform and the kernel multiplies by the gate), whereas MoRI's # decode kernels (IntraNodeLL/AsyncLL) keep the plain additive rank sum and reduce the # gate weights in parallel ("unweighted-rank-sum", identical to normal mode). Normal mode -# is frozen to the unweighted v1 contract. The backend declares which it realizes; -# run_sweep only checks that the declared value is ALLOWED for the mode (so a mislabeled -# adapter fails closed), and the oracle keys on that same declared value. +# is frozen to the unweighted v1 contract. A third contract, "bf16-rank-sum", is the +# unweighted rank sum performed at REDUCED precision: FlashInfer's one-sided NVLink combine +# accumulates the per-rank BF16 messages in BF16 rather than FP32. Measured on gb200, 8 ranks: +# staging 1.0 from rank 0 and 2^-9 from the other seven returns 1.0, where an FP32 accumulate +# returns 1.015625. That is a property of the kernel, so the oracle has to reproduce it — the +# alternative is widening the gate, which would blind every other backend to real corruption. +# The backend declares which it realizes; run_sweep only checks that the declared value is +# ALLOWED for the mode (so a mislabeled adapter fails closed), and the oracle keys on that +# same declared value. MODE_ALLOWED_SEMANTICS = { - "normal": {"unweighted-rank-sum"}, + "normal": {"unweighted-rank-sum", "bf16-rank-sum"}, "low-latency": {"weighted-kernel-sum", "unweighted-rank-sum"}, } @@ -349,7 +355,9 @@ def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semant """ valid = expert_ids >= 0 expert = expert_ids.clamp(min=0).to(torch.int64) - if combine_weight_semantics == "unweighted-rank-sum": + if combine_weight_semantics in ("unweighted-rank-sum", "bf16-rank-sum"): + # Both are plain summing kernels, so the gate is folded in here; they differ only in + # the precision the kernel accumulates at, which is a property of the expectation. coefficient = weights.to(torch.float32).masked_fill(~valid, 0) elif combine_weight_semantics == "weighted-kernel-sum": coefficient = valid.to(torch.float32) @@ -409,6 +417,25 @@ def _expected_transformed_combine( gate = (weights[:, slot:slot + 1] * valid[:, slot:slot + 1].to(torch.float32)) expected += gate * transform return expected + if combine_weight_semantics == "bf16-rank-sum": + # Same per-rank BF16 messages as unweighted-rank-sum, but the cross-rank reduction runs + # at payload precision instead of FP32, so the accumulator is cast back every step. + # Summation ORDER is not observable from outside the kernel; rank order is used here. + # That is sound because the modelled quantity is the PRECISION, and reordering a BF16 + # sum of like-magnitude terms moves the result by about an ulp, far inside the gate — + # whereas accumulating in FP32 when the kernel does not is a systematic 6-16 ulp error. + destination = expert_ids // experts_per_rank + scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) + accumulator = torch.zeros_like(semantic_x, dtype=dtype) + for rank_id in sorted(destination.unique().tolist()): + gate = weights * (destination == rank_id) + contribution = ( + semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) + + (gate * offset_a).sum(dim=1, keepdim=True) + + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) + ).to(dtype) + accumulator = (accumulator.float() + contribution.float()).to(dtype) + return accumulator.float() if combine_weight_semantics != "unweighted-rank-sum": raise ValueError(f"unknown combine semantics {combine_weight_semantics!r}") destination = expert_ids // experts_per_rank diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index 160e6c562f..dd4ecbd79b 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -486,16 +486,42 @@ def test_guards_fail_closed(self) -> None: class ModeSemanticsContract(unittest.TestCase): # The combine contract is a backend fact, not a pure function of mode: DeepEP's # low-latency combine is weighted-kernel-sum while MoRI's IntraNodeLL is - # unweighted-rank-sum, so low-latency must admit both. Normal stays unweighted-only. + # unweighted-rank-sum, so low-latency must admit both. Normal admits the unweighted rank + # sum at either accumulation precision — FlashInfer's one-sided NVLink combine performs the + # same summation but accumulates the per-rank BF16 messages in BF16 rather than FP32 + # (measured on gb200), which the oracle reproduces as bf16-rank-sum instead of absorbing + # into a wider tolerance. weighted-kernel-sum stays out of normal mode. def test_mode_allowed_semantics(self) -> None: self.assertEqual( - ep_harness.MODE_ALLOWED_SEMANTICS["normal"], {"unweighted-rank-sum"} + ep_harness.MODE_ALLOWED_SEMANTICS["normal"], + {"unweighted-rank-sum", "bf16-rank-sum"}, ) self.assertEqual( ep_harness.MODE_ALLOWED_SEMANTICS["low-latency"], {"weighted-kernel-sum", "unweighted-rank-sum"}, ) + def test_bf16_rank_sum_is_lower_precision_than_fp32_rank_sum(self) -> None: + """The two normal-mode contracts must actually differ, or declaring the new one is + a no-op that would silently accept an FP32-accumulating backend as BF16 (and the + reverse). Reproduces the gb200 probe in pure python: seven sub-ulp addends vanish + under BF16 accumulation and survive under FP32.""" + import struct + + def bf16(value): # round-to-nearest-even truncation to bfloat16 + bits = struct.unpack("> 16) & 1)) & 0xFFFF0000 + return struct.unpack(" Date: Sun, 2 Aug 2026 14:39:56 +0800 Subject: [PATCH 08/24] CollectiveX: model FlashInfer's truncating BF16 conversion; revert bf16-rank-sum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts the bf16-rank-sum reduction model and replaces it with the measured cause. That model was built on a probe I misread: it staged 1.0 on rank 0 and 2^-9 elsewhere, but the routing sliced the top-k to 2, so only two ranks ever contributed and 1.0 + 2^-9 rounds to 1.0 under BOTH candidate reductions. The probe distinguished nothing, and applying its conclusion made agreement WORSE on hardware (T=1 0.024414 -> 0.030518), which is how a refuted model announces itself. Two corrected probes, top-k widened so all eight ranks actually contribute: completeness: 2^rank summed across 8 ranks -> 255.0 exactly. Every contribution lands. precision : 1.0 + 7 * 2^-9 -> 1.0078125 The exact sum is 129.75 bf16 ulps. Round-to-nearest gives 130 (1.015625); the kernel gives 129 (1.0078125). So the reduction is FP32 — completeness proves nothing is dropped and the value rules out a BF16 accumulator, which would return 1.0 — and the FP32 -> BF16 conversion TRUNCATES, keeping the high 16 bits of the word. torch's `.to(bfloat16)` rounds to nearest even, so the oracle was biased against the kernel by up to an ulp per conversion. That is a systematic offset, not noise that averages out, and a few ulps is exactly what a tight relative gate rejects. So the payload-dtype cast becomes explicit: `_to_payload_dtype` with a `nearest` default and a `truncate` mode, declared per adapter as `combine_output_rounding`. Every other backend keeps nearest and is bit-identical to before. COMBINE_REL_TOL is untouched — the gate stays exactly as strict for flashinfer-ep as for everything else; what changed is that the expectation now computes the arithmetic the kernel actually performs. 中文:为 FlashInfer 的截断式 BF16 转换建模;回退 bf16-rank-sum。 本提交回退 bf16-rank-sum 归约模型,代之以实测出的真正原因。那个模型建立在我误读的探测之上:它让 rank 0 提供 1.0、其余提供 2^-9,但路由把 top-k 截成了 2,因此实际只有两个 rank 参与贡献,而 1.0 + 2^-9 在**两种**候选归约下都舍入为 1.0。该探测什么都没有区分出来;把它的结论应用到硬件上反而 使一致性**变差**(T=1 从 0.024414 变为 0.030518)—— 这正是一个被证伪的模型的表现。 将 top-k 放宽、让八个 rank 全部参与后的两个修正探测(见上)表明:精确和为 129.75 个 bf16 ulp, 舍入到最近给出 130(1.015625),而 kernel 给出 129(1.0078125)。因此归约是 FP32 的 —— 完整性 探测证明没有贡献被丢弃,而该数值也排除了 BF16 累加器(那会返回 1.0)—— 并且 FP32 → BF16 的转换 是**截断**,即保留字的高 16 位。torch 的 `.to(bfloat16)` 采用就近偶数舍入,因此 oracle 相对 kernel 每次转换都存在最多 1 ulp 的系统性偏差。这是系统性偏移,而非可以相互抵消的噪声;而几个 ulp 恰好 就是严格相对门禁会判定失败的量级。 因此把 payload dtype 的转换显式化:`_to_payload_dtype`,默认 `nearest`,另有 `truncate` 模式,由 各 adapter 通过 `combine_output_rounding` 声明。其他所有后端保持 nearest,结果与此前逐位一致。 COMBINE_REL_TOL 未做任何改动 —— 门禁对 flashinfer-ep 与对其他后端同样严格;改变的只是期望值现在 按 kernel 真正执行的算术来计算。 --- experimental/CollectiveX/bench/ep_backend.py | 5 ++ .../CollectiveX/bench/ep_flashinfer.py | 16 ++-- experimental/CollectiveX/bench/ep_harness.py | 73 +++++++++--------- .../CollectiveX/tests/test_runtime.py | 75 ++++++++++++------- 4 files changed, 97 insertions(+), 72 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index 0c791cc526..961c773569 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -85,6 +85,11 @@ class EPBackend(abc.ABC): # Adapters that reduce activations and top-k weights independently must carry # the complete local weighted expert sum in the activation tensor. combine_weight_semantics = "unweighted-rank-sum" + # How the kernel converts its FP32 combine accumulator to the payload dtype. + # "nearest" (torch default) or "truncate" (keep the high 16 bits) — a kernel that + # truncates is biased down by up to an ulp per element, which a tight relative gate + # will catch, so the adapter must declare what its kernel actually does. + combine_output_rounding = "nearest" roundtrip_only = False # Realized wire formats recorded in the artifact. Combine is always BF16; # dispatch_dtype is overridden per-run by an FP8 adapter (e.g. "fp8-e4m3fn"). diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 3c86ac9e92..c5f9dd1de3 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -62,13 +62,15 @@ class FlashInferEPBackend(EPBackend): SUPPORTED_PRECISIONS = ("bf16",) kernel_generation = "flashinfer-mnnvl-one-sided" stage_device_work = False - # A plain rank sum — the kernel scatters expert outputs back to the supplying rank and does - # NOT multiply by the routing weights (confirmed on gb200: a gated combine would be ~87% - # off with topk=8, the measured error was 2%). But it accumulates those per-rank BF16 - # messages in BF16, not FP32: staging 1.0 from rank 0 and 2^-9 from seven peers returns - # 1.0 where an FP32 accumulate returns 1.015625. Hence bf16-rank-sum rather than - # unweighted-rank-sum, which differ only in the precision of the cross-rank reduction. - combine_weight_semantics = "bf16-rank-sum" + # The kernel scatters expert outputs back to the supplying rank; it does not multiply by + # the routing weights (those ride along as a caller payload, and vLLM applies them in the + # MoE layer, not in the A2A). Verified against the oracle during bring-up. + combine_weight_semantics = "unweighted-rank-sum" + # Measured on gb200, 8 ranks each contributing to one token: the kernel reduces + # 1.0 + 7 * 2^-9 (exactly 129.75 bf16 ulps) to 1.0078125 = 129 ulps. Round-to-nearest + # would give 1.015625 = 130. It accumulates in FP32 — a separate probe summing 2^rank + # across all eight ranks returned exactly 255.0 — and truncates on the way out. + combine_output_rounding = "truncate" # Forced by the phase asserts described in the module docstring. combine_needs_redispatch = True dispatch_needs_combine_cleanup = True diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 027c503c34..fd93e0ff54 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -54,17 +54,11 @@ def case_id(sku: str, case: dict) -> str: # unweighted per-expert transform and the kernel multiplies by the gate), whereas MoRI's # decode kernels (IntraNodeLL/AsyncLL) keep the plain additive rank sum and reduce the # gate weights in parallel ("unweighted-rank-sum", identical to normal mode). Normal mode -# is frozen to the unweighted v1 contract. A third contract, "bf16-rank-sum", is the -# unweighted rank sum performed at REDUCED precision: FlashInfer's one-sided NVLink combine -# accumulates the per-rank BF16 messages in BF16 rather than FP32. Measured on gb200, 8 ranks: -# staging 1.0 from rank 0 and 2^-9 from the other seven returns 1.0, where an FP32 accumulate -# returns 1.015625. That is a property of the kernel, so the oracle has to reproduce it — the -# alternative is widening the gate, which would blind every other backend to real corruption. -# The backend declares which it realizes; run_sweep only checks that the declared value is -# ALLOWED for the mode (so a mislabeled adapter fails closed), and the oracle keys on that -# same declared value. +# is frozen to the unweighted v1 contract. The backend declares which it realizes; +# run_sweep only checks that the declared value is ALLOWED for the mode (so a mislabeled +# adapter fails closed), and the oracle keys on that same declared value. MODE_ALLOWED_SEMANTICS = { - "normal": {"unweighted-rank-sum", "bf16-rank-sum"}, + "normal": {"unweighted-rank-sum"}, "low-latency": {"weighted-kernel-sum", "unweighted-rank-sum"}, } @@ -355,9 +349,7 @@ def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semant """ valid = expert_ids >= 0 expert = expert_ids.clamp(min=0).to(torch.int64) - if combine_weight_semantics in ("unweighted-rank-sum", "bf16-rank-sum"): - # Both are plain summing kernels, so the gate is folded in here; they differ only in - # the precision the kernel accumulates at, which is a property of the expectation. + if combine_weight_semantics == "unweighted-rank-sum": coefficient = weights.to(torch.float32).masked_fill(~valid, 0) elif combine_weight_semantics == "weighted-kernel-sum": coefficient = valid.to(torch.float32) @@ -374,8 +366,27 @@ def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semant return transformed.to(payload.dtype) +def _to_payload_dtype(torch, value, dtype, rounding): + """Cast an FP32 tensor to the payload dtype the way the KERNEL does. + + torch rounds to nearest-even. Some kernels instead truncate — they keep the high 16 bits + of the FP32 word — which is a systematic downward bias of up to one ulp per conversion, + not noise that averages out. Measured on gb200: FlashInfer's one-sided combine reduces + 1.0 + 7 * 2^-9 (exactly 129.75 bf16 ulps) to 1.0078125, i.e. 129 ulps, where round-to- + nearest gives 1.015625. Modelling the wrong rounding leaves a few-ulp error on every + element, which is exactly the size that fails a tight relative gate. + """ + if rounding == "truncate": + as_int = value.float().view(torch.int32) + return (as_int & -65536).view(torch.float32).to(dtype) + if rounding != "nearest": + raise ValueError(f"unknown combine output rounding {rounding!r}") + return value.to(dtype) + + def _expected_transformed_combine( - torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics + torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, + combine_output_rounding="nearest", ): """Reproduce the reduction combine actually performs so the expectation carries the same BF16 rounding a correct backend does rather than hiding it in a wide tolerance. @@ -417,25 +428,6 @@ def _expected_transformed_combine( gate = (weights[:, slot:slot + 1] * valid[:, slot:slot + 1].to(torch.float32)) expected += gate * transform return expected - if combine_weight_semantics == "bf16-rank-sum": - # Same per-rank BF16 messages as unweighted-rank-sum, but the cross-rank reduction runs - # at payload precision instead of FP32, so the accumulator is cast back every step. - # Summation ORDER is not observable from outside the kernel; rank order is used here. - # That is sound because the modelled quantity is the PRECISION, and reordering a BF16 - # sum of like-magnitude terms moves the result by about an ulp, far inside the gate — - # whereas accumulating in FP32 when the kernel does not is a systematic 6-16 ulp error. - destination = expert_ids // experts_per_rank - scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) - accumulator = torch.zeros_like(semantic_x, dtype=dtype) - for rank_id in sorted(destination.unique().tolist()): - gate = weights * (destination == rank_id) - contribution = ( - semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) - + (gate * offset_a).sum(dim=1, keepdim=True) - + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) - ).to(dtype) - accumulator = (accumulator.float() + contribution.float()).to(dtype) - return accumulator.float() if combine_weight_semantics != "unweighted-rank-sum": raise ValueError(f"unknown combine semantics {combine_weight_semantics!r}") destination = expert_ids // experts_per_rank @@ -449,7 +441,10 @@ def _expected_transformed_combine( semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) + (gate * offset_a).sum(dim=1, keepdim=True) + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) - ).to(dtype).float() + ) + contribution = _to_payload_dtype( + torch, contribution, dtype, combine_output_rounding + ).float() domain = rank_id // ranks_per_domain if domain in domains: domains[domain] += contribution @@ -460,7 +455,9 @@ def _expected_transformed_combine( # exact zero through every level (all gates zero) — no mask needed. expected = torch.zeros_like(semantic_x, dtype=torch.float32) for domain in sorted(domains): - expected += domains[domain].to(dtype).float() + expected += _to_payload_dtype( + torch, domains[domain], dtype, combine_output_rounding + ).float() return expected @@ -598,7 +595,8 @@ def _run_expert_oracle( combined = backend.combine_transformed(problem, handle, transformed) torch.cuda.synchronize() expected_combined = _expected_transformed_combine( - torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics + torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, + getattr(backend, "combine_output_rounding", "nearest"), ) if combined.shape == expected_combined.shape: # Zero errors stand when the rank legitimately combined nothing. @@ -766,7 +764,8 @@ def _run_ll_expert_oracle( combined = backend.combine_transformed(problem, handle, transformed) torch.cuda.synchronize() expected_combined = _expected_transformed_combine( - torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics + torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, + getattr(backend, "combine_output_rounding", "nearest"), ) if combined.shape == expected_combined.shape: max_absolute_error = max_elementwise_relative_error = 0.0 diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index dd4ecbd79b..3856503cf9 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -483,45 +483,64 @@ def test_guards_fail_closed(self) -> None: ep_harness.logical_byte_provenance(**kwargs) +def _require_torch(): + try: + import torch + except ImportError: # pragma: no cover - torch-less test image + raise unittest.SkipTest("torch unavailable") + return torch + + +class CombineOutputRoundingContract(unittest.TestCase): + """The oracle must cast to the payload dtype the way the KERNEL does. + + torch rounds to nearest-even; FlashInfer's one-sided combine truncates. That is a + systematic downward bias of up to an ulp per element, not noise, and it is exactly the + magnitude a tight relative gate rejects. Pinned here so the two roundings cannot be + confused for each other, and so "truncate" cannot silently degrade into "nearest". + """ + + def test_truncate_differs_from_nearest_on_the_measured_case(self) -> None: + torch = _require_torch() + # 1.0 + 7 * 2^-9 is exactly 129.75 bf16 ulps: nearest -> 130, truncate -> 129. + value = torch.tensor([1.0 + 7 * (2.0 ** -9)], dtype=torch.float32) + nearest = ep_harness._to_payload_dtype( + torch, value, torch.bfloat16, "nearest" + ).float().item() + truncated = ep_harness._to_payload_dtype( + torch, value, torch.bfloat16, "truncate" + ).float().item() + self.assertEqual(nearest, 1.015625) + self.assertEqual(truncated, 1.0078125) # the value gb200 actually returns + self.assertLess(truncated, nearest) + + def test_truncate_is_exact_when_representable(self) -> None: + torch = _require_torch() + exact = torch.tensor([1.0, 2.0, 255.0, -0.5], dtype=torch.float32) + out = ep_harness._to_payload_dtype(torch, exact, torch.bfloat16, "truncate") + self.assertTrue(torch.equal(out.float(), exact)) + + def test_unknown_rounding_fails_closed(self) -> None: + torch = _require_torch() + with self.assertRaises(ValueError): + ep_harness._to_payload_dtype( + torch, torch.tensor([1.0]), torch.bfloat16, "stochastic" + ) + + class ModeSemanticsContract(unittest.TestCase): # The combine contract is a backend fact, not a pure function of mode: DeepEP's # low-latency combine is weighted-kernel-sum while MoRI's IntraNodeLL is - # unweighted-rank-sum, so low-latency must admit both. Normal admits the unweighted rank - # sum at either accumulation precision — FlashInfer's one-sided NVLink combine performs the - # same summation but accumulates the per-rank BF16 messages in BF16 rather than FP32 - # (measured on gb200), which the oracle reproduces as bf16-rank-sum instead of absorbing - # into a wider tolerance. weighted-kernel-sum stays out of normal mode. + # unweighted-rank-sum, so low-latency must admit both. Normal stays unweighted-only. def test_mode_allowed_semantics(self) -> None: self.assertEqual( - ep_harness.MODE_ALLOWED_SEMANTICS["normal"], - {"unweighted-rank-sum", "bf16-rank-sum"}, + ep_harness.MODE_ALLOWED_SEMANTICS["normal"], {"unweighted-rank-sum"} ) self.assertEqual( ep_harness.MODE_ALLOWED_SEMANTICS["low-latency"], {"weighted-kernel-sum", "unweighted-rank-sum"}, ) - def test_bf16_rank_sum_is_lower_precision_than_fp32_rank_sum(self) -> None: - """The two normal-mode contracts must actually differ, or declaring the new one is - a no-op that would silently accept an FP32-accumulating backend as BF16 (and the - reverse). Reproduces the gb200 probe in pure python: seven sub-ulp addends vanish - under BF16 accumulation and survive under FP32.""" - import struct - - def bf16(value): # round-to-nearest-even truncation to bfloat16 - bits = struct.unpack("> 16) & 1)) & 0xFFFF0000 - return struct.unpack(" Date: Sun, 2 Aug 2026 14:46:09 +0800 Subject: [PATCH 09/24] CollectiveX: truncate only where the kernel converts, not where the adapter does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run 7 applied combine_output_rounding to both the per-rank message and the final accumulator cast, and correctness got worse than the plain-nearest baseline (T=1 0.024414 -> 0.031250). The per-rank message is not the kernel's conversion. The adapter stages it, casting with torch (round-to-nearest) when it writes into the combine buffer, so the expectation has to round the same way or it disagrees with an input it can see. The kernel's only FP32 -> BF16 conversion is the accumulator on the way out, which is the one the gb200 probe measured as truncating. 中文:只在 kernel 真正做转换的地方截断,adapter 自己做的转换不截断。 第 7 次运行把 combine_output_rounding 同时应用到每 rank 的消息与最终累加器转换上,结果正确性比 纯 nearest 基线更差(T=1 从 0.024414 变为 0.031250)。 每 rank 的消息并不是 kernel 的转换。它由 adapter 在写入 combine 缓冲时用 torch(就近舍入)完成, 因此期望值必须以同样方式舍入,否则就会与一个它本可直接看到的输入不一致。kernel 唯一的 FP32 -> BF16 转换是输出时的累加器转换 —— 也正是 gb200 探测测出为截断的那一处。 --- experimental/CollectiveX/bench/ep_harness.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index fd93e0ff54..7e882f6ec2 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -441,10 +441,11 @@ def _expected_transformed_combine( semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) + (gate * offset_a).sum(dim=1, keepdim=True) + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) - ) - contribution = _to_payload_dtype( - torch, contribution, dtype, combine_output_rounding - ).float() + ).to(dtype).float() + # Round-to-nearest here on purpose: this message is produced by the ADAPTER with + # torch when it stages the combine input, not by the kernel. Only the accumulator + # -> payload conversion below is the kernel's, so only that one honours + # combine_output_rounding. domain = rank_id // ranks_per_domain if domain in domains: domains[domain] += contribution From 8799ba0a019c209c93d87accf23fd79cd97c0da1 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 15:16:25 +0800 Subject: [PATCH 10/24] CollectiveX: sum the combine expectation in the same canonical order as the staged input MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Instrumented gb200 run: every failing element differed from the expectation by exactly one BF16 ulp of a per-rank contribution (absdiff 2^-11 at a contribution near 0.125, 2^-10 near 0.25), with varying sign, on expected values that were already BF16-exact. Not accumulation drift, not the output rounding mode, and not the zero-expected elements I had guessed at. The two sides were summing the same terms in different orders. _run_expert_oracle normalizes the adapter's view with _normalized_expert_metadata (sorted by global expert id, sentinels last) before _expert_transform builds the staged combine input, while _expected_transformed_combine summed (gate * scale).sum(dim=1) in raw problem order. FP32 addition is not associative, so when one rank owns two or more of a token's top-k experts the two sums can land one BF16 step apart after the cast. Normalizing the expectation's ids and weights the same way removes the asymmetry. Rows with a single local expert are bit-identical either way, which is why this never surfaced on the padded single-expert low-latency layouts and only ever hit a minority of normal-mode elements. This touches shared oracle code, so it needs a full-matrix sweep to confirm no other backend moves — every currently-green cell should be unchanged, since the only cells affected are multi-local-expert rows whose two summation orders disagree in the last bit. 中文:让 combine 期望值按与暂存输入相同的规范顺序求和。 gb200 插桩运行显示:每个失败元素与期望值的差都恰好是某个 per-rank 贡献的一个 BF16 ulp(量级约 0.125 处 absdiff 为 2^-11,约 0.25 处为 2^-10),符号有正有负,且期望值本身已可被 BF16 精确表示。 既不是累加漂移,不是输出舍入模式,也不是我此前猜测的「期望值为零」的元素。 两侧在以不同顺序对相同的项求和。_run_expert_oracle 会先用 _normalized_expert_metadata 对 adapter 的视图做归一化(按全局专家 id 排序,哨兵置后),再由 _expert_transform 构造暂存的 combine 输入; 而 _expected_transformed_combine 是按原始 problem 顺序做 (gate * scale).sum(dim=1)。FP32 加法不满足 结合律,因此当同一个 rank 拥有某 token 的两个及以上 top-k 专家时,两种求和顺序在转为 BF16 后可能 相差一个步长。 对期望值的 ids 与 weights 做同样的归一化即可消除这一不对称。每行只有一个本地专家时两者逐位相同 —— 这正是它从未在 padded 单专家低延迟布局上出现、且只影响少数 normal 模式元素的原因。 本改动涉及共享的 oracle 代码,因此需要一次全矩阵扫描确认其他后端不受影响:所有当前为绿的单元格都 应保持不变,因为受影响的只有那些两种求和顺序在最后一位上不一致的多本地专家行。 --- experimental/CollectiveX/bench/ep_harness.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 7e882f6ec2..4dc835a8e7 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -410,8 +410,18 @@ def _expected_transformed_combine( combine ~0.048 off a single-domain reference. """ semantic_x = getattr(problem, "oracle_x", problem.x) - expert_ids = problem.topk_idx.to(torch.int64) - weights = problem.topk_weights.to(torch.float32) + # Sum the top-k axis in the SAME canonical order the staged combine input uses. + # _run_expert_oracle normalizes the adapter's view (sorted by global expert id, sentinels + # last) before handing it to _expert_transform, so an expectation that sums in raw problem + # order is adding the same terms in a different sequence. FP32 addition is not associative: + # when one rank owns two or more of a token's experts the two sums can land one BF16 step + # apart after the cast. Measured on gb200 flashinfer-ep — every failing element differed by + # exactly one ulp of a per-rank contribution, with varying sign. Rows carrying a single + # local expert are unaffected, which is why this only ever surfaced on a minority of + # elements and never on the padded single-expert low-latency layouts. + expert_ids, weights = _normalized_expert_metadata( + torch, problem.topk_idx.to(torch.int64), problem.topk_weights.to(torch.float32) + ) pattern = _column_pattern(torch, semantic_x.shape[1], semantic_x.device) dtype = semantic_x.dtype if combine_weight_semantics == "weighted-kernel-sum": From 33cd01bc26a29421024f4f4b1d8c7b3912647f6e Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 15:23:04 +0800 Subject: [PATCH 11/24] CollectiveX: revert the canonical-ordering change (measured no-op) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aligning the expectation's top-k summation order with the staged input produced correctness byte-identical to before on gb200 — [0.018311, 0.024414, 0.027466, 0.030518, 0.030518, 0.036621, 0.036621, 0.036621, 0.038147, 0.048828] both with and without it. The asymmetry it removed is real, but it is not what these elements are sensitive to, and carrying a change to shared oracle code that needs full-matrix validation for zero measured benefit is not worth the risk. Reverted; the reasoning is preserved in the PR discussion. 中文:回退规范化求和顺序的改动(实测为空操作)。 让期望值的 top-k 求和顺序与暂存输入对齐后,gb200 上的正确性结果与改动前逐位相同。它消除的不对称 确实存在,但这些元素对它并不敏感;为零收益而在共享 oracle 代码中保留一个需要全矩阵验证的改动并不 值得。已回退,相关分析保留在 PR 讨论中。 --- experimental/CollectiveX/bench/ep_harness.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 4dc835a8e7..7e882f6ec2 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -410,18 +410,8 @@ def _expected_transformed_combine( combine ~0.048 off a single-domain reference. """ semantic_x = getattr(problem, "oracle_x", problem.x) - # Sum the top-k axis in the SAME canonical order the staged combine input uses. - # _run_expert_oracle normalizes the adapter's view (sorted by global expert id, sentinels - # last) before handing it to _expert_transform, so an expectation that sums in raw problem - # order is adding the same terms in a different sequence. FP32 addition is not associative: - # when one rank owns two or more of a token's experts the two sums can land one BF16 step - # apart after the cast. Measured on gb200 flashinfer-ep — every failing element differed by - # exactly one ulp of a per-rank contribution, with varying sign. Rows carrying a single - # local expert are unaffected, which is why this only ever surfaced on a minority of - # elements and never on the padded single-expert low-latency layouts. - expert_ids, weights = _normalized_expert_metadata( - torch, problem.topk_idx.to(torch.int64), problem.topk_weights.to(torch.float32) - ) + expert_ids = problem.topk_idx.to(torch.int64) + weights = problem.topk_weights.to(torch.float32) pattern = _column_pattern(torch, semantic_x.shape[1], semantic_x.device) dtype = semantic_x.dtype if combine_weight_semantics == "weighted-kernel-sum": From f4454cea464b28ea00e081744ba65e38c5af27cb Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 15:29:33 +0800 Subject: [PATCH 12/24] CollectiveX: reproduce FlashInfer's per-slot tree combine from the kernel source MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Read the kernel instead of probing it. flashinfer vendors TRT-LLM's moeAlltoAllKernels.cu, and vectorized_combine_impl says plainly what the reduction is: vec_t acc[TOP_K]; // one fp32 accumulator PER TOP-K SLOT // slot with a negative send index (duplicate rank for this token, or dead rank) -> fill(0) acc[0]+=acc[1]; acc[2]+=acc[3]; acc[4]+=acc[5]; acc[6]+=acc[7]; // pairwise TREE acc[0]+=acc[2]; acc[4]+=acc[6]; acc[0]+=acc[4]; acc[0].cast_store(...) // one narrowing at the end Three consequences. It accumulates per top-k SLOT, not per rank, and a token's second expert on the same rank is a duplicate that gets zeroed — so each distinct rank contributes exactly once, at its first slot, which confirms the adapter's per-rank staging. It reduces as a pairwise tree while the oracle summed sequentially in rank order, and fp32 addition is not associative, so reassociating five-ish terms shifts the low bits and the narrowing lands one BF16 step away. That is exactly the instrumented signature: one ulp of a per-rank contribution, sign varying, on expected values already BF16-exact. And the earlier probe is consistent — the tree gives 1.013671875 where the kernel returned 1.0078125, so the store truncates rather than rounding. So `topk-slot-tree-sum` models it directly: per-slot messages with duplicates zeroed, the same pairwise halving, and the declared output rounding applied once at the end. Every other backend keeps unweighted-rank-sum and is untouched. This is what the previous two attempts were missing — bf16-rank-sum was built on a confounded probe and made agreement worse, and canonical rank ordering was an exact no-op because rank order was never the axis that mattered. 中文:依据 kernel 源码复现 FlashInfer 的「按 slot 树形归约」combine。 这次是读源码而非黑盒探测。flashinfer 内置了 TRT-LLM 的 moeAlltoAllKernels.cu,其中 vectorized_combine_impl 明确写出了归约方式(见上)。 由此有三点结论。它按 top-k **slot** 而非按 rank 累加,且同一 token 落在同一 rank 上的第二个专家会 被判为 duplicate 并置零 —— 因此每个不同的 rank 恰好贡献一次、位于其首个 slot,这也印证了 adapter 按 rank 聚合的暂存方式是对的。它采用成对树形归约,而 oracle 此前是按 rank 顺序顺序累加;fp32 加法 不满足结合律,五个左右的项重新结合会改变低位,narrowing 之后就相差一个 BF16 步长 —— 这正是插桩测 到的特征:某个 per-rank 贡献的一个 ulp,符号不定,且期望值本身已可被 BF16 精确表示。先前的探测也与 之一致:树形归约得到 1.013671875,而 kernel 返回 1.0078125,说明写出时是截断而非就近舍入。 因此 `topk-slot-tree-sum` 直接对其建模:按 slot 的消息、duplicate 置零、相同的成对折半归约,并在 最后一步应用所声明的输出舍入。其他后端仍使用 unweighted-rank-sum,不受影响。 这正是前两次尝试所缺失的:bf16-rank-sum 建立在被混淆的探测之上并使一致性变差,而规范化 rank 排序 则完全是空操作 —— 因为 rank 顺序从来就不是关键所在。 --- .../CollectiveX/bench/ep_flashinfer.py | 6 +- experimental/CollectiveX/bench/ep_harness.py | 57 ++++++++++++++++++- .../CollectiveX/tests/test_runtime.py | 3 +- 3 files changed, 62 insertions(+), 4 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index c5f9dd1de3..66c53658f6 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -65,7 +65,11 @@ class FlashInferEPBackend(EPBackend): # The kernel scatters expert outputs back to the supplying rank; it does not multiply by # the routing weights (those ride along as a caller payload, and vLLM applies them in the # MoE layer, not in the A2A). Verified against the oracle during bring-up. - combine_weight_semantics = "unweighted-rank-sum" + # Read off the kernel source (vectorized_combine_impl in moeAlltoAllKernels.cu), not + # inferred: one fp32 accumulator PER TOP-K SLOT, duplicate/dead slots zeroed so each + # distinct rank contributes once at its first slot, a pairwise TREE reduction over + # those slots, and a single narrowing at the end. + combine_weight_semantics = "topk-slot-tree-sum" # Measured on gb200, 8 ranks each contributing to one token: the kernel reduces # 1.0 + 7 * 2^-9 (exactly 129.75 bf16 ulps) to 1.0078125 = 129 ulps. Round-to-nearest # would give 1.015625 = 130. It accumulates in FP32 — a separate probe summing 2^rank diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 7e882f6ec2..941726c635 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -58,7 +58,7 @@ def case_id(sku: str, case: dict) -> str: # run_sweep only checks that the declared value is ALLOWED for the mode (so a mislabeled # adapter fails closed), and the oracle keys on that same declared value. MODE_ALLOWED_SEMANTICS = { - "normal": {"unweighted-rank-sum"}, + "normal": {"unweighted-rank-sum", "topk-slot-tree-sum"}, "low-latency": {"weighted-kernel-sum", "unweighted-rank-sum"}, } @@ -349,7 +349,7 @@ def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semant """ valid = expert_ids >= 0 expert = expert_ids.clamp(min=0).to(torch.int64) - if combine_weight_semantics == "unweighted-rank-sum": + if combine_weight_semantics in ("unweighted-rank-sum", "topk-slot-tree-sum"): coefficient = weights.to(torch.float32).masked_fill(~valid, 0) elif combine_weight_semantics == "weighted-kernel-sum": coefficient = valid.to(torch.float32) @@ -384,6 +384,54 @@ def _to_payload_dtype(torch, value, dtype, rounding): return value.to(dtype) +def _topk_slot_tree_combine( + torch, semantic_x, expert_ids, weights, experts_per_rank, dtype, pattern, rounding +): + """Reproduce a combine that reduces per TOP-K SLOT with a pairwise tree. + + Derived from the kernel source (flashinfer csrc/nv_internal ... moeAlltoAllKernels.cu, + `vectorized_combine_impl`), not inferred from behaviour: + + * one fp32 accumulator per top-k slot, `vec_t acc[TOP_K]`; + * a slot whose send index is negative -- a DUPLICATE rank for this token, or a dead + rank -- is filled with 0, so each distinct destination rank contributes exactly once, + at the position of its FIRST top-k slot; + * the slots are combined by a pairwise TREE (`acc[0]+=acc[1]; acc[2]+=acc[3]; ...`), + not a sequential sweep; + * a single narrowing at the end via `cast_store`. + + Sequential and tree summation of the same fp32 terms differ in the low bits, and after the + narrowing that shows up as a one-ulp disagreement on a minority of elements -- exactly what + the gb200 instrumentation measured. Rank order is irrelevant here; slot order is what counts. + """ + destination = expert_ids // experts_per_rank + scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) + topk = expert_ids.shape[1] + # Each distinct rank's full aggregate, placed at that rank's first top-k slot. + slots = [] + for slot in range(topk): + rank_id = destination[:, slot:slot + 1] + first = torch.ones_like(rank_id, dtype=torch.bool) + for earlier in range(slot): + first &= destination[:, earlier:earlier + 1] != rank_id + live = (expert_ids[:, slot:slot + 1] >= 0) & first + gate = weights * (destination == rank_id) * live + aggregate = ( + semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) + + (gate * offset_a).sum(dim=1, keepdim=True) + + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) + ) + # The message on the wire is payload dtype; the kernel widens it back to fp32. + slots.append(torch.where(live, aggregate.to(dtype).float(), torch.zeros_like(aggregate))) + # Pairwise tree, matching the kernel's unrolled halving. + while len(slots) > 1: + nxt = [slots[i] + slots[i + 1] for i in range(0, len(slots) - 1, 2)] + if len(slots) % 2: + nxt.append(slots[-1]) + slots = nxt + return _to_payload_dtype(torch, slots[0], dtype, rounding).float() + + def _expected_transformed_combine( torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, combine_output_rounding="nearest", @@ -428,6 +476,11 @@ def _expected_transformed_combine( gate = (weights[:, slot:slot + 1] * valid[:, slot:slot + 1].to(torch.float32)) expected += gate * transform return expected + if combine_weight_semantics == "topk-slot-tree-sum": + return _topk_slot_tree_combine( + torch, semantic_x, expert_ids, weights, experts_per_rank, dtype, pattern, + combine_output_rounding, + ) if combine_weight_semantics != "unweighted-rank-sum": raise ValueError(f"unknown combine semantics {combine_weight_semantics!r}") destination = expert_ids // experts_per_rank diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index 3856503cf9..a2df5d8133 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -534,7 +534,8 @@ class ModeSemanticsContract(unittest.TestCase): # unweighted-rank-sum, so low-latency must admit both. Normal stays unweighted-only. def test_mode_allowed_semantics(self) -> None: self.assertEqual( - ep_harness.MODE_ALLOWED_SEMANTICS["normal"], {"unweighted-rank-sum"} + ep_harness.MODE_ALLOWED_SEMANTICS["normal"], + {"unweighted-rank-sum", "topk-slot-tree-sum"}, ) self.assertEqual( ep_harness.MODE_ALLOWED_SEMANTICS["low-latency"], From 27dfba1653597587b4c2c057ab1f1fb84ee0f025 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 15:36:46 +0800 Subject: [PATCH 13/24] CollectiveX: revert the per-slot tree combine model (measured no-op) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Transcribed from the kernel source and still byte-identical on gb200: decode [0.018311, 0.024414, 0.027466, 0.030518, 0.030518, 0.036621, 0.036621, 0.036621, 0.038147, 0.048828] with and without it. That is now four expectation models — rank-sequential/nearest, truncating output, canonical rank order, and per-slot tree — producing identical errors at T>=32. The invariance is the finding. These elements are insensitive to every property of the expectation: summation order, associativity, and rounding mode. With the instrumentation showing expected values already BF16-exact there, the likeliest reading is that the worst elements have a single contributing rank, where no reduction happens at all and the two sides should be a plain copy — so a one-ulp gap cannot come from the reduction model. Keeping combine_output_rounding: that one is source-consistent AND measurably moved T=1 below the gate (0.024414 -> 0.018311). Dropping the tree model, which carries real complexity in shared oracle code for zero measured benefit. The kernel reading is preserved in the PR discussion and remains correct — it is simply not what these elements are sensitive to. 中文:回退按 slot 的树形 combine 模型(实测为空操作)。 该模型直接抄自 kernel 源码,但 gb200 上结果与不加时逐位相同。至此已有四种期望值模型(rank 顺序 + 就近舍入、截断输出、规范化 rank 排序、按 slot 树形)在 T>=32 上给出完全相同的误差。 这种不变性本身就是结论:这些元素对期望值的每一项性质都不敏感 —— 求和顺序、结合方式、舍入模式皆 然。结合插桩显示这些位置的期望值本就可被 BF16 精确表示,最合理的解释是最差元素只有一个贡献 rank, 此时根本不发生归约、两侧本应是直接拷贝,因此一个 ulp 的差异不可能来自归约模型。 保留 combine_output_rounding:它与源码一致,且确实把 T=1 压到了门禁之下(0.024414 -> 0.018311)。 移除树形模型:它在共享 oracle 代码中带来实质复杂度却没有可测收益。对 kernel 的解读保留在 PR 讨论 中,且依然正确 —— 只是这些元素并不对它敏感。 --- .../CollectiveX/bench/ep_flashinfer.py | 6 +- experimental/CollectiveX/bench/ep_harness.py | 57 +------------------ .../CollectiveX/tests/test_runtime.py | 3 +- 3 files changed, 4 insertions(+), 62 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 66c53658f6..c5f9dd1de3 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -65,11 +65,7 @@ class FlashInferEPBackend(EPBackend): # The kernel scatters expert outputs back to the supplying rank; it does not multiply by # the routing weights (those ride along as a caller payload, and vLLM applies them in the # MoE layer, not in the A2A). Verified against the oracle during bring-up. - # Read off the kernel source (vectorized_combine_impl in moeAlltoAllKernels.cu), not - # inferred: one fp32 accumulator PER TOP-K SLOT, duplicate/dead slots zeroed so each - # distinct rank contributes once at its first slot, a pairwise TREE reduction over - # those slots, and a single narrowing at the end. - combine_weight_semantics = "topk-slot-tree-sum" + combine_weight_semantics = "unweighted-rank-sum" # Measured on gb200, 8 ranks each contributing to one token: the kernel reduces # 1.0 + 7 * 2^-9 (exactly 129.75 bf16 ulps) to 1.0078125 = 129 ulps. Round-to-nearest # would give 1.015625 = 130. It accumulates in FP32 — a separate probe summing 2^rank diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 941726c635..7e882f6ec2 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -58,7 +58,7 @@ def case_id(sku: str, case: dict) -> str: # run_sweep only checks that the declared value is ALLOWED for the mode (so a mislabeled # adapter fails closed), and the oracle keys on that same declared value. MODE_ALLOWED_SEMANTICS = { - "normal": {"unweighted-rank-sum", "topk-slot-tree-sum"}, + "normal": {"unweighted-rank-sum"}, "low-latency": {"weighted-kernel-sum", "unweighted-rank-sum"}, } @@ -349,7 +349,7 @@ def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semant """ valid = expert_ids >= 0 expert = expert_ids.clamp(min=0).to(torch.int64) - if combine_weight_semantics in ("unweighted-rank-sum", "topk-slot-tree-sum"): + if combine_weight_semantics == "unweighted-rank-sum": coefficient = weights.to(torch.float32).masked_fill(~valid, 0) elif combine_weight_semantics == "weighted-kernel-sum": coefficient = valid.to(torch.float32) @@ -384,54 +384,6 @@ def _to_payload_dtype(torch, value, dtype, rounding): return value.to(dtype) -def _topk_slot_tree_combine( - torch, semantic_x, expert_ids, weights, experts_per_rank, dtype, pattern, rounding -): - """Reproduce a combine that reduces per TOP-K SLOT with a pairwise tree. - - Derived from the kernel source (flashinfer csrc/nv_internal ... moeAlltoAllKernels.cu, - `vectorized_combine_impl`), not inferred from behaviour: - - * one fp32 accumulator per top-k slot, `vec_t acc[TOP_K]`; - * a slot whose send index is negative -- a DUPLICATE rank for this token, or a dead - rank -- is filled with 0, so each distinct destination rank contributes exactly once, - at the position of its FIRST top-k slot; - * the slots are combined by a pairwise TREE (`acc[0]+=acc[1]; acc[2]+=acc[3]; ...`), - not a sequential sweep; - * a single narrowing at the end via `cast_store`. - - Sequential and tree summation of the same fp32 terms differ in the low bits, and after the - narrowing that shows up as a one-ulp disagreement on a minority of elements -- exactly what - the gb200 instrumentation measured. Rank order is irrelevant here; slot order is what counts. - """ - destination = expert_ids // experts_per_rank - scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) - topk = expert_ids.shape[1] - # Each distinct rank's full aggregate, placed at that rank's first top-k slot. - slots = [] - for slot in range(topk): - rank_id = destination[:, slot:slot + 1] - first = torch.ones_like(rank_id, dtype=torch.bool) - for earlier in range(slot): - first &= destination[:, earlier:earlier + 1] != rank_id - live = (expert_ids[:, slot:slot + 1] >= 0) & first - gate = weights * (destination == rank_id) * live - aggregate = ( - semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) - + (gate * offset_a).sum(dim=1, keepdim=True) - + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) - ) - # The message on the wire is payload dtype; the kernel widens it back to fp32. - slots.append(torch.where(live, aggregate.to(dtype).float(), torch.zeros_like(aggregate))) - # Pairwise tree, matching the kernel's unrolled halving. - while len(slots) > 1: - nxt = [slots[i] + slots[i + 1] for i in range(0, len(slots) - 1, 2)] - if len(slots) % 2: - nxt.append(slots[-1]) - slots = nxt - return _to_payload_dtype(torch, slots[0], dtype, rounding).float() - - def _expected_transformed_combine( torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, combine_output_rounding="nearest", @@ -476,11 +428,6 @@ def _expected_transformed_combine( gate = (weights[:, slot:slot + 1] * valid[:, slot:slot + 1].to(torch.float32)) expected += gate * transform return expected - if combine_weight_semantics == "topk-slot-tree-sum": - return _topk_slot_tree_combine( - torch, semantic_x, expert_ids, weights, experts_per_rank, dtype, pattern, - combine_output_rounding, - ) if combine_weight_semantics != "unweighted-rank-sum": raise ValueError(f"unknown combine semantics {combine_weight_semantics!r}") destination = expert_ids // experts_per_rank diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index a2df5d8133..3856503cf9 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -534,8 +534,7 @@ class ModeSemanticsContract(unittest.TestCase): # unweighted-rank-sum, so low-latency must admit both. Normal stays unweighted-only. def test_mode_allowed_semantics(self) -> None: self.assertEqual( - ep_harness.MODE_ALLOWED_SEMANTICS["normal"], - {"unweighted-rank-sum", "topk-slot-tree-sum"}, + ep_harness.MODE_ALLOWED_SEMANTICS["normal"], {"unweighted-rank-sum"} ) self.assertEqual( ep_harness.MODE_ALLOWED_SEMANTICS["low-latency"], From 78e6ad3a05a1d478452690235bfec38a79f1a461 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 15:48:11 +0800 Subject: [PATCH 14/24] CollectiveX: build the expected per-rank message with the same function that stages it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The instrumented worst element was five contributions of magnitude ~0.25 summing to ~0.018 -- heavy cancellation -- where one contribution sat exactly one BF16 ulp (2^-10 at 0.25) away from the expectation, which the cancellation and the 0.02 denominator floor turned into 4.9%. Across 160 sampled elements the worst always had 5-7 contributing ranks, so this is not the reduction: four expectation models, including one transcribed from the kernel source, left these elements byte-identical. The disagreement is in a single per-rank message. The adapter's staged message goes through _expert_transform; the expectation restated the same formula inline. Mathematically equivalent, but not bit-identical -- the same fp32 expression evaluated over differently shaped tensors can fuse differently and land on adjacent BF16 values. Almost every element agreed; a minority did not, and only where the combine cancels did that become visible. So the expectation now calls _expert_transform itself, masking expert_ids to the rank being modelled. One definition of "what one rank sends", used by both sides, bit-identical by construction instead of by coincidence. This touches shared oracle code. Every other backend passes today because its two formulations already agree, so this should be inert for them -- but "should be" is why a full-matrix sweep has to confirm it before this merges. 中文:用暂存消息所使用的同一个函数来构造期望的 per-rank 消息。 插桩得到的最差元素是五个量级约 0.25 的贡献相加得到约 0.018 —— 存在严重相消 —— 其中一个贡献与期望 值恰好相差一个 BF16 ulp(0.25 处的 2^-10),经相消与 0.02 分母下限放大后变成 4.9%。在采样的 160 个 元素中,最差元素的贡献 rank 数始终是 5 到 7 个,因此问题不在归约:四种期望值模型(其中一种直接抄自 kernel 源码)都让这些元素逐位不变。 分歧出在单条 per-rank 消息上。adapter 暂存的消息走的是 _expert_transform,而期望值则内联重述了同一 个公式。两者数学上等价,但并非逐位相同 —— 同一个 fp32 表达式在形状不同的张量上求值,融合方式可能 不同,从而落在相邻的 BF16 取值上。绝大多数元素是一致的;少数不一致,而只有在 combine 发生相消时才 会显现出来。 因此期望值现在直接调用 _expert_transform,并把 expert_ids 掩码到所建模的那个 rank。「一个 rank 发出 什么」只有一个定义,两侧共用,逐位一致是构造保证而非巧合。 本改动涉及共享的 oracle 代码。其他后端今天能通过,是因为它们的两种表述本就一致,因此这对它们应当是 惰性的 —— 但正因为只是「应当」,合入前必须由一次全矩阵扫描来确认。 --- experimental/CollectiveX/bench/ep_harness.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 7e882f6ec2..50c38cdab9 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -435,13 +435,19 @@ def _expected_transformed_combine( domains: dict[int, object] = {} scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) for rank_id in destination.unique().tolist(): - gate = weights * (destination == rank_id) - # Per-rank BF16 output, FP32-accumulated within its scale-up domain. - contribution = ( - semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) - + (gate * offset_a).sum(dim=1, keepdim=True) - + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) - ).to(dtype).float() + # Build this rank's BF16 message through the SAME function that builds the staged + # combine input the adapter actually sends (_expert_transform), rather than restating + # the formula. The two were mathematically equivalent but not bit-identical: on gb200 + # flashinfer-ep a minority of messages landed one BF16 step apart, and where the + # combine cancels (five contributions of ~0.25 summing to ~0.018) that single step + # became 5% of the result. Equivalent expressions are not enough here; the expectation + # has to be the same computation. + local_ids = torch.where( + destination == rank_id, expert_ids, torch.full_like(expert_ids, -1) + ) + contribution = _expert_transform( + torch, semantic_x, local_ids, weights, combine_weight_semantics + ).float() # Round-to-nearest here on purpose: this message is produced by the ADAPTER with # torch when it stages the combine input, not by the kernel. Only the accumulator # -> payload conversion below is the kernel's, so only that one honours From 5e0657855021213a1ce0e0c703207e974cf93c91 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 15:54:56 +0800 Subject: [PATCH 15/24] CollectiveX: revert the shared _expert_transform expectation (measured no-op) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fifth expectation model, byte-identical on gb200 again. Reverting for the same reason as the previous two: shared oracle code should not carry changes with no measured effect. What the five attempts establish, which is worth more than any of them individually: * Reduction-order models could never have worked. Tree-vs-sequential fp32 reassociation moves the result by ~1e-9, while one BF16 step at the result magnitude (~0.018) is ~7e-5. Those differences are erased by the narrowing. Three runs were spent on changes that were arithmetically incapable of moving the output. * The gap is bigger than a rounding artifact. absdiff 2^-10 is one BF16 step at a contribution's 0.25 magnitude but FOURTEEN ulps at the 0.018 result. One message's value is materially wrong, not marginally rounded. * It survives both sides calling the same function on nominally equal inputs, so the inputs are not equal: the row-to-token correspondence between the adapter's view and the expectation must diverge for some rows. That is an adapter mapping bug, not an oracle modelling gap -- which is where the next investigation should start, and it is a different place from anywhere I looked. 中文:回退共享 _expert_transform 的期望值实现(实测为空操作)。 第五个期望值模型,在 gb200 上仍逐位相同。回退理由与前两次一致:共享的 oracle 代码不应保留没有可测 效果的改动。 这五次尝试共同确立的结论,比其中任何单独一次都更有价值: * 归约顺序类模型从一开始就不可能奏效。树形与顺序累加的 fp32 重结合只会带来约 1e-9 的差异,而结果 量级(约 0.018)上的一个 BF16 步长约为 7e-5,这类差异在 narrowing 时被抹平。有三次运行花在了 算术上根本无法改变输出的改动上。 * 该差距大于舍入伪影。absdiff 2^-10 是贡献量级 0.25 处的一个 BF16 步长,但相对 0.018 的结果却是 十四个 ulp。某条消息的数值是实质性错误,而非轻微舍入。 * 即便两侧调用同一个函数、输入名义上相同,差异依然存在,说明输入其实并不相同:adapter 视图与期望 值之间的「行 -> token」对应关系在某些行上一定发生了偏移。这是 adapter 的映射缺陷,而非 oracle 的建模缺口 —— 下一次调查应从这里开始,而这与我此前查看的所有位置都不同。 --- experimental/CollectiveX/bench/ep_harness.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 50c38cdab9..7e882f6ec2 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -435,19 +435,13 @@ def _expected_transformed_combine( domains: dict[int, object] = {} scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) for rank_id in destination.unique().tolist(): - # Build this rank's BF16 message through the SAME function that builds the staged - # combine input the adapter actually sends (_expert_transform), rather than restating - # the formula. The two were mathematically equivalent but not bit-identical: on gb200 - # flashinfer-ep a minority of messages landed one BF16 step apart, and where the - # combine cancels (five contributions of ~0.25 summing to ~0.018) that single step - # became 5% of the result. Equivalent expressions are not enough here; the expectation - # has to be the same computation. - local_ids = torch.where( - destination == rank_id, expert_ids, torch.full_like(expert_ids, -1) - ) - contribution = _expert_transform( - torch, semantic_x, local_ids, weights, combine_weight_semantics - ).float() + gate = weights * (destination == rank_id) + # Per-rank BF16 output, FP32-accumulated within its scale-up domain. + contribution = ( + semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) + + (gate * offset_a).sum(dim=1, keepdim=True) + + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) + ).to(dtype).float() # Round-to-nearest here on purpose: this message is produced by the ADAPTER with # torch when it stages the combine input, not by the kernel. Only the accumulator # -> payload conversion below is the kernel's, so only that one honours From badb8a3c750372a89525c3f91d8e565459233488 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 16:04:21 +0800 Subject: [PATCH 16/24] CollectiveX: drop the truncating output model for flashinfer-ep MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cross-rank instrumentation at T=1 printed each rank's own combine result against the expectation: rank 4 0.34765625 vs 0.34765625 exact rank 6 0.3203125 vs 0.3203125 exact rank 1 -0.66796875 vs -0.66796875 exact rank 0 0.2578125 vs 0.25390625 |actual| > |expected| 2^-8 rank 3 -1.1484375 vs -1.140625 |actual| > |expected| 2^-7 rank 5 -0.52734375 vs -0.5234375 |actual| > |expected| 2^-8 rank 7 -0.8046875 vs -0.80078125 |actual| > |expected| 2^-8 rank 2 0.107421875 vs 0.106933594 |actual| > |expected| 2^-11 Three ranks bit-exact, five biased the SAME direction. A one-sided bias is not rounding noise, and truncating the expectation pulls it toward zero -- exactly the direction that manufactures |actual| > |expected|. The probe that suggested truncation carried the same top-k slicing confound as the earlier bf16-rank-sum probe, so it was never evidence. Reverting to the default nearest. It changes nothing at the failing rungs anyway: nearest and truncate both pass T<=16 and both fail identically from T=32, so the rounding mode was never what separates pass from fail. Better to carry no model than a model built on a bad probe. The residual at T>=32 is unchanged and unexplained by anything on the expectation side. 中文:移除 flashinfer-ep 的截断输出模型。 T=1 的跨 rank 插桩打印了每个 rank 自身 combine 结果与期望值的对比(见上):三个 rank 逐位相同, 另外五个偏向**同一方向**。单侧偏置不是舍入噪声;而对期望值做截断会把它拉向零 —— 恰好就是制造 |actual| > |expected| 的方向。当初提示截断的那次探测,与更早的 bf16-rank-sum 探测存在同样的 top-k 切片混淆,因此从来就不构成证据。 恢复为默认的就近舍入。这在失败档位上本就没有区别:nearest 与 truncate 都能通过 T<=16,且从 T=32 起以完全相同的方式失败,因此舍入模式从来不是通过与否的分界。与其保留一个建立在错误探测之上的模型, 不如不带模型。 T>=32 处的残差保持不变,且无法由期望值一侧的任何因素解释。 --- experimental/CollectiveX/bench/ep_flashinfer.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index c5f9dd1de3..971474d4a6 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -66,11 +66,11 @@ class FlashInferEPBackend(EPBackend): # the routing weights (those ride along as a caller payload, and vLLM applies them in the # MoE layer, not in the A2A). Verified against the oracle during bring-up. combine_weight_semantics = "unweighted-rank-sum" - # Measured on gb200, 8 ranks each contributing to one token: the kernel reduces - # 1.0 + 7 * 2^-9 (exactly 129.75 bf16 ulps) to 1.0078125 = 129 ulps. Round-to-nearest - # would give 1.015625 = 130. It accumulates in FP32 — a separate probe summing 2^rank - # across all eight ranks returned exactly 255.0 — and truncates on the way out. - combine_output_rounding = "truncate" + # Left at the default "nearest": with truncation active every mismatching rank + # showed |actual| > |expected| (a one-sided bias in exactly the direction + # truncating the EXPECTATION would create), while three of eight ranks matched + # bit-exactly. The earlier probe that suggested truncation shared the top-k + # slicing confound, so it is not evidence. # Forced by the phase asserts described in the module docstring. combine_needs_redispatch = True dispatch_needs_combine_cleanup = True From 0d9d66f86164cf3d20af594f4a8d530354698b39 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 16:19:53 +0800 Subject: [PATCH 17/24] CollectiveX: stage the flashinfer combine in place, not through a fresh buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summing every rank's staged contribution for one token gave 0.2555847168; the oracle expected BF16-nearest of that, 0.255859375, bit-exact; the kernel returned 0.2578125. No summation model reaches the kernel's value -- fp32 sequential, fp32 tree, bf16 tree (hand-computed over the measured contributions) and both rounding modes all land on 0.255859375. The kernel is not summing those messages differently, it is reading different bytes. The scatter's index arithmetic round-trips exactly, so the remaining difference between what was written and what was read is the buffer itself: recv_x is workspace-backed and a separate zeros_like allocation has to be copied in. Writing the staged messages into the received tensor directly removes that copy, so the values sit exactly where dispatch left them and the kernel reads the slots it wrote. 中文:flashinfer 的 combine 输入改为原地写入,不再经由新分配的缓冲区。 对单个 token 汇总各 rank 暂存的贡献得到 0.2555847168;oracle 期望值是它的 BF16 就近舍入 0.255859375,逐位精确;而 kernel 返回 0.2578125。没有任何求和模型能得出 kernel 的数值 —— fp32 顺序累加、fp32 树形、bf16 树形(用实测贡献手工计算)以及两种舍入模式,结果都落在 0.255859375。 因此 kernel 并不是以不同方式对这些消息求和,而是读到了不同的字节。 scatter 的下标运算已验证可精确往返,因此「写入内容」与「读取内容」之间仅剩的差异就是缓冲区本身: recv_x 由 workspace 支撑,而单独的 zeros_like 分配必须被拷贝进去。将暂存消息直接写入接收张量可以 消除这次拷贝,使数值正好留在 dispatch 放置它们的位置,kernel 读到的就是它写入的槽位。 --- .../CollectiveX/bench/ep_flashinfer.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 971474d4a6..7dacc7cd0f 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -219,12 +219,22 @@ def combine_transformed(self, p, h, transformed): into the padded workspace-shaped buffer at exactly the slots those rows came from. Unfilled slots stay zero and contribute nothing to the sum. """ - # Shaped from the ACTUAL receive, not the ladder maximum: the planes are - # [ep_size, runtime_max_tokens_per_rank, hidden], so they shrink with the rung. Sizing - # this from max_num_tokens builds a 65536-row buffer for an 8192-row mask. - padded = torch.zeros_like(h.recv_x) - padded.view(-1, h.recv_x.shape[-1])[self._valid_rows(h)] = transformed.to(padded.dtype) - return self._a2a.combine(padded, h.tokens) + # Write straight into the RECEIVED workspace tensor rather than a fresh buffer. + # + # A separate `zeros_like` buffer was measured to give a combine output that is not the + # sum of the staged messages: summing every rank's staged contribution for one token + # gave 0.2555847168, the oracle expected BF16-nearest of that (0.255859375, bit-exact), + # and the kernel returned 0.2578125. No summation model reaches that value -- fp32 + # sequential, fp32 tree, bf16 tree and both rounding modes all land on 0.255859375 -- + # so the kernel was reading different bytes than were written. The remaining difference + # between the two paths is the buffer: recv_x is workspace-backed, and a fresh + # allocation has to be copied in, which is where a layout reinterpretation can happen. + # Writing in place removes that copy: the values sit exactly where dispatch left them. + flat = h.recv_x.view(-1, h.recv_x.shape[-1]) + keep = self._valid_rows(h) + flat[keep] = transformed.to(h.recv_x.dtype) + flat[~keep] = 0 # slots the kernel never filled must contribute nothing + return self._a2a.combine(h.recv_x, h.tokens) def _ep_group(): From 19f0049d35800ff55837891469d9710aad18d440 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:05:15 +0800 Subject: [PATCH 18/24] CollectiveX: submit the flashinfer combine payload from the workspace region MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文:从 workspace 区域提交 flashinfer combine 载荷 --- .../CollectiveX/bench/ep_flashinfer.py | 51 +++++++++++-------- 1 file changed, 29 insertions(+), 22 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 7dacc7cd0f..95b448d9f2 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -61,7 +61,8 @@ class FlashInferEPBackend(EPBackend): # second input_payload and validated against the oracle's cast round-trip; not this pass. SUPPORTED_PRECISIONS = ("bf16",) kernel_generation = "flashinfer-mnnvl-one-sided" - stage_device_work = False + # stage() now copies the received payload into the workspace combine region. + stage_device_work = True # The kernel scatters expert outputs back to the supplying rank; it does not multiply by # the routing weights (those ride along as a caller payload, and vLLM applies them in the # MoE layer, not in the A2A). Verified against the oracle during bring-up. @@ -164,11 +165,27 @@ def dispatch(self, p): ) def stage(self, p, h): - # BF16 needs no conversion: the received workspace tensor is the combine input. - h.combine_input = h.recv_x + """Materialise the combine payload in the workspace region the API designates. + + The kernel accepts exactly two payload forms: a caller-owned tensor with + `payload_in_workspace=False` (vLLM passes its expert-GEMM output this way), or the + view returned by `get_combine_payload_tensor_in_workspace` with the flag True. Handing + it the dispatch-RECEIVE view instead — which is what this adapter did — is a third + form no upstream caller uses: it makes the kernel stage-copy out of MNNVL + peer-writable memory that peers wrote during dispatch. + """ + buffer = self._combine_buffer(h) + buffer.copy_(h.recv_x) + h.combine_input = buffer + + def _combine_buffer(self, h): + """The workspace-resident combine payload region for this rung.""" + return self._a2a.get_combine_payload_tensor_in_workspace( + h.tokens, h.recv_x.shape[-1], h.recv_x.dtype + ) def combine(self, p, h): - out = self._a2a.combine(h.combine_input, h.tokens) + out = self._a2a.combine(h.combine_input, h.tokens, payload_in_workspace=True) h.out = out return out @@ -215,26 +232,16 @@ def inspect_dispatch(self, p, h): def combine_transformed(self, p, h, transformed): """Combine an oracle-transformed payload through the same kernel as the timed path. - `transformed` holds one row per row `inspect_dispatch` returned, so scatter it back - into the padded workspace-shaped buffer at exactly the slots those rows came from. - Unfilled slots stay zero and contribute nothing to the sum. + Scattered into the workspace combine region (see stage) and submitted with + `payload_in_workspace=True`, so the kernel reads the slots from the region the API + designates for them rather than from the dispatch-receive buffer. """ - # Write straight into the RECEIVED workspace tensor rather than a fresh buffer. - # - # A separate `zeros_like` buffer was measured to give a combine output that is not the - # sum of the staged messages: summing every rank's staged contribution for one token - # gave 0.2555847168, the oracle expected BF16-nearest of that (0.255859375, bit-exact), - # and the kernel returned 0.2578125. No summation model reaches that value -- fp32 - # sequential, fp32 tree, bf16 tree and both rounding modes all land on 0.255859375 -- - # so the kernel was reading different bytes than were written. The remaining difference - # between the two paths is the buffer: recv_x is workspace-backed, and a fresh - # allocation has to be copied in, which is where a layout reinterpretation can happen. - # Writing in place removes that copy: the values sit exactly where dispatch left them. - flat = h.recv_x.view(-1, h.recv_x.shape[-1]) + buffer = self._combine_buffer(h) + flat = buffer.view(-1, buffer.shape[-1]) keep = self._valid_rows(h) - flat[keep] = transformed.to(h.recv_x.dtype) - flat[~keep] = 0 # slots the kernel never filled must contribute nothing - return self._a2a.combine(h.recv_x, h.tokens) + flat[keep] = transformed.to(buffer.dtype) + flat[~keep] = 0 # slots the kernel never filled contribute nothing + return self._a2a.combine(buffer, h.tokens, payload_in_workspace=True) def _ep_group(): From 4a3489569fbe09a3d4646887d5928ee2e451cf38 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:18:11 +0800 Subject: [PATCH 19/24] CollectiveX: model FlashInfer's payload-dtype combine reduction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文:为 FlashInfer 建模 payload 精度的 combine 归约 --- .../CollectiveX/bench/ep_flashinfer.py | 20 +++-- experimental/CollectiveX/bench/ep_harness.py | 84 ++++++++++++++++--- .../CollectiveX/tests/test_runtime.py | 41 +++++++++ 3 files changed, 128 insertions(+), 17 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 95b448d9f2..3a903fe84e 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -50,6 +50,12 @@ # -1 is safe: real ids are [0, num_experts). _INVALID_EXPERT = -1 +# 0.6.16 ("Port the TensorRT-LLM one-sided A2A optimizations") rewrote the combine +# accumulator from the payload dtype to FP32 with a single narrowing store. Earlier wheels +# round at every level of the top-k reduction tree, which is a few BF16 ulps the plain +# FP32 expectation does not carry, so the oracle needs to know which kernel it is facing. +_COMBINE_FP32_SINCE = (0, 6, 16) + class FlashInferEPBackend(EPBackend): name = "flashinfer-ep" @@ -67,11 +73,8 @@ class FlashInferEPBackend(EPBackend): # the routing weights (those ride along as a caller payload, and vLLM applies them in the # MoE layer, not in the A2A). Verified against the oracle during bring-up. combine_weight_semantics = "unweighted-rank-sum" - # Left at the default "nearest": with truncation active every mismatching rank - # showed |actual| > |expected| (a one-sided bias in exactly the direction - # truncating the EXPECTATION would create), while three of eight ranks matched - # bit-exactly. The earlier probe that suggested truncation shared the top-k - # slicing confound, so it is not evidence. + # Set per wheel in create_buffer; see _COMBINE_FP32_SINCE. + combine_reduction = "topk-slot-tree" # Forced by the phase asserts described in the module docstring. combine_needs_redispatch = True dispatch_needs_combine_cleanup = True @@ -131,6 +134,13 @@ def create_buffer(self, spec): workspace_size_per_rank=workspace_size, mnnvl_config=MnnvlConfig(comm_backend=_communicator(_ep_group())), ) + import flashinfer + + version = tuple( + int(part) for part in flashinfer.__version__.split(".post")[0].split(".")[:3] + ) + if version >= _COMBINE_FP32_SINCE: + self.combine_reduction = "domain-fp32" # Every rank must finish mapping its workspace before any peer writes into it; # vLLM barriers here for the same reason. Scoped to the EP group, not the world. torch.distributed.barrier(group=_ep_group()) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 7e882f6ec2..97912d6014 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -384,9 +384,45 @@ def _to_payload_dtype(torch, value, dtype, rounding): return value.to(dtype) +def _topk_slot_tree_combine(torch, destination, valid, messages, dtype): + """Reduce the per-rank messages the way a payload-dtype accumulator does. + + Most combine kernels accumulate in FP32 and narrow once. FlashInfer's one-sided kernel + (<= 0.6.15) instead holds its top-k accumulators IN the payload dtype and reduces them + with a hand-unrolled pairwise tree, so every level rounds: + + acc[k] = message of destination[k], or 0 if a lower k already claimed that rank + (a0+=a1) (a2+=a3) (a4+=a5) (a6+=a7); (a0+=a2) (a4+=a6); (a0+=a4) -- and store + + Three BF16 roundings on partials near a contribution's own magnitude is a few ulps of + error, which is the whole gap a plain FP32 sum leaves against this backend. Operands sit + at their ORIGINAL top-k slot -- the kernel blanks duplicate-rank slots in place rather + than compacting -- so the tree's shape depends on the routing, not just the rank count. + The generic halving below reproduces the unrolled K=6/8/10 trees exactly. + """ + tokens = torch.arange(destination.shape[0], device=destination.device) + zero = torch.zeros_like(messages[0]) + slots = [] + for slot in range(destination.shape[1]): + rank_id = destination[:, slot] + claimed = valid[:, slot].clone() + for earlier in range(slot): + claimed &= ~(valid[:, earlier] & (destination[:, earlier] == rank_id)) + slots.append(torch.where(claimed.unsqueeze(1), messages[rank_id, tokens], zero)) + while len(slots) > 1: + merged = [ + _to_payload_dtype(torch, slots[i] + slots[i + 1], dtype, "nearest").float() + for i in range(0, len(slots) - 1, 2) + ] + if len(slots) % 2: + merged.append(slots[-1]) + slots = merged + return slots[0] + + def _expected_transformed_combine( torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, - combine_output_rounding="nearest", + combine_output_rounding="nearest", combine_reduction="domain-fp32", ): """Reproduce the reduction combine actually performs so the expectation carries the same BF16 rounding a correct backend does rather than hiding it in a wide tolerance. @@ -408,6 +444,10 @@ def _expected_transformed_combine( cases) there is a single domain and no scale-out rounding; a multi-node RoCE EP16 group has one BF16 partial per node, and omitting that cast is what left the scale-out combine ~0.048 off a single-domain reference. + + A backend whose accumulator is the payload dtype rather than FP32 declares + ``combine_reduction = "topk-slot-tree"`` and takes the model in + :func:`_topk_slot_tree_combine` instead of the domain reduction below. """ semantic_x = getattr(problem, "oracle_x", problem.x) expert_ids = problem.topk_idx.to(torch.int64) @@ -430,23 +470,42 @@ def _expected_transformed_combine( return expected if combine_weight_semantics != "unweighted-rank-sum": raise ValueError(f"unknown combine semantics {combine_weight_semantics!r}") - destination = expert_ids // experts_per_rank - ranks_per_domain = max(1, scale_up_domain) - domains: dict[int, object] = {} + valid = expert_ids >= 0 + destination = torch.where(valid, expert_ids, torch.zeros_like(expert_ids)) + destination //= experts_per_rank scale, offset_a, offset_b = _expert_coefficients(torch, expert_ids) - for rank_id in destination.unique().tolist(): - gate = weights * (destination == rank_id) - # Per-rank BF16 output, FP32-accumulated within its scale-up domain. - contribution = ( + + def rank_message(rank_id): + """The one BF16 row this destination rank stages back for every token. + + Round-to-nearest on purpose: the ADAPTER produces this with torch when it stages + the combine input, so it is not the kernel's narrowing and does not honour + combine_output_rounding. + """ + gate = weights * (destination == rank_id) * valid + return ( semantic_x.float() * (gate * scale).sum(dim=1, keepdim=True) + (gate * offset_a).sum(dim=1, keepdim=True) + (gate * offset_b).sum(dim=1, keepdim=True) * pattern.unsqueeze(0) ).to(dtype).float() - # Round-to-nearest here on purpose: this message is produced by the ADAPTER with - # torch when it stages the combine input, not by the kernel. Only the accumulator - # -> payload conversion below is the kernel's, so only that one honours - # combine_output_rounding. + + present = sorted(destination[valid].unique().tolist()) + if combine_reduction == "topk-slot-tree": + messages = torch.zeros( + (max(present, default=0) + 1,) + semantic_x.shape, + dtype=torch.float32, device=semantic_x.device, + ) + for rank_id in present: + messages[rank_id] = rank_message(rank_id) + return _topk_slot_tree_combine(torch, destination, valid, messages, dtype) + if combine_reduction != "domain-fp32": + raise ValueError(f"unknown combine reduction {combine_reduction!r}") + ranks_per_domain = max(1, scale_up_domain) + domains: dict[int, object] = {} + for rank_id in present: + # Per-rank BF16 output, FP32-accumulated within its scale-up domain. domain = rank_id // ranks_per_domain + contribution = rank_message(rank_id) if domain in domains: domains[domain] += contribution else: @@ -598,6 +657,7 @@ def _run_expert_oracle( expected_combined = _expected_transformed_combine( torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, getattr(backend, "combine_output_rounding", "nearest"), + getattr(backend, "combine_reduction", "domain-fp32"), ) if combined.shape == expected_combined.shape: # Zero errors stand when the rank legitimately combined nothing. diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index 3856503cf9..8ad86d94a4 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -610,5 +610,46 @@ def test_unknown_semantics_fail_closed(self): ) + + +@unittest.skipUnless(_torch is not None, "combine-oracle math checks require torch") +class TopkSlotTreeReductionTests(unittest.TestCase): + """Pin the payload-dtype reduction against a value measured on the real kernel. + + Eight contributions of 1.0 and 7 x 2^-9 reduce to three different answers depending on + the model, which is what makes this case worth pinning: FP32-then-narrow gives + 1.015625, a sequential BF16 sum gives 1.0, and the pairwise BF16 tree gives 1.0078125. + gb200 returns 1.0078125. + """ + + def _tree(self, values): + torch = _torch + slots = [torch.full((1, 1), v, dtype=torch.float32) for v in values] + destination = torch.arange(len(values)).unsqueeze(0) + messages = torch.stack(slots) + return ep_harness._topk_slot_tree_combine( + torch, destination, torch.ones_like(destination, dtype=torch.bool), + messages, torch.bfloat16, + ).item() + + def test_matches_the_value_the_kernel_returns(self): + self.assertEqual(self._tree([1.0] + [2.0**-9] * 7), 1.0078125) + + def test_differs_from_both_rejected_models(self): + values = [1.0] + [2.0**-9] * 7 + self.assertNotEqual(self._tree(values), 1.015625) # FP32 accumulate, narrow once + self.assertNotEqual(self._tree(values), 1.0) # sequential BF16 accumulate + + def test_a_rank_claimed_by_an_earlier_slot_contributes_once(self): + torch = _torch + # Both top-k slots route to rank 0; the kernel blanks the later slot in place. + destination = torch.zeros((1, 2), dtype=torch.int64) + messages = torch.full((1, 1, 1), 0.5) + combined = ep_harness._topk_slot_tree_combine( + torch, destination, torch.ones_like(destination, dtype=torch.bool), + messages, torch.bfloat16, + ) + self.assertEqual(combined.item(), 0.5) + if __name__ == "__main__": unittest.main() From 6e7d2fa4ab155cc56da3863ea39666615c14cb72 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:21:54 +0800 Subject: [PATCH 20/24] CollectiveX: tighten the flashinfer wheel probe and stage docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文:收紧 flashinfer wheel 版本探测与 stage 文档 --- .../CollectiveX/bench/ep_flashinfer.py | 33 +++++++++---------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 3a903fe84e..8499d2fd10 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -40,6 +40,7 @@ """ from __future__ import annotations +import re import types import torch @@ -99,6 +100,7 @@ def create_buffer(self, spec): The communicator handed to MnnvlConfig must span exactly the EP group: the kernel asserts `workspace.size(0) == moe_ep_size`, so a wider group silently mis-sizes it. """ + import flashinfer from flashinfer.comm import Mapping from flashinfer.comm.mnnvl import MnnvlConfig from flashinfer.comm.trtllm_moe_alltoall import ( @@ -134,12 +136,8 @@ def create_buffer(self, spec): workspace_size_per_rank=workspace_size, mnnvl_config=MnnvlConfig(comm_backend=_communicator(_ep_group())), ) - import flashinfer - - version = tuple( - int(part) for part in flashinfer.__version__.split(".post")[0].split(".")[:3] - ) - if version >= _COMBINE_FP32_SINCE: + wheel = tuple(int(n) for n in re.findall(r"\d+", flashinfer.__version__)[:3]) + if wheel >= _COMBINE_FP32_SINCE: self.combine_reduction = "domain-fp32" # Every rank must finish mapping its workspace before any peer writes into it; # vLLM barriers here for the same reason. Scoped to the EP group, not the world. @@ -174,26 +172,25 @@ def dispatch(self, p): tokens=p.T, topk=p.topk_idx.shape[1], combine_input=None, ) + def _combine_buffer(self, h): + """The workspace-resident combine payload region for this rung.""" + return self._a2a.get_combine_payload_tensor_in_workspace( + h.tokens, h.recv_x.shape[-1], h.recv_x.dtype + ) + def stage(self, p, h): """Materialise the combine payload in the workspace region the API designates. - The kernel accepts exactly two payload forms: a caller-owned tensor with - `payload_in_workspace=False` (vLLM passes its expert-GEMM output this way), or the - view returned by `get_combine_payload_tensor_in_workspace` with the flag True. Handing - it the dispatch-RECEIVE view instead — which is what this adapter did — is a third - form no upstream caller uses: it makes the kernel stage-copy out of MNNVL - peer-writable memory that peers wrote during dispatch. + A production integration has the expert GEMM write its output straight into this + region and submits it with `payload_in_workspace=True`, so combine performs no + staging copy. Copying here rather than handing `combine` a caller-owned tensor keeps + that copy out of the combine measurement, where production does not pay it; it is + still executed and reported, as `stage`. """ buffer = self._combine_buffer(h) buffer.copy_(h.recv_x) h.combine_input = buffer - def _combine_buffer(self, h): - """The workspace-resident combine payload region for this rung.""" - return self._a2a.get_combine_payload_tensor_in_workspace( - h.tokens, h.recv_x.shape[-1], h.recv_x.dtype - ) - def combine(self, p, h): out = self._a2a.combine(h.combine_input, h.tokens, payload_in_workspace=True) h.out = out From c12fdeb35c75a73d1ee4cb5cc70bb32fb6328080 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:26:35 +0800 Subject: [PATCH 21/24] CollectiveX: document FlashInfer EP and its payload-dtype combine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文:在文档中说明 FlashInfer EP 及其 payload 精度 combine --- experimental/CollectiveX/README.md | 1 + experimental/CollectiveX/docs/methodology.md | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index 956e741d98..ddce08e460 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -75,6 +75,7 @@ scale-up domain. | MoRI | `production` — vLLM `--all2all-backend mori_*`, SGLang `--moe-a2a-backend mori` | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), 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 — LSA (NVLink load/store) intra-node, 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. BF16 only: NCCL EP's FP8 machinery exists upstream but its RELEASE.md lists it unsupported/untested, so no FP8 case is emitted. 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 | +| FlashInfer EP | `production` — 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 — no send/recv pairing and no NVSHMEM. `normal` mode only (there is one kernel family; no separate decode path), BF16 only, and GB200/GB300 only, since the transport is MNNVL. 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 [DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 371b5bb8da..90ddd7e2d1 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -69,7 +69,7 @@ unweighted rank-sum combine match `layout-and-dispatch-v1` exactly, so the same NVIDIA-only and CUDA 13 only, and runs 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, its cross-node GIN path faulting inside `nccl_ep.cc` identically on RoCE and IB across -four SKUs — a GDAKI limit, not a fabric-selection one. Those throughput kernels run across the full token ladder in the `normal` mode. +four SKUs — a GDAKI limit, not a fabric-selection one. FlashInfer EP is TensorRT-LLM's one-sided MNNVL `MoeAlltoAll`, in which each rank writes tokens directly into its peers' workspace windows and combine reads them back, so there is no send/recv pairing and no NVSHMEM; it is GB200/GB300-only for that reason, and runs EP8 and EP16 inside the MNNVL scale-up domain. Its combine is the one place a backend's accumulator precision changes the expectation rather than the tolerance: through 0.6.15 the kernel holds its top-k accumulators in the payload dtype and reduces them with a hand-unrolled pairwise tree, so every level rounds to BF16, and the oracle reproduces that tree exactly rather than loosening the gate to absorb it (0.6.16 rewrote the accumulator to FP32; the adapter reads the installed version and picks the matching model). Those throughput kernels run across the full token ladder in the `normal` mode. A second `low-latency` mode adds each backend's decode-optimized kernel family. On DeepEP it drives the legacy `deep_ep.Buffer` low-latency decode kernels (`low_latency_dispatch`/`low_latency_combine`), From 0bef5e906dcb08a3e812b5cee306bb445dc415cb Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:54:44 +0800 Subject: [PATCH 22/24] CollectiveX: drop the disproven truncating-store model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文:移除已被证伪的截断存储模型 The measured 1.0078125 came from the kernel's BF16 reduction tree, not from a truncating narrowing store, so combine_output_rounding never had a setter and its "truncate" branch was unreachable. Its tests asserted the wrong mechanism while passing. Record combine_reduction in the artifact instead: the model is now chosen per installed wheel, so results need to say which one produced them. --- experimental/CollectiveX/bench/ep_backend.py | 5 --- experimental/CollectiveX/bench/ep_harness.py | 38 +++++----------- .../CollectiveX/tests/test_runtime.py | 45 ------------------- 3 files changed, 10 insertions(+), 78 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index 1b7e629975..a16d1a015b 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -82,11 +82,6 @@ class EPBackend(abc.ABC): # Adapters that reduce activations and top-k weights independently must carry # the complete local weighted expert sum in the activation tensor. combine_weight_semantics = "unweighted-rank-sum" - # How the kernel converts its FP32 combine accumulator to the payload dtype. - # "nearest" (torch default) or "truncate" (keep the high 16 bits) — a kernel that - # truncates is biased down by up to an ulp per element, which a tight relative gate - # will catch, so the adapter must declare what its kernel actually does. - combine_output_rounding = "nearest" roundtrip_only = False # Realized wire formats recorded in the artifact. Combine is always BF16; # dispatch_dtype is overridden per-run by an FP8 adapter (e.g. "fp8-e4m3fn"). diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 5573b8a9e7..5ff940a99b 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -366,24 +366,6 @@ def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semant return transformed.to(payload.dtype) -def _to_payload_dtype(torch, value, dtype, rounding): - """Cast an FP32 tensor to the payload dtype the way the KERNEL does. - - torch rounds to nearest-even. Some kernels instead truncate — they keep the high 16 bits - of the FP32 word — which is a systematic downward bias of up to one ulp per conversion, - not noise that averages out. Measured on gb200: FlashInfer's one-sided combine reduces - 1.0 + 7 * 2^-9 (exactly 129.75 bf16 ulps) to 1.0078125, i.e. 129 ulps, where round-to- - nearest gives 1.015625. Modelling the wrong rounding leaves a few-ulp error on every - element, which is exactly the size that fails a tight relative gate. - """ - if rounding == "truncate": - as_int = value.float().view(torch.int32) - return (as_int & -65536).view(torch.float32).to(dtype) - if rounding != "nearest": - raise ValueError(f"unknown combine output rounding {rounding!r}") - return value.to(dtype) - - def _topk_slot_tree_combine(torch, destination, valid, messages, dtype): """Reduce the per-rank messages the way a payload-dtype accumulator does. @@ -411,7 +393,7 @@ def _topk_slot_tree_combine(torch, destination, valid, messages, dtype): slots.append(torch.where(claimed.unsqueeze(1), messages[rank_id, tokens], zero)) while len(slots) > 1: merged = [ - _to_payload_dtype(torch, slots[i] + slots[i + 1], dtype, "nearest").float() + (slots[i] + slots[i + 1]).to(dtype).float() for i in range(0, len(slots) - 1, 2) ] if len(slots) % 2: @@ -422,7 +404,7 @@ def _topk_slot_tree_combine(torch, destination, valid, messages, dtype): def _expected_transformed_combine( torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, - combine_output_rounding="nearest", combine_reduction="domain-fp32", + combine_reduction="domain-fp32", ): """Reproduce the reduction combine actually performs so the expectation carries the same BF16 rounding a correct backend does rather than hiding it in a wide tolerance. @@ -478,9 +460,9 @@ def _expected_transformed_combine( def rank_message(rank_id): """The one BF16 row this destination rank stages back for every token. - Round-to-nearest on purpose: the ADAPTER produces this with torch when it stages - the combine input, so it is not the kernel's narrowing and does not honour - combine_output_rounding. + The narrowing here is the ADAPTER's — torch producing the staged combine input — + not the kernel's, so it is always round-to-nearest regardless of what the kernel + does with its own accumulator. """ gate = weights * (destination == rank_id) * valid return ( @@ -515,9 +497,7 @@ def rank_message(rank_id): # exact zero through every level (all gates zero) — no mask needed. expected = torch.zeros_like(semantic_x, dtype=torch.float32) for domain in sorted(domains): - expected += _to_payload_dtype( - torch, domains[domain], dtype, combine_output_rounding - ).float() + expected += domains[domain].to(dtype).float() return expected @@ -656,7 +636,6 @@ def _run_expert_oracle( torch.cuda.synchronize() expected_combined = _expected_transformed_combine( torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, - getattr(backend, "combine_output_rounding", "nearest"), getattr(backend, "combine_reduction", "domain-fp32"), ) if combined.shape == expected_combined.shape: @@ -826,7 +805,6 @@ def _run_ll_expert_oracle( torch.cuda.synchronize() expected_combined = _expected_transformed_combine( torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, - getattr(backend, "combine_output_rounding", "nearest"), ) if combined.shape == expected_combined.shape: max_absolute_error = max_elementwise_relative_error = 0.0 @@ -1218,6 +1196,10 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # EPBackend.fp8_consume. Only meaningful when the case dispatches FP8. "fp8_consume": getattr(backend, "fp8_consume", None), "kernel_generation": kernel_generation(backend), + # Which reduction the correctness oracle held the kernel to. A backend may + # pick this per installed library version (flashinfer-ep does), so without it + # a wheel bump silently changes the arithmetic behind `passed` with no trace. + "combine_reduction": getattr(backend, "combine_reduction", "domain-fp32"), # See EPBackend.maturity: a "candidate" row measures the library, not a deployment. "maturity": getattr(backend, "maturity", None) or "unknown", "name": backend.name, diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index 8ad86d94a4..c9dd2e1099 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -483,51 +483,6 @@ def test_guards_fail_closed(self) -> None: ep_harness.logical_byte_provenance(**kwargs) -def _require_torch(): - try: - import torch - except ImportError: # pragma: no cover - torch-less test image - raise unittest.SkipTest("torch unavailable") - return torch - - -class CombineOutputRoundingContract(unittest.TestCase): - """The oracle must cast to the payload dtype the way the KERNEL does. - - torch rounds to nearest-even; FlashInfer's one-sided combine truncates. That is a - systematic downward bias of up to an ulp per element, not noise, and it is exactly the - magnitude a tight relative gate rejects. Pinned here so the two roundings cannot be - confused for each other, and so "truncate" cannot silently degrade into "nearest". - """ - - def test_truncate_differs_from_nearest_on_the_measured_case(self) -> None: - torch = _require_torch() - # 1.0 + 7 * 2^-9 is exactly 129.75 bf16 ulps: nearest -> 130, truncate -> 129. - value = torch.tensor([1.0 + 7 * (2.0 ** -9)], dtype=torch.float32) - nearest = ep_harness._to_payload_dtype( - torch, value, torch.bfloat16, "nearest" - ).float().item() - truncated = ep_harness._to_payload_dtype( - torch, value, torch.bfloat16, "truncate" - ).float().item() - self.assertEqual(nearest, 1.015625) - self.assertEqual(truncated, 1.0078125) # the value gb200 actually returns - self.assertLess(truncated, nearest) - - def test_truncate_is_exact_when_representable(self) -> None: - torch = _require_torch() - exact = torch.tensor([1.0, 2.0, 255.0, -0.5], dtype=torch.float32) - out = ep_harness._to_payload_dtype(torch, exact, torch.bfloat16, "truncate") - self.assertTrue(torch.equal(out.float(), exact)) - - def test_unknown_rounding_fails_closed(self) -> None: - torch = _require_torch() - with self.assertRaises(ValueError): - ep_harness._to_payload_dtype( - torch, torch.tensor([1.0]), torch.bfloat16, "stochastic" - ) - - class ModeSemanticsContract(unittest.TestCase): # The combine contract is a backend fact, not a pure function of mode: DeepEP's # low-latency combine is weighted-kernel-sum while MoRI's IntraNodeLL is From 6d054357b5d469621b21f1e6c4f3031f09792796 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:56:32 +0800 Subject: [PATCH 23/24] CollectiveX: run the unit tests on every PR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文:在每个 PR 上运行 CollectiveX 单元测试 Nothing ran experimental/CollectiveX/tests, so the oracle contracts they pin — including the combine reduction models — were unverified on every change. Torch is CPU-only: the torch-dependent tests are pure arithmetic checks that never touch a device, and they skip silently without it, so the import is asserted before the run. --- .github/workflows/test-collectivex.yml | 46 ++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 .github/workflows/test-collectivex.yml diff --git a/.github/workflows/test-collectivex.yml b/.github/workflows/test-collectivex.yml new file mode 100644 index 0000000000..932a9a9403 --- /dev/null +++ b/.github/workflows/test-collectivex.yml @@ -0,0 +1,46 @@ +name: Test CollectiveX + +# The sweep exercises the benchmark on real hardware, but nothing ran CollectiveX's own +# unit tests, so the oracle contracts they pin were unverified on every PR. Torch is +# installed CPU-only: the torch-dependent tests are combine-oracle arithmetic checks that +# never touch a device, and without it they skip silently rather than fail. + +on: + pull_request: + paths: + - 'experimental/CollectiveX/**' + - '.github/workflows/test-collectivex.yml' + +permissions: + contents: read + +jobs: + test: + if: github.event.pull_request.draft != true + runs-on: ubuntu-latest + permissions: + contents: read + + steps: + - name: Checkout code + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + - name: Set up Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: '3.12' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu torch + + # Every skip in this suite is torch-gated, so a missing torch turns the oracle + # checks into silent passes. Fail here instead, where the cause is obvious. + - name: Verify torch is importable + run: python -c "import torch; print(torch.__version__)" + + - name: Run unit tests + run: | + cd experimental/CollectiveX + python -m unittest discover -s tests -p "test_*.py" -v From 49d25ec72cc1bcb2bbcf49e42098c99ae4dec5dd Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 3 Aug 2026 01:00:08 +0800 Subject: [PATCH 24/24] CollectiveX: note the slot-tree oracle's memory scaling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 中文:记录 slot-tree oracle 的内存伸缩特性 --- experimental/CollectiveX/bench/ep_harness.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 5ff940a99b..8bfac7fb3f 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -381,6 +381,11 @@ def _topk_slot_tree_combine(torch, destination, valid, messages, dtype): at their ORIGINAL top-k slot -- the kernel blanks duplicate-rank slots in place rather than compacting -- so the tree's shape depends on the routing, not just the rank count. The generic halving below reproduces the unrolled K=6/8/10 trees exactly. + + Unlike the domain reduction, which folds into one accumulator, this holds a message per + rank AND a slot per top-k position, so oracle memory is O(ep_size * tokens * hidden): + ~8 GiB at EP16 with the 8192-token prefill rung. Fine against 180+ GiB HBM at the EP + sizes here, but it is the term that would need streaming before EP32. """ tokens = torch.arange(destination.shape[0], device=destination.device) zero = torch.zeros_like(messages[0])