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/.github/workflows/test-collectivex.yml b/.github/workflows/test-collectivex.yml new file mode 100644 index 0000000000..7e233ad363 --- /dev/null +++ b/.github/workflows/test-collectivex.yml @@ -0,0 +1,50 @@ +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: + # ready_for_review is not in the default set (opened/synchronize/reopened), so without + # it a PR developed as a draft has every run skipped by the draft gate below and gets + # no run at all on the transition — the suite would never see the merge candidate. + types: [opened, synchronize, reopened, ready_for_review] + 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 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/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py new file mode 100644 index 0000000000..8a7d33b9fc --- /dev/null +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -0,0 +1,314 @@ +#!/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 re +import types + +import torch + +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 + +# 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" + 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() 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. + combine_weight_semantics = "unweighted-rank-sum" + # 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 + 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.experts_per_rank = args.experts // world_size + + # ---- setup ------------------------------------------------------------------------------- + + def _topk_idx_dtype(self): + """int32 — what the kernel reads the routing plane as. + + Declaring it here means `make_problem` casts once, untimed, instead of `dispatch` + casting on every call inside the measured window. `topk_weights` is already FP32 + from `make_problem`, so both payloads reach the kernel with no conversion. + """ + return torch.int32 + + 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. + """ + import flashinfer + 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=_communicator(_ep_group())), + ) + 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. + torch.distributed.barrier(group=_ep_group()) + + # ---- transport contract ------------------------------------------------------------------ + + def dispatch(self, p): + """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. + """ + recv_x, recv_idx, recv_w = self._a2a.dispatch( + p.topk_idx, + [p.dispatch_x, p.topk_idx, p.topk_weights], + 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, + ) + + 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. + + 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(self, p, h): + out = self._a2a.combine(h.combine_input, h.tokens, payload_in_workspace=True) + h.out = out + return out + + def recv_tokens(self, h): + # 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 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 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] + 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=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): + """Combine an oracle-transformed payload through the same kernel as the timed path. + + 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. + """ + buffer = self._combine_buffer(h) + flat = buffer.view(-1, buffer.shape[-1]) + keep = self._valid_rows(h) + 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(): + """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 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 Get_size(self) -> int: + return self._group.size() + + def allgather(self, data): + gathered = [None] * self.Get_size() + dist.all_gather_object(gathered, data, group=self._group) + return gathered + + 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 barrier(self) -> None: + dist.barrier(group=self._group) + + 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) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index c52d30a70f..8bfac7fb3f 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -366,8 +366,50 @@ def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semant return transformed.to(payload.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. + + 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]) + 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 = [ + (slots[i] + slots[i + 1]).to(dtype).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 + torch, problem, experts_per_rank, scale_up_domain, combine_weight_semantics, + 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. @@ -389,6 +431,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) @@ -411,19 +457,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. + + 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 ( 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() + + 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: @@ -571,7 +640,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_reduction", "domain-fp32"), ) if combined.shape == expected_combined.shape: # Zero errors stand when the rank legitimately combined nothing. @@ -739,7 +809,7 @@ 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, ) if combined.shape == expected_combined.shape: max_absolute_error = max_elementwise_relative_error = 0.0 @@ -1131,6 +1201,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/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/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`), 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 00c3d5f30f..b422668f4f 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -41,6 +41,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 dc4becaada..b3c49b2315 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 diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index 160e6c562f..01f377b7a9 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -4,6 +4,7 @@ from __future__ import annotations import argparse +import ast import contextlib import io import json @@ -340,7 +341,8 @@ def _run_ep_parser() -> argparse.ArgumentParser: # Mirror of the parser bench/run_ep.py builds in main(). parser = argparse.ArgumentParser() parser.add_argument( - "--backend", required=True, choices=["deepep-v2", "mori", "uccl-ep", "nccl-ep"] + "--backend", required=True, + choices=["deepep-v2", "mori", "uccl-ep", "nccl-ep", "flashinfer-ep"], ) ep_harness.add_common_args(parser) return parser @@ -430,6 +432,52 @@ def test_nccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: self.assertEqual(args.case_id, nccl_case["case_id"]) self.assertEqual(args.out, "results/h200-dgxc_nccl-ep_bf16_decode_TS-c000.json") + def test_flashinfer_ep_case_round_trips_through_the_run_ep_parser(self) -> None: + # A flashinfer-ep case flows through the same generic codec; run_ep's --backend + # choices must accept "flashinfer-ep" and the result filename must carry the backend + # token so it never collides with the deepep-v2/nccl-ep legs of the same cell. + # The codec is SKU-agnostic, so this reuses the shared h200 fixture like its + # siblings; that flashinfer-ep is GB-only is a registry fact, pinned separately by + # test_matrix.test_flashinfer_ep_rollout_shape. + flashinfer_case = { + **self.CASE, + "backend": "flashinfer-ep", + "case_id": "h200-dgxc-flashinfer-ep-deepseek-v3-normal-decode-ep16-uniform-bf16", + } + argv = self._case_argv(["16", "2", "8", "8"], case=flashinfer_case) + args = self._run_ep_parser().parse_args(argv) + self.assertEqual(args.backend, "flashinfer-ep") + self.assertEqual(args.case_id, flashinfer_case["case_id"]) + self.assertEqual( + args.out, "results/h200-dgxc_flashinfer-ep_bf16_decode_TS-c000.json" + ) + + def test_mirrored_backend_choices_match_run_ep(self) -> None: + """The mirror is only worth having if it cannot drift from the real parser. + + Kept in sync by convention it has now drifted twice — a backend was added to + run_ep.py's choices while this fixture kept the old list, so a case_id that the real + CLI accepts raised SystemExit here. Read the real list out of the source (AST, not + import: importing run_ep pulls in torch and the vendor EP libraries) and compare. + """ + tree = ast.parse((BENCH / "run_ep.py").read_text()) + real = [ + [element.value for element in keyword.value.elts] + for node in ast.walk(tree) + if isinstance(node, ast.Call) + for argument in node.args + if isinstance(argument, ast.Constant) and argument.value == "--backend" + for keyword in node.keywords + if keyword.arg == "choices" + ] + self.assertEqual(len(real), 1, "expected exactly one --backend choices list") + mirrored = next( + action.choices + for action in self._run_ep_parser()._actions + if action.dest == "backend" + ) + self.assertEqual(sorted(real[0]), sorted(mirrored)) + # logical_byte_provenance is where FP8 changes MEASUREMENT semantics (asymmetric # per-direction byte counts), so its arithmetic and guards are pinned here on CPU. @@ -565,5 +613,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()