diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index 7715624d5..7f5738458 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -14,8 +14,10 @@ responsibility. The full measurement methodology is in [docs/methodology.md](doc The workload uses packed placement and one pinned `fixed-profile` resource configuration per backend/topology; there is no tuning sweep. Combine is always BF16; dispatch precision is a swept dimension — a BF16 control plus, on every backend whose FP8 dispatch is supported upstream -(DeepEP V2, MoRI, UCCL-EP), an FP8 dispatch, caller-prequantized in `normal` mode (in `low-latency` -the DeepEP and UCCL-EP kernels quantize internally from BF16; MoRI stays caller-prequantized). NCCL +(DeepEP V2, MoRI, UCCL-EP, FlashInfer EP), an FP8 dispatch, caller-prequantized in `normal` mode (in +`low-latency` the DeepEP and UCCL-EP kernels quantize internally from BF16; MoRI stays +caller-prequantized; FlashInfer has no `low-latency` path). That caller-side quantize is charged +inside the measured dispatch, because a production forward pass pays it on the critical path. NCCL EP is BF16-only this release, so it emits the control alone. Coverage is uniform routing only. Cases run in one of two modes: @@ -50,6 +52,16 @@ every position in the sequence; each iteration takes the cross-rank maximum befo p50/p90/p95/p99, and roundtrip p99 is the headline latency. A keyed BLAKE2b counter produces byte-identical routing and gate weights on every runtime. +`roundtrip` means dispatch then combine — the transport — in every row. Expert-output staging sits +outside it and is reported separately as `stage`; under FP8 that component is harness scaffolding +standing in for the expert GEMM, which in production consumes FP8 operands natively rather than +materialising a BF16 copy, so `stage` must not be summed into a total or compared between backends. +That was not always true: rows measured before this change carried the staging copy inside the chain +for MoRI BF16 and FlashInfer BF16, so `roundtrip` meant different things in different rows. The sweep +`version` stays 1 across the change, so `implementation.stage_excluded_from_roundtrip` and whether a +`stage` component is present are the only way to tell the two generations apart. +See [docs/methodology.md](docs/methodology.md) for the full contract. + Correctness is checked against an implementation-independent oracle that reproduces the backend's two-level reduction — intra-scale-up-domain FP32, then a BF16 cast of each domain's partial for the scale-out send. The combine gate is a tight max elementwise relative error below `8 * 2^-8` @@ -80,7 +92,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. EP16 is an unsupported coverage row on all three: the adapter pins `InterNodeV1` over 2x8 XGMI + RDMA, but its combine corrupts at the transport layer (ROCm/mori#475), so the registry ships `mori: [8]` and no EP16 case is dispatched. `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: `contrib/nccl_ep/RELEASE.md` says "No FP8 support", so no FP8 case is emitted. That note is worth re-testing rather than trusting — the C library at our pinned commit does read `inputs->scales` and switch on e4m3/e5m2, the two documented FP8 exclusions are expert-major layouts we do not use, and `NVIDIA/nccl` has not moved since 2026-06-11 while `NVIDIA/nccl-extensions` has replaced that row outright. NVIDIA-only and CUDA 13 only. EP8 scale-up on H100/H200/B200/B300 plus EP8 and EP16 on GB200/GB300, where EP16 stays inside the MNNVL scale-up domain. x86 EP16 scale-out is an unsupported coverage row: the cross-node GIN path faults inside `nccl_ep.cc` identically on RoCE and IB across four SKUs, a GDAKI limit rather than a fabric-selection one | -| 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 | +| 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), and GB200/GB300 only, since the transport is MNNVL. FP8 dispatch is caller-prequantized blockwise e4m3fn, carried as a fourth dispatch payload alongside its per-128-block FP32 scales, with the combine plane forced to BF16 — the C++ `toNvDataType` accepts only fp16/bf16/fp32 for combine, so an FP8 combine buffer would raise rather than corrupt. EP8 and EP16, both inside the scale-up domain. Unlike every other backend here, its combine accumulates in the PAYLOAD dtype rather than FP32: wheels before 0.6.16 reduce the top-k contributions with a pairwise BF16 tree that rounds at every level, so the oracle models that reduction directly (`combine_reduction = "topk-slot-tree"`) instead of widening the tolerance. 0.6.16 moved the accumulator to FP32, and the adapter switches models on the installed version | DeepEP V2 means the `ElasticBuffer` implementation introduced by [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 2741e719a..43b830f3f 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -110,8 +110,12 @@ class EPBackend(abc.ABC): # # `dequant` is a VERIFICATION HATCH, not a second metric: never a sweep axis, never a # default. It is retained because it costs nothing (BF16 needs the same staged-is-None - # branch) and because it reproduces historical numbers exactly for regression checks -- - # measured 302.0us against 302.5us in run 30177021271 at T=1. + # branch) and because it reproduces historical numbers for regression checks on the + # backends whose stage this repo has not since changed -- measured 302.0us against 302.5us + # in run 30177021271 at T=1 for deepep-v2/uccl-ep. It no longer reproduces MoRI fp8, whose + # stage now casts only the rows dispatch filled rather than the whole padded plane; and no + # env combination reproduces pre-hoist BF16 MoRI/FlashInfer roundtrips, because the hatch + # is fp8-only by design. # # A second measured mode is unnecessary because the mismatched-config cost is DERIVABLE # from what every run already emits: @@ -132,27 +136,100 @@ class EPBackend(abc.ABC): fp8_consume = os.environ.get("CX_FP8_CONSUME", "native") @property - def stages_fp8_natively(self) -> bool: - """Whether the chained roundtrip should skip the per-iteration `stage()`. - - Gated on precision, NOT on `stage_device_work` alone. The two are equivalent for - deepep-v2 and uccl, but MoRI sets `stage_device_work = self._fp8 or not - self._external_input`, so its scale-up kernels (IntraNode/IntraNodeLL) report True - for BF16 as well -- and their `stage()` really does run a device copy into the - registered combine-input buffer. That copy is not an fp8 dequant and nothing in this - change's evidence says it should leave the timed region, so BF16 keeps executing it - inline every iteration and its numbers are unmoved. - - (Whether MoRI's registered-buffer copy is a production cost or a harness artefact is - a real open question -- a native integration may have the expert GEMM write straight - into that buffer -- but it is a separate question from fp8 consumption, it applies to - both precisions equally, and it needs its own evidence.) + def stage_excluded_from_roundtrip(self) -> bool: + """Whether the chained roundtrip skips the per-iteration `stage()`. + + `roundtrip` must mean the same thing in every row, or it cannot be compared across + backends. It means dispatch -> combine: the transport, staging excluded. So the + answer is yes whenever `stage()` does device work, regardless of precision. + + This was previously gated on precision as well, which left `stage` inside the + roundtrip for exactly two configurations -- MoRI BF16 scale-up and FlashInfer BF16 -- + and transport-only for all 800+ other rows, so the headline compared two different + quantities and penalised those two. + + Gated on `stage_device_work` rather than applied blanket: where `stage()` is a bare + pointer assignment there is nothing to lift, and hoisting anyway would hand the + low-latency backends a VIEW into their double-buffered receive, whose parity flips on + each timed re-dispatch -- combine would then read the stale-parity buffer. + + `CX_FP8_CONSUME=dequant` still opts an fp8 run back into the inline stage, because + that switch exists to model a stack that really does dequantise between the two + collectives (see `fp8_consume`). """ - return ( - self.precision == "fp8" - and self.stage_device_work - and self.fp8_consume == "native" - ) + if not self.stage_device_work: + return False + return not (self.precision == "fp8" and self.fp8_consume == "dequant") + + def fused_quantize(self, eager): + """The fp8 quantize the TIMED dispatch should call, keyed on mode. + + Production quantises bf16->fp8 once per forward pass, with a single fused kernel, just + before the dispatch collective. This benchmark used to do it once per SHAPE in + `make_problem` -- outside every timed window -- so it omitted a real cost. Moving the + eager helper inside the window would have been worse than omitting it: the eager form is + a 9-launch composite (measured 19.2us on H100, 53.6us on MI300X) against ~1.5-4.9us for + the compiled single kernel, so charging it would publish this harness's kernel count + rather than production's cost, and flip fp8-vs-bf16 verdicts on that basis. + + LOW-LATENCY GETS THE EAGER FORM, unchanged and deliberately. Its dispatch kernel + quantises internally (`use_fp8`), so the cost is already inside the timed window, and + the oracle's payload gate compares the received bytes against `semantic_payload` -- which + must therefore keep matching the kernel's arithmetic, i.e. the eager helper's bits. + Swapping this globally would red every low-latency fp8 cell without touching LL timing. + + Compiled with `dynamic=False`: a dynamic-shape build of this same math measured 6.3x + slower once, and here that lands inside the timed window as silently inflated numbers. + The cache limit is raised because the quantize legitimately sees ~20+ shapes (one per + ladder rung, plus the oracle's receive-count shapes) against a default of 8, and + exceeding it makes dynamo fall back to EAGER SILENTLY -- the same failure class. + """ + if self.mode == "low-latency": + return eager + import torch + + torch._dynamo.config.cache_size_limit = 64 + if hasattr(torch._dynamo.config, "fail_on_recompile_limit_hit"): + # Prefer a loud failure over a silent eager fallback if the limit is ever hit. + torch._dynamo.config.fail_on_recompile_limit_hit = True + return torch.compile(eager, dynamic=False) + + def assert_quantize_identity(self, eager, fused, x) -> None: + """Fail loudly, untimed, if the compiled quantize is not the eager one bit-for-bit. + + The correctness oracle's payload gate is a `torch.equal`, and it compares the SENDER's + [T, hidden] quantize against the ORACLE's [receive_count, hidden] one -- so identity has + to hold per row, across batch sizes, not merely deterministically. Both properties were + verified on-metal for e4m3fn (H100) and e4m3fnuz (MI300X/MI325X) before this was enabled; + this check is what turns a future toolchain regression into a named failure here instead + of an unexplained payload mismatch across the whole fleet. + """ + if fused is eager: + return + import torch + + def bits(pair): + values, scales = pair + return values.view(torch.uint8), scales + + eager_values, eager_scales = bits(eager(x)) + fused_values, fused_scales = bits(fused(x)) + if not (torch.equal(eager_values, fused_values) + and torch.equal(eager_scales, fused_scales)): + raise RuntimeError( + "compiled fp8 quantize is not bitwise identical to the eager helper; the " + "oracle payload gate would fail fleet-wide" + ) + rows = min(int(x.shape[0]), 3) + if rows: + part_values, part_scales = bits(fused(x[:rows])) + whole_values, whole_scales = fused_values[:rows], fused_scales[:rows] + if not (torch.equal(part_values, whole_values) + and torch.equal(part_scales, whole_scales)): + raise RuntimeError( + "compiled fp8 quantize is not per-row invariant across batch sizes; the " + "oracle compares a different row count than the sender quantised" + ) def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) @@ -309,6 +386,18 @@ def _topk_idx_dtype(self): # ---- Timing template methods ----------------------------------------------------- + # Pairs issued back-to-back per `period` sample; 0 disables the component. A serving decode + # loop runs dispatch->combine->dispatch->combine without stopping, so its cost per layer is + # the pipeline's PERIOD, not the sum of separately-drained stages -- `roundtrip` measures the + # latter and overstates the former. Opt-in rather than default because the overlap is only + # sound where the backend tolerates rank drift: dispatch is a peer WRITE into another rank's + # buffer, and stream order on the receiver does not order the sender's remote writes. The + # collective bounds that drift to roughly one iteration (a dispatch cannot complete until + # every rank enters it), so a backend whose receive buffer is double-buffered per dispatch is + # safe and one with a single shared buffer is not. Enabling it for a backend that is not + # means a fast number over corrupted data, which is why the default is off. + pipeline_pairs = 0 + def timed_components(self): """Components measured for this backend: roundtrip always; the rest unless the backend exposes only a stateful paired round trip.""" @@ -317,32 +406,48 @@ def timed_components(self): components.extend(["dispatch", "combine"]) if self.stage_device_work: components.append("stage") + if self.pipeline_pairs > 1: + components.append("period") return components - def warm(self, problem, count): + def warm(self, problem, count, stage_every=False): """Untimed synchronized full round trips (fabric/clock warm-up; cold-jump-safe). Caches the dynamic receive cardinality once so adapters never read a device scalar during a timed trial (the count is stable for a fixed routing trace). + + `stage_every` re-materialises the combine input on every iteration. The default hoists it + after the first, mirroring `benchmark_roundtrip` on the same reasoning: where staging is + excluded from the chain the timed region stages nothing at all, so staging once per warm + iteration warms work the measurement never performs. For an FP8 dequant that is not free -- + it is ~247us against a 61us roundtrip, and at 32 warm iterations per component per trial it + was the largest single cost in the leg. `benchmark_stage` opts in, because there staging IS + the timed operation and its warm-up has to match. """ import torch + staged = None for _ in range(count): handle = self.dispatch(problem) if not hasattr(problem, "recv_tokens"): problem.recv_tokens = self.recv_tokens(handle) - self.stage(problem, handle) + if staged is None: + self.stage(problem, handle) + if not stage_every and self.stage_excluded_from_roundtrip: + staged = getattr(handle, self.combine_input_attr) + else: + setattr(handle, self.combine_input_attr, staged) self.combine(problem, handle) torch.cuda.synchronize() def run_roundtrip(self, problem, staged=None): """One chained round trip; returns combined activations. - `staged` supplies a pre-materialised combine input so the conversion pass stays out of - the timed region (see `fp8_consume`). When it is None the stage runs inline, which is - the `dequant` fp8 model and the BF16 path -- free for the adapters whose - received buffer is already the combine input, real device work for the ones that - must place it (mori, flashinfer-ep; both declare `stage_device_work`). + `staged` supplies a pre-materialised combine input so staging stays out of the timed + region, which is the default for every backend that does device work there (see + `stage_excluded_from_roundtrip`). It is None only when `stage()` is a bare pointer + assignment -- deepep-v2, uccl-ep and nccl-ep at BF16, where there is nothing to lift -- + or under the `CX_FP8_CONSUME=dequant` hatch, which wants the conversion back in the chain. """ handle = self.dispatch(problem) if staged is None: @@ -351,10 +456,41 @@ def run_roundtrip(self, problem, staged=None): setattr(handle, self.combine_input_attr, staged) return self.combine(problem, handle) + def benchmark_period(self, problem, warmup, iters): + """Steady-state cost per dispatch->combine pair, pairs issued back-to-back. + + `roundtrip` drains the GPU around every pair, so it reports the latency of an idle + pipeline and charges inter-rank entry stagger to whichever component the ranks entered + unevenly. A decode loop never stops between layers, so what it pays per layer is this + period. The two are different quantities, not competing estimates of one: quote + `roundtrip` for how long a single collective takes and `period` for what a continuous + stream costs, and never sum or compare them across backends -- only backends that opt in + via `pipeline_pairs` report it at all. + """ + import torch + + self.warm(problem, warmup) + pairs = max(2, int(self.pipeline_pairs)) + samples = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(pairs): + handle = self.dispatch(problem) + self.stage(problem, handle) + self.combine(problem, handle) + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / pairs) + return samples + def benchmark_component(self, component, problem, warmup, iters): """Measure one named component; every component gets the same warm-up first.""" if component == "roundtrip": return self.benchmark_roundtrip(problem, warmup, iters) + if component == "period": + return self.benchmark_period(problem, warmup, iters) if component == "dispatch": return self.benchmark_dispatch(problem, warmup, iters) if component == "stage": @@ -368,12 +504,19 @@ def benchmark_roundtrip(self, problem, warmup, iters): self.warm(problem, warmup) staged = None - if self.stages_fp8_natively: - # Materialise the expert-output stand-in ONCE, untimed. A native fp8 stack has no - # separate conversion between dispatch and combine, so the chained measurement - # must not contain one. Routing is fixed for a ladder point, so the same staged - # tensor is valid for every iteration (for MoRI it IS the registered combine - # buffer, already filled). + if self.stage_excluded_from_roundtrip: + # Materialise the expert-output stand-in ONCE, untimed, so the chained + # measurement is dispatch -> combine and nothing else. Routing is fixed for a + # ladder point, so the same staged tensor is valid for every iteration (for MoRI it is + # the dispatch output itself at BF16, or a fresh `[:rows]` BF16 cast under FP8; for + # FlashInfer it is the workspace combine region, which dispatch cannot clobber because + # that region sits past the end of every dispatch receive plane). + # + # Read the staged payload back through `combine_input_attr` rather than + # constructing one, so whatever the adapter put there round-trips unchanged. No + # current backend needs that -- nccl-ep is the one whose attribute holds a + # non-torch wrapper, and its `stage()` does no device work so it never reaches + # here -- but constructing a tensor instead would silently break it if it did. handle = self.dispatch(problem) self.stage(problem, handle) staged = getattr(handle, self.combine_input_attr) @@ -399,7 +542,8 @@ def finish_dispatch(hh, p=problem): def benchmark_stage(self, problem, warmup, iters): import torch - self.warm(problem, warmup) + # Staging is the timed operation here, so it must be warmed on every iteration. + self.warm(problem, warmup, stage_every=True) def prep_stage(p=problem): return self.dispatch(p) diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py index 672e0c658..b65796825 100644 --- a/experimental/CollectiveX/bench/ep_deepep_v2.py +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -26,6 +26,18 @@ # verifies the wheel's commit tag against the pin — it checks only that the loaded # deep_ep exposes ElasticBuffer (the from-source PR #605 capability). +# Low-latency receive sizing. These are deliberately two numbers, not one: _LL_BUFFER_CAP +# sizes the pre-allocated receive (and so fixes the transport footprint and the fp8 dequant +# volume), while _LL_LADDER_CAP bounds which token counts are measured. See `buffer_cap` for +# why the measured ladder stops below the buffer -- an upstream Blackwell combine defect at +# the 256 rung -- and `create_buffer` for why the buffer must not follow the ladder down. +_LL_BUFFER_CAP = 256 +_LL_LADDER_CAP = 128 +assert _LL_LADDER_CAP <= _LL_BUFFER_CAP <= 511, ( + "the LL receive cap must fit NVSHMEM_QP_DEPTH=1024 ((cap + 1) * 2 <= 1024 => cap <= 511) " + "and the measured ladder must fit inside the buffer" +) + def _fp8_cast_helpers(): """The pinned per-token FP8 cast pair (blockwise e4m3fn, per-128-block FP32 scale). @@ -47,8 +59,9 @@ def _ll_dequant_static(fp8, scales): ``[num_local_experts, cap*num_ranks, hidden]`` = (32, 2048, 7168) at EP8). The low-latency padded shape is constant on every dispatch, so a static (``dynamic=False``) compile fuses to one FP32 pass (~0.5 ms, 6.3x, bit-identical to the dynamic kernel on valid slots). The - dequant runs in every timed component's warmup and samples (~hundreds of thousands of - calls over the profile), so the dynamic kernel's per-call overhead overran the leg's + dequant runs in every timed `stage` sample, in `benchmark_stage`'s warm-up, and once per + other component's warm-up (it was every warm iteration until the staging hoist), so the + call count is large enough that the dynamic kernel's per-call overhead overran the leg's wall-clock budget (all ranks SIGKILLed ~22 min in, no result); the static form brings FP8 low-latency inside the budget BF16 already meets. Padding slots decode to NaN in both forms (FP8 padding bytes) — harmless, because combine is handle-indexed and never reads @@ -151,6 +164,10 @@ def __init__(self, args, rank, world_size, local_rank, device): # deep_ep.utils.math) so the timed stage() does no module lookup in the # measured region. self._to_fp8, self._cast_back = _fp8_cast_helpers() + # Normal/HT quantises inside the timed dispatch with the compiled single-kernel + # form; low-latency keeps the eager helper because its kernel quantises internally + # and the oracle gate is pinned to those bits. See EPBackend.fused_quantize. + self._quant = self.fused_quantize(self._to_fp8) if self.mode == "low-latency": # Legacy Buffer IBGDA decode path: a distinct kernel family whose combine # multiplies by the gate at the source (weighted), not an unweighted rank sum. @@ -166,11 +183,34 @@ def __init__(self, args, rank, world_size, local_rank, device): def buffer_cap(self, args): if self.mode == "low-latency": # LL pre-allocates a fixed [num_local_experts, cap * num_ranks, hidden] receive - # buffer, so cap is a hard per-rank dispatch-slot bound (the harness clamps the - # decode ladder to it and reports the dropped point). 256 sits well under the - # default NVSHMEM_QP_DEPTH ceiling ((cap + 1) * 2 <= 1024 => cap <= 511 with - # NVSHMEM_QP_DEPTH=1024) and is adjustable if the decode ladder needs more. - return 256 + # buffer, so the buffer cap is a hard per-rank dispatch-slot bound. The MEASURED + # ladder is clamped tighter than that buffer (the harness reports every dropped + # point, so the omission lands in the artifact rather than being silent) because + # DeepEP's low-latency combine corrupts at T=256 on every Blackwell SKU -- b200, + # gb200 and gb300, EP8 and EP16, both precisions, MNNVL and RDMA alike -- while + # Hopper stays clean. It is stochastic at roughly 1.5-3.3% per oracle invocation, + # so a passing leg proves nothing; the gate catches it as a 0.07-6.6 relative + # error against a 0.03125 tolerance. Tracked upstream as DeepEP issue #700. + # + # The likely fix already exists upstream and our pin simply predates it: + # PR #642 adds a CTA-scope `fence.proxy.async.shared::cta` before + # `mbarrier_arrive(empty_barriers[stage_idx])` in LOW_LATENCY_COMBINE_RECV, so the + # consumer's shared-memory reads retire before the stage is declared empty and the + # producer's next TMA load refills it. Signalling empty too early lets a row be + # assembled from two tiles, which is exactly the observed signature (one token row, + # norm preserved to 4 s.f., 16-40% of elements deviating). It closed #621, the same + # race reached from NVL72. COLLX_DEEPEP_V2_COMMIT is the head of PR #605, branched + # before #642 merged, so the fence is absent from our build; upstream main has both. + # An earlier on-metal test that appeared to rule fencing out used a device-scope + # __threadfence_system after the grid sync at internode_ll.cu:976 -- a different + # fence at a different site -- so it does not bear on #642. + # + # Raising this back to _LL_BUFFER_CAP is therefore gated on a pin bump, not on a + # new upstream release. That bump is deliberately NOT bundled here: it spans months + # of upstream change, re-baselines every deepep-v2 row including normal mode, and + # needs `rewrite_deepep_v2` made tolerant first (main already carries that + # 'libnccl' fix, so the rewrite's count(old) == 1 assertion would abort the stage). + return _LL_LADDER_CAP return None def create_buffer(self, spec): @@ -180,7 +220,28 @@ def create_buffer(self, spec): args, world_size = self.args, self.world_size self.max_tokens = spec.max_tokens_per_rank if self.mode == "low-latency": + # Size the LL buffer from the fixed cap, NOT from the clamped ladder. Deriving it + # from max(ladder) would halve the receive tensor the moment the ladder was + # clamped, and the receive footprint sets both the transport's memory traffic and + # the fp8 dequant volume (`_ll_recv_bf16` converts the whole padded receive) -- so + # every retained rung's numbers would shift and stop being comparable with the + # published series. Holding the buffer at 256 keeps them bit-comparable and leaves + # the top measured rung at half occupancy, which is the ladder/capacity decoupling + # the earlier capacity probe had to hand-roll. + if spec.max_tokens_per_rank > _LL_BUFFER_CAP: + raise RuntimeError( + f"low-latency ladder maximum {spec.max_tokens_per_rank} exceeds the LL " + f"buffer cap {_LL_BUFFER_CAP}" + ) + self.max_tokens = _LL_BUFFER_CAP self._create_ll_buffer(spec) + # The legacy LL receive is double-buffered with a parity that flips per dispatch, + # and a collective bounds rank drift to about one iteration (a dispatch cannot + # complete until every rank enters it), so two parities cover the worst overlap and + # pairs may be issued back-to-back. This is the pattern SGLang and vLLM use for + # two-micro-batch overlap. Only this path opts in; the ElasticBuffer normal-mode + # receive is not double-buffered that way. + self.pipeline_pairs = 8 return _require_runtime() jit_root = Path(os.environ["EP_JIT_CACHE_DIR"]) @@ -250,6 +311,23 @@ def _create_ll_buffer(self, spec): num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( self.max_tokens, args.hidden, world_size, args.experts ) + kwargs = {} + # On an MNNVL rack the scale-up fabric IS NVLink across trays, but the legacy Buffer + # defaults `allow_mnnvl=False` and a False there self-sets NVSHMEM_DISABLE_MNNVL -- + # so leaving it unset forced the low-latency kernels onto IBGDA on exactly the systems + # whose fast path is MNNVL, and measured the rack's slow path. Keyed on the topology + # the platform already reports rather than on the SKU name. Passed only when the pinned + # wheel accepts it, so an older deep_ep keeps working instead of raising on an unknown + # keyword. + if str(getattr(args, "scale_up_transport", "")) == "mnnvl": + import inspect + if "allow_mnnvl" in inspect.signature(deep_ep.Buffer.__init__).parameters: + kwargs["allow_mnnvl"] = True + else: + raise RuntimeError( + "MNNVL scale-up needs deep_ep.Buffer(allow_mnnvl=...); this wheel lacks it, " + "so the low-latency path would silently run over IBGDA" + ) self.buffer = deep_ep.Buffer( self.group, num_rdma_bytes=num_rdma_bytes, @@ -257,6 +335,7 @@ def _create_ll_buffer(self, spec): num_qps_per_rank=num_qps_per_rank, allow_nvlink_for_low_latency_mode=True, explicitly_destroy=True, + **kwargs, ) def _ll_recv_bf16(self, recv_x): @@ -286,7 +365,9 @@ def _topk_idx_dtype(self): def semantic_payload(self, x): if not self._fp8: return x - return self._cast_back(*self._to_fp8(x)) + # Same callable the wire uses, so the oracle cannot disagree with the sender by + # construction (low-latency: both are the eager helper, matching its kernel). + return self._cast_back(*self._quant(x)) def _encode_dispatch(self, x): if not self._fp8: @@ -296,8 +377,11 @@ def _encode_dispatch(self, x): # send x unquantized; expose the host round-trip as the oracle semantic so the # combine expectation models the FP8 transport (same as semantic_payload). return x, self._cast_back(*self._to_fp8(x)) - quantized = self._to_fp8(x) - return quantized, self._cast_back(*quantized) + # Normal/HT: send BF16 and quantise inside dispatch, where production pays it. The + # payload is therefore x itself; oracle_x is still the round trip, computed once here, + # untimed -- which also compiles this rung's shape before any timed region. + self.assert_quantize_identity(self._to_fp8, self._quant, x) + return x, self._cast_back(*self._quant(x)) def _ll_dispatch(self, p): # Verified pinned signature (legacy.py:553): @@ -322,8 +406,11 @@ def _ll_dispatch(self, p): def dispatch(self, p): if self.mode == "low-latency": return self._ll_dispatch(p) + # Quantise here, not in make_problem: production runs one fused bf16->fp8 kernel per + # forward pass immediately before this collective, so the timed window must contain it. + dispatch_x = self._quant(p.dispatch_x) if self._fp8 else p.dispatch_x recv_x, recv_topk_idx, recv_topk_weights, handle, _ = self.buffer.dispatch( - p.dispatch_x, + dispatch_x, topk_idx=p.topk_idx, topk_weights=p.topk_weights, num_experts=self.args.experts, diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 22a2eb1ac..6d605bf7a 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -58,15 +58,71 @@ _COMBINE_FP32_SINCE = (0, 6, 16) +def _wheel_has_fp32_combine(version: str) -> bool: + """Does this wheel accumulate combine in FP32, per `_COMBINE_FP32_SINCE`? + + Needs real version ordering, not a digit scrape: `0.6.16rc1` predates the rewrite, and + reading it as 0.6.16 models FP32 against a kernel that still rounds per level, which can + exceed COMBINE_REL_TOL and RED a correct run. The opposite error costs a few ulps, so + anything unparseable answers False. + """ + try: + from packaging.version import InvalidVersion, Version + except ImportError: # pragma: no cover - packaging ships with torch + digits = re.match(r"(?:\d+!)?(\d+(?:\.\d+){0,2})", version.strip()) + if digits is None: + return False + release = tuple(int(n) for n in digits.group(1).split(".")) + # No pre-release ordering available, so treat a marked pre-release as below its release. + if re.match(r"[._-]?(a|b|c|rc|alpha|beta|dev|pre)", version[digits.end():]): + return release > _COMBINE_FP32_SINCE + return release >= _COMBINE_FP32_SINCE + try: + return Version(version) >= Version(".".join(str(n) for n in _COMBINE_FP32_SINCE)) + except InvalidVersion: + return False + + +# FP8 block size, matching the DeepSeek-V3 recipe every other FP8 backend here uses. +_FP8_BLOCK = 128 + + +def _blockwise_cast_to_fp8(x): + """Per-128-channel e4m3 quantize: (values [m, n], FP32 scales [m, n//128]). + + A local pair rather than a shared one on purpose. deepep-v2 must use deep_ep's own helper + because its low-latency kernel quantises in-kernel and the oracle compares against those + bits; uccl-ep vendors a copy faithful to UCCL's. FlashInfer's combine is always BF16 and no + kernel here quantises, so nothing external pins the arithmetic -- the oracle only needs a + self-consistent round trip. Each copy is pinned to a different library's numerics, so this + is three contracts, not three copies of one. + """ + m, n = x.shape + blocks = x.view(m, -1, _FP8_BLOCK) + amax = blocks.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + values = (blocks * (448.0 / amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n) + return values, (amax / 448.0).view(m, -1) + + +def _blockwise_cast_back(values, scales): + """Inverse of _blockwise_cast_to_fp8, to BF16.""" + m, n = values.shape + blocks = values.view(m, -1, _FP8_BLOCK).float() + return (blocks * scales.view(m, -1, 1)).view(m, n).to(torch.bfloat16) + + 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",) + # FP8 is DISPATCH-side only: the scales ride as a fourth payload and combine stays BF16, + # so none of the 0.6.16+ combine-quant API is needed. The recipe is the same per-128-block + # e4m3 every other FP8 backend here uses, which keeps the fp8 axis comparable -- but note + # vLLM's integration accepts only nvfp4/mxfp8/bf16 on this transport, so an fp8 row here + # measures what the transport costs with the DeepSeek-V3 recipe rather than a configuration + # a deployment can currently select. `dispatch_dtype` records that per row. + SUPPORTED_PRECISIONS = ("bf16", "fp8") kernel_generation = "flashinfer-mnnvl-one-sided" # stage() now copies the received payload into the workspace combine region. stage_device_work = True @@ -83,6 +139,16 @@ class FlashInferEPBackend(EPBackend): def __init__(self, args, rank, world_size, local_rank, device): super().__init__(args, rank, world_size, local_rank, device) + self._fp8 = self.precision == "fp8" + if self._fp8: + # "-offpath" because vLLM cannot select this recipe on this transport today; the + # bytes and the block size match deepep-v2/uccl-ep so the axis stays comparable. + self.dispatch_dtype = "fp8-e4m3fn-blockwise-offpath" + self.dispatch_value_bytes = 1 + self.dispatch_scale_bytes_per_copy = ( + (args.hidden + _FP8_BLOCK - 1) // _FP8_BLOCK + ) * 4 + self._quant = self.fused_quantize(_blockwise_cast_to_fp8) self._a2a = None self._max_tokens = None self.experts_per_rank = args.experts // world_size @@ -98,6 +164,20 @@ def _topk_idx_dtype(self): """ return torch.int32 + def semantic_payload(self, x): + if not self._fp8: + return x + # Same callable the wire uses, so sender and oracle cannot disagree by construction. + return _blockwise_cast_back(*self._quant(x)) + + def _encode_dispatch(self, x): + if not self._fp8: + return x, None + # Send BF16 and quantise inside dispatch, where production pays it. oracle_x is the + # round trip, computed once here, untimed -- which also compiles this rung's shape. + self.assert_quantize_identity(_blockwise_cast_to_fp8, self._quant, x) + return x, _blockwise_cast_back(*self._quant(x)) + 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. @@ -122,7 +202,11 @@ def create_buffer(self, spec): 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 + dispatch_bytes = ( + hidden * self.dispatch_value_bytes + + self.dispatch_scale_bytes_per_copy + + top_k * 4 + top_k * 4 + ) combine_bytes = hidden * 2 workspace_size = moe_a2a_get_workspace_size_per_rank( ep_size=self.world_size, @@ -145,8 +229,8 @@ def create_buffer(self, spec): 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.library_version = flashinfer.__version__ + if _wheel_has_fp32_combine(flashinfer.__version__): 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. @@ -168,24 +252,60 @@ def dispatch(self, p): 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( + # Quantise here, not in make_problem: production runs one fused bf16->fp8 kernel per + # forward pass immediately before this collective. The scales then ride as their own + # payload, which shifts the expert ids to index 2 -- four payloads, and the kernel's + # kMaxPayloads is exactly 4, matching vLLM's own [values, scales, ids, weights] order. + if self._fp8: + values, scales = self._quant(p.dispatch_x) + payloads = [values, scales, p.topk_idx, p.topk_weights] + expert_id_index = 2 + else: + payloads = [p.dispatch_x, p.topk_idx, p.topk_weights] + expert_id_index = 1 + received = self._a2a.dispatch( p.topk_idx, - [p.dispatch_x, p.topk_idx, p.topk_weights], + payloads, p.T, invalid_token_expert_id=_INVALID_EXPERT, - expert_id_payload_index=1, + expert_id_payload_index=expert_id_index, ) + # One received tensor per payload. Under FP8 `recv_x` stays the VALUES plane and the + # scales ride beside it, so every shape-sensitive consumer keeps working on a tensor. + if self._fp8: + recv_x, recv_scales, recv_idx, recv_w = received + else: + (recv_x, recv_idx, recv_w), recv_scales = received, None return types.SimpleNamespace( - recv_x=recv_x, recv_idx=recv_idx, recv_w=recv_w, + recv_x=recv_x, recv_scales=recv_scales, 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.""" + """The workspace-resident combine payload region for this rung. + + Always BF16: combine carries BF16 whatever the dispatch precision was, so this cannot + key off `recv_x.dtype` -- under FP8 that would size the region for 1-byte values. + """ return self._a2a.get_combine_payload_tensor_in_workspace( - h.tokens, h.recv_x.shape[-1], h.recv_x.dtype + h.tokens, h.recv_x.shape[-1], torch.bfloat16 ) + def _filled_slot_index(self, p, h): + """Row indices of the receive slots dispatch actually filled, resolved once per rung. + + `_valid_rows` gives a boolean mask, and indexing with one needs the match count on the + host, which would put a device read inside the timed stage. Routing is fixed for a + ladder point, so resolve it to an integer index on first use -- which is always the + untimed `warm()` pass -- and cache it on the problem, the way `warm` already caches + `recv_tokens` for the same reason. + """ + index = getattr(p, "flashinfer_filled_slots", None) + if index is None: + index = self._valid_rows(h).nonzero(as_tuple=True)[0] + p.flashinfer_filled_slots = index + return index + def stage(self, p, h): """Materialise the combine payload in the workspace region the API designates. @@ -194,9 +314,26 @@ def stage(self, p, h): 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`. + + Only the filled slots are copied, which is what the kernel's own staging path does -- + `moeA2APrepareCombineKernel` returns early on `token_idx >= recv_counters[source]`. + Copying the whole plane moved 1/occupancy times too much: measured occupancy is + ~0.66 at EP8 and ~0.41 at EP16, i.e. 1.5x and 2.4x over-copy. The slots left + untouched are never read -- combine addresses peers exclusively through the + `topk_send_indices` recorded at dispatch, and an unfilled slot has none. """ buffer = self._combine_buffer(h) - buffer.copy_(h.recv_x) + filled = self._filled_slot_index(p, h) + hidden = h.recv_x.shape[-1] + flat_buffer = buffer.view(-1, buffer.shape[-1]) + source = h.recv_x.view(-1, hidden)[filled] + if self._fp8: + # Combine sends BF16, so the dequant lands here -- device work, which is why + # `stage` is a reported component for this backend under either precision. + source = _blockwise_cast_back( + source, h.recv_scales.view(-1, h.recv_scales.shape[-1])[filled] + ) + flat_buffer[filled] = source h.combine_input = buffer def combine(self, p, h): @@ -229,6 +366,12 @@ def inspect_dispatch(self, p, h): keep = self._valid_rows(h) hidden = h.recv_x.shape[-1] payload = h.recv_x.reshape(-1, hidden)[keep] + if self._fp8: + # The oracle compares a BF16 payload against semantic_payload's round trip, so the + # received FP8 slice is dequantised with the same pair that produced it. + payload = _blockwise_cast_back( + payload, h.recv_scales.reshape(-1, h.recv_scales.shape[-1])[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) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index ba34e135f..eb74620e8 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -962,6 +962,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> stage_pool = {T: [] for T in ladder} # measured only when stage launches device work comb_pool = {T: [] for T in ladder} # ... combine rt_pool = {T: [] for T in ladder} # independently measured round trip + period_pool = {T: [] for T in ladder} # steady-state per-pair cost, opt-in backends only spread_pool = {T: [] for T in ladder} # cross-rank (max-min) of the round trip, per iter # Cross-rank MIN per component. The LAST rank to enter a collective is the one that waited # least -- it started when its peers were already there -- so its duration is the closest @@ -977,7 +978,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # timed_components() encodes the roundtrip-only vs full-component contract # (and whether stage launches device work) once, in the base class. component_order = trial_order(backend.timed_components(), trial_index) - measured = {name: [] for name in ("dispatch", "stage", "combine", "roundtrip")} + measured = {name: [] for name in ("dispatch", "stage", "combine", "roundtrip", "period")} for component_name in component_order: # The base template gives every component the same synchronized # full-roundtrip warm-up before its timed trial and encodes the two @@ -993,6 +994,8 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> stage_pool[T] += _reduce_vec(torch, dist, device, measured["stage"], MAX) rt_max = _reduce_vec(torch, dist, device, measured["roundtrip"], MAX) rt_pool[T] += rt_max + if measured["period"]: + period_pool[T] += _reduce_vec(torch, dist, device, measured["period"], MAX) # Cross-rank SPREAD (max-min) of the same iterations. A collective cannot finish # before its slowest participant, so when ranks enter together every rank measures # nearly the same duration and the spread is small; a large spread means the ranks @@ -1066,6 +1069,20 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> backend.dispatch_value_bytes, backend.dispatch_scale_bytes_per_copy, ) combine_bytes = logical_byte_provenance(rstats["routed_copies"], args.hidden) + # Second byte basis, for the backends whose wire really does carry one copy per + # (token, expert) rather than per (token, dest-rank). Which basis applies is a property + # of the RECEIVE, not of the mode: deepep-v2, uccl-ep and nccl-ep low-latency receive + # per assignment, but MoRI's IntraNodeLL genuinely deduplicates, so keying this on + # "low-latency" would overstate MoRI by the dedup factor (~1.5x at EP8). Key it on the + # combine contract instead, which already encodes the distinction. `routed_copies` + # stays the canonical comparable basis; this is emitted alongside so a reader can + # convert a low-latency row onto the wire basis without guessing the factor. + assignment_copies = int(sum(rstats["expert_assignments_per_rank"])) + wire_basis = ( + "per-assignment" + if backend.combine_weight_semantics == "weighted-kernel-sum" + else "rank-deduplicated" + ) roundtrip_bytes = { field: dispatch_bytes[field] + combine_bytes[field] for field in dispatch_bytes } @@ -1077,6 +1094,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> "dispatch": _component(dp, len(d)), "isolated_sum": _component(isum, 0, derived=True), "roundtrip": _component(rtp, len(rt)), + "period": _component(_pcts(period_pool[T]), len(period_pool[T])), "stage": _component(sp, len(s)), }, # Skew-excluded companion to `components`: same iterations reduced with cross-rank @@ -1105,6 +1123,16 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> "roundtrip": roundtrip_bytes, "stage": stage_bytes, }, + # Copy counts behind the byte figures above, so a reader can rebase them. `routed` + # is the basis they use; `assignments` is the per-(token, expert) count; `wire` names + # which of the two this backend's kernels actually move. Kept OUT of + # `byte_provenance`, whose every value is a per-component byte breakdown -- a reader + # indexing it by component name must not meet a differently-shaped entry. + "logical_copies": { + "routed": int(rstats["routed_copies"]), + "assignments": assignment_copies, + "wire": wire_basis, + }, "receive": { "max": recv_max, "mean": recv_total / world_size, @@ -1181,6 +1209,15 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> }, "workload": { "cross_rank_consistent": routing_consistent, + # The ladder actually measured, plus any requested point the backend's cap + # excluded. This used to be a rank-0 stdout NOTE only, which made a clamped + # ladder invisible to anyone reading the artifact -- so a backend measuring a + # shorter ladder than the sweep requested looked identical to one that measured + # all of it. deepep-v2's low-latency mode clamps below its receive cap to avoid a + # rung upstream corrupts, and that omission has to be legible downstream. + "ladder_measured": list(ladder), + "ladder_dropped": list(dropped), + "ladder_cap": cap, }, "measurement": { "combine_dtype": backend.combine_dtype, @@ -1204,6 +1241,19 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # 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"), + # The library version the line above was decided FROM, where the backend knows it. + # Recording only the outcome leaves the decision unauditable: a reader can see that + # the oracle used the slot-tree model but not which wheel selected it, and so cannot + # tell a correct selection from a mis-parse. None where a backend does not report one. + "library_version": getattr(backend, "library_version", None), + # Whether `roundtrip` excludes expert-output staging. It always does now, unless + # the CX_FP8_CONSUME=dequant hatch is set, but it did not always: rows measured + # before that change carried the staging copy inside the chain for MoRI BF16 and + # FlashInfer BF16 only. Without this field those rows are indistinguishable from + # these ones while measuring a different quantity. + "stage_excluded_from_roundtrip": bool( + getattr(backend, "stage_excluded_from_roundtrip", False) + ), # 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/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index cb3c43951..eca73089e 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -54,9 +54,10 @@ def __init__(self, args, rank, world_size, local_rank, device): # argv. FP8 dispatch is caller-prequantized: MoRI's dispatch kernel keys purely # on the passed tensor dtype, so handing it an e4m3 tensor selects the FP8 # dispatch kernel with no in-kernel cast. Combine stays genuinely BF16 (quant_type - # "none"). With use_external_inp_buf False the launcher - # takes the zero-copy branch, EpCombineIntraNodeKernel_bf16_p2p; the _nop2p and - # _fp8cast variants both sit behind the external-buffer branch we never enter. + # "none"). With use_external_inp_buf True (the engines' setting, pinned unconditionally + # below) the launcher selects EpCombineIntraNodeKernel_bf16_nop2p; the _p2p variant is the + # registered-buffer path this adapter no longer takes, and _fp8cast stays unreachable + # because combine's quant_type is "none". self._fp8_dtype = None if self._fp8: arch = torch.cuda.get_device_properties(device).gcnArchName @@ -80,9 +81,15 @@ def __init__(self, args, rank, world_size, local_rank, device): # default; `kernel_type` kwarg omitted); scale-out EP16 uses InterNodeV1, whose # required enum member is an image-lineage check. # (kernel, generation label, (block_num, rdma_block_num, dispatch_warps, combine_warps)) + # Scale-up matches the config the engines pin: vLLM and SGLang both set block_num 80, + # rdma_block_num 0 and ONE warp_num_per_block of 16 that applies to dispatch and combine + # alike, and neither sets MORI_EP_LAUNCH_CONFIG_MODE, so production runs MANUAL with + # these numbers. 16 is also the kernel's hard ceiling (kMaxWarpGroups 8 x kWarpsPerGroup + # 2, with groupData[8] indexed by warpId/2 and no upstream guard), so it cannot go + # higher. The scale-out tuple is left as-is: EP16 is walled and never dispatched. kernel_name, self.kernel_generation, blocks = ( ("InterNodeV1", "inter-node-v1", (96, 64, 8, 8)) if scale_out - else ("IntraNode", "intranode", (80, 0, 16, 8)) + else ("IntraNode", "intranode", (80, 0, 16, 16)) ) if self.mode == "low-latency": # LOW-LATENCY (decode) mode: IntraNodeLL, the scale-up low-latency kernel. It is @@ -105,7 +112,7 @@ def __init__(self, args, rank, world_size, local_rank, device): "is out of scope; see platform_config ll_backends)" ) kernel_name, self.kernel_generation, blocks = ( - "IntraNodeLL", "intranode-ll", (80, 0, 16, 8) + "IntraNodeLL", "intranode-ll", (80, 0, 16, 16) ) self._kernel_type = None if kernel_name != "IntraNode": @@ -118,7 +125,12 @@ def __init__(self, args, rank, world_size, local_rank, device): self._inter_node = kernel_name == "InterNodeV1" self.num_qps = 1 self.block_num, self.rdma_block_num, self.dispatch_warps, self.combine_warps = blocks - self._external_input = self._inter_node + # External input buffer everywhere, as the engines run it (vLLM leaves MoRI's default, + # SGLang sets it explicitly). It must move together with `combine_warps`: MoRI's tuned + # tables key combine on `zero_copy`, so the engines' 16 warps belong to this mode and + # 4-8 to the registered one. The registered branches below stay, to keep the mode a + # one-line A/B; methodology.md carries the measured cost of mismatching them. + self._external_input = True # Registered-input MoRI copies expert output into a device-side symmetric buffer. External # input kernels consume the dispatch output directly, so their stage is not applicable. # Under FP8, stage also dequantizes the received fp8 payload to BF16 (device work) on @@ -127,6 +139,17 @@ def __init__(self, args, rank, world_size, local_rank, device): # Stash the __init__-only locals the moved create_buffer body reads back. self._gpus_per_node = gpus_per_node + def buffer_cap(self, args): + if self.mode == "low-latency": + # 256 tokens/rank, matching deepep-v2, uccl-ep and nccl-ep, so every backend's + # low-latency ladder ends at the same rung. MoRI imposes no bound of its own here + # -- this adapter previously allowed 512, which published a T=512 decode rung no + # other backend has, so the top of the ladder was unusable for any cross-vendor + # comparison. 256 is also vLLM's DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP, + # i.e. the capacity a deployment actually configures. + return 256 + return None + def create_buffer(self, spec): args, world_size, rank = self.args, self.world_size, self.rank gpus_per_node = self._gpus_per_node @@ -151,8 +174,13 @@ def create_buffer(self, spec): f"MoRI realized {realized_qps} QPs per PE; {self.num_qps} required" ) - # MoRI preallocates one communicator buffer for the case's entire ladder. - self._cap = max(512, spec.max_tokens_per_rank) + # MoRI preallocates one communicator buffer for the case's entire ladder. 256 matches + # the other three backends' low-latency cap and vLLM's own + # DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP; the previous 512 was this adapter's + # invention -- no MoRI bound requires it -- and it published a T=512 low-latency rung + # that no other backend has, so no cross-vendor comparison could use the top of the + # ladder. Normal mode still takes the ladder maximum, which is larger. + self._cap = max(256, spec.max_tokens_per_rank) # quant_type stays "none" for both precisions: dispatch precision is carried by # the passed tensor dtype (caller-prequantized e4m3 under FP8, BF16 otherwise), # and "none" keeps combine a genuine BF16 send. data_type is deprecated upstream @@ -222,8 +250,12 @@ def semantic_payload(self, x): def _encode_dispatch(self, x): if not self._fp8: return x, None - quantized = x.to(self._fp8_dtype) - return quantized, quantized.to(torch.bfloat16) + # Send BF16 and cast inside dispatch, where production pays it: vLLM and SGLang both run + # an aiter per-1x128 quant immediately before mori's dispatch, once per forward pass. + # MoRI's cast needs no compile -- it is already a single eager elementwise kernel, and + # elementwise is per-row invariant, so it carries none of the identity risk the blockwise + # backends do. oracle_x is the round trip, computed once here, untimed. + return x, x.to(self._fp8_dtype).to(torch.bfloat16) def make_problem(self, T, idx, weights, x): indices = idx.to(torch.int32) @@ -244,9 +276,14 @@ def make_problem(self, T, idx, weights, x): return problem def dispatch(self, p): + # See _encode_dispatch: the fp8 cast belongs in the timed window because production runs + # it per forward pass just before this collective. Low-latency is cast here too -- MoRI's + # IntraNodeLL takes a caller-prequantized tensor, unlike deepep-v2/uccl-ep whose LL + # kernels quantise internally, so there is no in-kernel cost already being charged. + dispatch_x = p.dispatch_x.to(self._fp8_dtype) if self._fp8 else p.dispatch_x dispatch_output, dispatch_weights, _scales, dispatch_indices, recv_num = ( self.op.dispatch( - p.dispatch_x, + dispatch_x, p.weights, p.scales, p.indices, @@ -267,16 +304,30 @@ def stage(self, p, h): rows = getattr(p, "recv_tokens", None) if not isinstance(rows, int) or rows < 0 or rows > h.dispatch_output.size(0): raise RuntimeError("MoRI receive count was not validated before staging") - # FP8: dispatch delivered an e4m3 payload; dequantize it to the BF16 combine sends. - h.combine_input = ( - h.dispatch_output.to(torch.bfloat16) if self._fp8 else h.dispatch_output - ) if self._external_input: + # The kernel's own staging copy is bounded by the receive count, not the buffer: with + # an external input buffer and combine quant_type "none" the launcher selects + # `EpCombineIntraNodeKernel_bf16_nop2p`, whose write-based staging loop is bounded by + # `tokenIdx < totalRecvTokenNum` over `args.inpTokenBuf`. (The `UseP2PRead` loop at + # intranode.hpp:542 carries the same bound but is compile-time dead on this path.) + # So it never reads past `rows` and only the filled rows need converting. + h.combine_input = ( + h.dispatch_output[:rows].to(torch.bfloat16) if self._fp8 else h.dispatch_output + ) return None + # Zero-copy path: combine only ever reads the `rows` slots dispatch filled, so cast + # and copy just those. `dispatch_output` is sized to the buffer cap times the world, + # independent of the token count, so casting all of it made this stage flat in T -- + # 61-114us across the whole decode ladder against 6-30us for BF16, ~99.8% of it + # padding at T=1. + source = h.dispatch_output[:rows] + if self._fp8: + # FP8: dispatch delivered an e4m3 payload; dequantize it to the BF16 combine sends. + source = source.to(torch.bfloat16) buffer = self.op.get_registered_combine_input_buffer( - torch.bfloat16, hidden_dim=h.combine_input.size(1) + torch.bfloat16, hidden_dim=h.dispatch_output.size(1) ) - buffer[:rows, :].copy_(h.combine_input[:rows, :]) + buffer[:rows, :].copy_(source) h.combine_input = buffer def combine(self, p, h): diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index 5a616add2..6c2d7aa4d 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -280,6 +280,13 @@ def _ensure_handle(self, p): ) if not self._ll: h.in_weights_t = self._t(p.topk_weights) + else: + # LL applies the gate in its COMBINE kernel, not on dispatch. Wrap the weights here + # with every other per-handle wrapper rather than per timed combine: building one + # costs a torch resolve, an np.asarray and a cybind allocation, and `time_us` + # charges host work inside the window, so a fresh wrapper per iteration was a + # per-call tax no other backend pays. + h.combine_weights_t = self._t(p.topk_weights) # combined output is restored to original token order: [num_tokens, hidden]. h.out = torch.empty((p.T, self.args.hidden), dtype=torch.bfloat16, device=self.device) h.out_t = self._t(h.out) @@ -309,7 +316,7 @@ def _ensure_handle(self, p): h.handle = self._handle torch.cuda.synchronize() if not self._ll: - h.count = int(h.recv_total.item()) + self._bind_ht_recv_count(h) self._bound = h else: h.handle = self._handle @@ -317,6 +324,26 @@ def _ensure_handle(self, p): p._nccl = h return h + def _bind_ht_recv_count(self, h): + """Read HT's received-token count and pre-wrap the combine input at that size. + + The combine call's staging copy is sized by the tensor it is HANDED -- upstream reads + `num_tokens = x->sizes[0]` and copies that many rows into the group's IPC staging -- + not by the group's buffer. Handing it the whole ladder-max receive plane therefore made + HT combine copy `max(ladder) * world` rows on every call regardless of T, putting a + rung-independent floor under it: ~55-80us on a decode leg (ladder max 512) against + ~470-1295us on a prefill leg (max 8192), while the per-token slope stayed within 12%. + Slicing is free (a contiguous leading-dim view) and is what upstream's own ep_test does + when it sizes the combine input to the actual receive count. + + Both callers are untimed -- handle creation and rebind, which only run on a shape change + -- so the `.item()` read never lands in a measured window. + """ + h.count = int(h.recv_total.item()) + # A rank that received nothing still needs a non-empty tensor for the shape checks; the + # routing map decides what combine reads, so the extra row cannot reach the output. + h.combine_in_t = self._t(self._recv_x[: max(h.count, 1)]) + def _rebind(self, h): """Point the single handle at h's routing (collective; untimed callers only). @@ -332,7 +359,7 @@ def _rebind(self, h): ) torch.cuda.synchronize() if not self._ll: - h.count = int(h.recv_total.item()) + self._bind_ht_recv_count(h) self._bound = h # ---- transport contract ------------------------------------------------------------------ @@ -375,8 +402,11 @@ def dispatch(self, p): def stage(self, p, h): # BF16 combine input is the received buffer itself; no device work (value correctness - # is exercised only through the oracle's combine_transformed path). - h.combine_input_t = self._recv_x_t + # is exercised only through the oracle's combine_transformed path). LL takes the full + # padded plane -- its kernel asserts that shape -- while HT takes only the rows this + # routing actually received, so its staging copy scales with T (see + # `_bind_ht_recv_count`). + h.combine_input_t = self._recv_x_t if self._ll else h.combine_in_t def combine(self, p, h): stream = self._stream() @@ -385,7 +415,7 @@ def combine(self, p, h): # source token's gate (CombineOutputs.topk_weights) before the FP32 accumulation. h.handle.combine( CombineInputs(tokens=h.combine_input_t), - CombineOutputs(tokens=h.out_t, topk_weights=self._t(p.topk_weights)), + CombineOutputs(tokens=h.out_t, topk_weights=h.combine_weights_t), config=self._combine_cfg, stream=stream, ) @@ -469,7 +499,7 @@ def _ll_combine_transformed(self, p, h, transformed): stream = self._stream() h.handle.combine( CombineInputs(tokens=self._t(combine_buf)), - CombineOutputs(tokens=h.out_t, topk_weights=self._t(p.topk_weights)), + CombineOutputs(tokens=h.out_t, topk_weights=h.combine_weights_t), config=self._combine_cfg, stream=stream, ) @@ -489,7 +519,8 @@ def combine_transformed(self, p, h, transformed): self._recv_x[: transformed.shape[0]].copy_(transformed.to(self._recv_x.dtype)) stream = self._stream() h.handle.combine( - CombineInputs(tokens=self._recv_x_t), + # Same sliced input the timed path uses, so the two cannot diverge in shape. + CombineInputs(tokens=h.combine_in_t), CombineOutputs(tokens=h.out_t), config=self._combine_cfg, stream=stream, diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 1a61821d0..632c87fd4 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -155,6 +155,10 @@ def __init__(self, args, rank, world_size, local_rank, device): ) self.dispatch_value_bytes = 1 self.dispatch_scale_bytes_per_copy = ((args.hidden + 127) // 128) * 4 + # Normal/HT quantises inside the timed dispatch with the compiled single-kernel + # form; low-latency keeps the eager helper, whose bits its in-kernel cast matches. + # See EPBackend.fused_quantize. + self._quant = self.fused_quantize(per_token_cast_to_fp8) if self.mode == "low-latency": # Legacy low-latency decode path: a distinct kernel family whose combine multiplies # by the gate at the source (weighted), not an unweighted rank sum. LL result tensors @@ -273,7 +277,8 @@ def _topk_idx_dtype(self): def semantic_payload(self, x): if not self._fp8: return x - return per_token_cast_back(*per_token_cast_to_fp8(x)) + # Same callable the wire uses, so sender and oracle cannot disagree by construction. + return per_token_cast_back(*self._quant(x)) def _encode_dispatch(self, x): if not self._fp8: @@ -282,11 +287,11 @@ def _encode_dispatch(self, x): # low_latency_dispatch takes BF16 x and casts to e4m3 inside the kernel, so send x # unquantized; expose the host round-trip as the oracle semantic. return x, per_token_cast_back(*per_token_cast_to_fp8(x)) - fp8, scales = per_token_cast_to_fp8(x) - # Column-major (TMA-compatible) scale layout the dispatch kernel expects, matching UCCL's - # own bench (`scales.T.contiguous().T`) and the LL scale-contiguity note below. - quantized = (fp8, scales.T.contiguous().T) - return quantized, per_token_cast_back(fp8, scales) + # Normal/HT: send BF16 and quantise inside dispatch, where production pays it. oracle_x + # is still the round trip, computed once here, untimed -- which also compiles this rung's + # shape before any timed region. + self.assert_quantize_identity(per_token_cast_to_fp8, self._quant, x) + return x, per_token_cast_back(*self._quant(x)) def _ll_recv_bf16(self, recv_x): """The padded per-expert receive as BF16 [num_local_experts, cap*num_ranks, hidden]. @@ -325,8 +330,17 @@ def dispatch(self, p): # it through so the same call serves both scopes. (num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, _) = self.buffer.get_dispatch_layout(p.topk_idx, self.args.experts) + # Quantise here, not in make_problem: production runs one fused bf16->fp8 kernel per + # forward pass immediately before this collective. The scales then need UCCL's + # column-major (TMA-compatible) layout, matching its own bench's `scales.T.contiguous().T`; + # production's kernel emits that layout directly, so charging the transpose here + # over-states by one small copy. + dispatch_x = p.dispatch_x + if self._fp8: + fp8, scales = self._quant(dispatch_x) + dispatch_x = (fp8, scales.T.contiguous().T) recv_x, recv_topk_idx, recv_topk_weights, _counts, handle, _event = self.buffer.dispatch( - x=p.dispatch_x, + x=dispatch_x, num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, is_token_in_rank=is_token_in_rank, diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 52de60ba9..c6f5d1fb5 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -17,7 +17,7 @@ "partition": "hpc-gpu-1", "account": "customer", "squash_dir": "/mnt/nfs/sa-shared/cx-squash", - "exclude_nodes": "hpc-gpu-1-0,hpc-gpu-1-1,hpc-gpu-1-4,hpc-gpu-1-5,hpc-gpu-1-7,hpc-gpu-1-8,hpc-gpu-1-13,hpc-gpu-1-16,hpc-gpu-1-19" + "exclude_nodes": "hpc-gpu-1-0,hpc-gpu-1-1,hpc-gpu-1-2,hpc-gpu-1-4,hpc-gpu-1-5,hpc-gpu-1-7,hpc-gpu-1-8,hpc-gpu-1-13,hpc-gpu-1-16,hpc-gpu-1-19" }, "network": { "socket_ifname": "eth0", @@ -61,7 +61,8 @@ "partition": "gpu-2", "account": "benchmark", "qos": "gpu-2_qos", - "squash_dir": "/home/sa-shared/containers" + "squash_dir": "/home/sa-shared/containers", + "exclude_nodes": "gpu-2-6" }, "network": { "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7", @@ -103,7 +104,7 @@ "scale_up_transport": "mnnvl", "launcher": "gb-nv", "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]}, - "ll_backends": {"nccl-ep": [8]}, + "ll_backends": {"deepep-v2": [8, 16], "nccl-ep": [8]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { "partition": "batch", @@ -123,7 +124,7 @@ "scale_up_transport": "mnnvl", "launcher": "gb-nv", "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]}, - "ll_backends": {"nccl-ep": [8]}, + "ll_backends": {"deepep-v2": [8, 16], "nccl-ep": [8]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { "partition": "batch_1", diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 8108c3fd4..0a4e4456b 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -22,9 +22,36 @@ It does not predict serving throughput without a separate correlation study. The implemented workload is `deepseek-v3`: hidden 7168, top-k 8, 256 routed experts, packed placement, and one pinned fixed resource profile per backend/topology. Combine is always BF16; dispatch precision is a swept dimension — a BF16 control and, on the backends whose FP8 dispatch is -supported upstream (DeepEP V2, MoRI, UCCL-EP), an FP8 dispatch (`bf16`, `fp8`), +supported upstream (DeepEP V2, MoRI, UCCL-EP, FlashInfer EP), an FP8 dispatch (`bf16`, `fp8`), caller-prequantized in `normal` mode (the `low-latency` kernels quantize FP8 internally from BF16 on -DeepEP and UCCL-EP, and stay caller-prequantized on MoRI). NCCL EP is BF16-only this release, so its +DeepEP and UCCL-EP, and stay caller-prequantized on MoRI). Because `normal`-mode FP8 is +caller-prequantized, that quantize is a cost a production forward pass pays on the critical path, so +it is charged **inside the measured dispatch** rather than prepared ahead of the timing window. On +DeepEP V2, UCCL-EP and FlashInfer EP it is issued as one fused kernel and guarded bitwise against its +eager reference; MoRI needs neither, because its quantize is a single plain dtype cast. This means an FP8 +`normal` dispatch number covers quantize-plus-transport while its BF16 control covers transport +alone, and it is why an FP8 `normal` row is not comparable to one measured before that change. The +sweep `version` deliberately stays 1 across it, so the version tag does NOT separate the two +generations — `implementation.stage_excluded_from_roundtrip` and the presence of a `stage` component +are the only discriminators, and a consumer comparing rows across that boundary has to key on them. + +Read that charge as a **fixed per-call cost, not a payload-proportional one**, or the FP8-versus-BF16 +comparison will be misread at the bottom of the ladder. Measured on DeepEP V2 decode, FP8 dispatch p50 +minus its BF16 control is roughly flat until the transport starts to dominate — 65us on h100, 59us on +b200, 27us on b300 and 57us on gb300 at T=1, holding within a microsecond or two through T=64, then +decaying, and by T=512 it has gone slightly negative on h100 and b300 where halved payload bytes more +than repay it. FlashInfer EP carries a larger one (~107us on gb300) because its codec is its own and +FP8 adds a fourth dispatch payload. At T=1 the FP8 path moves *fewer* bytes than BF16, so none of this +is transport: it is per-call work, and it exceeds the fused quantize's own device time (1.5-3.6us, +measured per SKU) by more than an order of magnitude, because the timing window has no host sync +before its start event (see below) and a near-idle stream at T=1 lets host-side launch cost land +inside it. Compiling the quantize reduces this rather than causing it — an eager quantize measures +33-39us worse per decode dispatch on h100 through the same window — but a production forward pass +issues one custom quantize op into a stream that is already busy, so the small-T end of an FP8 +`normal` row is the least production-representative number the suite emits. Compare FP8 and BF16 at +the top of the ladder, where the charge is repaid, rather than at T=1. +`low-latency` rows are unaffected: those kernels either quantize internally (nothing for the caller +to charge) or take pre-quantized input by API contract. NCCL EP is BF16-only this release, so its cells carry the control alone; the per-backend precision set lives in `sweep_matrix.py`'s `BACKEND_PRECISIONS` and a backend never emits a case for a precision it does not support. `normal`-mode cases use the @@ -35,6 +62,20 @@ cells carry the control alone; the per-backend precision set lives in `sweep_mat T=1..512 powers of two and prefill T=1024..8192 powers of two. Ladders are model-specific and live with the workload in `configs/sweep.json`. +A backend may clamp the ladder below that, and every clamped point is reported in the artifact +rather than dropped silently. Exactly one backend clamps today: DeepEP V2 in `low-latency` mode +pre-allocates a fixed receive, so its ladder cannot exceed that buffer, and it is currently held at +**T=128**, one rung below the 256-slot receive, because DeepEP's low-latency combine corrupts the +256 rung on every Blackwell SKU we run — B200, GB200 and GB300, EP8 and EP16, both precisions, +MNNVL and RDMA alike, while Hopper stays clean. It is stochastic at roughly 1.5-3.3% per +invocation and shows up as one wrong token row whose norm still matches, so it is a correctness +gate failure rather than a crash (upstream DeepEP issue #700). The receive stays sized at 256 +even though the ladder stops at 128: its footprint drives both the transport's memory traffic and +the FP8 dequant volume, so shrinking it with the ladder would move every retained rung and break +comparability with the published series. The likely upstream fix (DeepEP PR #642, a CTA-scope +fence so the combine consumer's shared-memory reads retire before the stage is recycled) landed +after the commit we pin, so raising this back to 256 is gated on a pin bump. + `sweep_matrix.py` materializes the requested SKUs, backends, EP sizes, and token ladders into a matrix document, then extracts strict per-shard controls. `--only-sku`, `--exclude-skus`, `--ep-sizes`, and `--precisions` select a subset; a subset produces a smaller matrix, not a @@ -57,7 +98,42 @@ request NCCL Device API LSA and fail closed unless the realized LSA team covers x86 EP16 scale-out uses the hybrid path with GIN and requires two logical scale-out domains represented by two physical RDMA ranks, with eight scale-up ranks per domain. GB EP16 remains MNNVL scale-up and uses LSA. MoRI EP8 uses the direct IntraNode kernel on every CDNA SKU; its EP16 InterNodeV1 path is -configured but unsupported (transport-layer combine corruption, ROCm/mori#475) and never dispatched. UCCL-EP is a drop-in, API-identical DeepEP replacement that keeps the legacy `Buffer` +configured but unsupported (transport-layer combine corruption, ROCm/mori#475) and never dispatched. +MoRI runs under its MANUAL launch mode with a pinned launch config, because that is what the engines +run: neither vLLM nor SGLang sets `MORI_EP_LAUNCH_CONFIG_MODE`, and for the BF16 and FP8 paths this +suite sweeps both pin block_num 80, +rdma_block_num 0, and `warp_num_per_block` 16 for the intra-node kernel, applied to dispatch and +combine alike (neither passes a per-call override, so combine inherits the 16). Both also run with an +external input buffer, which is MoRI's default and which SGLang sets explicitly. Those two settings +are pinned together deliberately: MoRI's own tuning tables key combine on `zero_copy`, selecting +roughly 16 warps for external input against 4-8 for a registered buffer, so taking the engines' warp +count while keeping a registered buffer would match neither, and the mismatch is not hypothetical in +either direction. Measured on MI300X, MI325X and MI355X: **in registered-buffer mode** 16 warps costs ++13-18% combine at T=128 and +61-78% at T=512 against 8, while **in the external-input mode the +engines actually run** the same 16 warps *wins* — 14-19% at T=128, 26-27% at T=256, and 9-14% at every +prefill rung including T=8192, the true top of the ladder. Below T=32 it gives up 0.2-2.5us, which is +the only rung range where 8 is ahead. All arms were correct at every rung, so the pairing is a +throughput question, not a correctness one. With an external input buffer the kernel does its own +staging copy, bounded by the receive count, so BF16 rows hand over the dispatch output unchanged: +their `stage` component is still declared, as an explicit unavailable marker with a null percentile +block and a zero sample count, the same way any backend whose staging is a bare pointer assignment +reports it. FP8 rows still stage for real, because the received payload has to be dequantized. These +numbers therefore describe the engine-integrated configuration, not MoRI's peak: its shipped tuning +tables reach a faster combine with per-shape block and warp counts no engine selects, and AUTO would +not reproduce them uniformly anyway: gfx950 ships no IntraNodeLL combine table and no BF16 rule for +normal-mode IntraNode dispatch, so AUTO falls back to hard-coded defaults for exactly those two and +couples the result to whichever MoRI revision is pinned. It would find genuinely tuned rules for the +other two paths, which is its own problem — a config that is tuned on some paths and defaulted on +others is not one number about the hardware. How far off peak is arch-dependent and only partly measured: comparing across +buffer modes on MI355X, with the registered mode's excluded BF16 stage added back so the comparison is +not flattered, a registered buffer at 8 warps is still 15% faster at T=512 decode and 8% at T=8192 +prefill than the shipped pairing. That margin did not reproduce as a clear win on gfx942, so treat +0-15% as the honest range rather than a single figure -- and note that the faster configuration is one +no engine runs, which is why this suite does not chase it. One +asymmetry is worth stating: the low-latency arm has no engine-integrated configuration to match at +all, because SGLang's low-latency path pins `AsyncLL` at 8 warps while this suite uses `IntraNodeLL` +(`AsyncLL` is split-phase and fails silently under a single-call harness), so the low-latency launch +config is inherited from the normal-mode tuple by choice rather than by precedent. UCCL-EP is a drop-in, API-identical DeepEP replacement that keeps the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum) but routes it over CPU-proxy GPUDirect RDMA on plain `libibverbs` — no NVSHMEM/IBGDA — with software message ordering, atomics, and flow control; its scale-up is single-node `cudaIpc` over NVLink/XGMI (so the scale-up domain is one physical node, @@ -115,9 +191,92 @@ Adapters may not generate routing or reinterpret one quantity as the other. Normal mode uses `layout-and-dispatch-v1`: dispatch timing includes layout plus communication, and combine returns activation payload through an unweighted rank-sum path. Expert-output staging is -outside isolated combine timing and inside the measured paired roundtrip. Each component declares +outside isolated combine timing AND outside the measured paired roundtrip, so `roundtrip` means +dispatch then combine — the transport — in every row. It is reported as its own `stage` component +wherever it does device work. The one exception is the `CX_FP8_CONSUME=dequant` verification hatch, +which puts the conversion back inside the chain on purpose. + +Under FP8, treat `stage` as **harness scaffolding rather than a phase a serving stack has**. Its work +is converting the received FP8 payload to the BF16 that combine sends, and in production nothing does +that as a separate step: the FP8 lands in the expert GEMM, which reads FP8 operands natively and emits +BF16, and that GEMM output is what combine receives. This suite deliberately does not run the expert +GEMM — it measures the collective, not the layer — so `stage` stands in for it. That is why `stage` is +excluded from `roundtrip`, and it is also why **`stage` must not be summed into a total or compared +between backends**: each adapter converts a different amount. DeepEP V2 and UCCL-EP convert only the +received rows in `normal` mode but the whole padded plane in `low-latency`, where the receive buffer is +`[experts, cap * ranks, hidden]` regardless of token count; MoRI converts only the received rows; +FlashInfer only the filled slots. So the same component name covers several different quantities, and +for two of the backends it covers a different one per mode. The one production path that *does* pay a separate materialised dequant is +a quant-format mismatch fallback (vLLM dequantises when `block_k` disagrees with DeepEP's block size); +`CX_FP8_CONSUME=dequant` exists to model exactly that case, and it is not the default because it is +not the fast path. + +Read `implementation.stage_excluded_from_roundtrip` as "there was device-work staging and it was +hoisted out of the chain", not as "this row's roundtrip is stage-free". It is gated on whether the +backend's `stage()` does device work at all, so it is `false` in two unrelated situations, and the +`stage` component is what separates them: **absent** means the backend has nothing to stage (the +staging is a bare pointer assignment, as for NCCL EP and for every BF16 row that hands the receive +buffer straight to combine), while **present alongside `false`** means the `dequant` hatch put the +conversion back inside the chain. A reader that treats `false` alone as "roundtrip includes staging" +will wrongly subtract a cost the row never paid. Each component declares availability, origin, and sample count. A paired-only API reports null isolated components. -`isolated_sum` is derived. The artifact records the mode so a reader can keep distinct measurement +`isolated_sum` is derived. + +Headline latency is the p99 of the per-iteration cross-rank MAX (`p50` is emitted alongside it, and +`summarize.py` prints both; the p99 is the figure the published cohorts rank on). That is not in +tension with the guidance below to rank by hand on p50: the published cohorts do not order cells by +raw p99, they group them into bootstrap equivalence bands, so a cell whose p99 is dominated by +worst-rank stalls rather than transport lands in a tie band instead of being declared a winner or a +loser. Reading a single pair of cells yourself has no such machinery, which is why the bracket below +is the manual procedure. MAX is the +reduction because a layer is not finished until +its slowest rank is, so MAX is the completion cost, and it charges inter-rank entry stagger to +whichever component the ranks entered unevenly. How much stagger there is depends on the code +path AND the precision, not only on the fleet: on identical h200 low-latency decode cells the +per-iteration spread is ~9.3 us for deepep-v2 and uccl-ep at BF16 (they share the legacy +`Buffer` path) against ~2.6 us for nccl-ep, and it collapses to ~2.8 us for those same two under +FP8, where the kernel quantises in-kernel and the heavier dispatch self-aligns the ranks. So the +term is not subtractable in any principled way, and MAX alone taxes some rows more than others. + +Some rows also carry a `period` component, and it answers a different question from `roundtrip`. +`roundtrip` drains the GPU around each pair, so it reports the latency of an idle pipeline. A decode +loop never stops between layers — the next dispatch is already in flight while the previous combine's +stragglers land — so what a serving stack pays per layer is the pipeline's PERIOD, which is smaller +than the sum of separately-drained stages and is also indifferent to how inter-rank entry stagger gets +attributed. Both are real; quote `roundtrip` for how long one collective takes and `period` for what a +continuous stream costs, and never sum them or treat one as a correction to the other. + +`period` is opt-in per backend (`pipeline_pairs`) rather than universal, because issuing pairs +back-to-back lets ranks drift, and dispatch is a peer WRITE into another rank's buffer — stream order +on the receiver does not order the sender's remote writes. A collective bounds that drift to roughly +one iteration, since a dispatch cannot complete until every rank enters it, so a receive buffer that +is double-buffered per dispatch is safe and one shared buffer is not. Today only DeepEP V2's +low-latency path opts in, which is the same two-micro-batch overlap SGLang and vLLM run. Enabling it +where the buffer cannot absorb the drift would produce a fast number over corrupted data, so it is off +by default and a row without the component simply did not measure it. + +Every row therefore also carries `cross_rank_min_us` (the same iterations reduced with MIN — the +skew-excluded floor) and `cross_rank_spread_us` (per-iteration MAX minus MIN). Read MAX and MIN +as a bracket. Two cells whose MAX gap is smaller than the larger contender's spread are not +separated by the data: rank on roundtrip p50 and call a winner only where MAX and MIN agree on +the ordering. Do not rank on p99 of MAX for multi-node decode cells, where it is dominated by +worst-rank stalls rather than transport — p99 of MIN is the synchronized-cost tail beside it. The +isolated components inherit the preceding operation's per-rank exit stagger, so treat them as +residual-wait diagnostics rather than per-operation costs; the paired roundtrip is the +comparable quantity. + +One backend's timed window omits a cost the others pay, deliberately. nccl-ep binds routing with +`ncclEpUpdateHandle`, a collective whose cost scales with the group's token capacity rather than with +the token count, so charging it per iteration would import a ladder-max-proportional term into +dispatch -- the same shape of artifact that sizing HT's combine input to the ladder maximum used to +put under combine. It is therefore bound during the untimed warm-up, which is also what NVIDIA's own +`ep_bench` does (CUDA events around dispatch and combine only, handle update outside the loop). In +low-latency mode there is nothing to exclude: `ncclEpUpdateHandle` returns immediately and the kernel +reads the cached routing inside the timed dispatch. Every other backend pays its layout per timed +call -- uccl-ep calls `get_dispatch_layout` inside dispatch; deepep-v2, MoRI and FlashInfer pass +routing on every call -- and those costs scale with tokens, so they belong in the window. + +The artifact records the mode so a reader can keep distinct measurement contracts separate. Every measured component uses one fixed timing profile, defined once in `configs/sweep.json` @@ -156,12 +315,16 @@ Logical payload bandwidth is: Payload bytes use rank-deduplicated token-rank activations and exclude expert metadata, padding, and backend buffer capacity. BF16 moves 2 bytes per value with no scale payload; an FP8 -dispatch moves 1 byte per value, plus per-128-block FP32 scales for DeepEP's and UCCL-EP's blockwise -codec (none for MoRI's plain e4m3 cast), while combine stays BF16 — so the dispatch and combine directions can carry +dispatch moves 1 byte per value, plus per-128-block FP32 scales for every blockwise codec here — +DeepEP V2, UCCL-EP and FlashInfer EP, which carries them as a fourth dispatch payload — and none for +MoRI's plain e4m3 cast, while combine stays BF16 — so the dispatch and combine directions can carry different byte counts and the roundtrip is their per-field sum. The rank-deduplicated count is exact -for the normal-mode layout; the low-latency layout sends one copy per (token, expert) assignment -rather than per (token, rank), so for a token whose experts share a destination rank this logical -count is a lower bound on the bytes the low-latency kernels actually move. Latency (the headline) is +for the normal-mode layout. It is also exact for a low-latency kernel that deduplicates per rank +(MoRI's `IntraNodeLL`, whose combine is an unweighted rank-sum). The low-latency kernels that apply +top-k weights inside combine instead send one copy per (token, expert) assignment rather than per +(token, rank), so for a token whose experts share a destination rank this logical count is a lower +bound on the bytes those kernels actually move. Each row states which basis it used in +`logical_copies`, so the two are never silently mixed. Latency (the headline) is measured directly and is unaffected. Algorithm bandwidth, bus bandwidth, wire utilization, and physical-link utilization are not emitted without a defined primitive model or transport counters. Logical bandwidth must never be labeled physical bandwidth. Payload and token diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh index eee35d1fa..1f891295f 100644 --- a/experimental/CollectiveX/launchers/launch_mi-amds.sh +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -51,7 +51,6 @@ if [ "$NODES" -gt 1 ]; then else export COLLX_TRANSPORT=xgmi fi -export COLLX_RUN_TIMEOUT="${COLLX_RUN_TIMEOUT:-1800}" collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_SQUASH_DIR COLLX_STAGE_DIR PARTITION="$COLLX_PARTITION"; SQUASH_DIR="$COLLX_SQUASH_DIR" diff --git a/experimental/CollectiveX/launchers/launch_single-slurm.sh b/experimental/CollectiveX/launchers/launch_single-slurm.sh index f30eef42e..b58eb0b63 100644 --- a/experimental/CollectiveX/launchers/launch_single-slurm.sh +++ b/experimental/CollectiveX/launchers/launch_single-slurm.sh @@ -114,17 +114,26 @@ for allocation_attempt in 1 2 3; do elif [ "$RUNNER" = b300 ] \ && ! collx_validate_cuda_context_on_job "$JOB_ID" "$NODES" "$GPN"; then validation_failure=cuda-context + elif ! collx_validate_gpu_health_on_job "$JOB_ID" "$NODES" "$GPN"; then + validation_failure=gpu-health else break fi retryable=0 [ "$RUNNER:$validation_failure" != h100-dgxc:network ] || retryable=1 [ "$RUNNER:$validation_failure" != b300:cuda-context ] || retryable=1 + # A throttled GPU is always someone else's node to fix, never this leg's to tolerate: one + # clamped device paces every rank, so retrying elsewhere is right on every SKU. + [ "$validation_failure" != gpu-health ] || retryable=1 if [ "$retryable" = 0 ] || [ "$allocation_attempt" = 3 ]; then if [ "$validation_failure" = network ]; then collx_log_tail "${COLLX_NETWORK_PROFILE_LOG:-}" collx_die "allocated nodes failed the network profile" fi + if [ "$validation_failure" = gpu-health ]; then + collx_log_tail "${COLLX_GPU_HEALTH_LOG:-}" + collx_die "allocated nodes hold a thermally throttled GPU" + fi collx_log_tail "$COLLX_CUDA_CONTEXT_LOG" collx_die "allocated nodes failed accelerator context validation" fi diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh index 24072604a..eafec3260 100644 --- a/experimental/CollectiveX/runtime/common.sh +++ b/experimental/CollectiveX/runtime/common.sh @@ -760,6 +760,27 @@ BASH printf '%s' "$sq" } +# Reject an allocation whose GPUs are throttled: collectives are barriers, so one clamped device +# paces every rank. `--gres` mirrors the cuda-context probe below rather than relying on a step +# inheriting the job's GRES by default, so the probe provably sees the devices it is judging -- +# a gate that sees none would pass everything while looking like protection. `--time` bounds the +# worst case: nvidia-smi can wedge uninterruptible on exactly the sick hardware this looks for, +# and Python's own timeout cannot reap a process in D-state. +collx_validate_gpu_health_on_job() { + local job_id="$1" nodes="$2" gpus_per_node="$3" log_label=gpu-health log + case "${COLLX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${COLLX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + log="$(collx_private_log_path "$log_label")" + export COLLX_GPU_HEALTH_LOG="$log" + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --gres=gpu:"$gpus_per_node" --time=5 --chdir=/tmp --input=all \ + --export="$(collx_host_exports)" python3 /dev/stdin gpu-health \ + < "$COLLX_RUNTIME_DIR/probe.py" >"$log" 2>&1 +} + # A clean nvidia-smi inventory does not prove that a prior cancelled workload # released every CUDA context. Retaining each primary context catches poisoned # allocations before a full shard spends time failing every case. @@ -906,7 +927,9 @@ collx_run_shard() { || { rm -f "$argv_file"; collx_die "case $ci produced no benchmark arguments"; } collx_log "EP${NGPUS}[$((ci + 1))/$expected_cases] $COLLX_BENCH" runtime_log="$(collx_private_log_path "runtime-c$(printf '%03d' "$ci")")" - if ! timeout -k 30 "${COLLX_RUN_TIMEOUT:-900}" \ + # A hang guard, NOT a work budget: at 900 it killed FP8 prefill cases that had already + # written complete, all-rungs-passed artifacts. 1800 is what the AMD launcher already used. + if ! timeout -k 30 "${COLLX_RUN_TIMEOUT:-1800}" \ srun --jobid="$JOB_ID" --nodes="$NODES" \ --ntasks="$NGPUS" --ntasks-per-node="$GPN" --chdir=/tmp \ --container-name="$container_name" --container-image="$SQUASH_FILE" \ diff --git a/experimental/CollectiveX/runtime/config.py b/experimental/CollectiveX/runtime/config.py index d1e6132bd..083f790ce 100644 --- a/experimental/CollectiveX/runtime/config.py +++ b/experimental/CollectiveX/runtime/config.py @@ -128,11 +128,12 @@ def _emit_argv(case: dict, version: object, runner: str, ts: str, index: int) -> iters, trials, warmup = str(case["timing"]).split(":") for flag, value in (("--iters", iters), ("--trials", trials), ("--warmup", warmup)): argv += [flag, value] - # precision is part of the filename so a cell's bf16 and fp8 legs (distinct shards - # sharing runner/backend/phase and each numbering cases from index 0) cannot collide - # when they land in the shared results/ dir under the same second-resolution ts. + # precision and mode are in the filename so a cell's legs -- distinct shards sharing + # runner/backend/phase, each numbering cases from index 0 -- cannot overwrite each other in + # the shared results/ dir when they share a second-resolution ts, as they do whenever + # several shards are driven from one loop rather than one CI job each. out = ( - f"results/{runner}_{case['backend']}_{case['precision']}_{case['phase']}" + f"results/{runner}_{case['backend']}_{case['precision']}_{case['mode']}_{case['phase']}" f"_{ts}-c{index:03d}.json" ) argv += ["--out", out] diff --git a/experimental/CollectiveX/runtime/probe.py b/experimental/CollectiveX/runtime/probe.py index 922a99682..073096c98 100644 --- a/experimental/CollectiveX/runtime/probe.py +++ b/experimental/CollectiveX/runtime/probe.py @@ -31,6 +31,109 @@ def validate_cuda_context(expected: int) -> None: raise SystemExit(1) +_GPU_HEALTH_FIELDS = ("index", "clocks_event_reasons.sw_thermal_slowdown", + "clocks_event_reasons.hw_thermal_slowdown", "temperature.gpu") + + +def gpu_health_faults(output: str, max_temperature_c: int = 90) -> list[str]: + """Throttled or overheating GPUs in an `nvidia-smi --format=csv,noheader` block. + + Split out from the I/O so the parsing is testable without hardware; see + tests/test_runtime.py::GpuHealthProbe. Returns [] for anything it cannot read, because the + caller treats an unreadable probe as healthy rather than blocking a leg on it. + """ + faults = [] + for line in output.splitlines(): + cells = [cell.strip() for cell in line.split(",")] + if len(cells) != len(_GPU_HEALTH_FIELDS): + continue + index, software, hardware, temperature = cells + # "Not Active" is the healthy reading, so compare exactly rather than searching for + # "Active" -- a substring test passes the fault straight through. + throttled = "Active" in (software, hardware) + try: + too_hot = int(temperature.split()[0]) > max_temperature_c + except (IndexError, ValueError): + too_hot = False + if throttled or too_hot: + faults.append( + f"gpu {index}: sw_thermal={software} hw_thermal={hardware} temp={temperature}" + ) + return faults + + +def gpu_temperature_spread(output: str) -> tuple[int, int, int] | None: + """`(hottest, median, spread)` GPU temperature, or None if unreadable. + + Reported, NOT gated on. The absolute threshold in `gpu_health_faults` is architecture + dependent and can be unreachable: an H100 engages its software thermal slowdown at ~86-87 C, + so a clamped H100 never crosses a 90 C limit even under load, and at pre-flight it is idle + anyway. In the one fault measured end to end, the only signal visible at pre-flight time was + the RELATIVE one -- the sick GPU idled at 55 C against ~30 C for every sibling, then clamped + under load. Recording the spread every run is how that becomes a gate with evidence behind its + threshold instead of a heuristic picked from a single incident. Healthy references so far: + 50-66 C under load on h100 (~16 C spread), 34-39 C on b200. + """ + temperatures = [] + for line in output.splitlines(): + cells = [cell.strip() for cell in line.split(",")] + if len(cells) != len(_GPU_HEALTH_FIELDS): + continue + try: + temperatures.append(int(cells[3].split()[0])) + except (IndexError, ValueError): + continue + if not temperatures: + return None + temperatures.sort() + median = temperatures[len(temperatures) // 2] + return temperatures[-1], median, temperatures[-1] - median + + +def validate_gpu_health(max_temperature_c: int = 90) -> None: + """Reject an allocation holding a thermally throttled GPU. + + Every collective is a barrier across all ranks, so one clamped device paces the whole leg. A + B200 with GPU 7 held at 120 MHz against 1965 MHz on its siblings ran a case 17x slower and was + killed by the wall-clock guard twice, at 900s and again at 1800s, looking exactly like a code + pathology. Rejecting the allocation costs seconds; diagnosing it cost two 30-minute burns. + + The signal is the throttle FLAG, not the clock: the allocation is idle at this point and an idle + B200 also reads 120 MHz, so a clock threshold cannot tell health from idleness. Temperature is a + second, independent signal because the flag can clear between samples while the fault persists. + + Fails OPEN on anything unexpected -- no `nvidia-smi`, non-zero exit, unparseable output. A check + that blocks legs when it cannot read the hardware is worse than the fault it looks for. + """ + import shutil + import subprocess + + if shutil.which("nvidia-smi") is None: + return + try: + output = subprocess.run( + ["nvidia-smi", f"--query-gpu={','.join(_GPU_HEALTH_FIELDS)}", + "--format=csv,noheader"], + capture_output=True, text=True, timeout=60, check=True, + ).stdout + except (OSError, subprocess.SubprocessError): + return + faults = gpu_health_faults(output, max_temperature_c) + for fault in faults: + _emit(f"gpu-health-fault {fault}") + if faults: + raise SystemExit(1) + # Positive control. Without it a gate that has gone BLIND -- no visible devices, or a driver + # old enough to spell these fields `clocks_throttle_reasons.*` -- writes an empty log and is + # indistinguishable from one that inspected eight healthy GPUs. Recording the count is what + # makes "the gate ran and saw N devices" checkable per cluster instead of assumed. + spread = gpu_temperature_spread(output) + detail = "" if spread is None else f" hottest={spread[0]}C median={spread[1]}C spread={spread[2]}C" + _emit( + f"gpu-health-checked gpus={sum(1 for line in output.splitlines() if line.strip())}{detail}" + ) + + def _emit(marker: str) -> None: # collx_validate_network_profile_on_job (runtime/common.sh) greps these exact strings # out of the per-node probe log to derive COLLX_SOCKET_IFNAME / COLLX_RDMA_LINK_LAYER and to @@ -113,11 +216,13 @@ def main() -> None: commands.add_parser("default-route-interface") command = commands.add_parser("prepare-cache"); command.add_argument("parent") command = commands.add_parser("cuda-context"); command.add_argument("expected", type=int) + commands.add_parser("gpu-health") command = commands.add_parser("network-profile"); command.add_argument("socket_names"); command.add_argument("rdma_devices"); command.add_argument("gid_index") args = parser.parse_args() if args.command == "default-route-interface": print(default_route_interface(), end="") elif args.command == "prepare-cache": print(prepare_cache(args.parent), end="") elif args.command == "cuda-context": validate_cuda_context(args.expected) + elif args.command == "gpu-health": validate_gpu_health() else: validate_network_profile(args.socket_names, args.rdma_devices, args.gid_index) diff --git a/experimental/CollectiveX/summarize.py b/experimental/CollectiveX/summarize.py index 5785d6757..ca6388375 100644 --- a/experimental/CollectiveX/summarize.py +++ b/experimental/CollectiveX/summarize.py @@ -43,11 +43,66 @@ def _identity(document: dict) -> tuple[str, str, str, str, str, str, int, str]: ) -def _headline(document: dict) -> tuple[int | str, float | str, float | str]: +def _topology(document: dict) -> str: + """Scale-up shape, because the same `ep` label is not the same hardware. + + GB200/GB300 run 4 GPUs per node inside a 72-GPU MNNVL domain, so their "EP8" spans two + trays over a rack fabric while every other SKU's EP8 is 8 GPUs in one node over NVLink or + XGMI. Nothing in the timing can fix that; printing it stops a reader comparing the two as + though they were the same configuration. + """ + topology = document.get("topology") or {} + per_node = topology.get("gpus_per_node") + domain = topology.get("scale_up_domain") + nodes = topology.get("nodes") + if per_node is None or domain is None: + return "-" + return f"{nodes}x{per_node}/d{domain}" + + +def _wire_basis(document: dict) -> str: + """Which copy basis this backend's kernels actually move. + + Low-latency deepep-v2/uccl-ep/nccl-ep receive one copy per (token, expert); MoRI's + IntraNodeLL deduplicates by destination rank. At the same token count that is ~1.5x + different combine traffic at EP8, so two low-latency rows are not doing equal work. + """ + rows = document["measurement"]["rows"] + copies = (rows[0] if rows else {}).get("logical_copies") or {} + return {"per-assignment": "assign", "rank-deduplicated": "dedup"}.get(copies.get("wire"), "-") + + +def _headline(document: dict) -> tuple: + """Headline row, with the skew bracket beside it. + + `p50`/`p99` are the cross-rank MAX: a layer is not finished until its slowest rank is, so + MAX is the completion cost. But entry stagger is charged to it, and how much is a property + of the BACKEND, not just the fleet — on identical h200 low-latency cells the per-iteration + spread is 9.2us for deepep-v2 and uccl-ep (which share a kernel family) against 2.0us for + nccl-ep, flat across the whole ladder. So MAX alone taxes some backends more than others, + and two cells whose MAX gap is smaller than that spread are not separated by the data. + `min50` is the same iterations reduced with MIN — the skew-excluded floor — and `skew` is + the per-iteration MAX-MIN. Read the pair as a bracket, not the two ends as rival metrics. + """ rows = document["measurement"]["rows"] + if not rows: + # This renderer validates nothing and must degrade rather than crash: a shard that + # reported an outcome but no measurement rows is malformed, not a reason to lose the + # whole table. + return ("-", "-", "-", "-", "-") row = next((item for item in rows if item["tokens_per_rank"] == 64), rows[len(rows) // 2]) latency = row["components"]["roundtrip"]["percentiles_us"] - return row["tokens_per_rank"], latency["p50"], latency["p99"] + + def percentile(block: str, name: str) -> float | str: + # Absent on rows measured before the skew diagnostics were emitted. + component = (row.get(block) or {}).get(name) or row.get(block) or {} + return (component.get("percentiles_us") or {}).get("p50", "-") + + return ( + row["tokens_per_rank"], latency["p50"], latency["p99"], + percentile("cross_rank_min_us", "roundtrip"), + percentile("cross_rank_spread_us", ""), + ) def render(documents: list[dict]) -> str: @@ -63,16 +118,18 @@ def render(documents: list[dict]) -> str: ) lines.append("") lines += [ - "| ver | sku | backend | mode | precision | suite | phase | routing | ep | outcome | T* | p50 us | p99 us |", - "|--:|---|---|---|---|---|---|---|--:|---|--:|--:|--:|", + "| ver | sku | backend | mode | precision | suite | phase | routing | ep | topo " + "| wire | outcome | T* | p50 us | p99 us | min50 us | skew us |", + "|--:|---|---|---|---|---|---|---|--:|---|---|---|--:|--:|--:|--:|--:|", ] for document in documents: sku, backend, suite, routing, mode, phase, ep, precision = _identity(document) - token, p50, p99 = _headline(document) + token, p50, p99, min50, skew = _headline(document) + topo, wire = _topology(document), _wire_basis(document) lines.append( f"| {document['version']} | {sku} | `{backend}` | {mode} | {precision} | {suite} | " - f"{phase} | {routing} | {ep} | " - f"{document['outcome']['status']} | {token} | {p50} | {p99} |" + f"{phase} | {routing} | {ep} | {topo} | {wire} | " + f"{document['outcome']['status']} | {token} | {p50} | {p99} | {min50} | {skew} |" ) if not documents: lines.append("\n> No valid native outcome documents found.") diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index 11bab352c..2b6dbc370 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -44,7 +44,9 @@ def _load_config(name: str) -> dict[str, Any]: # 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",), + # FP8 is dispatch-side only here (scales as a fourth payload, combine stays BF16), and + # uses the same per-128-block e4m3 recipe as deepep-v2/uccl-ep so the axis is comparable. + "flashinfer-ep": ("bf16", "fp8"), } # 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_ep_nccl_handle.py b/experimental/CollectiveX/tests/test_ep_nccl_handle.py index 943625ee7..f68bb75a2 100644 --- a/experimental/CollectiveX/tests/test_ep_nccl_handle.py +++ b/experimental/CollectiveX/tests/test_ep_nccl_handle.py @@ -96,6 +96,10 @@ def backend(ll=True): b.args = types.SimpleNamespace(hidden=16) b._t = lambda x: x b._stream = lambda: 0 + # create_buffer always runs before the first _ensure_handle, so the HT receive plane + # exists by then: `_bind_ht_recv_count` slices it to the received-token count. A list + # stands in for the tensor because `_t` is identity here and only the slice is exercised. + b._recv_x = list(range(64)) return b @@ -146,6 +150,44 @@ def test_ll_never_passes_layout_info_on_rebind(self): b._ensure_handle(problem(2)) self.assertEqual([info for _, info in b._ep_group.handle.updates], [None]) + def test_ll_gate_wrapper_is_built_once_per_handle(self): + """LL applies the gate in combine, so its weights wrapper must be cached like the rest. + + Building one per timed combine costs a torch resolve, an np.asarray and a cybind + allocation, and `time_us` charges host work inside the window — a per-call tax no other + backend pays. HT never needs it: FWD forbids input weights on its combine. + """ + ll = backend(ll=True) + pa = problem(1) + h = ll._ensure_handle(pa) + self.assertTrue(hasattr(h, "combine_weights_t")) + self.assertEqual(h.combine_weights_t, "w1") + # Re-entering the SAME problem -- the timed loop's steady state -- reuses the handle and + # therefore the wrapper; a fresh problem object legitimately builds its own. + self.assertIs(ll._ensure_handle(pa).combine_weights_t, h.combine_weights_t) + + ht = backend(ll=False) + self.assertFalse(hasattr(ht._ensure_handle(problem(1)), "combine_weights_t")) + + def test_ht_combine_input_is_sliced_to_the_received_count(self): + """HT combine's staging copy is sized by the tensor it is handed, not by the group. + + Passing the whole ladder-max receive plane made combine copy max(ladder) * world rows + on every call regardless of T, which put a rung-independent floor under it. LL must keep + the full padded plane -- its kernel asserts that shape. + """ + b = backend(ll=False) + h = b._ensure_handle(problem(1)) + # 7 is what the stubbed `torch.zeros(...).item()` reports as the received count. + self.assertEqual(h.count, 7) + self.assertEqual(h.combine_in_t, list(range(7))) + # The point of the fix: the slice, not the whole 64-row plane. + self.assertLess(len(h.combine_in_t), len(b._recv_x)) + + ll = backend(ll=True) + ll_h = ll._ensure_handle(problem(1)) + self.assertFalse(hasattr(ll_h, "combine_in_t")) + def test_ht_rebind_carries_that_problems_counters(self): """HT re-runs the metadata exchange into the rebound problem's own counter tensors.""" b = backend(ll=False) diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index f89035b79..583df4856 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -242,7 +242,11 @@ def test_flashinfer_ep_rollout_shape(self): 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"}) + # FP8 is dispatch-side only: scales ride as a fourth payload (the kernel's kMaxPayloads + # is exactly 4) and combine stays BF16, so none of the 0.6.16+ combine-quant API is + # needed. Same per-128-block e4m3 recipe as deepep-v2/uccl-ep, so the axis stays + # comparable across backends. + self.assertEqual({item["case"]["precision"] for item in cases}, {"bf16", "fp8"}) # 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(): diff --git a/experimental/CollectiveX/tests/test_roundtrip_staging.py b/experimental/CollectiveX/tests/test_roundtrip_staging.py index b6338fc50..5a8071d40 100644 --- a/experimental/CollectiveX/tests/test_roundtrip_staging.py +++ b/experimental/CollectiveX/tests/test_roundtrip_staging.py @@ -1,8 +1,11 @@ #!/usr/bin/env python3 """Contract for what the chained roundtrip measures. -`stage` exists only for FP8 (`stage_device_work = self._fp8`), so charging it to the chained -roundtrip compares FP8 and BF16 through structurally different pipelines. Real stacks decide +`stage` is not an FP8-only cost: deepep-v2, uccl-ep and MoRI all set +`stage_device_work = self._fp8` (MoRI's `self._fp8 or not self._external_input` collapses to that, +now that it always uses an external input buffer), but FlashInfer sets it unconditionally, so its +BF16 rows do real device work there. Charging it to the chained roundtrip therefore made +`roundtrip` mean different things in different rows. Real stacks decide this on quant-format match: SGLang's DeepEP dispatcher contains no dequant at all, and vLLM returns the dispatched fp8 + scales untouched when `block_k == DEEPEP_QUANT_BLOCK_SIZE`, dequantising only as a mismatch fallback. These tests pin both models. @@ -64,8 +67,9 @@ def test_staged_input_keeps_the_conversion_out_of_the_chain(self): self.assertNotIn("stage", b.calls) def test_without_staged_input_the_stage_runs_inline(self): - # The fp8 `dequant` model takes this path, and so does BF16 — free for the adapters - # whose receive buffer is already the combine input, real work for mori/flashinfer-ep. + # The fp8 `dequant` model takes this path, as does any backend whose `stage` is a bare + # pointer assignment. mori/flashinfer-ep no longer do: their real device copy is + # hoisted, so the chained roundtrip is dispatch -> combine for every row. b = _StubBackend(stage_device_work=True, fp8_consume="dequant") b.run_roundtrip(object()) self.assertEqual(b.calls, ["dispatch", "stage", "combine(staged-by-stage)"]) @@ -84,45 +88,140 @@ def test_default_models_the_native_path(self): self.assertEqual(ep_backend.EPBackend.fp8_consume, "native") -class NativeStagingGate(unittest.TestCase): - """`stage_device_work` does NOT imply fp8, so the gate must check precision. +class RoundtripStagingGate(unittest.TestCase): + """`roundtrip` must mean dispatch -> combine in EVERY row, or it is not comparable. - MoRI sets `stage_device_work = self._fp8 or not self._external_input`, so its scale-up - kernels report True for BF16 too, and their stage() does a real copy into the registered - combine-input buffer. Gating on stage_device_work alone silently lifted that copy out of - the BF16 timed region -- a precision-asymmetric change of exactly the kind this file - exists to prevent. + The gate was previously precision-dependent, which left `stage` inside the roundtrip for + exactly two configurations -- MoRI BF16 scale-up and FlashInfer BF16 -- and transport-only + for every other, so the headline compared two different quantities. It is now gated on + `stage_device_work` alone, with `CX_FP8_CONSUME=dequant` as the sole opt-out. """ - def test_bf16_with_device_staging_keeps_the_stage_inline(self): + def test_bf16_with_device_staging_now_lifts_the_copy_out(self): + # MoRI BF16 scale-up and FlashInfer BF16: a real device copy, previously charged to + # the chained roundtrip and to nothing else's. mori_intranode_bf16 = _StubBackend( stage_device_work=True, fp8_consume="native", precision="bf16" ) - self.assertFalse(mori_intranode_bf16.stages_fp8_natively) - mori_intranode_bf16.run_roundtrip(object()) + self.assertTrue(mori_intranode_bf16.stage_excluded_from_roundtrip) + mori_intranode_bf16.run_roundtrip(object(), staged="pre-materialised") self.assertEqual( - mori_intranode_bf16.calls, ["dispatch", "stage", "combine(staged-by-stage)"] + mori_intranode_bf16.calls, ["dispatch", "combine(pre-materialised)"] ) def test_fp8_with_device_staging_lifts_the_conversion_out(self): self.assertTrue( _StubBackend( stage_device_work=True, fp8_consume="native", precision="fp8" - ).stages_fp8_natively + ).stage_excluded_from_roundtrip ) - def test_the_hatch_and_no_op_stages_never_take_the_fast_path(self): - self.assertFalse( # CX_FP8_CONSUME=dequant restores the inline stage + def test_a_no_op_stage_is_never_hoisted(self): + # deepep-v2 / uccl-ep / nccl-ep at BF16: `stage` is a pointer assignment, so there is + # nothing to lift -- and hoisting anyway would hand the low-latency backends a view + # into their double-buffered receive, whose parity flips on each re-dispatch. + self.assertFalse( _StubBackend( - stage_device_work=True, fp8_consume="dequant", precision="fp8" - ).stages_fp8_natively + stage_device_work=False, fp8_consume="native", precision="bf16" + ).stage_excluded_from_roundtrip ) - self.assertFalse( # nccl-ep / bf16 deepep-v2: nothing to lift + self.assertFalse( _StubBackend( stage_device_work=False, fp8_consume="native", precision="fp8" - ).stages_fp8_natively + ).stage_excluded_from_roundtrip ) + def test_the_dequant_hatch_restores_the_inline_stage(self): + # CX_FP8_CONSUME=dequant models a stack that really does convert between the two + # collectives, so that run wants the stage back inside the chain. + backend = _StubBackend( + stage_device_work=True, fp8_consume="dequant", precision="fp8" + ) + self.assertFalse(backend.stage_excluded_from_roundtrip) + backend.run_roundtrip(object()) + self.assertEqual(backend.calls, ["dispatch", "stage", "combine(staged-by-stage)"]) + + def test_the_hatch_does_not_apply_to_bf16(self): + # The hatch is about fp8 consumption; a BF16 row has no conversion to model. + self.assertTrue( + _StubBackend( + stage_device_work=True, fp8_consume="dequant", precision="bf16" + ).stage_excluded_from_roundtrip + ) + +class WarmStaging(unittest.TestCase): + """Warm-up must not rehearse work the timed region skips. + + Where staging is excluded from the chain, the timed roundtrip stages nothing, so staging on + every warm iteration warms a path the measurement never takes. For an FP8 dequant that was + the largest single cost in the leg (~247us x 32 iterations x every component x every trial). + `benchmark_stage` is the exception: staging is its timed operation. + """ + + @staticmethod + def _warm(backend, count, **kwargs): + # `warm` imports torch for one synchronize; a stub keeps this runnable without a GPU. + fake = types.ModuleType("torch") + fake.cuda = types.SimpleNamespace(synchronize=lambda: None) + saved = sys.modules.get("torch") + sys.modules["torch"] = fake + try: + backend.warm(types.SimpleNamespace(), count, **kwargs) + finally: + if saved is None: + del sys.modules["torch"] + else: + sys.modules["torch"] = saved + + def test_stages_once_when_the_chain_excludes_staging(self): + b = _StubBackend(stage_device_work=True, fp8_consume="native") + self._warm(b, 5) + self.assertEqual(b.calls.count("dispatch"), 5) + self.assertEqual(b.calls.count("stage"), 1) + # Every later iteration still hands combine the staged payload, not a stale None. + self.assertEqual(b.calls.count("combine(staged-by-stage)"), 5) + + def test_stage_every_rehearses_it_on_every_iteration(self): + b = _StubBackend(stage_device_work=True, fp8_consume="native") + self._warm(b, 5, stage_every=True) + self.assertEqual(b.calls.count("stage"), 5) + + def test_a_chain_that_includes_staging_keeps_warming_it(self): + # The `dequant` hatch puts the conversion back in the timed chain, so warm-up must match. + b = _StubBackend(stage_device_work=True, fp8_consume="dequant") + self._warm(b, 5) + self.assertEqual(b.calls.count("stage"), 5) + +class SteadyStatePeriod(unittest.TestCase): + """`period` is opt-in, because the overlap it measures is only sound for some backends. + + A decode loop never stops between layers, so its per-layer cost is the pipeline's period, + not the sum of separately-drained stages. Measuring that means issuing pairs back-to-back, + which lets ranks drift — and dispatch is a peer WRITE into another rank's buffer, so stream + order on the receiver does not order the sender's remote writes. A double-buffered receive + covers the ~one iteration of drift a collective permits; a single shared buffer does not. + Defaulting this on would produce a fast number over corrupted data. + """ + + def test_off_by_default_so_no_backend_pipelines_accidentally(self): + b = _StubBackend(stage_device_work=False, fp8_consume="native") + self.assertEqual(b.pipeline_pairs, 0) + self.assertNotIn("period", b.timed_components()) + + def test_declaring_pairs_adds_the_component(self): + b = _StubBackend(stage_device_work=False, fp8_consume="native") + b.pipeline_pairs = 8 + self.assertIn("period", b.timed_components()) + # and it never displaces the drained latency measurement + self.assertIn("roundtrip", b.timed_components()) + + def test_a_single_pair_is_not_a_pipeline(self): + # pipeline_pairs = 1 measures exactly what roundtrip already does, so it must not + # advertise a second name for the same quantity. + b = _StubBackend(stage_device_work=False, fp8_consume="native") + b.pipeline_pairs = 1 + self.assertNotIn("period", b.timed_components()) + if __name__ == "__main__": unittest.main() diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index 01f377b7a..19f9ef8c9 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -27,6 +27,7 @@ import config # noqa: E402 import stage # noqa: E402 import ep_harness # noqa: E402 (stdlib-only at module top) +import ep_backend # noqa: E402 (torch is imported lazily inside its methods) # configs/platform_config.json is shared by matrix scheduling, operator/network @@ -379,7 +380,7 @@ def test_case_args_round_trips_through_the_run_ep_parser(self) -> None: self.assertEqual(args.version, 1) self.assertEqual(args.seed, self.CASE["seed"]) self.assertEqual((args.iters, args.trials, args.warmup), (8, 256, 32)) - self.assertEqual(args.out, "results/h200-dgxc_deepep-v2_bf16_decode_TS-c000.json") + self.assertEqual(args.out, "results/h200-dgxc_deepep-v2_bf16_normal_decode_TS-c000.json") def test_case_args_fails_closed_on_placement_mismatch(self) -> None: with self.assertRaises(subprocess.CalledProcessError): @@ -401,6 +402,15 @@ def test_low_latency_case_round_trips_through_the_run_ep_parser(self) -> None: args = self._run_ep_parser().parse_args(argv) self.assertEqual((args.mode, args.phase, args.scope), ("low-latency", "decode", "scale-up")) self.assertEqual(args.case_id, ll_case["case_id"]) + # The filename must carry the mode. Without it this case and its normal-mode sibling + # produce byte-identical paths (same runner, backend, precision, phase and index), so + # driving both under one timestamp silently overwrites one artifact with the other -- + # which cost two on-metal runs before it was noticed. CI hides it by giving every leg + # its own job and therefore its own ts. + # ... which is what test_case_args_round_trips pins as `..._bf16_normal_decode_...`. + self.assertEqual( + args.out, "results/h200-dgxc_deepep-v2_bf16_low-latency_decode_TS-c000.json" + ) def test_uccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: # A uccl-ep case flows through the same generic codec; run_ep's --backend choices @@ -415,7 +425,7 @@ def test_uccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: args = self._run_ep_parser().parse_args(argv) self.assertEqual(args.backend, "uccl-ep") self.assertEqual(args.case_id, uccl_case["case_id"]) - self.assertEqual(args.out, "results/h200-dgxc_uccl-ep_bf16_decode_TS-c000.json") + self.assertEqual(args.out, "results/h200-dgxc_uccl-ep_bf16_normal_decode_TS-c000.json") def test_nccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: # A nccl-ep case flows through the same generic codec; run_ep's --backend choices must @@ -430,7 +440,7 @@ def test_nccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: args = self._run_ep_parser().parse_args(argv) self.assertEqual(args.backend, "nccl-ep") self.assertEqual(args.case_id, nccl_case["case_id"]) - self.assertEqual(args.out, "results/h200-dgxc_nccl-ep_bf16_decode_TS-c000.json") + self.assertEqual(args.out, "results/h200-dgxc_nccl-ep_bf16_normal_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 @@ -449,7 +459,7 @@ def test_flashinfer_ep_case_round_trips_through_the_run_ep_parser(self) -> None: 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" + args.out, "results/h200-dgxc_flashinfer-ep_bf16_normal_decode_TS-c000.json" ) def test_mirrored_backend_choices_match_run_ep(self) -> None: @@ -654,5 +664,289 @@ def test_a_rank_claimed_by_an_earlier_slot_contributes_once(self): ) self.assertEqual(combined.item(), 0.5) + def test_the_first_slot_claiming_a_rank_is_the_one_that_survives(self): + torch = _torch + # Which duplicate the kernel blanks is invisible when the repeated payload is the only + # value in play, so the case above cannot tell keep-first from keep-last. These + # contributions cancel instead: blanking the LATER slot reduces (256+1)->256 and + # (-256+0)->-256 to 0.0, while blanking the EARLIER one gives (0+1)->1 and + # (-256+256)->0, i.e. 1.0. The 257 -> 256 step is the BF16 rounding that makes the two + # models disagree at all. + destination = torch.tensor([[0, 1, 2, 0]]) + messages = torch.tensor([[[256.0]], [[1.0]], [[-256.0]]]) + combined = ep_harness._topk_slot_tree_combine( + torch, destination, torch.ones_like(destination, dtype=torch.bool), + messages, torch.bfloat16, + ) + self.assertEqual(combined.item(), 0.0) + + + +@unittest.skipUnless(_torch is not None, "quantize-identity checks require torch") +class FusedQuantizeGate(unittest.TestCase): + """The fp8 quantize moved inside the timed dispatch, so its identity is load-bearing. + + The oracle's payload gate is a `torch.equal` that compares the SENDER's [T, hidden] quantize + against the ORACLE's [receive_count, hidden] one, so the callable must be bit-identical to + the eager helper AND per-row invariant across batch sizes. Verified on-metal for e4m3fn + (H100: 9 kernels/19.2us eager vs 1/1.51us fused) and e4m3fnuz (MI300X, the arch with no + prior precedent for a compiled fp32->fp8 convert). These tests pin the guard, not the + compiler. + """ + + @staticmethod + def _fuse(mode, eager): + # Both methods read only `self.mode`, so call them unbound rather than instantiating an + # abstract backend; that also documents that they depend on nothing else. + return ep_backend.EPBackend.fused_quantize(types.SimpleNamespace(mode=mode), eager) + + @staticmethod + def _check(eager, fused, x): + return ep_backend.EPBackend.assert_quantize_identity( + types.SimpleNamespace(mode="normal"), eager, fused, x + ) + + def test_low_latency_keeps_the_eager_helper(self): + # Its dispatch kernel quantises internally and the oracle gate is pinned to those bits, + # so swapping in a compiled callable would red every LL fp8 cell without touching timing. + def eager(x): + return x, x + self.assertIs(self._fuse("low-latency", eager), eager) + + def test_normal_mode_wraps_the_helper(self): + def eager(x): + return x, x + self.assertIsNot(self._fuse("normal", eager), eager) + + def test_identity_check_accepts_an_equivalent_callable(self): + torch = _torch + x = torch.randn(8, 256, dtype=torch.bfloat16) + + def eager(t): + return t.to(torch.float8_e4m3fn), t.float().abs().amax(dim=1) + + self._check(eager, lambda t: eager(t), x) # must not raise + + def test_identity_check_rejects_a_divergent_callable(self): + torch = _torch + x = torch.randn(8, 256, dtype=torch.bfloat16) + + def eager(t): + return t.to(torch.float8_e4m3fn), t.float().abs().amax(dim=1) + + def divergent(t): + values, scales = eager(t) + return values, scales + 1 # one differing scale is enough to red a cell + with self.assertRaises(RuntimeError): + self._check(eager, divergent, x) + + def test_identity_check_rejects_a_shape_dependent_callable(self): + # A callable that is deterministic but NOT per-row invariant still breaks the gate, + # because the oracle quantises a different row count than the sender did. + torch = _torch + x = torch.randn(8, 256, dtype=torch.bfloat16) + + def eager(t): + return t.to(torch.float8_e4m3fn), t.float().abs().amax(dim=1) + + def shape_dependent(t): + values, scales = eager(t) + return values, scales * float(t.shape[0]) + with self.assertRaises(RuntimeError): + self._check(eager, shape_dependent, x) + +class GpuHealthProbe(unittest.TestCase): + """Reject an allocation holding a throttled GPU before it burns the wall-clock guard. + + Collectives are barriers, so one clamped device paces every rank. A B200 with GPU 7 held at + 120 MHz against 1965 MHz on its siblings ran a case 17x slower and was killed twice, at 900s + and again at 1800s, looking exactly like a code pathology. The flag is the signal, not the + clock: the allocation is idle when this runs and an idle B200 also reads 120 MHz. + """ + + HEALTHY = "\n".join(f"{i}, Not Active, Not Active, 3{i} " for i in range(8)) + + def _swap(self, line_in: str, line_out: str) -> str: + self.assertIn(line_in, self.HEALTHY) # guard the fixture against silent drift + return self.HEALTHY.replace(line_in, line_out) + + def test_healthy_allocation_passes(self): + self.assertEqual(probe.gpu_health_faults(self.HEALTHY), []) + + def test_a_thermally_throttled_gpu_is_rejected(self): + output = self._swap("7, Not Active, Not Active, 37 ", "7, Active, Active, 93 ") + faults = probe.gpu_health_faults(output) + self.assertEqual(len(faults), 1) + self.assertIn("gpu 7", faults[0]) + + def test_either_throttle_flag_alone_is_enough(self): + for cells in ("7, Active, Not Active, 88 ", "7, Not Active, Active, 88 "): + with self.subTest(cells=cells): + output = self._swap("7, Not Active, Not Active, 37 ", cells) + self.assertEqual(len(probe.gpu_health_faults(output)), 1) + + def test_not_active_is_not_read_as_active(self): + # A substring search for "Active" matches "Not Active" and passes every fault through, + # which is the whole failure mode this probe exists to avoid. + self.assertEqual(probe.gpu_health_faults(self.HEALTHY), []) + + def test_temperature_is_an_independent_signal(self): + # The flag can clear between samples while the fault persists, so heat alone rejects. + output = self._swap("3, Not Active, Not Active, 33 ", "3, Not Active, Not Active, 95 ") + faults = probe.gpu_health_faults(output) + self.assertEqual(len(faults), 1) + self.assertIn("gpu 3", faults[0]) + + def test_unreadable_output_fails_open(self): + # Blocking legs when the hardware cannot be read is worse than the fault being looked for. + for output in ("", "nonsense\n", "1, Not Active\n", self.HEALTHY.replace("32 ", "[N/A] ")): + with self.subTest(output=output[:20]): + self.assertEqual(probe.gpu_health_faults(output), []) + + def _run_validate(self, csv: str, has_smi: bool = True): + """Drive validate_gpu_health with a stubbed nvidia-smi; returns (exit_code, stdout).""" + import shutil + real_which = shutil.which + shutil.which = (lambda name: "/usr/bin/nvidia-smi") if has_smi else (lambda name: None) + + class FakeSubprocess: + SubprocessError = subprocess.SubprocessError + + @staticmethod + def run(*args, **kwargs): + return types.SimpleNamespace(stdout=csv) + + sys.modules["subprocess"] = FakeSubprocess + captured = io.StringIO() + try: + with contextlib.redirect_stdout(captured): + probe.validate_gpu_health() + code = 0 + except SystemExit as exit_: + code = exit_.code + finally: + sys.modules["subprocess"] = subprocess + shutil.which = real_which + return code, captured.getvalue() + + def test_a_healthy_check_records_how_many_gpus_it_saw(self): + # Without this marker a gate that went BLIND -- no visible devices, or a driver spelling + # the fields clocks_throttle_reasons.* -- writes an empty log and is indistinguishable + # from one that inspected eight healthy GPUs. + code, out = self._run_validate(self.HEALTHY) + self.assertEqual(code, 0) + self.assertIn("gpu-health-checked gpus=8", out) + + def test_a_fault_exits_nonzero_and_names_the_gpu(self): + code, out = self._run_validate( + self._swap("7, Not Active, Not Active, 37 ", "7, Active, Active, 93 ") + ) + self.assertEqual(code, 1) + self.assertIn("gpu-health-fault gpu 7", out) + self.assertNotIn("gpu-health-checked", out) + + def test_a_missing_nvidia_smi_is_silent_and_passes(self): + code, out = self._run_validate(self.HEALTHY, has_smi=False) + self.assertEqual(code, 0) + self.assertEqual(out, "") + def test_the_temperature_spread_is_reported_for_the_signal_no_gate_can_see(self): + # An H100 engages software thermal slowdown at ~86-87 C, so a clamped one never crosses the + # 90 C limit; and at pre-flight it is idle, so the flag is clear too. The measured fault was + # visible ONLY as an idle outlier (55 C against ~30 C siblings), so the spread is recorded + # every run to build the evidence a relative gate would need. + sick = self._swap("3, Not Active, Not Active, 33 ", "3, Not Active, Not Active, 55 ") + self.assertEqual(probe.gpu_temperature_spread(sick), (55, 35, 20)) + hottest, median, spread = probe.gpu_temperature_spread(self.HEALTHY) + self.assertEqual((hottest, median), (37, 34)) # 8 temps -> median is index 4 + self.assertLess(spread, 10) + + def test_the_spread_appears_in_the_healthy_marker(self): + sick = self._swap("3, Not Active, Not Active, 33 ", "3, Not Active, Not Active, 55 ") + code, out = self._run_validate(sick) + self.assertEqual(code, 0) # reported, deliberately NOT gated on + self.assertIn("spread=20C", out) + + def test_the_spread_is_none_when_unreadable(self): + for output in ("", "nonsense\n", self.HEALTHY.replace("33 ", "[N/A] ")): + with self.subTest(output=output[:16]): + result = probe.gpu_temperature_spread(output) + self.assertTrue(result is None or result[2] < 10) + + +class LowLatencyCapDecoupling(unittest.TestCase): + """The LL receive size and the measured ladder are two numbers and must stay two. + + The measured ladder stops below the receive cap to skip a token count DeepEP's low-latency + combine corrupts on Blackwell (upstream #700; the fix is #642, which our pin predates). The + receive must NOT follow the ladder down: its footprint sets the transport's memory traffic + and the fp8 dequant volume, so sizing it from `max(ladder)` would shift every retained + rung and break comparability with the published series. Asserted over the source because + importing the adapter needs a built deep_ep. + """ + + @classmethod + def setUpClass(cls): + cls.tree = ast.parse((BENCH / "ep_deepep_v2.py").read_text()) + cls.consts = { + t.id: node.value.value + for node in cls.tree.body + if isinstance(node, ast.Assign) + for t in node.targets + if isinstance(t, ast.Name) and isinstance(node.value, ast.Constant) + } + + def _func(self, name): + for node in ast.walk(self.tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + self.fail(f"{name} not found in ep_deepep_v2.py") + + def test_both_caps_exist_and_the_ladder_stops_strictly_below_the_buffer(self): + buf, ladder = self.consts.get("_LL_BUFFER_CAP"), self.consts.get("_LL_LADDER_CAP") + self.assertIsInstance(buf, int) + self.assertIsInstance(ladder, int) + # Strictly below: equality is the configuration that corrupts, and it also puts the top + # measured rung at 100% occupancy, which is what made capacity and last-rung + # indistinguishable in the original investigation. + self.assertLess(ladder, buf) + # NVSHMEM_QP_DEPTH=1024 asserts nvshmem_qp_depth >= (cap + 1) * 2 at construction. + self.assertLessEqual(buf, 511) + + def test_buffer_cap_clamps_the_ladder_by_the_constant_not_a_literal(self): + returns = [ + n.value for n in ast.walk(self._func("buffer_cap")) + if isinstance(n, ast.Return) and n.value is not None + ] + names = {n.id for n in returns if isinstance(n, ast.Name)} + self.assertIn("_LL_LADDER_CAP", names) + # A bare literal here would drift out of step with the constants above. + self.assertEqual([n for n in returns if isinstance(n, ast.Constant) and n.value is not None], []) + + def test_the_low_latency_receive_is_sized_from_the_cap_not_the_ladder(self): + # Guards the regression that would silently re-baseline every LL row. + assigned = [ + node.value for node in ast.walk(self._func("create_buffer")) + if isinstance(node, ast.Assign) + for t in node.targets + if isinstance(t, ast.Attribute) and t.attr == "max_tokens" + ] + self.assertTrue( + any(isinstance(v, ast.Name) and v.id == "_LL_BUFFER_CAP" for v in assigned), + "create_buffer must set self.max_tokens = _LL_BUFFER_CAP on the low-latency path", + ) + + def test_a_clamped_ladder_is_recorded_in_the_artifact_not_only_on_stdout(self): + # A clamped ladder was previously visible only as a rank-0 stdout NOTE, so an artifact + # from a backend that measured 8 rungs was indistinguishable from one that measured 9. + # The emitted record must carry what ran and what was excluded. + harness = ast.parse((BENCH / "ep_harness.py").read_text()) + keys = { + k.value for node in ast.walk(harness) if isinstance(node, ast.Dict) + for k in node.keys if isinstance(k, ast.Constant) and isinstance(k.value, str) + } + for required in ("ladder_measured", "ladder_dropped", "ladder_cap"): + self.assertIn(required, keys, f"the emitted record must include {required}") + + if __name__ == "__main__": unittest.main()