From e8fc87acc49b1397f809ec16a296009f8871f7c9 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:35:18 +0800 Subject: [PATCH 01/20] CollectiveX: graph-replay every capturable EP backend, keeping rank alignment and the chained period Graph mode now keeps the measurement contract eager mode has: - each timed replay starts behind a device-side rank barrier, with the timing events as graph nodes, so host launch skew (~75us across b200 EP16 nodes) no longer lands in the cross-rank MAX - the chained family is captured as one graph of unrolled pairs per sibling, so pair_period, chain floors, chain health and the chained oracle are published again under replay - graphed rows carry a -cudagraph kernel_generation suffix so the store never pools them with eager samples Graphed backends: MoRI (both modes), uccl-ep (low-latency), and deepep-v2 normal-mode decode, which now runs ElasticBuffer as vLLM's graphed deepep_v2 decode does (do_cpu_sync=False, worst-case receive; kernel generation v2-elastic-buffer-nosync). deepep-v2 prefill and uccl-ep normal stay eager. --- experimental/CollectiveX/bench/ep_backend.py | 144 +++++++++++++----- .../CollectiveX/bench/ep_deepep_v2.py | 23 ++- experimental/CollectiveX/bench/ep_harness.py | 90 ++++++----- experimental/CollectiveX/bench/ep_mori.py | 5 + experimental/CollectiveX/bench/ep_uccl.py | 5 + experimental/CollectiveX/docs/methodology.md | 54 ++++--- .../CollectiveX/tests/test_backends.py | 53 +++++++ experimental/CollectiveX/tests/test_chain.py | 53 ++++++- 8 files changed, 325 insertions(+), 102 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index e1e28ee6c..ed0aa2393 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -173,7 +173,7 @@ def stage_excluded_from_roundtrip(self) -> bool: @property def cuda_graph_supported(self) -> bool: """Whether this realized backend/mode has a graph-safe fixed-shape roundtrip.""" - return self.mode in self.CUDA_GRAPH_MODES + return getattr(self, "mode", None) in self.CUDA_GRAPH_MODES @property def cuda_graph_enabled(self) -> bool: @@ -393,6 +393,65 @@ def _topk_idx_dtype(self): import torch return torch.int64 + # ---- CUDA graph capture ---------------------------------------------------------------- + + # Spin after the alignment all-reduce so every rank's host has enqueued its replay before the + # stream reaches it; ~50us at 2GHz, far above a graph launch, so the replay start is set by + # the barrier release on every rank rather than by host launch latency. + _GRAPH_ALIGN_SPIN_CYCLES = 100_000 + + def _graph_align(self): + """Enqueue a device-side rank barrier on the current stream, without a host sync.""" + import torch + import torch.distributed as dist + + token = getattr(self, "_graph_align_token", None) + if token is None: + token = self._graph_align_token = torch.zeros(1, device=self.device) + dist.all_reduce(token) + torch.cuda._sleep(self._GRAPH_ALIGN_SPIN_CYCLES) + + def _capture_pairs(self, problem, staged, pairs, marks): + """Capture `pairs` back-to-back dispatch -> combine pairs into one graph. + + `marks` selects which windows get external event nodes: "pair" (the whole pair), + "dispatch", "combine". Event records are graph nodes, so they cost the stream nothing on + the host -- the six-events-per-pair defect the eager chain splits around does not exist + here. Returns (graph, {mark: (starts, ends)}, last combined output). + """ + import torch + import torch.distributed as dist + + def events(): + return [torch.cuda.Event(enable_timing=True, external=True) for _ in range(pairs)] + + stamps = {mark: (events(), events()) for mark in marks} + + def record(mark, edge, i): + if mark in stamps: + stamps[mark][edge][i].record() + + dist.barrier() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + combined = None + with torch.cuda.graph(graph, capture_error_mode="relaxed"): + for i in range(pairs): + record("pair", 0, i) + record("dispatch", 0, i) + handle = self.dispatch(problem) + record("dispatch", 1, i) + if staged is None: + self.stage(problem, handle) + else: + handle.combine_input = staged + record("combine", 0, i) + combined = self.combine(problem, handle) + record("combine", 1, i) + record("pair", 1, i) + torch.cuda.synchronize() + return graph, stamps, combined + # ---- Timing template methods ----------------------------------------------------- def timed_components(self): @@ -494,6 +553,8 @@ def benchmark_chain(self, problem, warmup, iters, drop): staged = handle.combine_input self.combine(problem, handle) # drain the pair backends require torch.cuda.synchronize() + if self.cuda_graph_enabled: + return self._benchmark_chain_graph(problem, staged, iters, drop) # Events are allocated BEFORE the loops: an allocation between two record() calls is host # work inside a window meant to belong to the stream, a measurable fraction of the period # at the bottom of the ladder. @@ -547,6 +608,43 @@ def series(starts, ends): "combined": combined.clone(), } + def _benchmark_chain_graph(self, problem, staged, iters, drop): + """The chained family under capture: each chain is ONE graph of `iters` unrolled pairs. + + That is the shape a serving decode graph has -- every layer's dispatch -> combine back to + back inside a single replay -- so the period keeps its eager meaning (free-running pairs, + entry skew amortised across the chain) with launch overhead removed. Same two siblings as + the eager chain and the same returned series, so `run_sweep` reduces both identically. + Each graph replays once untimed (first-launch upload), then once aligned and timed. + """ + import torch + + floors, floor_stamps, _ = self._capture_pairs( + problem, staged, iters, ("dispatch", "combine") + ) + period, period_stamps, combined = self._capture_pairs(problem, staged, iters, ("pair",)) + for graph in (floors, period): + graph.replay() + torch.cuda.synchronize() + self._graph_align() + graph.replay() + torch.cuda.synchronize() + + def series(starts, ends): + return [ + start.elapsed_time(end) * 1000.0 # ms -> us + for start, end in zip(starts[drop:], ends[drop:]) + ] + + pair_start, pair_end = period_stamps["pair"] + return { + "pair": series(pair_start, pair_end), + "start_to_start": series(pair_start[:-1], pair_start[1:]), + "dispatch": series(*floor_stamps["dispatch"]), + "combine": series(*floor_stamps["combine"]), + "combined": combined.clone(), + } + def benchmark_component(self, component, problem, warmup, iters): """Measure one named component; every component gets the same warm-up first.""" if self.cuda_graph_enabled: @@ -582,42 +680,16 @@ def benchmark_roundtrip(self, problem, warmup, iters, graph_component="roundtrip self.combine(problem, handle) # drain the pair backends require torch.cuda.synchronize() if self.cuda_graph_enabled: - # Capture replaces the existing roundtrip callable in place. Capture and its warmup - # are excluded; the ordinary time_us event pipeline measures replay directly. - import torch.distributed as dist - - dist.barrier() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - interval = ( - ( - torch.cuda.Event(enable_timing=True, external=True), - torch.cuda.Event(enable_timing=True, external=True), - ) - if graph_component != "roundtrip" else None + # One captured pair, bracketed by external events for the timed component. Capture + # and its warm-up are excluded; each timed replay starts behind a device-side rank + # barrier (`_graph_align`) so the cross-rank MAX is the operation, not launch skew. + mark = "pair" if graph_component == "roundtrip" else graph_component + graph, stamps, combined = self._capture_pairs(problem, staged, 1, (mark,)) + starts, ends = stamps[mark] + samples = time_cuda_graph_phase_us( + torch, graph.replay, warmup, iters, (starts[0], ends[0]), + align=self._graph_align, ) - with torch.cuda.graph(graph, capture_error_mode="relaxed"): - if graph_component == "dispatch": - interval[0].record() - handle = self.dispatch(problem) - if graph_component == "dispatch": - interval[1].record() - if staged is None: - self.stage(problem, handle) - else: - handle.combine_input = staged - if graph_component == "combine": - interval[0].record() - combined = self.combine(problem, handle) - if graph_component == "combine": - interval[1].record() - torch.cuda.synchronize() - if interval is None: - samples = time_us(torch, graph.replay, warmup, iters) - else: - samples = time_cuda_graph_phase_us( - torch, graph.replay, warmup, iters, interval - ) # Prove replay, rather than capture, writes the output used by the correctness gate. combined.fill_(float("nan")) diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py index 179c7baed..de3564211 100644 --- a/experimental/CollectiveX/bench/ep_deepep_v2.py +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -137,8 +137,9 @@ class DeepEPV2Backend(EPBackend): kernel_generation = "v2-elastic-buffer" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # ElasticBuffer normal dispatch performs a host synchronization; the legacy decode kernels - # are explicitly graph compatible. + # The legacy decode kernels are explicitly graph compatible. ElasticBuffer normal mode is + # graph compatible only without its host sync, which this adapter drops for the decode phase + # alone (see `_normal_cpu_sync` and `cuda_graph_supported`). CUDA_GRAPH_MODES = ("low-latency",) stage_device_work = False requires_fresh_pair = False @@ -168,6 +169,16 @@ def __init__(self, args, rank, world_size, local_rank, device): # Normal/HT quantises inside the timed dispatch with the compiled form; low-latency # keeps the eager helper, whose bits its in-kernel quantise matches. See fused_quantize. self._quant = self.fused_quantize(self._to_fp8) + # Normal-mode decode runs ElasticBuffer the way vLLM's deepep_v2 decode path does + # (prepare_finalize/deepep_v2.py, use_cudagraph=True): do_expand=False, do_cpu_sync=False, + # receive sized to the worst case (num_max_tokens_per_rank * num_ranks) with the valid + # prefix read from the handle's device-side psum. That is the graph-capturable contract, + # and it lands one row per (token, destination rank) with an unweighted rank-sum combine: + # the rank-major shape. Prefill keeps the host sync that sizes the receive exactly, as + # vLLM's (uncaptured) prefill does. + self._normal_cpu_sync = self.mode == "normal" and args.phase != "decode" + if self.mode == "normal" and not self._normal_cpu_sync: + self.kernel_generation = "v2-elastic-buffer-nosync" 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. @@ -180,6 +191,12 @@ def __init__(self, args, rank, world_size, local_rank, device): # dispatch must be drained by its combine. self.requires_fresh_pair = True + @property + def cuda_graph_supported(self) -> bool: + if self.mode == "normal": + return not getattr(self, "_normal_cpu_sync", True) + return super().cuda_graph_supported + def buffer_cap(self, args): if self.mode == "low-latency": # LL pre-allocates a fixed [num_local_experts, cap * num_ranks, hidden] receive @@ -377,7 +394,7 @@ def dispatch(self, p): num_qps=self.num_qps, async_with_compute_stream=False, do_handle_copy=True, - do_cpu_sync=True, + do_cpu_sync=self._normal_cpu_sync, do_expand=False, ) return types.SimpleNamespace( diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 4c688ac55..1757109e8 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -318,14 +318,24 @@ def sample(): def time_cuda_graph_phase_us( - torch, fn, warmup: int, iters: int, interval + torch, fn, warmup: int, iters: int, interval, align=None ) -> list[float]: - """Time one event-record interval captured inside graph replay.""" + """Time one event-record interval captured inside graph replay. + + `interval` is a pair of external events recorded as graph nodes, so the host's replay launch + never lands in the window. `align()` runs before each replay with no host sync between the + two: it enqueues a device-side rank barrier, so every rank's replay starts when that barrier + releases instead of when its own host got round to launching the graph. Without it each + sample restarts from the preceding synchronize and ranks enter ~75us apart across nodes + (b200 EP16), which the cross-rank MAX then reports as latency. + """ for _ in range(max(0, warmup)): fn() torch.cuda.synchronize() samples = [] for _ in range(iters): + if align is not None: + align() fn() torch.cuda.synchronize() samples.append(interval[0].elapsed_time(interval[1]) * 1000.0) @@ -333,8 +343,16 @@ def time_cuda_graph_phase_us( def kernel_generation(backend) -> str: - """Return the adapter's declared kernel family.""" - return getattr(backend, "kernel_generation", None) or "n-a" + """Return the adapter's declared kernel family, suffixed when timed under graph replay. + + Replay removes launch overhead that eager timing pays, so the two regimes are different + series: the suffix keeps the durable store from pooling a graphed row with the eager rows + published under the same kernel family. + """ + family = getattr(backend, "kernel_generation", None) or "n-a" + if getattr(backend, "cuda_graph_enabled", False): + return f"{family}-cudagraph" + return family def _reduce_vec(torch, dist, device, vals, op): @@ -1027,8 +1045,8 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # `chain_health` as "unavailable", indistinguishable from a backend that cannot be chained. # Requiring two kept pairs here is what lets Pass 2b compute the health scalars # unconditionally and Pass 3 assert the chained oracle ran. - if (not cuda_graph and (min(args.chain_iters, args.chain_trials) <= 0 - or not 0 <= args.chain_drop <= args.chain_iters - 2)): + if (min(args.chain_iters, args.chain_trials) <= 0 + or not 0 <= args.chain_drop <= args.chain_iters - 2): if rank == 0: print(f"ERROR: chain iters/trials must be positive and 0 <= drop <= iters - 2; got " f"{args.chain_iters}:{args.chain_trials}:{args.chain_drop}") @@ -1164,7 +1182,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # (every FP8 adapter by default, since stage_device_work IS the fp8 flag) the staged # stand-in is decoupled from each pair's dispatch, so chained and drained are not # comparable -- see the call site for the measurement that established this. - chain_output_applicable = not cuda_graph and not backend.stage_excluded_from_roundtrip + chain_output_applicable = not backend.stage_excluded_from_roundtrip cuda_graph_output_applicable = cuda_graph and not backend.stage_excluded_from_roundtrip # ---- Pass 2: every backend uses the same rotated point order. @@ -1235,7 +1253,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # already yields chain_iters free-running pairs, so a handful of trials out-samples the # fresh-entry components' 256 for a fraction of the wall clock. Ladder order still rotates # per trial, as above. ---- - for trial_index in range(0 if cuda_graph else args.chain_trials): + for trial_index in range(args.chain_trials): final_chain_trial = trial_index == args.chain_trials - 1 for T in trial_order(list(ladder), trial_index): chained = backend.benchmark_chain( @@ -1324,14 +1342,12 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> ) pre = gate[T]["oracle_pre"] chain_oracle = gate[T]["oracle_chain"] - if cuda_graph: - chain_ok = True - chain_max_rel = 0.0 - else: - # The eager chained oracle is required whenever that pipeline was measured. - assert chain_oracle is not None, "chained oracle missing despite a validated budget" - chain_ok = bool(chain_oracle["passed"]) - chain_max_rel = chain_oracle["max_elementwise_relative_error"] or 0.0 + # The chained ORACLE is ANDed in like the other two, so a chained-regime failure reds the + # leg. The budget gate rejects chain_trials=0 up front, so a missing chained oracle is a + # harness bug, not a configuration. + assert chain_oracle is not None, "chained oracle missing despite a validated budget" + chain_ok = bool(chain_oracle["passed"]) + chain_max_rel = chain_oracle["max_elementwise_relative_error"] or 0.0 # The chained-OUTPUT check gates again, on a measured magnitude rather than a verdict. # It was briefly demoted on the theory its tolerance was too tight for FP8; probe # 31180411148 (h100, deepep-v2, EP8, low-latency) falsified that: @@ -1385,23 +1401,24 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> recv_max = _reduce_int(torch, dist, device, g["recv_local"], MAX) recv_min = _reduce_int(torch, dist, device, g["recv_local"], MIN) global_ok = _reduce_int(torch, dist, device, g["local_ok"], MIN) - if cuda_graph: - post_chain_state_passed = None - chain_last_output_passed = None - chain_output_error = None - else: - # Agreed across ranks like `passed`, not rank 0's local view. - post_chain_state_passed = bool( - _reduce_int(torch, dist, device, g["chain_local_ok"], MIN) - ) - chain_last_output_passed = bool( - _reduce_int(torch, dist, device, g["chain_output_local_ok"], MIN) - ) - chain_output_error = _reduce_vec( - torch, dist, device, [g["chain_output_error"]], MAX - )[0] - if not chain_output_applicable: - chain_last_output_passed, chain_output_error = None, None + # Agreed across ranks like `passed`, not rank 0's local view. + post_chain_state_passed = bool( + _reduce_int(torch, dist, device, g["chain_local_ok"], MIN) + ) + # null where the check does not apply (staging hoisted): the artifact says "not + # asked", never a bare False that a reader would mistake for a failed comparison. + # The reduce still runs on every rank so the collective stays aligned. + chain_last_output_passed = bool( + _reduce_int(torch, dist, device, g["chain_output_local_ok"], MIN) + ) + # Published whether or not the verdict passed. Without it the artifact records THAT the + # chained output differed but never BY HOW MUCH, which is the difference between a + # transport corruption and a tolerance set too tight for a backend's accumulator. + chain_output_error = _reduce_vec( + torch, dist, device, [g["chain_output_error"]], MAX + )[0] + if not chain_output_applicable: + chain_last_output_passed, chain_output_error = None, None if cuda_graph: cuda_graph_output_rewritten = bool( _reduce_int(torch, dist, device, g["cuda_graph_output_rewritten"], MIN) @@ -1706,9 +1723,10 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> "stage_excluded_from_roundtrip": bool( getattr(backend, "stage_excluded_from_roundtrip", False) ), - # Graph mode replaces the eager component/chain pipeline in place. Existing component - # fields contain replay samples; no parallel graph component exists. - "chained_period": not cuda_graph, + # Whether this document's rows carry the chained family. Consumers key the headline on + # presence, as for `stage_excluded_from_roundtrip`. Graph mode keeps it: the chain is + # captured as one graph of unrolled pairs (EPBackend._benchmark_chain_graph). + "chained_period": True, "cuda_graph_replay": cuda_graph, "cuda_graph_supported": bool( getattr(backend, "cuda_graph_supported", False) diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index 36cd3b889..a2aa3c0b5 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,6 +42,11 @@ class MoRIBackend(EPBackend): maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") + # Both kernel families launch with host-built args only (no per-call host read of counts; + # the reset moved on-device in ROCm/mori#86 for vLLM's graphs), and SGLang captures AsyncLL + # decode split-phase inside its decode graph. `stage` slices by the untimed, per-rung + # `recv_tokens`, which is fixed for a rung's routing and so safe to bake into a capture. + CUDA_GRAPH_MODES = ("normal", "low-latency") requires_fresh_pair = True def __init__(self, args, rank, world_size, local_rank, device): diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 76eefb401..799747495 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -133,6 +133,11 @@ class UCCLEPBackend(EPBackend): kernel_generation = "uccl-legacy-buffer" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") + # The low-latency kernels are plain launches whose only host state is the double-buffer + # toggle, baked into a capture exactly as in DeepEP; every capture holds whole pairs (an even + # call count), so the toggle lands where it started. Normal mode host-syncs on its receive + # counters unless dispatched with `num_worst_tokens`, which this adapter does not do. + CUDA_GRAPH_MODES = ("low-latency",) stage_device_work = False requires_fresh_pair = False receive_layout = "token-rank" diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index d98e78f07..1db5ef001 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -282,9 +282,8 @@ Reading `false` alone as "roundtrip includes staging" subtracts a cost the row n availability, origin, and sample count. A paired-only API reports null isolated components. `isolated_sum` is derived. -Headline latency is `components.roundtrip` for default CUDA-graph rows and the **chained pair -period** (`components.pair_period`, defined under Chained Pair Period below) for eager rows that -carry one. The earlier eager flip shipped **held** while the +Headline latency is the **chained pair period** (`components.pair_period`, defined under Chained +Pair Period below) for every row that carries one, graph-replayed or eager. The earlier eager flip shipped **held** while the six-events-per-pair chain described below, whose inner records inflated small-T periods fleet-wide, was replaced by the two-pass chain, and was released on 2026-08-06 once the b200, h200 and gb200 hand references were confirmed against two-pass fleet artifacts (runs 31092783122 and @@ -309,23 +308,35 @@ rather than per-operation costs. The paired roundtrip is the comparable quantity ### CUDA Graph Replay -Graph-compatible backend/mode pairs capture the existing fixed-shape -dispatch→stage→combine roundtrip and measure `CUDAGraph.replay()` by default. Capture and replay -warmup are excluded. The result is published directly as `components.roundtrip`, with origin -`cuda-graph-replay`. Its capture has no internal timing nodes. Separate dispatch and combine -invocations each recapture the roundtrip with one event pair around only the requested phase, -preserving those existing component fields without charging their instrumentation to roundtrip. -There is no `graph_*` component or separate graph output path. `stage`, `pair_period`, chain -floors, and chain health are unavailable, so a graph-mode document contains only graph-derived -latency values. `isolated_sum` remains the derived sum of dispatch and combine. The ordinary -cross-rank MAX/MIN/spread reductions still apply to the replay samples. - -`COLLX_CUDA_GRAPH=0` restores the eager pipeline unchanged, including isolated components and the -chained pair period below. Modes not declared graph-compatible by their adapter also remain eager. +Serving engines capture their decode step, so graph-compatible backend/mode pairs are measured +under `CUDAGraph.replay()` by default. The graphed set follows what each library supports without +changing its contract: nccl-ep (both modes), flashinfer-ep (normal), MoRI (both modes), uccl-ep +(low-latency), and deepep-v2 low-latency plus normal-mode **decode**, which runs ElasticBuffer as +vLLM's graphed `deepep_v2` decode does (`do_cpu_sync=False`, worst-case receive, valid prefix read +from the handle on device; kernel generation `v2-elastic-buffer-nosync`). deepep-v2 normal prefill +keeps the host sync that sizes its receive exactly and stays eager, as does uccl-ep normal mode, +whose dispatch host-syncs unless padded to `num_worst_tokens`. + +Every family keeps its eager meaning under replay; only the launch mechanism changes: + +- **Fresh-entry components** (`roundtrip`, `dispatch`, `combine`) capture one pair with external + event nodes around the timed window, so host launch cost never enters it. Each timed replay + starts behind a device-side rank barrier (an all-reduce followed by a fixed spin, with no host + sync before the replay), so ranks enter together instead of from their own host's last + synchronize: without it, b200 EP16 ranks entered ~75 µs apart and the cross-rank MAX reported + the stagger as latency. Component origin is `cuda-graph-replay`. +- **The chained family** (`pair_period`, chain floors, chain health) captures each sibling chain as + ONE graph of `chain_iters` unrolled pairs, the shape a decode graph has, and replays it once + untimed and once behind the barrier. Event records are graph nodes and cost the host nothing, + so the floors sibling carries op windows without the eager six-events-per-pair defect. The + chained oracle and the chained-output check apply unchanged. + +`stage` is not separately timed under replay. `COLLX_CUDA_GRAPH=0` restores the eager pipeline. Every captured output is poisoned after timing and replayed once more; a finite rewrite gates the case, and where staging is not hoisted that replay is also compared with an untimed drained pair. -The artifact records `implementation.cuda_graph_replay`, `cuda_graph_supported`, and -`chained_period`, while the component origin makes the measurement visible at row granularity. +A graphed row's `kernel_generation` carries a `-cudagraph` suffix, so the durable store never +pools graph-replayed and eager samples of one kernel family into a single series. The artifact +also records `implementation.cuda_graph_replay` and `cuda_graph_supported`. ### Chained Pair Period @@ -647,7 +658,7 @@ One raw case document carries `record_type: "case-attempt"`, the single `version `combine_reduction` and `library_version` (which reduction the oracle held the kernel to, and the installed library that selected it), and two generation discriminators: `stage_excluded_from_roundtrip` (whether `roundtrip` excludes expert-output staging, discussed - above), `chained_period` (whether this document's rows carry the eager chained family), + above), `chained_period` (whether this document's rows carry the chained family), `cuda_graph_supported` (whether the adapter declares this mode graph-safe), and `cuda_graph_replay` (whether the existing measurement pipeline used replay). - `topology`: requested SKU/product, placement, `gpus_per_node`, nodes, scale-up domain, `scope`, @@ -660,9 +671,8 @@ One raw case document carries `record_type: "case-attempt"`, the single `version - `provenance`: the mounted image tag and source SHA, and - `outcome`: `status` (`success` or `invalid`) and `reasons`. -Each `rows` entry carries point latency (graph replay in `components.roundtrip` by default where -supported, otherwise the eager `components` plus `components.pair_period`, `chain_floor_us` and -`chain_health` (see Chained Pair Period)), byte +Each `rows` entry carries point latency (`components`, graph-replayed by default where supported, +plus `components.pair_period`, `chain_floor_us` and `chain_health` (see Chained Pair Period)), byte accounting, token rate, correctness, load, and fanout, while per-point statistics are summarized in place, not emitted as separate documents. Each dispatched case writes exactly this one raw result document, while unsupported or never-run cells produce no diff --git a/experimental/CollectiveX/tests/test_backends.py b/experimental/CollectiveX/tests/test_backends.py index 36d35cdbf..fe6d1f990 100644 --- a/experimental/CollectiveX/tests/test_backends.py +++ b/experimental/CollectiveX/tests/test_backends.py @@ -550,5 +550,58 @@ def test_ht_combine_input_is_sliced_to_the_received_count(self): self.assertEqual(h.combine_in_t, list(range(7))) self.assertLess(len(h.combine_in_t), len(b._recv_x)) + +def _deepep_v2_stubs(): + """Fake torch / deep_ep so `import ep_deepep_v2` succeeds without the benchmark image.""" + torch = types.ModuleType("torch") + torch.compile = lambda *a, **k: (lambda fn: fn) + dist = types.ModuleType("torch.distributed") + dist.group = types.SimpleNamespace(WORLD="world") + torch.distributed = dist + deep_ep = types.ModuleType("deep_ep") + deep_ep.ElasticBuffer = type("ElasticBuffer", (), {}) + deep_ep.Buffer = type("Buffer", (), {}) + return {"torch": torch, "torch.distributed": dist, "deep_ep": deep_ep} + + +class DeepEPV2GraphContract(unittest.TestCase): + """Normal-mode decode drops ElasticBuffer's host sync -- vLLM's graphed deepep_v2 decode + contract -- and only that makes it graph-capturable; prefill keeps the exact-size sync.""" + + def _backend(self, **updates): + with mock.patch.dict(sys.modules, _deepep_v2_stubs()): + sys.modules.pop("ep_deepep_v2", None) + import ep_deepep_v2 + sys.modules.pop("ep_deepep_v2", None) + return ep_deepep_v2.DeepEPV2Backend(args(**updates), 0, 8, 0, "cpu") + + def _dispatched_cpu_sync(self, backend): + calls = [] + + def dispatch(*_args, **kwargs): + calls.append(kwargs["do_cpu_sync"]) + return "recv_x", "recv_idx", "recv_w", "handle", None + + backend.buffer = types.SimpleNamespace(dispatch=dispatch) + backend.max_tokens, backend.num_sms, backend.num_qps = 8, 1, 1 + backend.dispatch(types.SimpleNamespace(dispatch_x="x", topk_idx="i", topk_weights="w")) + return calls[0] + + def test_normal_decode_is_the_no_sync_graphed_contract(self): + backend = self._backend(mode="normal", phase="decode") + self.assertIs(self._dispatched_cpu_sync(backend), False) + self.assertTrue(backend.cuda_graph_supported) + self.assertEqual(backend.kernel_generation, "v2-elastic-buffer-nosync") + + def test_normal_prefill_keeps_the_host_sync_and_stays_eager(self): + backend = self._backend(mode="normal", phase="prefill") + self.assertIs(self._dispatched_cpu_sync(backend), True) + self.assertFalse(backend.cuda_graph_supported) + self.assertEqual(backend.kernel_generation, "v2-elastic-buffer") + + def test_low_latency_stays_graphed(self): + backend = self._backend(mode="low-latency", phase="decode") + self.assertTrue(backend.cuda_graph_supported) + if __name__ == "__main__": unittest.main() diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index 965e75ba2..70f3d4ce1 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -113,6 +113,7 @@ def graph_context(_graph, **kwargs): CUDAGraph=lambda: _TraceGraph(clock, log), graph=graph_context, synchronize=lambda *args, **kwargs: log.append("sync"), + _sleep=lambda *args, **kwargs: log.append("align_spin"), current_stream=lambda *args, **kwargs: types.SimpleNamespace( synchronize=lambda: log.append("sync") ), @@ -320,6 +321,47 @@ def test_each_graph_component_uses_its_own_capture(self): self.assertTrue(problem._cuda_graph_output.cloned) self.assertTrue(problem._cuda_graph_output_rewritten) + def test_every_timed_replay_starts_behind_a_device_side_rank_barrier(self): + # Without the barrier each replay restarts from the preceding synchronize and ranks enter + # ~75us apart across nodes; the barrier must sit between the sync and the replay, with no + # host sync in between, on every timed replay and on none of the warm-ups. + backend = _ChainBackend(stage_device_work=False, fp8_consume="native", precision="bf16") + backend.mode = "normal" + backend.CUDA_GRAPH_MODES = ("normal",) + with mock.patch.dict(os.environ, {}, clear=True), \ + trace_torch(backend.clock, backend.calls): + backend.benchmark_component("roundtrip", new_problem(), warmup=2, iters=3) + replays = [i for i, call in enumerate(backend.calls) if call == "graph_replay"] + warmups, timed = replays[:2], replays[2:5] + for index in warmups: + self.assertNotEqual(backend.calls[index - 1], "align_spin") + for index in timed: + self.assertEqual(backend.calls[index - 2:index], ["all_reduce", "align_spin"]) + + def test_the_graph_chain_is_one_capture_of_unrolled_pairs_per_sibling(self): + iters, drop = 5, 1 + backend = _ChainBackend(stage_device_work=False, fp8_consume="native", precision="bf16") + backend.mode = "normal" + backend.CUDA_GRAPH_MODES = ("normal",) + with mock.patch.dict(os.environ, {}, clear=True), \ + trace_torch(backend.clock, backend.calls): + series = backend.benchmark_chain(new_problem(), 0, iters, drop) + # Floors and period siblings: one capture each, holding every pair of the chain. + self.assertEqual(backend.calls.count("capture_begin"), 2) + begin = [i for i, call in enumerate(backend.calls) if call == "capture_begin"] + end = [i for i, call in enumerate(backend.calls) if call == "capture_end"] + for lo, hi in zip(begin, end): + self.assertEqual( + ops_only(backend.calls[lo:hi]), ["dispatch", "stage", "combine"] * iters + ) + # Each graph replays once untimed, then once behind the barrier. + self.assertEqual(backend.calls.count("graph_replay"), 4) + self.assertEqual(backend.calls.count("align_spin"), 2) + for key in ("pair", "dispatch", "combine"): + self.assertEqual(len(series[key]), iters - drop) + self.assertEqual(len(series["start_to_start"]), iters - drop - 1) + self.assertTrue(series["combined"].cloned) + def test_external_switch_restores_the_eager_component_pipeline(self): backend = _ChainBackend() backend.mode = "normal" @@ -860,8 +902,9 @@ def setUpClass(cls): def test_existing_component_fields_carry_graph_measurements(self): self.assertEqual(self.swept.rc, 0) self.assertIs(self.swept.doc["implementation"]["cuda_graph_replay"], True) - self.assertIs(self.swept.doc["implementation"]["chained_period"], False) - self.assertFalse(any(event[0] == "chain" for event in self.swept.events)) + # Graph mode keeps the chained family: the chain is captured, not dropped. + self.assertIs(self.swept.doc["implementation"]["chained_period"], True) + self.assertTrue(any(event[0] == "chain" for event in self.swept.events)) for row in self.swept.rows: with self.subTest(tokens=row["tokens_per_rank"]): roundtrip = row["components"]["roundtrip"] @@ -873,11 +916,11 @@ def test_existing_component_fields_carry_graph_measurements(self): self.assertEqual(row["components"]["dispatch"]["origin"], "cuda-graph-replay") self.assertEqual(row["components"]["combine"]["origin"], "cuda-graph-replay") self.assertEqual(row["components"]["isolated_sum"]["percentiles_us"]["p50"], 24.0) - for name in ("stage", "pair_period"): - self.assertIsNone(row["components"][name]["percentiles_us"]) + self.assertIsNone(row["components"]["stage"]["percentiles_us"]) + self.assertIsNotNone(row["components"]["pair_period"]["percentiles_us"]) self.assertIs(row["correctness"]["cuda_graph_output_rewritten"], True) self.assertIs(row["correctness"]["cuda_graph_last_output_passed"], True) - self.assertIsNone(row["correctness"]["post_chain_state_passed"]) + self.assertIs(row["correctness"]["post_chain_state_passed"], True) def test_a_replay_that_does_not_rewrite_its_output_reds_the_case(self): swept = drive(backend_factory=_BrokenGraphSweepBackend) From a86227ffd4f42528d60977179c83822e4767affa Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:42:08 +0800 Subject: [PATCH 02/20] CollectiveX: selectable nccl-ep LL layout (expert-major alongside rank-major); keep MoRI AsyncLL eager COLLX_NCCL_LL_LAYOUT=expert-major restores the weighted per-expert LL contract deepep-v2 LL uses, under its original kernel generation, as the like-for-like row against the DeepEP-API backends; rank-major stays the default. MoRI AsyncLL trips a recv-copy device assertion under capture on mi300x-tw, so only MoRI normal mode is graphed. --- experimental/CollectiveX/bench/ep_mori.py | 12 +- experimental/CollectiveX/bench/ep_nccl.py | 112 ++++++++++++++++-- .../CollectiveX/tests/test_backends.py | 33 ++++++ 3 files changed, 143 insertions(+), 14 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index a2aa3c0b5..c7c221710 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,11 +42,13 @@ class MoRIBackend(EPBackend): maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # Both kernel families launch with host-built args only (no per-call host read of counts; - # the reset moved on-device in ROCm/mori#86 for vLLM's graphs), and SGLang captures AsyncLL - # decode split-phase inside its decode graph. `stage` slices by the untimed, per-rung - # `recv_tokens`, which is fixed for a rung's routing and so safe to bake into a capture. - CUDA_GRAPH_MODES = ("normal", "low-latency") + # Normal-mode kernels launch with host-built args only (no per-call host read of counts; the + # reset moved on-device in ROCm/mori#86 for vLLM's graphs). `stage` slices by the untimed, + # per-rung `recv_tokens`, fixed for a rung's routing and so safe to bake into a capture. + # AsyncLL stays eager: captured, its recv-copy kernel trips `(pe >= 0) && (pe < worldSize)` + # (low_latency_async.cpp:375) on mi300x-tw at bf16 and fp8, and upstream has no AsyncLL graph + # test to say what state the split phase expects between replays. + CUDA_GRAPH_MODES = ("normal",) requires_fresh_pair = True def __init__(self, args, rank, world_size, local_rank, device): diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index c0b08f59f..15f79cd02 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -104,6 +104,7 @@ class NCCLEPBackend(EPBackend): receive_layout = "token-rank" combine_weight_semantics = "unweighted-rank-sum" zero_copy = True + _ll_expert_major = False def __init__(self, args, rank, world_size, local_rank, device): super().__init__(args, rank, world_size, local_rank, device) @@ -119,9 +120,21 @@ def __init__(self, args, rank, world_size, local_rank, device): self.num_local_experts = self.experts_per_rank self._internode = world_size > int(args.scale_up_domain) self._ll = self.mode == "low-latency" + # LL layout: rank-major (TensorRT-LLM's NCCL EP contract, the default) or expert-major + # (DeepEP LL's contract, as vLLM/SGLang decode consume it). Both are native LL layouts. + layout = os.environ.get("COLLX_NCCL_LL_LAYOUT", "rank-major") + if layout not in ("rank-major", "expert-major"): + raise ValueError(f"COLLX_NCCL_LL_LAYOUT must be rank-major or expert-major, got {layout!r}") + self._ll_expert_major = self._ll and layout == "expert-major" # LL rank-major follows the inference-framework contract. Direct windows are scale-up only. - self.zero_copy = not self._ll or not self._internode - if self._ll: + self.zero_copy = (not self._ll or not self._internode) and not self._ll_expert_major + if self._ll_expert_major: + # Weighted source-side combine over a per-expert padded receive: deepep-v2 LL's + # contract, so this is the like-for-like row against the DeepEP-API backends. + self.kernel_generation = "nccl-ep-v02-ll" + self.receive_layout = "token-expert" + self.combine_weight_semantics = "weighted-kernel-sum" + elif self._ll: self.kernel_generation = ( "nccl-ep-v02-ll-rm-zc" if self.zero_copy else "nccl-ep-v02-ll-rm" ) @@ -135,7 +148,10 @@ def __init__(self, args, rank, world_size, local_rank, device): # low-latency Buffer — no timed component needs a fresh dispatch or a draining combine; # both modes keep requires_fresh_pair False. self._algorithm = Algorithm.LOW_LATENCY if self._ll else Algorithm.HIGH_THROUGHPUT - self._layout = Layout.RANK_MAJOR if self._ll else Layout.FLAT + if self._ll: + self._layout = Layout.EXPERT_MAJOR if self._ll_expert_major else Layout.RANK_MAJOR + else: + self._layout = Layout.FLAT # send_only=0 runs each dispatch/combine as a complete SEND|RECV operation. # FWD pass carries top-k weights on dispatch (HT) and forbids them on the HT combine # input (the combine is a plain rank sum). @@ -234,7 +250,20 @@ def create_buffer(self, spec): self._ep_group = nccl_ep.Group.create(self._comm, config) dev = self.device - if self._ll: + if self._ll_expert_major: + # EXPERT_MAJOR recv: [num_local_experts, max_dispatch*num_ranks, hidden]. + slots = self.max_dispatch * self.world_size + self._recv_x = torch.empty( + (self.num_local_experts, slots, hidden), dtype=torch.bfloat16, device=dev + ) + # Per-local-expert received-token counts, written by NCCL EP during dispatch. + self._recv_count = torch.empty( + (self.num_local_experts,), dtype=torch.int32, device=dev + ) + # Zeroed scratch the combine oracle scatters the transformed rows into. + self._combine_scratch = torch.empty_like(self._recv_x) + self._recv_count_t = self._t(self._recv_count) + elif self._ll: # RANK_MAJOR receive: [source rank, source slot, hidden]. self._recv_x = nccl_core.torch.empty( (self.world_size, self.max_dispatch, hidden), dtype=torch.bfloat16, device=dev @@ -329,10 +358,12 @@ def _ensure_handle(self, p): in_tokens_t=self._t(p.dispatch_x), topk_idx_t=topk_idx_t, ) - if not self._ll: - h.in_weights_t = self._t(p.topk_weights) + if self._ll_expert_major: + # Expert-major applies the gate in its combine kernel, not on dispatch. Wrapped once + # per handle: `time_us` charges the wrapper's host work to the window. + h.combine_weights_t = self._t(p.topk_weights) else: - # LL rank-major transports weights with dispatch. + # HT carries weights on dispatch; LL rank-major transports them with dispatch too. h.in_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) @@ -419,7 +450,20 @@ def dispatch(self, p): # read here: the bound problem's counters are deterministic and already read # (_bind_ht_recv_count) in the untimed rebind. h.handle.update(h.topk_idx_t, layout_info=h.layout_info, stream=stream) - if self._ll: + if self._ll_expert_major: + # LL EXPERT_MAJOR: tokens in, 3D per-expert padded tokens out, per-expert recv + # counts written into expert_counters. No weights on the dispatch (the gate is + # applied by the combine kernel at the source). + h.handle.dispatch( + DispatchInputs(tokens=h.in_tokens_t), + DispatchOutputs(tokens=self._recv_x_t), + layout_info=LayoutInfo(expert_counters=self._recv_count_t), + config=self._dispatch_cfg, + stream=stream, + ) + h.recv_x = self._recv_x + h.recv_count = self._recv_count + elif self._ll: # LL RANK_MAJOR returns one plane per source rank. h.handle.dispatch( DispatchInputs(tokens=h.in_tokens_t, topk_weights=h.in_weights_t), @@ -462,7 +506,17 @@ def stage(self, p, h): def combine(self, p, h): stream = self._stream() - # Both layouts use an unweighted rank-sum combine. + if self._ll_expert_major: + # Weighted combine: the kernel multiplies each expert contribution by the source + # token's gate before the FP32 accumulation. + h.handle.combine( + CombineInputs(tokens=h.combine_input), + CombineOutputs(tokens=h.out_t, topk_weights=h.combine_weights_t), + config=self._combine_cfg, + stream=stream, + ) + return h.out + # HT and LL rank-major use an unweighted rank-sum combine. h.handle.combine( CombineInputs(tokens=h.combine_input), CombineOutputs(tokens=h.out_t), @@ -502,7 +556,30 @@ def _ll_inspect_dispatch(self, p, h): ), ) + def _ll_em_inspect_dispatch(self, p, h): + """Flat per-slot view over the EXPERT_MAJOR padded receive (mirror of + ep_deepep_v2._ll_inspect_dispatch): each local expert's valid tokens are packed at the + front [0:recv_count[e]] of its slot dimension. Flatten to the oracle's compact + (expert, slot) row-major contract and keep the coordinates for the combine scatter.""" + recv_bf16 = h.recv_x # [E, S, hidden] BF16 + num_slots = recv_bf16.shape[1] + counts = h.recv_count.to(torch.int64) # [E] + slot_valid = ( + torch.arange(num_slots, device=recv_bf16.device).unsqueeze(0) < counts.unsqueeze(1) + ) + slot_expert, slot_j = slot_valid.nonzero(as_tuple=True) + h.slot_expert = slot_expert + h.slot_j = slot_j + local_lo = self.rank * self.num_local_experts + return types.SimpleNamespace( + payload=recv_bf16[slot_expert, slot_j], + expert_ids=local_lo + slot_expert.to(torch.int64), + local_expert_counts=counts, + ) + def inspect_dispatch(self, p, h): + if self._ll_expert_major: + return self._ll_em_inspect_dispatch(p, h) if self._ll: return self._ll_inspect_dispatch(p, h) # HT FLAT normal recv: front-packed to recv_total_counter, one row per received token. @@ -546,7 +623,24 @@ def _ll_combine_transformed(self, p, h, transformed): ) return h.out[: p.T] + def _ll_em_combine_transformed(self, p, h, transformed): + """Scatter the oracle-transformed rows back into a zeroed EXPERT_MAJOR combine buffer at + the (expert, slot) coordinates inspect read them from, then run the weighted combine; + the kernel applies p.topk_weights, so the staged transform is unweighted.""" + combine_buf = self._combine_scratch + combine_buf.zero_() + combine_buf[h.slot_expert, h.slot_j] = transformed.to(combine_buf.dtype) + h.handle.combine( + CombineInputs(tokens=self._t(combine_buf)), + CombineOutputs(tokens=h.out_t, topk_weights=h.combine_weights_t), + config=self._combine_cfg, + stream=self._stream(), + ) + return h.out[: p.T] + def combine_transformed(self, p, h, transformed): + if self._ll_expert_major: + return self._ll_em_combine_transformed(p, h, transformed) if self._ll: return self._ll_combine_transformed(p, h, transformed) # `transformed` is the oracle's per-received-token combine input [count, hidden] diff --git a/experimental/CollectiveX/tests/test_backends.py b/experimental/CollectiveX/tests/test_backends.py index fe6d1f990..df4baf938 100644 --- a/experimental/CollectiveX/tests/test_backends.py +++ b/experimental/CollectiveX/tests/test_backends.py @@ -251,6 +251,39 @@ def base_init(instance, options, rank, world_size, local_rank, device): self.assertEqual(ll.combine_reduction, "rank-fp32") self.assertEqual(getattr(ht, "combine_reduction", "domain-fp32"), "domain-fp32") + def test_ll_layout_selector_restores_the_expert_major_contract(self): + module = self._module() + module.dist.group = types.SimpleNamespace(WORLD=object()) + + def base_init(instance, options, rank, world_size, local_rank, device): + instance.args = options + instance.mode = options.mode + + common = dict(experts=384, hidden=7168, topk=6, scale_up_domain=8) + with mock.patch.object(module.EPBackend, "__init__", base_init): + with mock.patch.dict(os.environ, {"COLLX_NCCL_LL_LAYOUT": "expert-major"}): + em = module.NCCLEPBackend( + types.SimpleNamespace(mode="low-latency", **common), 0, 8, 0, "cuda:0" + ) + rm = module.NCCLEPBackend( + types.SimpleNamespace(mode="low-latency", **common), 0, 8, 0, "cuda:0" + ) + with mock.patch.dict(os.environ, {"COLLX_NCCL_LL_LAYOUT": "flat"}), \ + self.assertRaisesRegex(ValueError, "COLLX_NCCL_LL_LAYOUT"): + module.NCCLEPBackend( + types.SimpleNamespace(mode="low-latency", **common), 0, 8, 0, "cuda:0" + ) + + self.assertEqual(em._layout, module.Layout.EXPERT_MAJOR) + self.assertEqual( + (em.kernel_generation, em.receive_layout, em.combine_weight_semantics), + ("nccl-ep-v02-ll", "token-expert", "weighted-kernel-sum"), + ) + self.assertFalse(em.zero_copy) + self.assertEqual(getattr(em, "combine_reduction", "domain-fp32"), "domain-fp32") + self.assertEqual(rm._layout, module.Layout.RANK_MAJOR) + self.assertEqual(rm.combine_reduction, "rank-fp32") + def test_ladder_cap_drops_only_oversized_measurement_points(self): module = self._module() backend = self._backend(module, low_latency=True) From 828c7fb2f5f751e035309581d7f743da5b63401c Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:14:24 +0800 Subject: [PATCH 03/20] CollectiveX: keep MoRI eager -- ROCm torch rejects external events inside a graph capture --- experimental/CollectiveX/bench/ep_mori.py | 12 +++++------- experimental/CollectiveX/docs/methodology.md | 7 ++++--- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index c7c221710..f7238b7b3 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,13 +42,11 @@ class MoRIBackend(EPBackend): maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # Normal-mode kernels launch with host-built args only (no per-call host read of counts; the - # reset moved on-device in ROCm/mori#86 for vLLM's graphs). `stage` slices by the untimed, - # per-rung `recv_tokens`, fixed for a rung's routing and so safe to bake into a capture. - # AsyncLL stays eager: captured, its recv-copy kernel trips `(pe >= 0) && (pe < worldSize)` - # (low_latency_async.cpp:375) on mi300x-tw at bf16 and fp8, and upstream has no AsyncLL graph - # test to say what state the split phase expects between replays. - CUDA_GRAPH_MODES = ("normal",) + # Eager in both modes. The kernels themselves are capturable (host-built args only; the reset + # moved on-device in ROCm/mori#86 for vLLM's graphs), but ROCm torch rejects the external + # events the graph timing records as graph nodes ("External events are disallowed in rocm"), + # so replay could not be timed with the same windows as every CUDA backend. + CUDA_GRAPH_MODES = () requires_fresh_pair = True def __init__(self, args, rank, world_size, local_rank, device): diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 1db5ef001..5f1a88cf9 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -310,12 +310,13 @@ rather than per-operation costs. The paired roundtrip is the comparable quantity Serving engines capture their decode step, so graph-compatible backend/mode pairs are measured under `CUDAGraph.replay()` by default. The graphed set follows what each library supports without -changing its contract: nccl-ep (both modes), flashinfer-ep (normal), MoRI (both modes), uccl-ep -(low-latency), and deepep-v2 low-latency plus normal-mode **decode**, which runs ElasticBuffer as +changing its contract: nccl-ep (both modes), flashinfer-ep (normal), uccl-ep (low-latency), and +deepep-v2 low-latency plus normal-mode **decode**, which runs ElasticBuffer as vLLM's graphed `deepep_v2` decode does (`do_cpu_sync=False`, worst-case receive, valid prefix read from the handle on device; kernel generation `v2-elastic-buffer-nosync`). deepep-v2 normal prefill keeps the host sync that sizes its receive exactly and stays eager, as does uccl-ep normal mode, -whose dispatch host-syncs unless padded to `num_worst_tokens`. +whose dispatch host-syncs unless padded to `num_worst_tokens`. MoRI stays eager in both modes: +ROCm torch rejects the external events the replay windows are recorded with. Every family keeps its eager meaning under replay; only the launch mechanism changes: From b25ca14a8307f28b0cafdfa2c409b328f9ac5977 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Fri, 25 Sep 2026 12:08:53 -0500 Subject: [PATCH 04/20] fix: route MI325X CollectiveX sweeps through the live Slurm pool --- .github/workflows/collectivex-sweep.yml | 6 +++--- .../CollectiveX/configs/platform_config.json | 17 ++--------------- 2 files changed, 5 insertions(+), 18 deletions(-) diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 3d68f67cb..5cfe0ba1d 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -151,7 +151,7 @@ jobs: vars.NODE_SLOT_SCHEDULER_ENABLED == 'true' && format( '["self-hosted",{0},{1},{2},{3}]', - toJSON(matrix.sku), + toJSON(matrix.sku == 'mi325x' && 'cluster:mi325x-amds' || matrix.sku), toJSON(format('nodes:{0}', matrix.nodes)), toJSON(format( 'ci-job-{0}-{1}', @@ -162,7 +162,7 @@ jobs: ) || format( '["self-hosted",{0},{1},{2},{3}]', - toJSON(matrix.sku), + toJSON(matrix.sku == 'mi325x' && 'cluster:mi325x-amds' || matrix.sku), toJSON(format('nodes:{0}', matrix.nodes)), toJSON(format( 'ci-job-{0}-{1}', @@ -172,7 +172,7 @@ jobs: toJSON(format('ci-attempt-{0}', github.run_attempt)) ) ) || - format('[{0}]', toJSON(matrix.sku)) + format('[{0}]', toJSON(matrix.sku == 'mi325x' && 'cluster:mi325x-amds' || matrix.sku)) ) }} name: p${{ needs.setup.outputs.priority }} | ${{ matrix.sku }} ${{ matrix.backend }} shard ${{ matrix.id }} timeout-minutes: 350 diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index e2c1ed3c3..7e24bc989 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -129,19 +129,6 @@ "exclude_nodes": "im-gb300-r01-c003,im-gb300-r01-c005" } }, - "mi325x-tw": { - "arch": "gfx942", - "product": "mi325x", - "image": "rocm/sgl-dev:sglang-0.5.14-rocm720-mi35x-mori-0701", - "image_platform": "linux/amd64", - "gpus_per_node": 8, - "scale_up_domain": 8, - "scale_up_transport": "xgmi", - "launcher": "mi-tw", - "backends": {"mori": [8], "uccl-ep": [8]}, - "ll_backends": {"mori": [8]}, - "fabric": {"nic": "n/a (single-node scale-up)", "switch": "n/a (single-node scale-up)"} - }, "mi300x-tw": { "arch": "gfx942", "product": "mi300x", @@ -184,8 +171,8 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-amds", - "backends": {}, - "ll_backends": {}, + "backends": {"mori": [8], "uccl-ep": [8]}, + "ll_backends": {"mori": [8]}, "fabric": { "nic": "n/a (single-node scale-up)", "switch": "n/a (single-node scale-up)" From 15796e5aab766b3265f7fde7ba581956951bf37d Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 01:59:52 +0800 Subject: [PATCH 05/20] CollectiveX: keep nccl-ep HT eager -- graphed zero-copy HT corrupts intermittently across nodes on x86 --- experimental/CollectiveX/bench/ep_nccl.py | 6 +++++- experimental/CollectiveX/docs/methodology.md | 6 ++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index 15f79cd02..a8742878e 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -98,7 +98,11 @@ class NCCLEPBackend(EPBackend): kernel_generation = "nccl-ep-v02-ht-routed-zc" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16",) - CUDA_GRAPH_MODES = ("normal", "low-latency") + # LL replays; HT stays eager. Graphed zero-copy HT failed the combine oracle intermittently + # across nodes on x86 (b200 EP16 T=128, h200 EP16 T=32 and prefill T=1024; runs 35994093313, + # 36113759089) while eager zero-copy HT passed every cell and was as fast or faster + # (run 36114371399), so eager is the better HT configuration on every pool measured. + CUDA_GRAPH_MODES = ("low-latency",) stage_device_work = False requires_fresh_pair = False receive_layout = "token-rank" diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 5f1a88cf9..2f2fff2ce 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -310,13 +310,15 @@ rather than per-operation costs. The paired roundtrip is the comparable quantity Serving engines capture their decode step, so graph-compatible backend/mode pairs are measured under `CUDAGraph.replay()` by default. The graphed set follows what each library supports without -changing its contract: nccl-ep (both modes), flashinfer-ep (normal), uccl-ep (low-latency), and +changing its contract: nccl-ep (low-latency), flashinfer-ep (normal), uccl-ep (low-latency), and deepep-v2 low-latency plus normal-mode **decode**, which runs ElasticBuffer as vLLM's graphed `deepep_v2` decode does (`do_cpu_sync=False`, worst-case receive, valid prefix read from the handle on device; kernel generation `v2-elastic-buffer-nosync`). deepep-v2 normal prefill keeps the host sync that sizes its receive exactly and stays eager, as does uccl-ep normal mode, whose dispatch host-syncs unless padded to `num_worst_tokens`. MoRI stays eager in both modes: -ROCm torch rejects the external events the replay windows are recorded with. +ROCm torch rejects the external events the replay windows are recorded with. nccl-ep normal (HT) +stays eager because graphed zero-copy HT failed the combine oracle intermittently across nodes on +x86, while eager zero-copy HT was correct everywhere and no slower. Every family keeps its eager meaning under replay; only the launch mechanism changes: From b0d7acf176da9db48999384b66ca5f5afd8112b6 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 02:51:20 +0800 Subject: [PATCH 06/20] CollectiveX: address the per-library review of graph replay - nccl-ep HT combine takes the full static receive plane the FLAT contract requires (the count slice broke it in both regimes); kernel generation -> nccl-ep-v02-ht-routed-zc-static - value-check every graphed row: an untimed capture with stage inside the graph, dispatch outputs and result poisoned before its replay, compared with an eager drained pair - alignment spin sized to wall time per GPU (calibrated), not a cycle count; graphed fresh-entry p90/p95/p99 withheld until alignment is clean - deepep-v2 no-sync decode sizes its receive to the next power of two of T, as vLLM does - per-case graph gates: flashinfer decode only; uccl-ep LL intranode only, not b200 FP8, not with the proxy's adaptive sleep - MoRI graphed in both modes, capturing timing events via hipEventRecordWithFlags on ROCm - expert-major nccl-ep LL kernel generation -> nccl-ep-v02-ll-em (no complete() in its windows) --- experimental/CollectiveX/bench/ep_backend.py | 159 ++++++++++++++++-- .../CollectiveX/bench/ep_deepep_v2.py | 16 +- .../CollectiveX/bench/ep_flashinfer.py | 7 + experimental/CollectiveX/bench/ep_harness.py | 51 ++++-- experimental/CollectiveX/bench/ep_mori.py | 11 +- experimental/CollectiveX/bench/ep_nccl.py | 29 ++-- experimental/CollectiveX/bench/ep_uccl.py | 20 +++ experimental/CollectiveX/docs/methodology.md | 71 +++++--- .../CollectiveX/tests/test_backends.py | 87 ++++++++-- experimental/CollectiveX/tests/test_chain.py | 65 ++++++- 10 files changed, 432 insertions(+), 84 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index ed0aa2393..d2c067335 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -395,10 +395,30 @@ def _topk_idx_dtype(self): # ---- CUDA graph capture ---------------------------------------------------------------- - # Spin after the alignment all-reduce so every rank's host has enqueued its replay before the - # stream reaches it; ~50us at 2GHz, far above a graph launch, so the replay start is set by - # the barrier release on every rank rather than by host launch latency. - _GRAPH_ALIGN_SPIN_CYCLES = 100_000 + # Wall time the stream spins after the alignment all-reduce, so every rank's host has enqueued + # its replay before the stream reaches it and the replay start is set by the barrier release, + # not by host launch latency. Converted to cycles per GPU by `_calibrate_align_spin`: + # `torch.cuda._sleep` counts SM cycles, so a fixed cycle count spun 48-70us across the clock + # range and left ~15us of cross-rank skew on gb200 (a fast-clocked rank released early). + _GRAPH_ALIGN_SPIN_US = 100.0 + # Attributes of a dispatch handle that hold what dispatch wrote; the replay check poisons them + # so a replay that skipped (or stalely reused) the dispatch cannot reproduce a valid output. + _DISPATCH_OUTPUT_FIELDS = ("recv_x", "recv_scales", "dispatch_output") + + def _calibrate_align_spin(self): + """Measure this GPU's current spin rate and size the alignment spin to a wall time.""" + import torch + + probe = 200_000 + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda._sleep(probe // 10) # ramp clocks before the measured spin + start.record() + torch.cuda._sleep(probe) + end.record() + torch.cuda.synchronize() + elapsed_us = max(start.elapsed_time(end) * 1000.0, 1e-3) + self._graph_align_cycles = max(1, int(probe * self._GRAPH_ALIGN_SPIN_US / elapsed_us)) def _graph_align(self): """Enqueue a device-side rank barrier on the current stream, without a host sync.""" @@ -408,33 +428,94 @@ def _graph_align(self): token = getattr(self, "_graph_align_token", None) if token is None: token = self._graph_align_token = torch.zeros(1, device=self.device) + if getattr(self, "_graph_align_cycles", None) is None: + self._calibrate_align_spin() dist.all_reduce(token) - torch.cuda._sleep(self._GRAPH_ALIGN_SPIN_CYCLES) + torch.cuda._sleep(self._graph_align_cycles) + + def _graph_event(self): + """A timing event whose record() can become a node of a graph being captured. + + CUDA torch does this with `external=True`. ROCm torch before 2.13 rejects external events + ("External events are disallowed in rocm") although HIP >= 7 supports them, so there the + event is created normally, recorded once outside capture so it exists and counts as + recorded, and captured with `hipEventRecordWithFlags(..., hipEventRecordExternal)` -- + the call torch 2.13 itself makes (pytorch#178264). + """ + import torch + + if not getattr(torch.version, "hip", None): + return torch.cuda.Event(enable_timing=True, external=True) + event = torch.cuda.Event(enable_timing=True) + event.record() + event._collx_hip_external = True + return event + + @staticmethod + def _record_graph_event(event): + import torch + + if not getattr(event, "_collx_hip_external", False): + event.record() + return + import ctypes + + hip = EPBackend._hip_runtime() + rc = hip.hipEventRecordWithFlags( + ctypes.c_void_p(event.cuda_event), + ctypes.c_void_p(torch.cuda.current_stream().cuda_stream), + ctypes.c_uint(0x1), # hipEventRecordExternal + ) + if rc != 0: + raise RuntimeError(f"hipEventRecordWithFlags(external) failed with hipError {rc}") + + @staticmethod + def _hip_runtime(): + lib = getattr(EPBackend, "_hip_lib", None) + if lib is None: + import ctypes + import os as _os + + import torch + + candidates = ["libamdhip64.so", _os.path.join(_os.path.dirname(torch.__file__), "lib", + "libamdhip64.so")] + for name in candidates: + try: + lib = ctypes.CDLL(name) + break + except OSError: + continue + if lib is None: + raise RuntimeError("libamdhip64.so not loadable for graph event capture") + lib.hipEventRecordWithFlags.restype = ctypes.c_int + EPBackend._hip_lib = lib + return lib def _capture_pairs(self, problem, staged, pairs, marks): """Capture `pairs` back-to-back dispatch -> combine pairs into one graph. - `marks` selects which windows get external event nodes: "pair" (the whole pair), - "dispatch", "combine". Event records are graph nodes, so they cost the stream nothing on - the host -- the six-events-per-pair defect the eager chain splits around does not exist - here. Returns (graph, {mark: (starts, ends)}, last combined output). + `marks` selects which windows get event nodes: "pair" (the whole pair), "dispatch", + "combine". Event records are graph nodes, so they cost the stream nothing on the host -- + the six-events-per-pair defect the eager chain splits around does not exist here. + Returns (graph, {mark: (starts, ends)}, last combined output, last dispatch handle). """ import torch import torch.distributed as dist def events(): - return [torch.cuda.Event(enable_timing=True, external=True) for _ in range(pairs)] + return [self._graph_event() for _ in range(pairs)] stamps = {mark: (events(), events()) for mark in marks} def record(mark, edge, i): if mark in stamps: - stamps[mark][edge][i].record() + self._record_graph_event(stamps[mark][edge][i]) dist.barrier() torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() - combined = None + combined = handle = None with torch.cuda.graph(graph, capture_error_mode="relaxed"): for i in range(pairs): record("pair", 0, i) @@ -450,7 +531,48 @@ def record(mark, edge, i): record("combine", 1, i) record("pair", 1, i) torch.cuda.synchronize() - return graph, stamps, combined + return graph, stamps, combined, handle + + @staticmethod + def _poison(tensor): + """Overwrite a tensor with 0xFF bytes: NaN for bf16/fp16/fp32/fp8-e4m3, -1 for ints.""" + import torch + + if tensor is None: + return + if isinstance(tensor, (tuple, list)): + for part in tensor: + EPBackend._poison(part) + return + if not isinstance(tensor, torch.Tensor) or not tensor.numel(): + return + try: + tensor.view(torch.uint8).fill_(0xFF) + except RuntimeError: + tensor.fill_(float("nan") if tensor.is_floating_point() else -1) + + def graph_replay_output(self, problem): + """The value check for graph replay: one untimed capture with `stage` INSIDE the graph. + + The timed captures hoist staging where `stage` does device work, so their output never + depends on that replay's dispatch and a stale replay would still look correct. This one + stages per pair, then poisons what dispatch wrote and the combined output before its only + replay: the result is valid only if the replay itself re-ran dispatch, stage and combine. + The caller compares it with an eager drained pair through the same code path. + """ + import torch + + self.warm(problem, 1) + graph, _, combined, handle = self._capture_pairs(problem, None, 1, ()) + graph.replay() # first launch uploads the graph; its output is discarded + torch.cuda.synchronize() + for field in self._DISPATCH_OUTPUT_FIELDS: + self._poison(getattr(handle, field, None)) + self._poison(combined) + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + return combined.clone() # ---- Timing template methods ----------------------------------------------------- @@ -619,10 +741,13 @@ def _benchmark_chain_graph(self, problem, staged, iters, drop): """ import torch - floors, floor_stamps, _ = self._capture_pairs( + floors, floor_stamps, _, _ = self._capture_pairs( problem, staged, iters, ("dispatch", "combine") ) - period, period_stamps, combined = self._capture_pairs(problem, staged, iters, ("pair",)) + period, period_stamps, combined, _ = self._capture_pairs( + problem, staged, iters, ("pair",) + ) + self._calibrate_align_spin() for graph in (floors, period): graph.replay() torch.cuda.synchronize() @@ -684,7 +809,9 @@ def benchmark_roundtrip(self, problem, warmup, iters, graph_component="roundtrip # and its warm-up are excluded; each timed replay starts behind a device-side rank # barrier (`_graph_align`) so the cross-rank MAX is the operation, not launch skew. mark = "pair" if graph_component == "roundtrip" else graph_component - graph, stamps, combined = self._capture_pairs(problem, staged, 1, (mark,)) + graph, stamps, combined, _ = self._capture_pairs(problem, staged, 1, (mark,)) + # Re-measure the spin rate per timed series: clocks move with load and temperature. + self._calibrate_align_spin() starts, ends = stamps[mark] samples = time_cuda_graph_phase_us( torch, graph.replay, warmup, iters, (starts[0], ends[0]), diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py index de3564211..c3065df9a 100644 --- a/experimental/CollectiveX/bench/ep_deepep_v2.py +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -197,6 +197,20 @@ def cuda_graph_supported(self) -> bool: return not getattr(self, "_normal_cpu_sync", True) return super().cuda_graph_supported + def _dispatch_capacity(self, tokens): + """Per-call `num_max_tokens_per_rank`. + + With the host sync the receive is sized exactly, so the buffer maximum is only a bound. + Without it DeepEP allocates `num_max_tokens_per_rank * num_ranks` receive rows + (elastic/buffer.hpp "allocate with the worst case"), so passing the ladder maximum made a + T=1 dispatch receive 4096 rows at EP8 and the FP8 stage dequantize all of them (stage + 59 -> 212us, b200 EP8). vLLM's graphed decode (prepare_finalize/deepep_v2.py) passes the + next power of two of the batch, which bounds both the receive and the JIT variants. + """ + if self._normal_cpu_sync: + return self.max_tokens + return min(self.max_tokens, 1 << max(0, int(tokens) - 1).bit_length()) + def buffer_cap(self, args): if self.mode == "low-latency": # LL pre-allocates a fixed [num_local_experts, cap * num_ranks, hidden] receive @@ -388,7 +402,7 @@ def dispatch(self, p): topk_idx=p.topk_idx, topk_weights=p.topk_weights, num_experts=self.args.experts, - num_max_tokens_per_rank=self.max_tokens, + num_max_tokens_per_rank=self._dispatch_capacity(p.T), expert_alignment=1, num_sms=self.num_sms, num_qps=self.num_qps, diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 8fe654cba..61a2768e3 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -120,6 +120,13 @@ class FlashInferEPBackend(EPBackend): # Forced by the phase asserts described in the module docstring. requires_fresh_pair = True + @property + def cuda_graph_supported(self) -> bool: + # Decode only: graph replay removes host-bound library overhead at decode sizes (gb200 EP8 + # T=1 pair period 151 -> 43us) but changes nothing at prefill (T=8192 1400 vs 1387us), and + # engines capture decode, not prefill, so prefill keeps the eager series it already has. + return super().cuda_graph_supported and getattr(self.args, "phase", None) == "decode" + def __init__(self, args, rank, world_size, local_rank, device): super().__init__(args, rank, world_size, local_rank, device) self._fp8 = self.precision == "fp8" diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index 1757109e8..c21edd000 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -213,6 +213,20 @@ def _pcts(xs): CUDA_GRAPH_ORIGIN = "cuda-graph-replay" +def _published_tails(percentiles, graph_replay): + """Fresh-entry percentiles as published: under graph replay only the median. + + A graphed fresh-entry sample starts behind the alignment barrier, but a rank whose host is + late to launch its replay still stalls the others, and those stalls own the tail (gb200 + flashinfer T=1: roundtrip p99 858us against a 30us combine p99). Until that alignment is + clean, the p90/p95/p99 of these series describe host jitter, so they are withheld (null) + rather than published as operation tails. The chained family is unaffected. + """ + if not graph_replay or percentiles is None: + return percentiles + return {key: (value if key == "p50" else None) for key, value in percentiles.items()} + + def _component(percentiles, count, *, derived=False, origin=None): """One component block: availability, the reduction behind it, percentiles, sample count. @@ -1183,7 +1197,9 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # stand-in is decoupled from each pair's dispatch, so chained and drained are not # comparable -- see the call site for the measurement that established this. chain_output_applicable = not backend.stage_excluded_from_roundtrip - cuda_graph_output_applicable = cuda_graph and not backend.stage_excluded_from_roundtrip + # Every graphed row is value-checked: `graph_replay_output` stages inside its own capture, so + # the comparison is defined even where the timed captures hoist staging. + cuda_graph_output_applicable = cuda_graph # ---- Pass 2: every backend uses the same rotated point order. # Per-iteration cross-rank MAX samples are pooled across trials. ---- @@ -1230,20 +1246,20 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> samples[T].dispatch_min += _reduce_vec(torch, dist, device, measured["dispatch"], MIN) samples[T].combine_min += _reduce_vec(torch, dist, device, measured["combine"], MIN) - # The existing roundtrip measurement is graph replay in graph mode. Verify that a replay - # overwrote its poisoned output, and, where staging was not hoisted, compare that output with - # an ordinary drained pair. These checks are untimed and add no parallel measurement path. + # Graph mode: verify that the timed replays overwrote their poisoned output, then value-check + # replay itself -- a capture with staging inside it, dispatch outputs and result poisoned + # before its only replay -- against an ordinary drained pair. Untimed; collective in ladder + # order on every rank. if cuda_graph: for T in ladder: problem = problems[T] rewritten = bool(getattr(problem, "_cuda_graph_output_rewritten", False)) gate[T]["cuda_graph_output_rewritten"] &= int(rewritten) if cuda_graph_output_applicable: + replayed = backend.graph_replay_output(problem) drained = backend.run_roundtrip(problem) torch.cuda.synchronize() - output_ok, output_error = _chain_output_matches( - problem._cuda_graph_output, drained - ) + output_ok, output_error = _chain_output_matches(replayed, drained) gate[T]["cuda_graph_output_local_ok"] &= int(output_ok) gate[T]["cuda_graph_output_error"] = max( gate[T]["cuda_graph_output_error"], output_error @@ -1389,6 +1405,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> rstats = g["rstats"] d, s, c, rt = samples[T].dispatch, samples[T].stage, samples[T].combine, samples[T].roundtrip dp, sp, cp, rtp = _pcts(d), _pcts(s), _pcts(c), _pcts(rt) + pub = lambda pcts: _published_tails(pcts, cuda_graph) # noqa: E731 # isolated_sum = SUM of the isolated dispatch+stage+combine percentiles. Stage contributes # zero when it is explicitly not applicable. This is NOT a measured chained operation # (can't reveal shared sync / launch amortization / overlap) — do NOT use for throughput @@ -1438,8 +1455,8 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> max_rel = _reduce_vec(torch, dist, device, [g["max_rel"]], MAX)[0] point_ok = bool(global_ok) and recv_total > 0 throughput = { - percentile_name: gt / (latency_us * 1e-6) - for percentile_name, latency_us in rtp.items() + percentile_name: (gt / (latency_us * 1e-6) if latency_us is not None else None) + for percentile_name, latency_us in pub(rtp).items() } # Canonical LOGICAL payload bytes come from the routing trace (NOT backend recv # tensors): one copy per unique (token, dest-rank) pair. Dispatch carries the @@ -1494,18 +1511,18 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> rows.append({ "components": { "combine": _component( - cp, len(c), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None + pub(cp), len(c), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None ), "dispatch": _component( - dp, len(d), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None + pub(dp), len(d), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None ), - "isolated_sum": _component(isum, 0, derived=True), + "isolated_sum": _component(pub(isum), 0, derived=True), # What a serving decode loop pays per MoE layer: the steady-state period of # back-to-back dispatch->combine pairs, every backend, cross-rank median. Not # `roundtrip` (drained around every pair, an idle-pipeline latency). Do not sum it. "pair_period": _component(chainp, len(chain), origin=CHAIN_PERIOD_ORIGIN), "roundtrip": _component( - rtp, len(rt), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None + pub(rtp), len(rt), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None ), "stage": _component(sp, len(s)), }, @@ -1531,15 +1548,15 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # never the operation getting faster. "cross_rank_min_us": { "combine": _component( - _pcts(samples[T].combine_min), len(samples[T].combine_min), + pub(_pcts(samples[T].combine_min)), len(samples[T].combine_min), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None, ), "dispatch": _component( - _pcts(samples[T].dispatch_min), len(samples[T].dispatch_min), + pub(_pcts(samples[T].dispatch_min)), len(samples[T].dispatch_min), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None, ), "roundtrip": _component( - _pcts(samples[T].roundtrip_min), len(samples[T].roundtrip_min), + pub(_pcts(samples[T].roundtrip_min)), len(samples[T].roundtrip_min), origin=CUDA_GRAPH_ORIGIN if cuda_graph else None, ), }, @@ -1774,6 +1791,8 @@ def _point_summary(row): percentiles = row["components"]["dispatch"]["percentiles_us"] if not percentiles: return f"T={row['tokens_per_rank']}:n/a{period_summary}" + if percentiles.get("p99") is None: + return f"T={row['tokens_per_rank']}:disp_p50={percentiles['p50']:.1f}us{period_summary}" return (f"T={row['tokens_per_rank']}:disp_p99={percentiles['p99']:.1f}us" f"{period_summary}") diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index f7238b7b3..b3c72adcc 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,11 +42,12 @@ class MoRIBackend(EPBackend): maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # Eager in both modes. The kernels themselves are capturable (host-built args only; the reset - # moved on-device in ROCm/mori#86 for vLLM's graphs), but ROCm torch rejects the external - # events the graph timing records as graph nodes ("External events are disallowed in rocm"), - # so replay could not be timed with the same windows as every CUDA backend. - CUDA_GRAPH_MODES = () + # Both kernel families launch with host-built args only (no per-call host read of counts; the + # reset moved on-device in ROCm/mori#86 for vLLM's graphs), and MoRI's own benchmark captures + # IntraNode dispatch/combine and N-pair graphs. `stage` slices by the untimed per-rung + # `recv_tokens`, fixed for a rung's routing and so safe to bake into a capture. The timing + # events are captured through hipEventRecordWithFlags on ROCm (EPBackend._graph_event). + CUDA_GRAPH_MODES = ("normal", "low-latency") requires_fresh_pair = True def __init__(self, args, rank, world_size, local_rank, device): diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index a8742878e..9bde74e7c 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -95,7 +95,9 @@ class NCCLEPBackend(EPBackend): # per-row discriminator that change lacked. "v02" marks the nccl-extensions v0.2 mover # (new kernels: LL combine fence, B200 EP16 fix, HT gains) so pre-upgrade rows never # pool with post-upgrade rows. - kernel_generation = "nccl-ep-v02-ht-routed-zc" + # "-static" marks the combine input bound to the full static receive plane (see + # `_bind_ht_recv_count`); "-zc" rows before it sliced that input to the received count. + kernel_generation = "nccl-ep-v02-ht-routed-zc-static" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16",) # LL replays; HT stays eager. Graphed zero-copy HT failed the combine oracle intermittently @@ -135,7 +137,9 @@ def __init__(self, args, rank, world_size, local_rank, device): if self._ll_expert_major: # Weighted source-side combine over a per-expert padded receive: deepep-v2 LL's # contract, so this is the like-for-like row against the DeepEP-API backends. - self.kernel_generation = "nccl-ep-v02-ll" + # "-em" separates this from the pre-#3370 "nccl-ep-v02-ll" rows, whose timed windows + # also carried a handle.complete() per op; v0.2 needs complete() only after send_only. + self.kernel_generation = "nccl-ep-v02-ll-em" self.receive_layout = "token-expert" self.combine_weight_semantics = "weighted-kernel-sum" elif self._ll: @@ -407,18 +411,19 @@ def _ensure_handle(self, p): return h def _bind_ht_recv_count(self, h): - """Read HT's received-token count and pre-wrap the combine input at that size. - - Upstream sizes the combine staging copy from the tensor it is handed (`num_tokens = - x->sizes[0]`), not from the group's buffer, so handing it the whole ladder-max plane put a - rung-independent floor under HT combine -- ~470-1295us on a prefill leg (ladder max 8192). - Slicing is a free leading-dim view and matches upstream's own ep_test. Both callers are - untimed (handle creation and rebind), so the `.item()` read never lands in a window. + """Read HT's received-token count and bind the combine input to the full receive plane. + + The FLAT contract (ep_enums.h, NCCL_EP_LAYOUT_FLAT) gives combine the SAME + `[num_recv_slots, hidden]` shape as the dispatch output, static at the group's + `max_recv_tokens_per_rank` -- "Required under CUDA Graph capture"; sizing it to the + received count is only valid for a group created with `max_recv_tokens_per_rank = + NCCL_EP_AUTO`, which this one is not. An earlier revision sliced it to the count, which + broke that contract in both regimes. The slice existed to dodge a whole-plane staging + copy (~470-1295us on prefill); zero-copy HT elides that staging, so the full plane costs + nothing. The count is still read here (untimed) for `recv_tokens` and the oracle. """ 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._window_t(self._recv_x[: max(h.count, 1)]) + h.combine_in_t = self._recv_x_t def _rebind(self, h): """Point the single handle at h's routing (collective; untimed callers only). diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 799747495..732d0a8a3 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -28,6 +28,7 @@ """ from __future__ import annotations +import os import sys import types @@ -143,6 +144,25 @@ class UCCLEPBackend(EPBackend): receive_layout = "token-rank" combine_weight_semantics = "unweighted-rank-sum" + # (product, precision) cases measured faster eager than graphed: b200 FP8 low-latency pair + # period 0.98x baseline eager vs 1.04x graphed, every rung (runs 36114371399, 36113759089). + _EAGER_CASES = frozenset({("b200", "fp8")}) + + @property + def cuda_graph_supported(self) -> bool: + if not super().cuda_graph_supported: + return False + args = self.args + # Intranode only: at EP8 the LL kernels take the IPC path. Scale-out LL runs through the + # CPU proxy, which a replay reaches only via GPU-written queues -- never validated under + # capture -- and whose adaptive sleeper is woken only by a host call replay skips. + if self.world_size > int(getattr(args, "scale_up_domain", self.world_size)): + return False + if os.environ.get("UCCL_RDMA_ADAPTIVE_SLEEP", "0") not in ("", "0"): + return False + product = str(getattr(args, "runner", "")).split("-")[0] + return (product, self.precision) not in self._EAGER_CASES + def __init__(self, args, rank, world_size, local_rank, device): super().__init__(args, rank, world_size, local_rank, device) self.group = dist.group.WORLD diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 2f2fff2ce..aca228512 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -309,34 +309,56 @@ rather than per-operation costs. The paired roundtrip is the comparable quantity ### CUDA Graph Replay Serving engines capture their decode step, so graph-compatible backend/mode pairs are measured -under `CUDAGraph.replay()` by default. The graphed set follows what each library supports without -changing its contract: nccl-ep (low-latency), flashinfer-ep (normal), uccl-ep (low-latency), and -deepep-v2 low-latency plus normal-mode **decode**, which runs ElasticBuffer as -vLLM's graphed `deepep_v2` decode does (`do_cpu_sync=False`, worst-case receive, valid prefix read -from the handle on device; kernel generation `v2-elastic-buffer-nosync`). deepep-v2 normal prefill -keeps the host sync that sizes its receive exactly and stays eager, as does uccl-ep normal mode, -whose dispatch host-syncs unless padded to `num_worst_tokens`. MoRI stays eager in both modes: -ROCm torch rejects the external events the replay windows are recorded with. nccl-ep normal (HT) -stays eager because graphed zero-copy HT failed the combine oracle intermittently across nodes on -x86, while eager zero-copy HT was correct everywhere and no slower. +under `CUDAGraph.replay()` by default. The graphed set is each library's best measured +configuration that passed every check, without changing its contract: + +- **nccl-ep** low-latency. HT stays eager: graphed HT failed the combine oracle intermittently + across x86 nodes while eager HT was correct and no slower. +- **flashinfer-ep** decode. Prefill stays eager; graphs change nothing there. +- **uccl-ep** low-latency, intranode only, except b200 FP8, which measured faster eager. Normal + mode host-syncs unless padded to `num_worst_tokens` and stays eager. +- **deepep-v2** low-latency and normal-mode **decode**. Normal decode runs ElasticBuffer as + vLLM's graphed `deepep_v2` decode does: `do_cpu_sync=False`, receive sized to the next power of + two of T, valid prefix read from the handle on device. Its kernel generation is + `v2-elastic-buffer-nosync`. Normal prefill keeps the host sync that sizes its receive exactly + and stays eager. +- **MoRI** both modes. ROCm torch before 2.13 rejects external events, so the timing events are + captured with `hipEventRecordWithFlags(..., hipEventRecordExternal)`, the call newer torch makes. Every family keeps its eager meaning under replay; only the launch mechanism changes: -- **Fresh-entry components** (`roundtrip`, `dispatch`, `combine`) capture one pair with external - event nodes around the timed window, so host launch cost never enters it. Each timed replay - starts behind a device-side rank barrier (an all-reduce followed by a fixed spin, with no host - sync before the replay), so ranks enter together instead of from their own host's last - synchronize: without it, b200 EP16 ranks entered ~75 µs apart and the cross-rank MAX reported - the stagger as latency. Component origin is `cuda-graph-replay`. +- **Fresh-entry components** (`roundtrip`, `dispatch`, `combine`) capture one pair with event + nodes around the timed window, so host launch cost never enters it. + - Each timed replay starts behind a device-side rank barrier: an all-reduce, then a spin of fixed + **wall time**, with no host sync before the replay. The spin is converted to cycles by + measuring each GPU's spin rate per timed series, because `torch.cuda._sleep` counts SM cycles + and a fixed count left ~15 µs of skew between differently clocked ranks. + - Without the barrier, b200 EP16 ranks entered ~75 µs apart and the cross-rank MAX reported the + stagger as latency. A rank whose host is late to launch still stalls the others, and those + stalls own the tail. + - These series therefore publish **only p50** under replay: p90/p95/p99 are null and so is the + matching token rate. Component origin is `cuda-graph-replay`. - **The chained family** (`pair_period`, chain floors, chain health) captures each sibling chain as ONE graph of `chain_iters` unrolled pairs, the shape a decode graph has, and replays it once - untimed and once behind the barrier. Event records are graph nodes and cost the host nothing, - so the floors sibling carries op windows without the eager six-events-per-pair defect. The - chained oracle and the chained-output check apply unchanged. + untimed and once behind the barrier. + - Event records are graph nodes and cost the host nothing, so the floors sibling carries op + windows without the eager six-events-per-pair defect. + - The chained oracle and the chained-output check apply unchanged, and the period keeps its + tails. `stage` is not separately timed under replay. `COLLX_CUDA_GRAPH=0` restores the eager pipeline. -Every captured output is poisoned after timing and replayed once more; a finite rewrite gates the -case, and where staging is not hoisted that replay is also compared with an untimed drained pair. + +The value check runs on every graphed row: + +- Each timed capture's output is poisoned after timing and replayed once more; a finite rewrite + gates the case. +- A separate untimed capture then stages INSIDE the graph (the timed captures hoist staging where + it does device work). +- Before its only timed replay, what dispatch wrote and the combined output are overwritten with + 0xFF bytes (NaN for every float payload). +- The replay's result must match an eager drained pair, so a replay that skipped or stalely reused + its dispatch fails `cuda_graph_last_output_passed`. + A graphed row's `kernel_generation` carries a `-cudagraph` suffix, so the durable store never pools graph-replayed and eager samples of one kernel family into a single series. The artifact also records `implementation.cuda_graph_replay` and `cuda_graph_supported`. @@ -440,8 +462,11 @@ it (as NVIDIA's own `ep_bench` does: CUDA events around dispatch and combine onl outside the loop) on the argument that its capacity-proportional cost would import a ladder-max term into dispatch; that argument describes exactly what production pays, since engines size the handle to their max token capacity and update it per step. The timed window now includes the -update; rows carry `kernel_generation` `nccl-ep-v02-ht-routed-zc` -(`nccl-ep-v02-ll-rm-zc` for scale-up low-latency and `nccl-ep-v02-ll-rm` for scale-out). +update; rows carry `kernel_generation` `nccl-ep-v02-ht-routed-zc-static` +(`nccl-ep-v02-ll-rm-zc` for scale-up low-latency and `nccl-ep-v02-ll-rm` for scale-out; +`nccl-ep-v02-ll-em` under `COLLX_NCCL_LL_LAYOUT=expert-major`). `-static` marks HT combine taking +the full static receive plane the FLAT contract requires; earlier `-zc` rows sliced it to the +received count. The `v02` component discriminates the `nccl-extensions` v0.2 mover from earlier wheels, and pre-change `nccl-ep-ht`/`nccl-ep-ht-routed` rows are a different measurement contract or mover — the per-row discriminator the earliest NCCL changes lacked. HT uses zero-copy. LL uses the diff --git a/experimental/CollectiveX/tests/test_backends.py b/experimental/CollectiveX/tests/test_backends.py index df4baf938..01434d57e 100644 --- a/experimental/CollectiveX/tests/test_backends.py +++ b/experimental/CollectiveX/tests/test_backends.py @@ -277,7 +277,7 @@ def base_init(instance, options, rank, world_size, local_rank, device): self.assertEqual(em._layout, module.Layout.EXPERT_MAJOR) self.assertEqual( (em.kernel_generation, em.receive_layout, em.combine_weight_semantics), - ("nccl-ep-v02-ll", "token-expert", "weighted-kernel-sum"), + ("nccl-ep-v02-ll-em", "token-expert", "weighted-kernel-sum"), ) self.assertFalse(em.zero_copy) self.assertEqual(getattr(em, "combine_reduction", "domain-fp32"), "domain-fp32") @@ -543,6 +543,7 @@ def backend(ll=True): # create_buffer always runs before the first _ensure_handle, so the HT receive plane exists # by then; a list stands in for the tensor because `_t` is identity here. b._recv_x = list(range(64)) + b._recv_x_t = ("window", "full-plane") return b @@ -573,15 +574,14 @@ def test_ll_rank_major_weight_wrapper_is_built_once_per_handle(self): self.assertIs(ll._ensure_handle(pa).in_weights_t, first_weights) ll._t.assert_not_called() - 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: the whole ladder-max - receive plane put a rung-independent floor under it. LL keeps the full padded plane.""" + def test_ht_combine_input_is_the_full_static_receive_plane(self): + """The FLAT contract gives combine the dispatch output's static [num_recv_slots, hidden] + shape (required under graph capture); a count-sized slice needs an AUTO-sized group.""" 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))) - self.assertLess(len(h.combine_in_t), len(b._recv_x)) + self.assertIs(h.combine_in_t, b._recv_x_t) def _deepep_v2_stubs(): @@ -608,18 +608,33 @@ def _backend(self, **updates): sys.modules.pop("ep_deepep_v2", None) return ep_deepep_v2.DeepEPV2Backend(args(**updates), 0, 8, 0, "cpu") - def _dispatched_cpu_sync(self, backend): + def _dispatch_kwargs(self, backend, tokens=3): calls = [] def dispatch(*_args, **kwargs): - calls.append(kwargs["do_cpu_sync"]) + calls.append(kwargs) return "recv_x", "recv_idx", "recv_w", "handle", None backend.buffer = types.SimpleNamespace(dispatch=dispatch) - backend.max_tokens, backend.num_sms, backend.num_qps = 8, 1, 1 - backend.dispatch(types.SimpleNamespace(dispatch_x="x", topk_idx="i", topk_weights="w")) + backend.max_tokens, backend.num_sms, backend.num_qps = 512, 1, 1 + backend.dispatch(types.SimpleNamespace( + T=tokens, dispatch_x="x", topk_idx="i", topk_weights="w", + )) return calls[0] + def _dispatched_cpu_sync(self, backend): + return self._dispatch_kwargs(backend)["do_cpu_sync"] + + def test_no_sync_decode_sizes_the_receive_to_the_next_power_of_two(self): + # Worst-case sizing is num_max_tokens_per_rank * num_ranks rows; the ladder maximum made + # every rung receive (and FP8-dequantize) the T=512 plane. vLLM rounds the batch up. + backend = self._backend(mode="normal", phase="decode") + for tokens, capacity in ((1, 1), (3, 4), (64, 64), (65, 128), (512, 512)): + kwargs = self._dispatch_kwargs(backend, tokens) + self.assertEqual(kwargs["num_max_tokens_per_rank"], capacity) + prefill = self._backend(mode="normal", phase="prefill") + self.assertEqual(self._dispatch_kwargs(prefill, 3)["num_max_tokens_per_rank"], 512) + def test_normal_decode_is_the_no_sync_graphed_contract(self): backend = self._backend(mode="normal", phase="decode") self.assertIs(self._dispatched_cpu_sync(backend), False) @@ -636,5 +651,57 @@ def test_low_latency_stays_graphed(self): backend = self._backend(mode="low-latency", phase="decode") self.assertTrue(backend.cuda_graph_supported) + +class PerCaseGraphGates(unittest.TestCase): + """Graph replay is each adapter's default only where it was measured best and safe.""" + + def _load(self, name, extra): + torch = types.ModuleType("torch") + dist = types.ModuleType("torch.distributed") + dist.group = types.SimpleNamespace(WORLD="world") + torch.distributed = dist + torch.compile = lambda *a, **k: (lambda fn: fn) + modules = {"torch": torch, "torch.distributed": dist, **extra} + with mock.patch.dict(sys.modules, modules): + sys.modules.pop(name, None) + module = __import__(name) + sys.modules.pop(name, None) + return module + + def _instance(self, cls, mode, world_size=8, **fields): + backend = object.__new__(cls) + backend.mode, backend.world_size = mode, world_size + backend.precision = fields.pop("precision", "bf16") + backend.args = types.SimpleNamespace(scale_up_domain=8, runner="h200-dgxc", **fields) + return backend + + def test_uccl_low_latency_graphs_intranode_except_b200_fp8(self): + deep_ep = types.ModuleType("deep_ep") + deep_ep.Buffer, deep_ep.Config = object, object + module = self._load("ep_uccl", {"deep_ep": deep_ep}) + cls = module.UCCLEPBackend + with mock.patch.dict(os.environ, {}, clear=True): + self.assertTrue(self._instance(cls, "low-latency").cuda_graph_supported) + self.assertFalse(self._instance(cls, "normal").cuda_graph_supported) + self.assertFalse( + self._instance(cls, "low-latency", world_size=16).cuda_graph_supported + ) + b200_fp8 = self._instance(cls, "low-latency", precision="fp8") + b200_fp8.args.runner = "b200-nscale" + self.assertFalse(b200_fp8.cuda_graph_supported) + b200_fp8.precision = "bf16" + self.assertTrue(b200_fp8.cuda_graph_supported) + with mock.patch.dict(os.environ, {"UCCL_RDMA_ADAPTIVE_SLEEP": "1"}): + self.assertFalse(self._instance(cls, "low-latency").cuda_graph_supported) + + def test_flashinfer_graphs_decode_only(self): + module = self._load("ep_flashinfer", {}) + cls = module.FlashInferEPBackend if hasattr(module, "FlashInferEPBackend") else next( + value for value in vars(module).values() + if isinstance(value, type) and issubclass(value, EPBackend) and value is not EPBackend + ) + self.assertTrue(self._instance(cls, "normal", phase="decode").cuda_graph_supported) + self.assertFalse(self._instance(cls, "normal", phase="prefill").cuda_graph_supported) + if __name__ == "__main__": unittest.main() diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index 70f3d4ce1..49b3728da 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -119,6 +119,7 @@ def graph_context(_graph, **kwargs): ), ), distributed=dist, + version=types.SimpleNamespace(hip=None, cuda="13.0"), zeros=tensor, ones=tensor, empty=tensor, full=tensor, tensor=tensor, float32="float32", float64="float64", bfloat16="bfloat16", int32="int32", isfinite=lambda _value: types.SimpleNamespace( @@ -356,7 +357,11 @@ def test_the_graph_chain_is_one_capture_of_unrolled_pairs_per_sibling(self): ) # Each graph replays once untimed, then once behind the barrier. self.assertEqual(backend.calls.count("graph_replay"), 4) - self.assertEqual(backend.calls.count("align_spin"), 2) + aligned = [ + i for i, call in enumerate(backend.calls) + if call == "graph_replay" and backend.calls[i - 2:i] == ["all_reduce", "align_spin"] + ] + self.assertEqual(len(aligned), 2) for key in ("pair", "dispatch", "combine"): self.assertEqual(len(series[key]), iters - drop) self.assertEqual(len(series["start_to_start"]), iters - drop - 1) @@ -379,6 +384,51 @@ def test_external_switch_restores_the_eager_component_pipeline(self): backend.timed_components() +class GraphAlignmentAndValueCheck(unittest.TestCase): + def test_the_alignment_spin_is_sized_to_wall_time_not_a_cycle_count(self): + # A GPU spinning twice as fast (the probe spin takes half the time) must be handed twice + # the cycles, so every rank releases after the same wall time whatever its SM clock. + spins = [] + fake = types.SimpleNamespace(cuda=types.SimpleNamespace( + _sleep=spins.append, synchronize=lambda: None, + Event=lambda **kwargs: types.SimpleNamespace( + record=lambda: None, elapsed_time=lambda other: probe_ms, + ), + )) + backend = _ChainBackend() + results = {} + for probe_ms in (0.1, 0.05): # the 200k-cycle probe took 100us, then 50us + with mock.patch.dict(sys.modules, {"torch": fake}): + backend._calibrate_align_spin() + results[probe_ms] = backend._graph_align_cycles + target = ep_backend.EPBackend._GRAPH_ALIGN_SPIN_US + self.assertEqual(results[0.1], int(200_000 * target / 100.0)) + self.assertEqual(results[0.05], 2 * results[0.1]) + + def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(self): + backend = _ChainBackend(stage_device_work=True, fp8_consume="native", precision="fp8") + order = [] + handle = types.SimpleNamespace(recv_x="recv", recv_scales=None, combine_input=None) + combined = _Combined(1.0) + graph = types.SimpleNamespace(replay=lambda: order.append("replay")) + backend.warm = lambda problem, count: order.append("warm") + backend._capture_pairs = lambda problem, staged, pairs, marks: ( + order.append(("capture", staged, marks)) or (graph, {}, combined, handle) + ) + backend._poison = lambda tensor: order.append(("poison", tensor)) + fake = types.SimpleNamespace(cuda=types.SimpleNamespace(synchronize=lambda: None)) + with mock.patch.dict(sys.modules, {"torch": fake}): + result = backend.graph_replay_output(new_problem()) + # Staging runs INSIDE the capture (staged=None), and the poison lands between the upload + # replay and the replay whose output is returned. + self.assertEqual(order[1], ("capture", None, ())) + first, last = order.index("replay"), len(order) - 1 - order[::-1].index("replay") + poisoned = [entry for entry in order[first:last] if isinstance(entry, tuple)] + self.assertIn(("poison", "recv"), poisoned) + self.assertIn(("poison", combined), poisoned) + self.assertTrue(result.cloned) + + class EventPlacement(unittest.TestCase): """Which events each sibling chain may carry. The stub charges host work nothing, so these assert record placement in the trace rather than window values.""" @@ -665,6 +715,10 @@ def combine_transformed(self, problem, handle, transformed): class _GraphSweepBackend(_SweepBackend): CUDA_GRAPH_MODES = ("normal",) + def graph_replay_output(self, problem): + self.events.append(("graph-value-check", problem.T)) + return f"graph-{problem.T}" + def benchmark_component(self, component, problem, warmup, iters): self.events.append(("graph", component, problem.T)) problem._cuda_graph_output = f"graph-{problem.T}" @@ -921,6 +975,15 @@ def test_existing_component_fields_carry_graph_measurements(self): self.assertIs(row["correctness"]["cuda_graph_output_rewritten"], True) self.assertIs(row["correctness"]["cuda_graph_last_output_passed"], True) self.assertIs(row["correctness"]["post_chain_state_passed"], True) + # Graphed fresh-entry tails are withheld; the chained period keeps its tails. + for name in ("roundtrip", "dispatch", "combine", "isolated_sum"): + tails = row["components"][name]["percentiles_us"] + self.assertIsNotNone(tails["p50"]) + self.assertEqual([tails[k] for k in ("p90", "p95", "p99")], [None] * 3) + self.assertIsNone(row["cross_rank_min_us"]["roundtrip"]["percentiles_us"]["p99"]) + self.assertIsNotNone(row["components"]["pair_period"]["percentiles_us"]["p99"]) + checked = [event[1] for event in self.swept.events if event[0] == "graph-value-check"] + self.assertEqual(sorted(checked), sorted(LADDER)) def test_a_replay_that_does_not_rewrite_its_output_reds_the_case(self): swept = drive(backend_factory=_BrokenGraphSweepBackend) From 901fbcd61325bac2e2fa6538f23cb97725d03d68 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 03:26:30 +0800 Subject: [PATCH 07/20] CollectiveX: render withheld graph-replay tails as '-' in the step summaries --- experimental/CollectiveX/bandwidth.py | 11 +++++++---- experimental/CollectiveX/summarize.py | 3 ++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/experimental/CollectiveX/bandwidth.py b/experimental/CollectiveX/bandwidth.py index b927a580c..93f5d14ce 100644 --- a/experimental/CollectiveX/bandwidth.py +++ b/experimental/CollectiveX/bandwidth.py @@ -76,7 +76,7 @@ def _ep(document: dict) -> int: def _algbw_per_gpu(total_logical_bytes: int, latency_us: float, ep: int) -> float | None: """Per-GPU effective GB/s, or None when the latency cannot yield a rate. Bytes are the AGGREGATE world payload (routed_copies, routing.py), hence the divide by EP size.""" - if latency_us <= 0: + if latency_us is None or latency_us <= 0: return None return total_logical_bytes / (latency_us * 1e-6) / 1e9 / ep @@ -156,9 +156,12 @@ def _cell(row: dict, component: str, ep: int) -> str: return f"{component}=n/a" nbytes = _wire_bytes(row, component) p50 = _algbw_per_gpu(nbytes, percentiles["p50"], ep) - p99 = _algbw_per_gpu(nbytes, percentiles["p99"], ep) - return f"{component}=n/a" if p50 is None or p99 is None \ - else f"{component}={p50:6.1f}/{p99:<6.1f}" + p99 = _algbw_per_gpu(nbytes, percentiles.get("p99"), ep) + if p50 is None: + return f"{component}=n/a" + # Graph-replayed fresh-entry rows publish p50 only (tails withheld, see methodology). + tail = "-" if p99 is None else f"{p99:<6.1f}" + return f"{component}={p50:6.1f}/{tail}" def _sort_key(document: dict): diff --git a/experimental/CollectiveX/summarize.py b/experimental/CollectiveX/summarize.py index 47e58826b..ca498da42 100644 --- a/experimental/CollectiveX/summarize.py +++ b/experimental/CollectiveX/summarize.py @@ -102,7 +102,8 @@ def percentile(block: str, name: str) -> float | str: return (component.get("percentiles_us") or {}).get("p50", "-") return ( - row["tokens_per_rank"], latency["p50"], latency["p99"], + row["tokens_per_rank"], latency["p50"], + "-" if latency.get("p99") is None else latency["p99"], percentile("cross_rank_min_us", "roundtrip"), percentile("cross_rank_spread_us", ""), period is not None, From 3f4a146724a9f395369d9be4f04c38ab97e87157 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 03:28:00 +0800 Subject: [PATCH 08/20] CollectiveX: publish per-pass oracle verdicts so a failure can be placed before or after the measured regimes --- experimental/CollectiveX/bench/ep_harness.py | 9 +++++++++ experimental/CollectiveX/tests/test_chain.py | 5 +++++ 2 files changed, 14 insertions(+) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index c21edd000..b91297d96 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -1418,6 +1418,14 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> recv_max = _reduce_int(torch, dist, device, g["recv_local"], MAX) recv_min = _reduce_int(torch, dist, device, g["recv_local"], MIN) global_ok = _reduce_int(torch, dist, device, g["local_ok"], MIN) + # Which oracle pass failed, agreed across ranks. `max_relative_error` folds all three, so + # without these a failure cannot be placed before (Pass 1, before any timing or capture) + # or after the measured regimes -- the distinction that attributes it to them or not. + oracle_verdicts = { + name: bool(_reduce_int(torch, dist, device, int(bool(g[key]["passed"])), MIN)) + for name, key in (("pre", "oracle_pre"), ("chained", "oracle_chain"), + ("post", "oracle_post")) + } # Agreed across ranks like `passed`, not rank 0's local view. post_chain_state_passed = bool( _reduce_int(torch, dist, device, g["chain_local_ok"], MIN) @@ -1593,6 +1601,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # Max elementwise relative error (COMBINE_MAG_FLOOR-clamped) # against the BF16-faithful expected combine. "max_relative_error": max_rel, + "oracle_passed": oracle_verdicts, "passed": point_ok, }, "global_tokens": gt, diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index 49b3728da..4fd9fbecf 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -855,6 +855,11 @@ def test_the_drained_oracles_still_red_the_case_on_their_own(self): self.assertIs(row["correctness"]["passed"], False) self.assertIs(row["correctness"]["post_chain_state_passed"], True) self.assertIs(row["correctness"]["chain_last_output_passed"], True) + # The per-pass verdicts place the failure in the pass that produced it. + self.assertEqual( + row["correctness"]["oracle_passed"], + {"pre": phase != "pre", "chained": True, "post": phase != "post"}, + ) def test_the_output_check_is_skipped_where_staging_is_hoisted(self): # Under the hoist the chain captures one warm-up dispatch's staged stand-in and reuses From 858dc2d0209c92c074b8dba3ba95dfbe6220f7d7 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 05:03:49 +0800 Subject: [PATCH 09/20] CollectiveX: barrier the graph replay value check after poisoning; report a non-finite output mismatch as inf Zero-copy and RDMA transports write into peers' receive buffers, so without the barrier a fast rank's checked replay landed data in a slow peer's buffer before that peer poisoned it (false failures on nccl-ep LL-zc EP8 and MoRI EP16, error published as 0.0 because NaN vanished from the cross-rank MAX). --- experimental/CollectiveX/bench/ep_backend.py | 7 +++++++ experimental/CollectiveX/bench/ep_harness.py | 4 ++++ experimental/CollectiveX/tests/test_chain.py | 18 ++++++++++++++++-- 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index d2c067335..874e84d86 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -561,6 +561,7 @@ def graph_replay_output(self, problem): The caller compares it with an eager drained pair through the same code path. """ import torch + import torch.distributed as dist self.warm(problem, 1) graph, _, combined, handle = self._capture_pairs(problem, None, 1, ()) @@ -570,6 +571,12 @@ def graph_replay_output(self, problem): self._poison(getattr(handle, field, None)) self._poison(combined) torch.cuda.synchronize() + # Every rank must finish poisoning before ANY rank replays: zero-copy and RDMA transports + # write straight into peers' receive buffers, so a fast rank's replay would otherwise land + # its payload in a slow peer's buffer before that peer poisoned it, and the poison would + # overwrite fresh data (seen as NaN output on nccl-ep LL-zc EP8 and MoRI EP16). + dist.barrier() + torch.cuda.synchronize() graph.replay() torch.cuda.synchronize() return combined.clone() diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index b91297d96..fe2b514ff 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -725,6 +725,10 @@ def _chain_output_matches(chained, drained): error = (chained.float() - drained.float()).abs() relative = error / drained.float().abs().clamp_min(COMBINE_MAG_FLOOR) worst = float(relative.max().item()) + # torch.max propagates NaN, so a non-finite element lands here. Report it as inf: NaN would + # vanish from the cross-rank MAX and publish a failed check beside an error of 0.0. + if not math.isfinite(worst): + return False, float("inf") return worst < COMBINE_REL_TOL, worst diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index 4fd9fbecf..11cdb5282 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -416,8 +416,11 @@ def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(sel order.append(("capture", staged, marks)) or (graph, {}, combined, handle) ) backend._poison = lambda tensor: order.append(("poison", tensor)) - fake = types.SimpleNamespace(cuda=types.SimpleNamespace(synchronize=lambda: None)) - with mock.patch.dict(sys.modules, {"torch": fake}): + dist = types.SimpleNamespace(barrier=lambda: order.append("barrier")) + fake = types.SimpleNamespace( + cuda=types.SimpleNamespace(synchronize=lambda: None), distributed=dist, + ) + with mock.patch.dict(sys.modules, {"torch": fake, "torch.distributed": dist}): result = backend.graph_replay_output(new_problem()) # Staging runs INSIDE the capture (staged=None), and the poison lands between the upload # replay and the replay whose output is returned. @@ -426,6 +429,11 @@ def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(sel poisoned = [entry for entry in order[first:last] if isinstance(entry, tuple)] self.assertIn(("poison", "recv"), poisoned) self.assertIn(("poison", combined), poisoned) + # Every rank finishes poisoning before any rank replays: peers write into each other's + # receive buffers, so an unbarriered replay races a slow peer's poison. + last_poison = max(i for i, entry in enumerate(order) if isinstance(entry, tuple) + and entry[0] == "poison") + self.assertIn("barrier", order[last_poison:last]) self.assertTrue(result.cloned) @@ -1091,6 +1099,12 @@ def test_the_verdict_and_the_magnitude_together(self): self.assertIs(got, ok) self.assertAlmostEqual(error, expected_error) + def test_a_non_finite_output_is_an_unbounded_mismatch_not_zero_error(self): + # NaN would vanish from the cross-rank MAX and publish "failed, error 0.0". + got, error = ep_harness._chain_output_matches(_Vec([float("nan"), 1.0]), _Vec([1.0, 1.0])) + self.assertIs(got, False) + self.assertEqual(error, float("inf")) + def test_near_zero_elements_are_judged_against_the_magnitude_floor(self): # Relative error against a denominator of 1e-6 would be huge; the floor keeps # numerically-tiny elements from redding a healthy chain. From 94f43b536880aed954facaa5e7fbe1d05a0cbf68 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 08:02:47 +0800 Subject: [PATCH 10/20] CollectiveX: drop MoRI low-latency on mi325x -- it fails eager on that gfx942 pool (run 36189510167), independent of graph replay --- experimental/CollectiveX/configs/platform_config.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 014bdf4f9..7b7293358 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -159,7 +159,7 @@ "scale_up_transport": "xgmi", "launcher": "mi-amds", "backends": {"mori": [8], "uccl-ep": [8]}, - "ll_backends": {"mori": [8]}, + "ll_backends": {}, "fabric": { "nic": "n/a (single-node scale-up)", "switch": "n/a (single-node scale-up)" From 7e7bc343485c740f3111e74f5034e1abc36a8da6 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:22:06 +0800 Subject: [PATCH 11/20] CollectiveX: publish which oracle sub-checks failed, per pass --- experimental/CollectiveX/bench/ep_harness.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index fe2b514ff..ebaac4092 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -1425,11 +1425,19 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # Which oracle pass failed, agreed across ranks. `max_relative_error` folds all three, so # without these a failure cannot be placed before (Pass 1, before any timing or capture) # or after the measured regimes -- the distinction that attributes it to them or not. - oracle_verdicts = { - name: bool(_reduce_int(torch, dist, device, int(bool(g[key]["passed"])), MIN)) - for name, key in (("pre", "oracle_pre"), ("chained", "oracle_chain"), - ("post", "oracle_post")) - } + oracle_verdicts, oracle_failed_checks = {}, {} + for name, key in (("pre", "oracle_pre"), ("chained", "oracle_chain"), + ("post", "oracle_post")): + report = g[key] + oracle_verdicts[name] = bool( + _reduce_int(torch, dist, device, int(bool(report["passed"])), MIN) + ) + # Which sub-checks failed on ANY rank (dispatch payload/metadata/counts vs the combine + # values): the first thing a failed pass has to answer, collective on every rank. + oracle_failed_checks[name] = [ + check for check in _ORACLE_CHECKS + if not _reduce_int(torch, dist, device, int(bool(report["checks"][check])), MIN) + ] # Agreed across ranks like `passed`, not rank 0's local view. post_chain_state_passed = bool( _reduce_int(torch, dist, device, g["chain_local_ok"], MIN) @@ -1606,6 +1614,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # against the BF16-faithful expected combine. "max_relative_error": max_rel, "oracle_passed": oracle_verdicts, + "oracle_failed_checks": oracle_failed_checks, "passed": point_ok, }, "global_tokens": gt, From d143198090c2ddab84a37711bee926866cf7815f Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:35:03 +0800 Subject: [PATCH 12/20] CollectiveX: move MoRI graph replay to its own PR; MoRI stays eager here --- experimental/CollectiveX/bench/ep_backend.py | 58 ++------------------ experimental/CollectiveX/bench/ep_mori.py | 9 +-- experimental/CollectiveX/docs/methodology.md | 4 +- 3 files changed, 10 insertions(+), 61 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index 874e84d86..ee7c1ab66 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -433,64 +433,16 @@ def _graph_align(self): dist.all_reduce(token) torch.cuda._sleep(self._graph_align_cycles) - def _graph_event(self): - """A timing event whose record() can become a node of a graph being captured. - - CUDA torch does this with `external=True`. ROCm torch before 2.13 rejects external events - ("External events are disallowed in rocm") although HIP >= 7 supports them, so there the - event is created normally, recorded once outside capture so it exists and counts as - recorded, and captured with `hipEventRecordWithFlags(..., hipEventRecordExternal)` -- - the call torch 2.13 itself makes (pytorch#178264). - """ - import torch - - if not getattr(torch.version, "hip", None): - return torch.cuda.Event(enable_timing=True, external=True) - event = torch.cuda.Event(enable_timing=True) - event.record() - event._collx_hip_external = True - return event - @staticmethod - def _record_graph_event(event): + def _graph_event(): + """A timing event whose record() becomes a node of the graph being captured.""" import torch - if not getattr(event, "_collx_hip_external", False): - event.record() - return - import ctypes - - hip = EPBackend._hip_runtime() - rc = hip.hipEventRecordWithFlags( - ctypes.c_void_p(event.cuda_event), - ctypes.c_void_p(torch.cuda.current_stream().cuda_stream), - ctypes.c_uint(0x1), # hipEventRecordExternal - ) - if rc != 0: - raise RuntimeError(f"hipEventRecordWithFlags(external) failed with hipError {rc}") + return torch.cuda.Event(enable_timing=True, external=True) @staticmethod - def _hip_runtime(): - lib = getattr(EPBackend, "_hip_lib", None) - if lib is None: - import ctypes - import os as _os - - import torch - - candidates = ["libamdhip64.so", _os.path.join(_os.path.dirname(torch.__file__), "lib", - "libamdhip64.so")] - for name in candidates: - try: - lib = ctypes.CDLL(name) - break - except OSError: - continue - if lib is None: - raise RuntimeError("libamdhip64.so not loadable for graph event capture") - lib.hipEventRecordWithFlags.restype = ctypes.c_int - EPBackend._hip_lib = lib - return lib + def _record_graph_event(event): + event.record() def _capture_pairs(self, problem, staged, pairs, marks): """Capture `pairs` back-to-back dispatch -> combine pairs into one graph. diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index b3c72adcc..f34776526 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,12 +42,9 @@ class MoRIBackend(EPBackend): maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # Both kernel families launch with host-built args only (no per-call host read of counts; the - # reset moved on-device in ROCm/mori#86 for vLLM's graphs), and MoRI's own benchmark captures - # IntraNode dispatch/combine and N-pair graphs. `stage` slices by the untimed per-rung - # `recv_tokens`, fixed for a rung's routing and so safe to bake into a capture. The timing - # events are captured through hipEventRecordWithFlags on ROCm (EPBackend._graph_event). - CUDA_GRAPH_MODES = ("normal", "low-latency") + # Eager here; graph replay for MoRI is its own change (ROCm torch before 2.13 rejects the + # external events the replay windows are recorded with, so it needs its own event path). + CUDA_GRAPH_MODES = () requires_fresh_pair = True def __init__(self, args, rank, world_size, local_rank, device): diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index aca228512..93deff4f2 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -322,8 +322,8 @@ configuration that passed every check, without changing its contract: two of T, valid prefix read from the handle on device. Its kernel generation is `v2-elastic-buffer-nosync`. Normal prefill keeps the host sync that sizes its receive exactly and stays eager. -- **MoRI** both modes. ROCm torch before 2.13 rejects external events, so the timing events are - captured with `hipEventRecordWithFlags(..., hipEventRecordExternal)`, the call newer torch makes. +- **MoRI** stays eager: ROCm torch before 2.13 rejects the external events the replay windows are + recorded with. Every family keeps its eager meaning under replay; only the launch mechanism changes: From fd9431afaf8a477f996e1baf737d061927373058 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:35:05 +0800 Subject: [PATCH 13/20] CollectiveX: graph-replay MoRI, capturing timing events via hipEventRecordWithFlags on ROCm ROCm torch before 2.13 rejects Event(external=True) ('External events are disallowed in rocm') although HIP >= 7 supports external event records in a capture. Create the event normally, record it once outside capture, and capture its record with hipEventRecordWithFlags(..., hipEventRecordExternal) -- the call torch 2.13 makes (pytorch#178264). MoRI's kernels launch with host-built args only, so both kernel families replay. --- experimental/CollectiveX/bench/ep_backend.py | 58 ++++++++++++++++++-- experimental/CollectiveX/bench/ep_mori.py | 9 ++- experimental/CollectiveX/docs/methodology.md | 4 +- 3 files changed, 61 insertions(+), 10 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index ee7c1ab66..874e84d86 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -433,16 +433,64 @@ def _graph_align(self): dist.all_reduce(token) torch.cuda._sleep(self._graph_align_cycles) - @staticmethod - def _graph_event(): - """A timing event whose record() becomes a node of the graph being captured.""" + def _graph_event(self): + """A timing event whose record() can become a node of a graph being captured. + + CUDA torch does this with `external=True`. ROCm torch before 2.13 rejects external events + ("External events are disallowed in rocm") although HIP >= 7 supports them, so there the + event is created normally, recorded once outside capture so it exists and counts as + recorded, and captured with `hipEventRecordWithFlags(..., hipEventRecordExternal)` -- + the call torch 2.13 itself makes (pytorch#178264). + """ import torch - return torch.cuda.Event(enable_timing=True, external=True) + if not getattr(torch.version, "hip", None): + return torch.cuda.Event(enable_timing=True, external=True) + event = torch.cuda.Event(enable_timing=True) + event.record() + event._collx_hip_external = True + return event @staticmethod def _record_graph_event(event): - event.record() + import torch + + if not getattr(event, "_collx_hip_external", False): + event.record() + return + import ctypes + + hip = EPBackend._hip_runtime() + rc = hip.hipEventRecordWithFlags( + ctypes.c_void_p(event.cuda_event), + ctypes.c_void_p(torch.cuda.current_stream().cuda_stream), + ctypes.c_uint(0x1), # hipEventRecordExternal + ) + if rc != 0: + raise RuntimeError(f"hipEventRecordWithFlags(external) failed with hipError {rc}") + + @staticmethod + def _hip_runtime(): + lib = getattr(EPBackend, "_hip_lib", None) + if lib is None: + import ctypes + import os as _os + + import torch + + candidates = ["libamdhip64.so", _os.path.join(_os.path.dirname(torch.__file__), "lib", + "libamdhip64.so")] + for name in candidates: + try: + lib = ctypes.CDLL(name) + break + except OSError: + continue + if lib is None: + raise RuntimeError("libamdhip64.so not loadable for graph event capture") + lib.hipEventRecordWithFlags.restype = ctypes.c_int + EPBackend._hip_lib = lib + return lib def _capture_pairs(self, problem, staged, pairs, marks): """Capture `pairs` back-to-back dispatch -> combine pairs into one graph. diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index f34776526..b3c72adcc 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,9 +42,12 @@ class MoRIBackend(EPBackend): maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # Eager here; graph replay for MoRI is its own change (ROCm torch before 2.13 rejects the - # external events the replay windows are recorded with, so it needs its own event path). - CUDA_GRAPH_MODES = () + # Both kernel families launch with host-built args only (no per-call host read of counts; the + # reset moved on-device in ROCm/mori#86 for vLLM's graphs), and MoRI's own benchmark captures + # IntraNode dispatch/combine and N-pair graphs. `stage` slices by the untimed per-rung + # `recv_tokens`, fixed for a rung's routing and so safe to bake into a capture. The timing + # events are captured through hipEventRecordWithFlags on ROCm (EPBackend._graph_event). + CUDA_GRAPH_MODES = ("normal", "low-latency") requires_fresh_pair = True def __init__(self, args, rank, world_size, local_rank, device): diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 93deff4f2..aca228512 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -322,8 +322,8 @@ configuration that passed every check, without changing its contract: two of T, valid prefix read from the handle on device. Its kernel generation is `v2-elastic-buffer-nosync`. Normal prefill keeps the host sync that sizes its receive exactly and stays eager. -- **MoRI** stays eager: ROCm torch before 2.13 rejects the external events the replay windows are - recorded with. +- **MoRI** both modes. ROCm torch before 2.13 rejects external events, so the timing events are + captured with `hipEventRecordWithFlags(..., hipEventRecordExternal)`, the call newer torch makes. Every family keeps its eager meaning under replay; only the launch mechanism changes: From 2e0f0923705534e5da3fb275a6dbbabf4c6e5504 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:38:04 +0800 Subject: [PATCH 14/20] CollectiveX: test the ROCm graph-event record path --- experimental/CollectiveX/bench/ep_backend.py | 3 +- experimental/CollectiveX/tests/test_chain.py | 45 ++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index 874e84d86..fa7b7d205 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -433,7 +433,8 @@ def _graph_align(self): dist.all_reduce(token) torch.cuda._sleep(self._graph_align_cycles) - def _graph_event(self): + @staticmethod + def _graph_event(): """A timing event whose record() can become a node of a graph being captured. CUDA torch does this with `external=True`. ROCm torch before 2.13 rejects external events diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index 11cdb5282..85154b5c1 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -437,6 +437,51 @@ def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(sel self.assertTrue(result.cloned) +class RocmGraphEventRecord(unittest.TestCase): + """ROCm torch < 2.13 rejects Event(external=True); the capture records through HIP instead.""" + + def _torch(self, hip, records): + class Event: + def __init__(self, **kwargs): + if kwargs.get("external"): + raise RuntimeError("External events are disallowed in rocm") + self.cuda_event = 0xE0 + + def record(self): + records.append("record") + + return types.SimpleNamespace( + version=types.SimpleNamespace(hip=hip), + cuda=types.SimpleNamespace( + Event=Event, + current_stream=lambda: types.SimpleNamespace(cuda_stream=0x5), + ), + ) + + def test_rocm_events_are_recorded_once_outside_then_captured_through_hip(self): + records, hip_calls = [], [] + fake_hip = types.SimpleNamespace( + hipEventRecordWithFlags=lambda event, stream, flags: hip_calls.append( + (event.value, stream.value, flags.value) + ) or 0 + ) + with mock.patch.dict(sys.modules, {"torch": self._torch("7.2", records)}), \ + mock.patch.object(ep_backend.EPBackend, "_hip_runtime", lambda: fake_hip): + event = ep_backend.EPBackend._graph_event() + self.assertEqual(records, ["record"]) # materialized outside capture + ep_backend.EPBackend._record_graph_event(event) + self.assertEqual(records, ["record"]) # the captured record went through HIP + self.assertEqual(hip_calls, [(0xE0, 0x5, 0x1)]) # hipEventRecordExternal + + def test_a_failed_hip_record_raises(self): + fake_hip = types.SimpleNamespace(hipEventRecordWithFlags=lambda *args: 1) + with mock.patch.dict(sys.modules, {"torch": self._torch("7.2", [])}), \ + mock.patch.object(ep_backend.EPBackend, "_hip_runtime", lambda: fake_hip): + event = ep_backend.EPBackend._graph_event() + with self.assertRaisesRegex(RuntimeError, "hipEventRecordWithFlags"): + ep_backend.EPBackend._record_graph_event(event) + + class EventPlacement(unittest.TestCase): """Which events each sibling chain may carry. The stub charges host work nothing, so these assert record placement in the trace rather than window values.""" From b7ae179fc3064a8b9c9da7b3a206adea8be4f89a Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 18:27:34 +0800 Subject: [PATCH 15/20] CollectiveX: fence the nccl-ep HT oracle's combine-input write across ranks The oracle writes its transformed rows into the zero-copy receive window and then combines; nothing ordered that local write against peers' combine. After graph replay shifted rank timing the combine read a peer's half-written input (combine_values only, EP16, h100/h200 3/3 runs). Fenced, the same graphed cells pass 4/4. Graphed HT stays off: correct now, but 1.03-1.11x eager's pair period. --- experimental/CollectiveX/bench/ep_nccl.py | 17 +++++++++++++---- experimental/CollectiveX/docs/methodology.md | 5 +++-- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index 9bde74e7c..778c07aba 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -100,10 +100,10 @@ class NCCLEPBackend(EPBackend): kernel_generation = "nccl-ep-v02-ht-routed-zc-static" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16",) - # LL replays; HT stays eager. Graphed zero-copy HT failed the combine oracle intermittently - # across nodes on x86 (b200 EP16 T=128, h200 EP16 T=32 and prefill T=1024; runs 35994093313, - # 36113759089) while eager zero-copy HT passed every cell and was as fast or faster - # (run 36114371399), so eager is the better HT configuration on every pool measured. + # LL replays; HT stays eager. Graphed HT's oracle failures were the oracle's own unfenced + # write into the zero-copy window (fixed in `combine_transformed`); with that fixed, graphed + # HT is correct but 1.03-1.11x eager's pair period on h100/h200 EP16 (runs 36231927003.. + # 36231931954 vs 36176100175), so eager is HT's best configuration. CUDA_GRAPH_MODES = ("low-latency",) stage_device_work = False requires_fresh_pair = False @@ -660,6 +660,15 @@ def combine_transformed(self, p, h, transformed): # destination ranks back to each token's home rank. self._recv_x.zero_() self._recv_x[: transformed.shape[0]].copy_(transformed.to(self._recv_x.dtype)) + # Fence the write across ranks. `_recv_x` is the zero-copy window peers access directly, + # and nothing orders this rank's local write against a peer's combine touching it: after + # graph replay shifted rank timing, the combine read a peer's half-written input + # (combine_values failed, dispatch checks clean, EP16 only; h100/h200 3/3 runs). With this + # fence the same cells passed 4/4 (runs 36231927003..36231931954). The timed path writes + # nothing between dispatch and combine, so it needs no fence; this is the oracle's write. + torch.cuda.synchronize() + dist.barrier() + torch.cuda.synchronize() stream = self._stream() h.handle.combine( # Same sliced input the timed path uses, so the two cannot diverge in shape. diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 93deff4f2..cd842302e 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -312,8 +312,9 @@ Serving engines capture their decode step, so graph-compatible backend/mode pair under `CUDAGraph.replay()` by default. The graphed set is each library's best measured configuration that passed every check, without changing its contract: -- **nccl-ep** low-latency. HT stays eager: graphed HT failed the combine oracle intermittently - across x86 nodes while eager HT was correct and no slower. +- **nccl-ep** low-latency. HT stays eager: graphed HT is correct but 3-11% slower than eager on + h100/h200 EP16. (Its earlier oracle failures were the oracle writing its combine input into the + zero-copy window without a cross-rank fence; that write is now fenced.) - **flashinfer-ep** decode. Prefill stays eager; graphs change nothing there. - **uccl-ep** low-latency, intranode only, except b200 FP8, which measured faster eager. Normal mode host-syncs unless padded to `num_worst_tokens` and stays eager. From 394075ed36e085077e5b24bee094db6b3b02ba11 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:50:19 +0800 Subject: [PATCH 16/20] CollectiveX: name a pool's runners in the platform registry instead of the workflow Each matrix cell carries runs-on as `runner`: the platform's `runner_label`, else the SKU. MI325X declares cluster:mi325x-amds there, replacing the SKU special case in collectivex-sweep.yml. --- .github/workflows/collectivex-sweep.yml | 6 +++--- experimental/CollectiveX/configs/platform_config.json | 1 + experimental/CollectiveX/swap_matrix.py | 1 + experimental/CollectiveX/sweep_matrix.py | 8 ++++++++ experimental/CollectiveX/tests/test_matrix.py | 7 +++++++ experimental/CollectiveX/tests/test_swap_matrix.py | 5 +++++ 6 files changed, 25 insertions(+), 3 deletions(-) diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 5cfe0ba1d..b5f03502e 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -151,7 +151,7 @@ jobs: vars.NODE_SLOT_SCHEDULER_ENABLED == 'true' && format( '["self-hosted",{0},{1},{2},{3}]', - toJSON(matrix.sku == 'mi325x' && 'cluster:mi325x-amds' || matrix.sku), + toJSON(matrix.runner), toJSON(format('nodes:{0}', matrix.nodes)), toJSON(format( 'ci-job-{0}-{1}', @@ -162,7 +162,7 @@ jobs: ) || format( '["self-hosted",{0},{1},{2},{3}]', - toJSON(matrix.sku == 'mi325x' && 'cluster:mi325x-amds' || matrix.sku), + toJSON(matrix.runner), toJSON(format('nodes:{0}', matrix.nodes)), toJSON(format( 'ci-job-{0}-{1}', @@ -172,7 +172,7 @@ jobs: toJSON(format('ci-attempt-{0}', github.run_attempt)) ) ) || - format('[{0}]', toJSON(matrix.sku == 'mi325x' && 'cluster:mi325x-amds' || matrix.sku)) + format('[{0}]', toJSON(matrix.runner)) ) }} name: p${{ needs.setup.outputs.priority }} | ${{ matrix.sku }} ${{ matrix.backend }} shard ${{ matrix.id }} timeout-minutes: 350 diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 7b7293358..598dc3c67 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -158,6 +158,7 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-amds", + "runner_label": "cluster:mi325x-amds", "backends": {"mori": [8], "uccl-ep": [8]}, "ll_backends": {}, "fabric": { diff --git a/experimental/CollectiveX/swap_matrix.py b/experimental/CollectiveX/swap_matrix.py index 040376c15..121083bfb 100644 --- a/experimental/CollectiveX/swap_matrix.py +++ b/experimental/CollectiveX/swap_matrix.py @@ -23,6 +23,7 @@ def build_matrix(platforms: dict, only_sku: str, exclude_skus: str) -> dict: { "id": f"swap-{sku}", "sku": sku, + "runner": platform.get("runner_label", sku), "backend": "swap-blocks", "nodes": 1, "gpus_per_node": 1, diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index 919e88283..1e672037f 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -28,6 +28,13 @@ def _load_config(name: str) -> dict[str, Any]: PLATFORMS = _load_config("platform_config.json")["platforms"] # Per-backend production/candidate map for the matrix and docs; see EPBackend.maturity. BACKEND_MATURITY = _load_config("platform_config.json")["backend_maturity"] + + +def _runner_label(sku: str) -> str: + """The runs-on label for a pool: the SKU itself unless the registry names its runners.""" + return PLATFORMS[sku].get("runner_label", sku) + + SWEEP_BACKENDS = tuple(dict.fromkeys( backend for platform in PLATFORMS.values() for backend in platform["backends"] )) @@ -252,6 +259,7 @@ def resolve_matrix( shards_by_sku.setdefault(sku, []).append({ "id": f"{sku}-{target}{mode_segment}-{precision}-n{nodes}", "sku": sku, + "runner": _runner_label(sku), "backend": target, "mode": mode, "launcher": PLATFORMS[sku]["launcher"], diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index 6f143c54c..5416c5285 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -31,6 +31,13 @@ def test_every_shard_has_an_exact_positive_node_request(self): {shard["nodes"]}, ) + def test_shards_run_on_the_registry_runner_label_else_the_sku(self): + for shard in matrix(backend="all")["include"]: + with self.subTest(shard=shard["id"]): + platform = sweep_matrix.PLATFORMS[shard["sku"]] + self.assertEqual(shard["runner"], platform.get("runner_label", shard["sku"])) + self.assertEqual(sweep_matrix._runner_label("mi325x"), "cluster:mi325x-amds") + def test_only_real_platform_cells_are_unsupported(self): platform = { "product": "test-gpu", "gpus_per_node": 8, "scale_up_domain": 8, diff --git a/experimental/CollectiveX/tests/test_swap_matrix.py b/experimental/CollectiveX/tests/test_swap_matrix.py index a9aa8480a..d8ad34fe1 100644 --- a/experimental/CollectiveX/tests/test_swap_matrix.py +++ b/experimental/CollectiveX/tests/test_swap_matrix.py @@ -20,6 +20,7 @@ def test_selection_preserves_vendor_and_single_gpu_allocations(self): { "id": "swap-amd-test", "sku": "amd-test", + "runner": "amd-test", "backend": "swap-blocks", "nodes": 1, "gpus_per_node": 1, @@ -33,6 +34,10 @@ def test_selection_preserves_vendor_and_single_gpu_allocations(self): self.assertEqual( build_matrix(platforms, "cuda-test", "")["include"][0]["vendor"], "nvidia" ) + labelled = {"amd-test": {"arch": "gfx942", "runner_label": "cluster:amd-pool"}} + self.assertEqual( + build_matrix(labelled, "amd-test", "")["include"][0]["runner"], "cluster:amd-pool" + ) for only, exclude in [ ("missing", ""), ("", "missing"), From a4367bf001376f4701c175c060e3b4f66034a90d Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:50:20 +0800 Subject: [PATCH 17/20] CollectiveX: inline the graph-alignment spin as a calibration argument --- experimental/CollectiveX/bench/ep_backend.py | 19 ++++++++++--------- experimental/CollectiveX/tests/test_chain.py | 3 +-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index ee7c1ab66..dee1eea69 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -395,18 +395,19 @@ def _topk_idx_dtype(self): # ---- CUDA graph capture ---------------------------------------------------------------- - # Wall time the stream spins after the alignment all-reduce, so every rank's host has enqueued - # its replay before the stream reaches it and the replay start is set by the barrier release, - # not by host launch latency. Converted to cycles per GPU by `_calibrate_align_spin`: - # `torch.cuda._sleep` counts SM cycles, so a fixed cycle count spun 48-70us across the clock - # range and left ~15us of cross-rank skew on gb200 (a fast-clocked rank released early). - _GRAPH_ALIGN_SPIN_US = 100.0 # Attributes of a dispatch handle that hold what dispatch wrote; the replay check poisons them # so a replay that skipped (or stalely reused) the dispatch cannot reproduce a valid output. _DISPATCH_OUTPUT_FIELDS = ("recv_x", "recv_scales", "dispatch_output") - def _calibrate_align_spin(self): - """Measure this GPU's current spin rate and size the alignment spin to a wall time.""" + def _calibrate_align_spin(self, spin_us=100.0): + """Size the post-barrier alignment spin to `spin_us` of wall time on this GPU. + + The stream spins after the alignment all-reduce so every rank's host has enqueued its + replay before the stream reaches it: the replay start is set by the barrier release, not + by host launch latency. `torch.cuda._sleep` counts SM cycles, so a fixed cycle count spun + 48-70us across the clock range and left ~15us of cross-rank skew on gb200 (a fast-clocked + rank released early); measuring the spin rate converts the wall time to cycles per GPU. + """ import torch probe = 200_000 @@ -418,7 +419,7 @@ def _calibrate_align_spin(self): end.record() torch.cuda.synchronize() elapsed_us = max(start.elapsed_time(end) * 1000.0, 1e-3) - self._graph_align_cycles = max(1, int(probe * self._GRAPH_ALIGN_SPIN_US / elapsed_us)) + self._graph_align_cycles = max(1, int(probe * spin_us / elapsed_us)) def _graph_align(self): """Enqueue a device-side rank barrier on the current stream, without a host sync.""" diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index 11cdb5282..30eb63fef 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -401,8 +401,7 @@ def test_the_alignment_spin_is_sized_to_wall_time_not_a_cycle_count(self): with mock.patch.dict(sys.modules, {"torch": fake}): backend._calibrate_align_spin() results[probe_ms] = backend._graph_align_cycles - target = ep_backend.EPBackend._GRAPH_ALIGN_SPIN_US - self.assertEqual(results[0.1], int(200_000 * target / 100.0)) + self.assertEqual(results[0.1], 200_000) # the default 100us spin at 2 cycles/ns self.assertEqual(results[0.05], 2 * results[0.1]) def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(self): From 6fa861ceb086d0bd00774b565e2f72b50b38063e Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sat, 26 Sep 2026 22:50:06 +0800 Subject: [PATCH 18/20] CollectiveX: replay nccl-ep HT at decode, as a captured decode step runs it HT prefill stays eager. Graphed HT decode measures 1.03-1.11x eager's pair period on h100/h200 EP16; the cost is NCCL's per-replay host-callback node on the captured cross-node routing all-gather, which a captured decode step pays too. --- experimental/CollectiveX/bench/ep_nccl.py | 20 ++++++++++++++----- experimental/CollectiveX/docs/methodology.md | 15 +++++++++----- .../CollectiveX/tests/test_backends.py | 9 +++++++++ 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index 778c07aba..2ecf6128f 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -100,11 +100,9 @@ class NCCLEPBackend(EPBackend): kernel_generation = "nccl-ep-v02-ht-routed-zc-static" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16",) - # LL replays; HT stays eager. Graphed HT's oracle failures were the oracle's own unfenced - # write into the zero-copy window (fixed in `combine_transformed`); with that fixed, graphed - # HT is correct but 1.03-1.11x eager's pair period on h100/h200 EP16 (runs 36231927003.. - # 36231931954 vs 36176100175), so eager is HT's best configuration. - CUDA_GRAPH_MODES = ("low-latency",) + # HT replays at decode only, the regime an engine's captured decode step would run it in + # (see `cuda_graph_supported`); LL replays everywhere. + CUDA_GRAPH_MODES = ("normal", "low-latency") stage_device_work = False requires_fresh_pair = False receive_layout = "token-rank" @@ -112,6 +110,18 @@ class NCCLEPBackend(EPBackend): zero_copy = True _ll_expert_major = False + @property + def cuda_graph_supported(self) -> bool: + # HT prefill stays eager: engines run prefill uncaptured. HT decode replays because a + # captured decode step would run it that way, even though it is slower there: 1.03-1.11x + # eager's pair period on h100/h200 EP16 (runs 36231927003..36231931954 vs 36176100175). + # Captured, the per-step routing ncclAllGather is a proxy-driven cross-node collective, + # which NCCL fronts with a host-callback node on every replay (enqueue.cc, persistent + # plans), ~+50us of dispatch at EP16 and nothing within one node. + if self.mode == "normal" and getattr(self.args, "phase", None) != "decode": + return False + return super().cuda_graph_supported + def __init__(self, args, rank, world_size, local_rank, device): super().__init__(args, rank, world_size, local_rank, device) # NCCL EP group creation requires the NCCL Device API (LSA symmetric memory), which NCCL diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index cd842302e..193c2169a 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -310,11 +310,16 @@ rather than per-operation costs. The paired roundtrip is the comparable quantity Serving engines capture their decode step, so graph-compatible backend/mode pairs are measured under `CUDAGraph.replay()` by default. The graphed set is each library's best measured -configuration that passed every check, without changing its contract: - -- **nccl-ep** low-latency. HT stays eager: graphed HT is correct but 3-11% slower than eager on - h100/h200 EP16. (Its earlier oracle failures were the oracle writing its combine input into the - zero-copy window without a cross-rank fence; that write is now fenced.) +configuration that passed every check, without changing its contract. The one exception is nccl-ep HT +decode, below: + +- **nccl-ep** low-latency and HT **decode**. HT decode replays because a captured decode step + would run it that way, although it measures 3-11% slower than eager on h100/h200 EP16: captured, + the per-step routing `ncclAllGather` is a cross-node collective driven by NCCL's proxy, which + NCCL fronts with a host-callback node on every replay (about +50us of dispatch at EP16, nothing + within one node). HT prefill stays eager, as engines run prefill uncaptured. (HT's earlier graphed + oracle failures were the oracle writing its combine input into the zero-copy window without a + cross-rank fence; that write is now fenced.) - **flashinfer-ep** decode. Prefill stays eager; graphs change nothing there. - **uccl-ep** low-latency, intranode only, except b200 FP8, which measured faster eager. Normal mode host-syncs unless padded to `num_worst_tokens` and stays eager. diff --git a/experimental/CollectiveX/tests/test_backends.py b/experimental/CollectiveX/tests/test_backends.py index 01434d57e..997ea7321 100644 --- a/experimental/CollectiveX/tests/test_backends.py +++ b/experimental/CollectiveX/tests/test_backends.py @@ -703,5 +703,14 @@ def test_flashinfer_graphs_decode_only(self): self.assertTrue(self._instance(cls, "normal", phase="decode").cuda_graph_supported) self.assertFalse(self._instance(cls, "normal", phase="prefill").cuda_graph_supported) + def test_nccl_graphs_low_latency_and_ht_decode_only(self): + with mock.patch.dict(sys.modules, _stub_modules()): + import importlib + import ep_nccl + cls = importlib.reload(ep_nccl).NCCLEPBackend + self.assertTrue(self._instance(cls, "normal", phase="decode").cuda_graph_supported) + self.assertFalse(self._instance(cls, "normal", phase="prefill").cuda_graph_supported) + self.assertTrue(self._instance(cls, "low-latency", phase="decode").cuda_graph_supported) + if __name__ == "__main__": unittest.main() From 7fc7c2ab44c68d164fd35617323afd5cb973f352 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 27 Sep 2026 12:05:00 +0800 Subject: [PATCH 19/20] CollectiveX: trim the graph-replay change Fold the per-backend graph-gate tests into table-driven cases, shorten comments and docstrings to what the code needs (rationale stays in methodology), compress the CUDA Graph Replay section, and inline the event wrappers only the MoRI change needs. No behavior change. --- experimental/CollectiveX/bench/ep_backend.py | 61 ++---- .../CollectiveX/bench/ep_deepep_v2.py | 23 +-- .../CollectiveX/bench/ep_flashinfer.py | 4 +- experimental/CollectiveX/bench/ep_harness.py | 45 ++--- experimental/CollectiveX/bench/ep_nccl.py | 29 +-- experimental/CollectiveX/bench/ep_uccl.py | 10 +- experimental/CollectiveX/docs/methodology.md | 84 +++----- experimental/CollectiveX/sweep_matrix.py | 8 +- .../CollectiveX/tests/test_backends.py | 183 +++++++----------- experimental/CollectiveX/tests/test_chain.py | 84 ++++---- experimental/CollectiveX/tests/test_matrix.py | 8 +- 11 files changed, 189 insertions(+), 350 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index dee1eea69..694deb8b5 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -395,18 +395,14 @@ def _topk_idx_dtype(self): # ---- CUDA graph capture ---------------------------------------------------------------- - # Attributes of a dispatch handle that hold what dispatch wrote; the replay check poisons them - # so a replay that skipped (or stalely reused) the dispatch cannot reproduce a valid output. + # Handle attributes dispatch writes; the replay check poisons them (see graph_replay_output). _DISPATCH_OUTPUT_FIELDS = ("recv_x", "recv_scales", "dispatch_output") def _calibrate_align_spin(self, spin_us=100.0): """Size the post-barrier alignment spin to `spin_us` of wall time on this GPU. - The stream spins after the alignment all-reduce so every rank's host has enqueued its - replay before the stream reaches it: the replay start is set by the barrier release, not - by host launch latency. `torch.cuda._sleep` counts SM cycles, so a fixed cycle count spun - 48-70us across the clock range and left ~15us of cross-rank skew on gb200 (a fast-clocked - rank released early); measuring the spin rate converts the wall time to cycles per GPU. + The spin lets every host enqueue its replay before the stream reaches it. `_sleep` counts + SM cycles, so a fixed count left ~15us of skew between differently clocked gb200 ranks. """ import torch @@ -434,36 +430,23 @@ def _graph_align(self): dist.all_reduce(token) torch.cuda._sleep(self._graph_align_cycles) - @staticmethod - def _graph_event(): - """A timing event whose record() becomes a node of the graph being captured.""" - import torch - - return torch.cuda.Event(enable_timing=True, external=True) - - @staticmethod - def _record_graph_event(event): - event.record() - def _capture_pairs(self, problem, staged, pairs, marks): """Capture `pairs` back-to-back dispatch -> combine pairs into one graph. - `marks` selects which windows get event nodes: "pair" (the whole pair), "dispatch", - "combine". Event records are graph nodes, so they cost the stream nothing on the host -- - the six-events-per-pair defect the eager chain splits around does not exist here. - Returns (graph, {mark: (starts, ends)}, last combined output, last dispatch handle). + `marks` picks the windows that get event nodes ("pair", "dispatch", "combine"). Returns + (graph, {mark: (starts, ends)}, last combined output, last dispatch handle). """ import torch import torch.distributed as dist def events(): - return [self._graph_event() for _ in range(pairs)] + return [torch.cuda.Event(enable_timing=True, external=True) for _ in range(pairs)] stamps = {mark: (events(), events()) for mark in marks} def record(mark, edge, i): if mark in stamps: - self._record_graph_event(stamps[mark][edge][i]) + stamps[mark][edge][i].record() dist.barrier() torch.cuda.synchronize() @@ -505,13 +488,10 @@ def _poison(tensor): tensor.fill_(float("nan") if tensor.is_floating_point() else -1) def graph_replay_output(self, problem): - """The value check for graph replay: one untimed capture with `stage` INSIDE the graph. + """Graph replay's value check: an untimed capture with `stage` INSIDE the graph. - The timed captures hoist staging where `stage` does device work, so their output never - depends on that replay's dispatch and a stale replay would still look correct. This one - stages per pair, then poisons what dispatch wrote and the combined output before its only - replay: the result is valid only if the replay itself re-ran dispatch, stage and combine. - The caller compares it with an eager drained pair through the same code path. + Dispatch's output and the result are poisoned before the only replay, so the returned + output is valid only if that replay re-ran dispatch, stage and combine. """ import torch import torch.distributed as dist @@ -524,10 +504,8 @@ def graph_replay_output(self, problem): self._poison(getattr(handle, field, None)) self._poison(combined) torch.cuda.synchronize() - # Every rank must finish poisoning before ANY rank replays: zero-copy and RDMA transports - # write straight into peers' receive buffers, so a fast rank's replay would otherwise land - # its payload in a slow peer's buffer before that peer poisoned it, and the poison would - # overwrite fresh data (seen as NaN output on nccl-ep LL-zc EP8 and MoRI EP16). + # Peers write straight into each other's receive buffers: without the barrier a fast + # rank's replay lands before a slow peer's poison, which then overwrites fresh data. dist.barrier() torch.cuda.synchronize() graph.replay() @@ -691,13 +669,9 @@ def series(starts, ends): } def _benchmark_chain_graph(self, problem, staged, iters, drop): - """The chained family under capture: each chain is ONE graph of `iters` unrolled pairs. - - That is the shape a serving decode graph has -- every layer's dispatch -> combine back to - back inside a single replay -- so the period keeps its eager meaning (free-running pairs, - entry skew amortised across the chain) with launch overhead removed. Same two siblings as - the eager chain and the same returned series, so `run_sweep` reduces both identically. - Each graph replays once untimed (first-launch upload), then once aligned and timed. + """The chained family under capture: each sibling is ONE graph of `iters` unrolled pairs, + the shape of a decode graph. Returns the eager chain's series; each graph replays once + untimed (upload), then once aligned and timed. """ import torch @@ -765,9 +739,8 @@ def benchmark_roundtrip(self, problem, warmup, iters, graph_component="roundtrip self.combine(problem, handle) # drain the pair backends require torch.cuda.synchronize() if self.cuda_graph_enabled: - # One captured pair, bracketed by external events for the timed component. Capture - # and its warm-up are excluded; each timed replay starts behind a device-side rank - # barrier (`_graph_align`) so the cross-rank MAX is the operation, not launch skew. + # One captured pair with event nodes around the timed component; each timed replay + # starts behind `_graph_align` so the cross-rank MAX is the operation, not launch skew. mark = "pair" if graph_component == "roundtrip" else graph_component graph, stamps, combined, _ = self._capture_pairs(problem, staged, 1, (mark,)) # Re-measure the spin rate per timed series: clocks move with load and temperature. diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py index c3065df9a..c9ea6438c 100644 --- a/experimental/CollectiveX/bench/ep_deepep_v2.py +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -137,9 +137,7 @@ class DeepEPV2Backend(EPBackend): kernel_generation = "v2-elastic-buffer" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # The legacy decode kernels are explicitly graph compatible. ElasticBuffer normal mode is - # graph compatible only without its host sync, which this adapter drops for the decode phase - # alone (see `_normal_cpu_sync` and `cuda_graph_supported`). + # Legacy decode kernels are graph compatible; ElasticBuffer normal only without its host sync. CUDA_GRAPH_MODES = ("low-latency",) stage_device_work = False requires_fresh_pair = False @@ -169,13 +167,8 @@ def __init__(self, args, rank, world_size, local_rank, device): # Normal/HT quantises inside the timed dispatch with the compiled form; low-latency # keeps the eager helper, whose bits its in-kernel quantise matches. See fused_quantize. self._quant = self.fused_quantize(self._to_fp8) - # Normal-mode decode runs ElasticBuffer the way vLLM's deepep_v2 decode path does - # (prepare_finalize/deepep_v2.py, use_cudagraph=True): do_expand=False, do_cpu_sync=False, - # receive sized to the worst case (num_max_tokens_per_rank * num_ranks) with the valid - # prefix read from the handle's device-side psum. That is the graph-capturable contract, - # and it lands one row per (token, destination rank) with an unweighted rank-sum combine: - # the rank-major shape. Prefill keeps the host sync that sizes the receive exactly, as - # vLLM's (uncaptured) prefill does. + # Normal decode runs ElasticBuffer as vLLM's graphed deepep_v2 decode does + # (do_cpu_sync=False, valid prefix read on device); prefill keeps the exact-size sync. self._normal_cpu_sync = self.mode == "normal" and args.phase != "decode" if self.mode == "normal" and not self._normal_cpu_sync: self.kernel_generation = "v2-elastic-buffer-nosync" @@ -198,14 +191,8 @@ def cuda_graph_supported(self) -> bool: return super().cuda_graph_supported def _dispatch_capacity(self, tokens): - """Per-call `num_max_tokens_per_rank`. - - With the host sync the receive is sized exactly, so the buffer maximum is only a bound. - Without it DeepEP allocates `num_max_tokens_per_rank * num_ranks` receive rows - (elastic/buffer.hpp "allocate with the worst case"), so passing the ladder maximum made a - T=1 dispatch receive 4096 rows at EP8 and the FP8 stage dequantize all of them (stage - 59 -> 212us, b200 EP8). vLLM's graphed decode (prepare_finalize/deepep_v2.py) passes the - next power of two of the batch, which bounds both the receive and the JIT variants. + """Per-call `num_max_tokens_per_rank`: without the host sync DeepEP receives that many + rows per rank, so pass the next power of two of the batch, as vLLM's graphed decode does. """ if self._normal_cpu_sync: return self.max_tokens diff --git a/experimental/CollectiveX/bench/ep_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 61a2768e3..d927053f7 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -122,9 +122,7 @@ class FlashInferEPBackend(EPBackend): @property def cuda_graph_supported(self) -> bool: - # Decode only: graph replay removes host-bound library overhead at decode sizes (gb200 EP8 - # T=1 pair period 151 -> 43us) but changes nothing at prefill (T=8192 1400 vs 1387us), and - # engines capture decode, not prefill, so prefill keeps the eager series it already has. + # Decode only: replay cuts gb200 EP8 T=1 pair period 151 -> 43us but leaves prefill flat. return super().cuda_graph_supported and getattr(self.args, "phase", None) == "decode" def __init__(self, args, rank, world_size, local_rank, device): diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py index ebaac4092..c4e776d4c 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -216,11 +216,8 @@ def _pcts(xs): def _published_tails(percentiles, graph_replay): """Fresh-entry percentiles as published: under graph replay only the median. - A graphed fresh-entry sample starts behind the alignment barrier, but a rank whose host is - late to launch its replay still stalls the others, and those stalls own the tail (gb200 - flashinfer T=1: roundtrip p99 858us against a 30us combine p99). Until that alignment is - clean, the p90/p95/p99 of these series describe host jitter, so they are withheld (null) - rather than published as operation tails. The chained family is unaffected. + A rank whose host is late to launch its replay stalls the others, so graphed fresh-entry + tails measure host jitter (gb200 flashinfer T=1 roundtrip p99 858us vs combine p99 30us). """ if not graph_replay or percentiles is None: return percentiles @@ -336,12 +333,8 @@ def time_cuda_graph_phase_us( ) -> list[float]: """Time one event-record interval captured inside graph replay. - `interval` is a pair of external events recorded as graph nodes, so the host's replay launch - never lands in the window. `align()` runs before each replay with no host sync between the - two: it enqueues a device-side rank barrier, so every rank's replay starts when that barrier - releases instead of when its own host got round to launching the graph. Without it each - sample restarts from the preceding synchronize and ranks enter ~75us apart across nodes - (b200 EP16), which the cross-rank MAX then reports as latency. + `align()` enqueues a device-side rank barrier before each replay, so replays start together + rather than ~75us apart (b200 EP16), which the cross-rank MAX would report as latency. """ for _ in range(max(0, warmup)): fn() @@ -357,12 +350,7 @@ def time_cuda_graph_phase_us( def kernel_generation(backend) -> str: - """Return the adapter's declared kernel family, suffixed when timed under graph replay. - - Replay removes launch overhead that eager timing pays, so the two regimes are different - series: the suffix keeps the durable store from pooling a graphed row with the eager rows - published under the same kernel family. - """ + """Return the adapter's kernel family; `-cudagraph` keeps replayed rows a separate series.""" family = getattr(backend, "kernel_generation", None) or "n-a" if getattr(backend, "cuda_graph_enabled", False): return f"{family}-cudagraph" @@ -725,8 +713,7 @@ def _chain_output_matches(chained, drained): error = (chained.float() - drained.float()).abs() relative = error / drained.float().abs().clamp_min(COMBINE_MAG_FLOOR) worst = float(relative.max().item()) - # torch.max propagates NaN, so a non-finite element lands here. Report it as inf: NaN would - # vanish from the cross-rank MAX and publish a failed check beside an error of 0.0. + # NaN would vanish from the cross-rank MAX and publish "failed, error 0.0". if not math.isfinite(worst): return False, float("inf") return worst < COMBINE_REL_TOL, worst @@ -1201,8 +1188,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> # stand-in is decoupled from each pair's dispatch, so chained and drained are not # comparable -- see the call site for the measurement that established this. chain_output_applicable = not backend.stage_excluded_from_roundtrip - # Every graphed row is value-checked: `graph_replay_output` stages inside its own capture, so - # the comparison is defined even where the timed captures hoist staging. + # `graph_replay_output` stages inside its own capture, so every graphed row is comparable. cuda_graph_output_applicable = cuda_graph # ---- Pass 2: every backend uses the same rotated point order. @@ -1250,10 +1236,8 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> samples[T].dispatch_min += _reduce_vec(torch, dist, device, measured["dispatch"], MIN) samples[T].combine_min += _reduce_vec(torch, dist, device, measured["combine"], MIN) - # Graph mode: verify that the timed replays overwrote their poisoned output, then value-check - # replay itself -- a capture with staging inside it, dispatch outputs and result poisoned - # before its only replay -- against an ordinary drained pair. Untimed; collective in ladder - # order on every rank. + # Graph mode: the timed replays must rewrite their poisoned output, and a poisoned replay + # (graph_replay_output) must match a drained pair. Untimed, in ladder order on every rank. if cuda_graph: for T in ladder: problem = problems[T] @@ -1422,9 +1406,8 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> recv_max = _reduce_int(torch, dist, device, g["recv_local"], MAX) recv_min = _reduce_int(torch, dist, device, g["recv_local"], MIN) global_ok = _reduce_int(torch, dist, device, g["local_ok"], MIN) - # Which oracle pass failed, agreed across ranks. `max_relative_error` folds all three, so - # without these a failure cannot be placed before (Pass 1, before any timing or capture) - # or after the measured regimes -- the distinction that attributes it to them or not. + # Which oracle pass failed (before, within or after the measured regimes), and which of + # its sub-checks, agreed across ranks. `max_relative_error` folds all three passes. oracle_verdicts, oracle_failed_checks = {}, {} for name, key in (("pre", "oracle_pre"), ("chained", "oracle_chain"), ("post", "oracle_post")): @@ -1432,8 +1415,6 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> oracle_verdicts[name] = bool( _reduce_int(torch, dist, device, int(bool(report["passed"])), MIN) ) - # Which sub-checks failed on ANY rank (dispatch payload/metadata/counts vs the combine - # values): the first thing a failed pass has to answer, collective on every rank. oracle_failed_checks[name] = [ check for check in _ORACLE_CHECKS if not _reduce_int(torch, dist, device, int(bool(report["checks"][check])), MIN) @@ -1762,9 +1743,7 @@ def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> "stage_excluded_from_roundtrip": bool( getattr(backend, "stage_excluded_from_roundtrip", False) ), - # Whether this document's rows carry the chained family. Consumers key the headline on - # presence, as for `stage_excluded_from_roundtrip`. Graph mode keeps it: the chain is - # captured as one graph of unrolled pairs (EPBackend._benchmark_chain_graph). + # Whether rows carry the chained family; graph mode captures it too. "chained_period": True, "cuda_graph_replay": cuda_graph, "cuda_graph_supported": bool( diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index 2ecf6128f..26529c313 100644 --- a/experimental/CollectiveX/bench/ep_nccl.py +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -100,8 +100,6 @@ class NCCLEPBackend(EPBackend): kernel_generation = "nccl-ep-v02-ht-routed-zc-static" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16",) - # HT replays at decode only, the regime an engine's captured decode step would run it in - # (see `cuda_graph_supported`); LL replays everywhere. CUDA_GRAPH_MODES = ("normal", "low-latency") stage_device_work = False requires_fresh_pair = False @@ -112,12 +110,8 @@ class NCCLEPBackend(EPBackend): @property def cuda_graph_supported(self) -> bool: - # HT prefill stays eager: engines run prefill uncaptured. HT decode replays because a - # captured decode step would run it that way, even though it is slower there: 1.03-1.11x - # eager's pair period on h100/h200 EP16 (runs 36231927003..36231931954 vs 36176100175). - # Captured, the per-step routing ncclAllGather is a proxy-driven cross-node collective, - # which NCCL fronts with a host-callback node on every replay (enqueue.cc, persistent - # plans), ~+50us of dispatch at EP16 and nothing within one node. + # HT replays at decode only, as a captured decode step runs it (engines run prefill + # uncaptured). It is slower there at EP16: see methodology, CUDA Graph Replay. if self.mode == "normal" and getattr(self.args, "phase", None) != "decode": return False return super().cuda_graph_supported @@ -423,14 +417,9 @@ def _ensure_handle(self, p): def _bind_ht_recv_count(self, h): """Read HT's received-token count and bind the combine input to the full receive plane. - The FLAT contract (ep_enums.h, NCCL_EP_LAYOUT_FLAT) gives combine the SAME - `[num_recv_slots, hidden]` shape as the dispatch output, static at the group's - `max_recv_tokens_per_rank` -- "Required under CUDA Graph capture"; sizing it to the - received count is only valid for a group created with `max_recv_tokens_per_rank = - NCCL_EP_AUTO`, which this one is not. An earlier revision sliced it to the count, which - broke that contract in both regimes. The slice existed to dodge a whole-plane staging - copy (~470-1295us on prefill); zero-copy HT elides that staging, so the full plane costs - nothing. The count is still read here (untimed) for `recv_tokens` and the oracle. + FLAT combine takes the dispatch output's static `[num_recv_slots, hidden]` shape + (ep_enums.h; a count-sized slice needs an NCCL_EP_AUTO group). Zero-copy elides the + staging copy the old slice avoided. The count is read here, untimed, for the oracle. """ h.count = int(h.recv_total.item()) h.combine_in_t = self._recv_x_t @@ -670,12 +659,8 @@ def combine_transformed(self, p, h, transformed): # destination ranks back to each token's home rank. self._recv_x.zero_() self._recv_x[: transformed.shape[0]].copy_(transformed.to(self._recv_x.dtype)) - # Fence the write across ranks. `_recv_x` is the zero-copy window peers access directly, - # and nothing orders this rank's local write against a peer's combine touching it: after - # graph replay shifted rank timing, the combine read a peer's half-written input - # (combine_values failed, dispatch checks clean, EP16 only; h100/h200 3/3 runs). With this - # fence the same cells passed 4/4 (runs 36231927003..36231931954). The timed path writes - # nothing between dispatch and combine, so it needs no fence; this is the oracle's write. + # `_recv_x` is the zero-copy window peers read directly: fence this write across ranks + # or a peer's combine can read it half-written (EP16). The timed path writes nothing here. torch.cuda.synchronize() dist.barrier() torch.cuda.synchronize() diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 732d0a8a3..12ad1d596 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -134,10 +134,8 @@ class UCCLEPBackend(EPBackend): kernel_generation = "uccl-legacy-buffer" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # The low-latency kernels are plain launches whose only host state is the double-buffer - # toggle, baked into a capture exactly as in DeepEP; every capture holds whole pairs (an even - # call count), so the toggle lands where it started. Normal mode host-syncs on its receive - # counters unless dispatched with `num_worst_tokens`, which this adapter does not do. + # LL's only host state is the double-buffer toggle; captures hold whole pairs, so it returns + # to where it started. Normal mode host-syncs on its receive counters. CUDA_GRAPH_MODES = ("low-latency",) stage_device_work = False requires_fresh_pair = False @@ -153,9 +151,7 @@ def cuda_graph_supported(self) -> bool: if not super().cuda_graph_supported: return False args = self.args - # Intranode only: at EP8 the LL kernels take the IPC path. Scale-out LL runs through the - # CPU proxy, which a replay reaches only via GPU-written queues -- never validated under - # capture -- and whose adaptive sleeper is woken only by a host call replay skips. + # Intranode only: scale-out LL runs through the CPU proxy, unvalidated under capture. if self.world_size > int(getattr(args, "scale_up_domain", self.world_size)): return False if os.environ.get("UCCL_RDMA_ADAPTIVE_SLEEP", "0") not in ("", "0"): diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 193c2169a..3ce3da987 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -309,65 +309,39 @@ rather than per-operation costs. The paired roundtrip is the comparable quantity ### CUDA Graph Replay Serving engines capture their decode step, so graph-compatible backend/mode pairs are measured -under `CUDAGraph.replay()` by default. The graphed set is each library's best measured -configuration that passed every check, without changing its contract. The one exception is nccl-ep HT -decode, below: - -- **nccl-ep** low-latency and HT **decode**. HT decode replays because a captured decode step - would run it that way, although it measures 3-11% slower than eager on h100/h200 EP16: captured, - the per-step routing `ncclAllGather` is a cross-node collective driven by NCCL's proxy, which - NCCL fronts with a host-callback node on every replay (about +50us of dispatch at EP16, nothing - within one node). HT prefill stays eager, as engines run prefill uncaptured. (HT's earlier graphed - oracle failures were the oracle writing its combine input into the zero-copy window without a - cross-rank fence; that write is now fenced.) -- **flashinfer-ep** decode. Prefill stays eager; graphs change nothing there. -- **uccl-ep** low-latency, intranode only, except b200 FP8, which measured faster eager. Normal - mode host-syncs unless padded to `num_worst_tokens` and stays eager. -- **deepep-v2** low-latency and normal-mode **decode**. Normal decode runs ElasticBuffer as - vLLM's graphed `deepep_v2` decode does: `do_cpu_sync=False`, receive sized to the next power of - two of T, valid prefix read from the handle on device. Its kernel generation is - `v2-elastic-buffer-nosync`. Normal prefill keeps the host sync that sizes its receive exactly - and stays eager. -- **MoRI** stays eager: ROCm torch before 2.13 rejects the external events the replay windows are - recorded with. - -Every family keeps its eager meaning under replay; only the launch mechanism changes: +under `CUDAGraph.replay()` by default: each library's best measured configuration that passed +every check, without changing its contract. + +- **nccl-ep** low-latency and HT **decode**. HT decode replays as a captured decode step would + run it, although it is 3-11% slower than eager at h100/h200 EP16: captured, the per-step routing + `ncclAllGather` is a proxy-driven cross-node collective that NCCL fronts with a host-callback + node on every replay (about +50 µs of dispatch; nothing within one node). HT prefill stays eager. +- **flashinfer-ep** decode; prefill stays eager (graphs change nothing there). +- **uccl-ep** low-latency, intranode only, except b200 FP8 (faster eager). Normal mode host-syncs. +- **deepep-v2** low-latency and normal **decode**, run as vLLM's graphed `deepep_v2` decode runs + ElasticBuffer (`do_cpu_sync=False`, receive sized to the next power of two of T; kernel + generation `v2-elastic-buffer-nosync`). Normal prefill keeps its host sync and stays eager. +- **MoRI** stays eager: ROCm torch before 2.13 rejects the external timing events. + +Only the launch mechanism changes; every family keeps its eager meaning: - **Fresh-entry components** (`roundtrip`, `dispatch`, `combine`) capture one pair with event - nodes around the timed window, so host launch cost never enters it. - - Each timed replay starts behind a device-side rank barrier: an all-reduce, then a spin of fixed - **wall time**, with no host sync before the replay. The spin is converted to cycles by - measuring each GPU's spin rate per timed series, because `torch.cuda._sleep` counts SM cycles - and a fixed count left ~15 µs of skew between differently clocked ranks. - - Without the barrier, b200 EP16 ranks entered ~75 µs apart and the cross-rank MAX reported the - stagger as latency. A rank whose host is late to launch still stalls the others, and those - stalls own the tail. - - These series therefore publish **only p50** under replay: p90/p95/p99 are null and so is the - matching token rate. Component origin is `cuda-graph-replay`. -- **The chained family** (`pair_period`, chain floors, chain health) captures each sibling chain as - ONE graph of `chain_iters` unrolled pairs, the shape a decode graph has, and replays it once - untimed and once behind the barrier. - - Event records are graph nodes and cost the host nothing, so the floors sibling carries op - windows without the eager six-events-per-pair defect. - - The chained oracle and the chained-output check apply unchanged, and the period keeps its - tails. + nodes around the timed window. Each timed replay starts behind a device-side rank barrier (an + all-reduce, then a spin of fixed wall time, calibrated per GPU because `torch.cuda._sleep` + counts SM cycles); without it b200 EP16 ranks entered ~75 µs apart. A rank whose host launches + late still stalls the others, so these series publish **only p50** (p90/p95/p99 and the matching + token rate are null). Component origin is `cuda-graph-replay`. +- **The chained family** (`pair_period`, floors, health) captures each sibling chain as one graph + of `chain_iters` unrolled pairs, the shape of a decode graph, replayed once untimed and once + behind the barrier. Its oracle, output check and tails apply unchanged. `stage` is not separately timed under replay. `COLLX_CUDA_GRAPH=0` restores the eager pipeline. - -The value check runs on every graphed row: - -- Each timed capture's output is poisoned after timing and replayed once more; a finite rewrite - gates the case. -- A separate untimed capture then stages INSIDE the graph (the timed captures hoist staging where - it does device work). -- Before its only timed replay, what dispatch wrote and the combined output are overwritten with - 0xFF bytes (NaN for every float payload). -- The replay's result must match an eager drained pair, so a replay that skipped or stalely reused - its dispatch fails `cuda_graph_last_output_passed`. - -A graphed row's `kernel_generation` carries a `-cudagraph` suffix, so the durable store never -pools graph-replayed and eager samples of one kernel family into a single series. The artifact -also records `implementation.cuda_graph_replay` and `cuda_graph_supported`. +Every graphed row is value-checked: each timed capture's output is poisoned and must be rewritten +by a further replay, and a separate capture with staging inside it has dispatch's output and its +result overwritten with 0xFF bytes before its only replay, whose output must then match an eager +drained pair (`cuda_graph_last_output_passed`). Graphed rows carry a `-cudagraph` suffix on +`kernel_generation`, so the durable store never pools them with eager rows; the artifact also +records `implementation.cuda_graph_replay` and `cuda_graph_supported`. ### Chained Pair Period diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index 1e672037f..e00bea9c5 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -30,11 +30,6 @@ def _load_config(name: str) -> dict[str, Any]: BACKEND_MATURITY = _load_config("platform_config.json")["backend_maturity"] -def _runner_label(sku: str) -> str: - """The runs-on label for a pool: the SKU itself unless the registry names its runners.""" - return PLATFORMS[sku].get("runner_label", sku) - - SWEEP_BACKENDS = tuple(dict.fromkeys( backend for platform in PLATFORMS.values() for backend in platform["backends"] )) @@ -259,7 +254,8 @@ def resolve_matrix( shards_by_sku.setdefault(sku, []).append({ "id": f"{sku}-{target}{mode_segment}-{precision}-n{nodes}", "sku": sku, - "runner": _runner_label(sku), + # runs-on label: the SKU unless the registry names the pool's runners. + "runner": PLATFORMS[sku].get("runner_label", sku), "backend": target, "mode": mode, "launcher": PLATFORMS[sku]["launcher"], diff --git a/experimental/CollectiveX/tests/test_backends.py b/experimental/CollectiveX/tests/test_backends.py index 997ea7321..64ddbe88b 100644 --- a/experimental/CollectiveX/tests/test_backends.py +++ b/experimental/CollectiveX/tests/test_backends.py @@ -2,6 +2,8 @@ """EPBackend contracts: ladder/spec construction, the staging-vs-roundtrip gate, and the NCCL EP handle.""" from __future__ import annotations +import contextlib +import importlib import os import sys import types @@ -584,133 +586,96 @@ def test_ht_combine_input_is_the_full_static_receive_plane(self): self.assertIs(h.combine_in_t, b._recv_x_t) -def _deepep_v2_stubs(): - """Fake torch / deep_ep so `import ep_deepep_v2` succeeds without the benchmark image.""" +@contextlib.contextmanager +def _stubbed(name, extra=None): + """Import one adapter module against a fake torch (plus `extra` fake modules).""" torch = types.ModuleType("torch") torch.compile = lambda *a, **k: (lambda fn: fn) dist = types.ModuleType("torch.distributed") dist.group = types.SimpleNamespace(WORLD="world") torch.distributed = dist - deep_ep = types.ModuleType("deep_ep") - deep_ep.ElasticBuffer = type("ElasticBuffer", (), {}) - deep_ep.Buffer = type("Buffer", (), {}) - return {"torch": torch, "torch.distributed": dist, "deep_ep": deep_ep} + with mock.patch.dict(sys.modules, {"torch": torch, "torch.distributed": dist, **(extra or {})}): + sys.modules.pop(name, None) + yield __import__(name) + sys.modules.pop(name, None) -class DeepEPV2GraphContract(unittest.TestCase): - """Normal-mode decode drops ElasticBuffer's host sync -- vLLM's graphed deepep_v2 decode - contract -- and only that makes it graph-capturable; prefill keeps the exact-size sync.""" +def _deep_ep(*classes): + module = types.ModuleType("deep_ep") + for cls in classes: + setattr(module, cls, type(cls, (), {})) + return module - def _backend(self, **updates): - with mock.patch.dict(sys.modules, _deepep_v2_stubs()): - sys.modules.pop("ep_deepep_v2", None) - import ep_deepep_v2 - sys.modules.pop("ep_deepep_v2", None) - return ep_deepep_v2.DeepEPV2Backend(args(**updates), 0, 8, 0, "cpu") - def _dispatch_kwargs(self, backend, tokens=3): - calls = [] +def _gate(cls, mode, world_size=8, precision="bf16", runner="h200-dgxc", **fields): + """An adapter carrying only what `cuda_graph_supported` reads.""" + backend = object.__new__(cls) + backend.mode, backend.world_size, backend.precision = mode, world_size, precision + backend.args = types.SimpleNamespace(scale_up_domain=8, runner=runner, **fields) + return backend + - def dispatch(*_args, **kwargs): - calls.append(kwargs) - return "recv_x", "recv_idx", "recv_w", "handle", None - - backend.buffer = types.SimpleNamespace(dispatch=dispatch) - backend.max_tokens, backend.num_sms, backend.num_qps = 512, 1, 1 - backend.dispatch(types.SimpleNamespace( - T=tokens, dispatch_x="x", topk_idx="i", topk_weights="w", - )) - return calls[0] - - def _dispatched_cpu_sync(self, backend): - return self._dispatch_kwargs(backend)["do_cpu_sync"] - - def test_no_sync_decode_sizes_the_receive_to_the_next_power_of_two(self): - # Worst-case sizing is num_max_tokens_per_rank * num_ranks rows; the ladder maximum made - # every rung receive (and FP8-dequantize) the T=512 plane. vLLM rounds the batch up. - backend = self._backend(mode="normal", phase="decode") - for tokens, capacity in ((1, 1), (3, 4), (64, 64), (65, 128), (512, 512)): - kwargs = self._dispatch_kwargs(backend, tokens) - self.assertEqual(kwargs["num_max_tokens_per_rank"], capacity) - prefill = self._backend(mode="normal", phase="prefill") - self.assertEqual(self._dispatch_kwargs(prefill, 3)["num_max_tokens_per_rank"], 512) - - def test_normal_decode_is_the_no_sync_graphed_contract(self): - backend = self._backend(mode="normal", phase="decode") - self.assertIs(self._dispatched_cpu_sync(backend), False) - self.assertTrue(backend.cuda_graph_supported) - self.assertEqual(backend.kernel_generation, "v2-elastic-buffer-nosync") - - def test_normal_prefill_keeps_the_host_sync_and_stays_eager(self): - backend = self._backend(mode="normal", phase="prefill") - self.assertIs(self._dispatched_cpu_sync(backend), True) - self.assertFalse(backend.cuda_graph_supported) - self.assertEqual(backend.kernel_generation, "v2-elastic-buffer") - - def test_low_latency_stays_graphed(self): - backend = self._backend(mode="low-latency", phase="decode") - self.assertTrue(backend.cuda_graph_supported) - - -class PerCaseGraphGates(unittest.TestCase): +class GraphReplayDefaults(unittest.TestCase): """Graph replay is each adapter's default only where it was measured best and safe.""" - def _load(self, name, extra): - torch = types.ModuleType("torch") - dist = types.ModuleType("torch.distributed") - dist.group = types.SimpleNamespace(WORLD="world") - torch.distributed = dist - torch.compile = lambda *a, **k: (lambda fn: fn) - modules = {"torch": torch, "torch.distributed": dist, **extra} - with mock.patch.dict(sys.modules, modules): - sys.modules.pop(name, None) - module = __import__(name) - sys.modules.pop(name, None) - return module - - def _instance(self, cls, mode, world_size=8, **fields): - backend = object.__new__(cls) - backend.mode, backend.world_size = mode, world_size - backend.precision = fields.pop("precision", "bf16") - backend.args = types.SimpleNamespace(scale_up_domain=8, runner="h200-dgxc", **fields) - return backend + def test_deepep_v2_normal_decode_drops_the_host_sync_and_rounds_the_receive_up(self): + calls = [] + with _stubbed("ep_deepep_v2", {"deep_ep": _deep_ep("ElasticBuffer", "Buffer")}) as module: + def make(phase, mode="normal"): + backend = module.DeepEPV2Backend(args(mode=mode, phase=phase), 0, 8, 0, "cpu") + backend.buffer = types.SimpleNamespace( + dispatch=lambda *a, **k: calls.append(k) or ("x", "i", "w", "h", None) + ) + backend.max_tokens, backend.num_sms, backend.num_qps = 512, 1, 1 + return backend - def test_uccl_low_latency_graphs_intranode_except_b200_fp8(self): - deep_ep = types.ModuleType("deep_ep") - deep_ep.Buffer, deep_ep.Config = object, object - module = self._load("ep_uccl", {"deep_ep": deep_ep}) - cls = module.UCCLEPBackend - with mock.patch.dict(os.environ, {}, clear=True): - self.assertTrue(self._instance(cls, "low-latency").cuda_graph_supported) - self.assertFalse(self._instance(cls, "normal").cuda_graph_supported) - self.assertFalse( - self._instance(cls, "low-latency", world_size=16).cuda_graph_supported - ) - b200_fp8 = self._instance(cls, "low-latency", precision="fp8") - b200_fp8.args.runner = "b200-nscale" - self.assertFalse(b200_fp8.cuda_graph_supported) - b200_fp8.precision = "bf16" - self.assertTrue(b200_fp8.cuda_graph_supported) - with mock.patch.dict(os.environ, {"UCCL_RDMA_ADAPTIVE_SLEEP": "1"}): - self.assertFalse(self._instance(cls, "low-latency").cuda_graph_supported) + decode, prefill, ll = make("decode"), make("prefill"), make("decode", "low-latency") - def test_flashinfer_graphs_decode_only(self): - module = self._load("ep_flashinfer", {}) - cls = module.FlashInferEPBackend if hasattr(module, "FlashInferEPBackend") else next( - value for value in vars(module).values() - if isinstance(value, type) and issubclass(value, EPBackend) and value is not EPBackend + def dispatched(backend, tokens): + backend.dispatch(types.SimpleNamespace( + T=tokens, dispatch_x="x", topk_idx="i", topk_weights="w", + )) + return calls[-1]["num_max_tokens_per_rank"], calls[-1]["do_cpu_sync"] + + # vLLM's graphed decode passes the next power of two of the batch; prefill syncs exactly. + for tokens, capacity in ((1, 1), (3, 4), (65, 128), (512, 512)): + self.assertEqual(dispatched(decode, tokens), (capacity, False)) + self.assertEqual(dispatched(prefill, 3), (512, True)) + self.assertEqual( + (decode.cuda_graph_supported, decode.kernel_generation), + (True, "v2-elastic-buffer-nosync"), ) - self.assertTrue(self._instance(cls, "normal", phase="decode").cuda_graph_supported) - self.assertFalse(self._instance(cls, "normal", phase="prefill").cuda_graph_supported) + self.assertEqual( + (prefill.cuda_graph_supported, prefill.kernel_generation), (False, "v2-elastic-buffer") + ) + self.assertTrue(ll.cuda_graph_supported) - def test_nccl_graphs_low_latency_and_ht_decode_only(self): + def test_uccl_graphs_intranode_low_latency_except_b200_fp8(self): + with _stubbed("ep_uccl", {"deep_ep": _deep_ep("Buffer", "Config")}) as module: + cls = module.UCCLEPBackend + with mock.patch.dict(os.environ, {}, clear=True): + for gate, expected in ( + (_gate(cls, "low-latency"), True), + (_gate(cls, "normal"), False), + (_gate(cls, "low-latency", world_size=16), False), + (_gate(cls, "low-latency", precision="fp8", runner="b200-nscale"), False), + (_gate(cls, "low-latency", runner="b200-nscale"), True), + ): + self.assertIs(gate.cuda_graph_supported, expected) + with mock.patch.dict(os.environ, {"UCCL_RDMA_ADAPTIVE_SLEEP": "1"}): + self.assertFalse(_gate(cls, "low-latency").cuda_graph_supported) + + def test_flashinfer_and_nccl_ht_graph_decode_only(self): + with _stubbed("ep_flashinfer") as module: + flashinfer = module.FlashInferEPBackend with mock.patch.dict(sys.modules, _stub_modules()): - import importlib import ep_nccl - cls = importlib.reload(ep_nccl).NCCLEPBackend - self.assertTrue(self._instance(cls, "normal", phase="decode").cuda_graph_supported) - self.assertFalse(self._instance(cls, "normal", phase="prefill").cuda_graph_supported) - self.assertTrue(self._instance(cls, "low-latency", phase="decode").cuda_graph_supported) + nccl = importlib.reload(ep_nccl).NCCLEPBackend + for cls in (flashinfer, nccl): + self.assertTrue(_gate(cls, "normal", phase="decode").cuda_graph_supported) + self.assertFalse(_gate(cls, "normal", phase="prefill").cuda_graph_supported) + self.assertTrue(_gate(nccl, "low-latency", phase="decode").cuda_graph_supported) + if __name__ == "__main__": unittest.main() diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index 30eb63fef..777b93716 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -322,48 +322,42 @@ def test_each_graph_component_uses_its_own_capture(self): self.assertTrue(problem._cuda_graph_output.cloned) self.assertTrue(problem._cuda_graph_output_rewritten) - def test_every_timed_replay_starts_behind_a_device_side_rank_barrier(self): - # Without the barrier each replay restarts from the preceding synchronize and ranks enter - # ~75us apart across nodes; the barrier must sit between the sync and the replay, with no - # host sync in between, on every timed replay and on none of the warm-ups. + def _graph_backend(self): backend = _ChainBackend(stage_device_work=False, fp8_consume="native", precision="bf16") - backend.mode = "normal" - backend.CUDA_GRAPH_MODES = ("normal",) + backend.mode, backend.CUDA_GRAPH_MODES = "normal", ("normal",) + return backend + + def _aligned_replays(self, calls): + """Replays preceded by the device-side rank barrier (all-reduce, then the spin).""" + return [i for i, call in enumerate(calls) + if call == "graph_replay" and calls[i - 2:i] == ["all_reduce", "align_spin"]] + + def test_every_timed_replay_starts_behind_a_device_side_rank_barrier(self): + # Without it each replay restarts from the preceding sync and ranks enter ~75us apart. + backend = self._graph_backend() with mock.patch.dict(os.environ, {}, clear=True), \ trace_torch(backend.clock, backend.calls): backend.benchmark_component("roundtrip", new_problem(), warmup=2, iters=3) replays = [i for i, call in enumerate(backend.calls) if call == "graph_replay"] - warmups, timed = replays[:2], replays[2:5] - for index in warmups: - self.assertNotEqual(backend.calls[index - 1], "align_spin") - for index in timed: - self.assertEqual(backend.calls[index - 2:index], ["all_reduce", "align_spin"]) + self.assertEqual(self._aligned_replays(backend.calls), replays[2:5]) def test_the_graph_chain_is_one_capture_of_unrolled_pairs_per_sibling(self): iters, drop = 5, 1 - backend = _ChainBackend(stage_device_work=False, fp8_consume="native", precision="bf16") - backend.mode = "normal" - backend.CUDA_GRAPH_MODES = ("normal",) + backend = self._graph_backend() with mock.patch.dict(os.environ, {}, clear=True), \ trace_torch(backend.clock, backend.calls): series = backend.benchmark_chain(new_problem(), 0, iters, drop) - # Floors and period siblings: one capture each, holding every pair of the chain. - self.assertEqual(backend.calls.count("capture_begin"), 2) - begin = [i for i, call in enumerate(backend.calls) if call == "capture_begin"] - end = [i for i, call in enumerate(backend.calls) if call == "capture_end"] + calls = backend.calls + begin = [i for i, call in enumerate(calls) if call == "capture_begin"] + end = [i for i, call in enumerate(calls) if call == "capture_end"] + # Floors and period siblings: one capture each holding every pair, each replayed once + # untimed and once behind the barrier. + self.assertEqual(len(begin), 2) for lo, hi in zip(begin, end): - self.assertEqual( - ops_only(backend.calls[lo:hi]), ["dispatch", "stage", "combine"] * iters - ) - # Each graph replays once untimed, then once behind the barrier. - self.assertEqual(backend.calls.count("graph_replay"), 4) - aligned = [ - i for i, call in enumerate(backend.calls) - if call == "graph_replay" and backend.calls[i - 2:i] == ["all_reduce", "align_spin"] - ] - self.assertEqual(len(aligned), 2) - for key in ("pair", "dispatch", "combine"): - self.assertEqual(len(series[key]), iters - drop) + self.assertEqual(ops_only(calls[lo:hi]), ["dispatch", "stage", "combine"] * iters) + self.assertEqual(calls.count("graph_replay"), 4) + self.assertEqual(len(self._aligned_replays(calls)), 2) + self.assertEqual({len(series[k]) for k in ("pair", "dispatch", "combine")}, {iters - drop}) self.assertEqual(len(series["start_to_start"]), iters - drop - 1) self.assertTrue(series["combined"].cloned) @@ -406,33 +400,27 @@ def test_the_alignment_spin_is_sized_to_wall_time_not_a_cycle_count(self): def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(self): backend = _ChainBackend(stage_device_work=True, fp8_consume="native", precision="fp8") - order = [] + order, combined = [], _Combined(1.0) handle = types.SimpleNamespace(recv_x="recv", recv_scales=None, combine_input=None) - combined = _Combined(1.0) graph = types.SimpleNamespace(replay=lambda: order.append("replay")) backend.warm = lambda problem, count: order.append("warm") backend._capture_pairs = lambda problem, staged, pairs, marks: ( - order.append(("capture", staged, marks)) or (graph, {}, combined, handle) + order.append(("capture", staged)) or (graph, {}, combined, handle) ) backend._poison = lambda tensor: order.append(("poison", tensor)) dist = types.SimpleNamespace(barrier=lambda: order.append("barrier")) - fake = types.SimpleNamespace( - cuda=types.SimpleNamespace(synchronize=lambda: None), distributed=dist, - ) + fake = types.SimpleNamespace(cuda=types.SimpleNamespace(synchronize=lambda: None), + distributed=dist) with mock.patch.dict(sys.modules, {"torch": fake, "torch.distributed": dist}): result = backend.graph_replay_output(new_problem()) - # Staging runs INSIDE the capture (staged=None), and the poison lands between the upload - # replay and the replay whose output is returned. - self.assertEqual(order[1], ("capture", None, ())) - first, last = order.index("replay"), len(order) - 1 - order[::-1].index("replay") - poisoned = [entry for entry in order[first:last] if isinstance(entry, tuple)] - self.assertIn(("poison", "recv"), poisoned) - self.assertIn(("poison", combined), poisoned) - # Every rank finishes poisoning before any rank replays: peers write into each other's - # receive buffers, so an unbarriered replay races a slow peer's poison. - last_poison = max(i for i, entry in enumerate(order) if isinstance(entry, tuple) - and entry[0] == "poison") - self.assertIn("barrier", order[last_poison:last]) + # Staging runs INSIDE the capture (staged=None). Between the upload replay and the + # returned one, dispatch's output and the result are poisoned, then every rank barriers: + # peers write into each other's buffers, so an unbarriered replay races a slow poison. + self.assertEqual(order, [ + "warm", ("capture", None), "replay", + ("poison", "recv"), ("poison", None), ("poison", None), ("poison", combined), + "barrier", "replay", + ]) self.assertTrue(result.cloned) diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index 5416c5285..73d838278 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -32,11 +32,9 @@ def test_every_shard_has_an_exact_positive_node_request(self): ) def test_shards_run_on_the_registry_runner_label_else_the_sku(self): - for shard in matrix(backend="all")["include"]: - with self.subTest(shard=shard["id"]): - platform = sweep_matrix.PLATFORMS[shard["sku"]] - self.assertEqual(shard["runner"], platform.get("runner_label", shard["sku"])) - self.assertEqual(sweep_matrix._runner_label("mi325x"), "cluster:mi325x-amds") + runners = {shard["sku"]: shard["runner"] for shard in matrix(backend="all")["include"]} + self.assertEqual(runners["mi325x"], "cluster:mi325x-amds") + self.assertEqual(runners["h200-dgxc"], "h200-dgxc") def test_only_real_platform_cells_are_unsupported(self): platform = { From cbdac219e126eab5bc35c6a9cc445dda05132f03 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Sun, 27 Sep 2026 23:54:39 +0800 Subject: [PATCH 20/20] CollectiveX: trim the MoRI graph-replay change --- experimental/CollectiveX/bench/ep_backend.py | 51 +++++++------------- experimental/CollectiveX/bench/ep_mori.py | 7 +-- experimental/CollectiveX/tests/test_chain.py | 48 +++++++----------- 3 files changed, 37 insertions(+), 69 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index 1a028cdf4..2a8b2c34c 100644 --- a/experimental/CollectiveX/bench/ep_backend.py +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -3,6 +3,7 @@ from __future__ import annotations import abc +import functools import os import types from dataclasses import dataclass, field @@ -432,14 +433,9 @@ def _graph_align(self): @staticmethod def _graph_event(): - """A timing event whose record() can become a node of a graph being captured. - - CUDA torch does this with `external=True`. ROCm torch before 2.13 rejects external events - ("External events are disallowed in rocm") although HIP >= 7 supports them, so there the - event is created normally, recorded once outside capture so it exists and counts as - recorded, and captured with `hipEventRecordWithFlags(..., hipEventRecordExternal)` -- - the call torch 2.13 itself makes (pytorch#178264). - """ + """A timing event whose record() can be captured into a graph. ROCm torch < 2.13 rejects + `external=True`, so there the event is recorded once to exist and its captured record goes + through hipEventRecordWithFlags, as torch 2.13 does (pytorch#178264).""" import torch if not getattr(torch.version, "hip", None): @@ -451,15 +447,13 @@ def _graph_event(): @staticmethod def _record_graph_event(event): + import ctypes + import torch if not getattr(event, "_collx_hip_external", False): - event.record() - return - import ctypes - - hip = EPBackend._hip_runtime() - rc = hip.hipEventRecordWithFlags( + return event.record() + rc = EPBackend._hip_runtime().hipEventRecordWithFlags( ctypes.c_void_p(event.cuda_event), ctypes.c_void_p(torch.cuda.current_stream().cuda_stream), ctypes.c_uint(0x1), # hipEventRecordExternal @@ -468,27 +462,16 @@ def _record_graph_event(event): raise RuntimeError(f"hipEventRecordWithFlags(external) failed with hipError {rc}") @staticmethod + @functools.cache def _hip_runtime(): - lib = getattr(EPBackend, "_hip_lib", None) - if lib is None: - import ctypes - import os as _os - - import torch - - candidates = ["libamdhip64.so", _os.path.join(_os.path.dirname(torch.__file__), "lib", - "libamdhip64.so")] - for name in candidates: - try: - lib = ctypes.CDLL(name) - break - except OSError: - continue - if lib is None: - raise RuntimeError("libamdhip64.so not loadable for graph event capture") - lib.hipEventRecordWithFlags.restype = ctypes.c_int - EPBackend._hip_lib = lib - return lib + import ctypes + + import torch + + try: + return ctypes.CDLL("libamdhip64.so") + except OSError: # pip ROCm torch bundles the runtime in torch/lib + return ctypes.CDLL(os.path.join(os.path.dirname(torch.__file__), "lib", "libamdhip64.so")) def _capture_pairs(self, problem, staged, pairs, marks): """Capture `pairs` back-to-back dispatch -> combine pairs into one graph. diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py index b3c72adcc..107c2456e 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,11 +42,8 @@ class MoRIBackend(EPBackend): maturity = "production" # vLLM --all2all-backend mori_*; SGLang --moe-a2a-backend mori SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") - # Both kernel families launch with host-built args only (no per-call host read of counts; the - # reset moved on-device in ROCm/mori#86 for vLLM's graphs), and MoRI's own benchmark captures - # IntraNode dispatch/combine and N-pair graphs. `stage` slices by the untimed per-rung - # `recv_tokens`, fixed for a rung's routing and so safe to bake into a capture. The timing - # events are captured through hipEventRecordWithFlags on ROCm (EPBackend._graph_event). + # Kernels launch with host-built args only (reset on-device since ROCm/mori#86); `stage` + # slices by the per-rung `recv_tokens`, fixed for a rung, so it is safe to capture. CUDA_GRAPH_MODES = ("normal", "low-latency") requires_fresh_pair = True diff --git a/experimental/CollectiveX/tests/test_chain.py b/experimental/CollectiveX/tests/test_chain.py index da4244a83..8a2a3c88b 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -427,46 +427,34 @@ def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(sel class RocmGraphEventRecord(unittest.TestCase): """ROCm torch < 2.13 rejects Event(external=True); the capture records through HIP instead.""" - def _torch(self, hip, records): + def _record(self, rc): + records, calls = [], [] + class Event: + cuda_event = 0xE0 + def __init__(self, **kwargs): - if kwargs.get("external"): - raise RuntimeError("External events are disallowed in rocm") - self.cuda_event = 0xE0 + assert not kwargs.get("external"), "External events are disallowed in rocm" def record(self): records.append("record") - return types.SimpleNamespace( - version=types.SimpleNamespace(hip=hip), - cuda=types.SimpleNamespace( - Event=Event, - current_stream=lambda: types.SimpleNamespace(cuda_stream=0x5), - ), - ) + fake = types.SimpleNamespace(version=types.SimpleNamespace(hip="7.2"), cuda=types.SimpleNamespace( + Event=Event, current_stream=lambda: types.SimpleNamespace(cuda_stream=0x5))) + hip = types.SimpleNamespace(hipEventRecordWithFlags=lambda e, s, f: calls.append( + (e.value, s.value, f.value)) or rc) + with mock.patch.dict(sys.modules, {"torch": fake}), \ + mock.patch.object(ep_backend.EPBackend, "_hip_runtime", lambda: hip): + ep_backend.EPBackend._record_graph_event(ep_backend.EPBackend._graph_event()) + return records, calls def test_rocm_events_are_recorded_once_outside_then_captured_through_hip(self): - records, hip_calls = [], [] - fake_hip = types.SimpleNamespace( - hipEventRecordWithFlags=lambda event, stream, flags: hip_calls.append( - (event.value, stream.value, flags.value) - ) or 0 - ) - with mock.patch.dict(sys.modules, {"torch": self._torch("7.2", records)}), \ - mock.patch.object(ep_backend.EPBackend, "_hip_runtime", lambda: fake_hip): - event = ep_backend.EPBackend._graph_event() - self.assertEqual(records, ["record"]) # materialized outside capture - ep_backend.EPBackend._record_graph_event(event) - self.assertEqual(records, ["record"]) # the captured record went through HIP - self.assertEqual(hip_calls, [(0xE0, 0x5, 0x1)]) # hipEventRecordExternal + # One record materializes the event outside capture; the captured one goes through HIP. + self.assertEqual(self._record(0), (["record"], [(0xE0, 0x5, 0x1)])) def test_a_failed_hip_record_raises(self): - fake_hip = types.SimpleNamespace(hipEventRecordWithFlags=lambda *args: 1) - with mock.patch.dict(sys.modules, {"torch": self._torch("7.2", [])}), \ - mock.patch.object(ep_backend.EPBackend, "_hip_runtime", lambda: fake_hip): - event = ep_backend.EPBackend._graph_event() - with self.assertRaisesRegex(RuntimeError, "hipEventRecordWithFlags"): - ep_backend.EPBackend._record_graph_event(event) + with self.assertRaisesRegex(RuntimeError, "hipEventRecordWithFlags"): + self._record(1) class EventPlacement(unittest.TestCase):