diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 3d68f67cbc..b5f03502ed 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.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), + 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)) + 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/bandwidth.py b/experimental/CollectiveX/bandwidth.py index b927a580ca..93f5d14cec 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/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py index e1e28ee6c1..694deb8b5d 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,125 @@ def _topk_idx_dtype(self): import torch return torch.int64 + # ---- CUDA graph capture ---------------------------------------------------------------- + + # 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 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 + + 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 * spin_us / elapsed_us)) + + 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) + if getattr(self, "_graph_align_cycles", None) is None: + self._calibrate_align_spin() + dist.all_reduce(token) + torch.cuda._sleep(self._graph_align_cycles) + + def _capture_pairs(self, problem, staged, pairs, marks): + """Capture `pairs` back-to-back dispatch -> combine pairs into one graph. + + `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 [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 = handle = 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, 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): + """Graph replay's value check: an untimed capture with `stage` INSIDE the graph. + + 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 + + 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() + # 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() + torch.cuda.synchronize() + return combined.clone() + # ---- Timing template methods ----------------------------------------------------- def timed_components(self): @@ -494,6 +613,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 +668,42 @@ def series(starts, ends): "combined": combined.clone(), } + def _benchmark_chain_graph(self, problem, staged, iters, drop): + """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 + + floors, floor_stamps, _, _ = self._capture_pairs( + problem, staged, iters, ("dispatch", "combine") + ) + 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() + 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 +739,17 @@ 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 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. + self._calibrate_align_spin() + 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 179c7baed4..c9ea6438cd 100644 --- a/experimental/CollectiveX/bench/ep_deepep_v2.py +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -137,8 +137,7 @@ 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. + # 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 @@ -168,6 +167,11 @@ 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 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" 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 +184,20 @@ 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 _dispatch_capacity(self, tokens): + """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 + 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 @@ -371,13 +389,13 @@ 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, 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_flashinfer.py b/experimental/CollectiveX/bench/ep_flashinfer.py index 8fe654cba2..d927053f72 100644 --- a/experimental/CollectiveX/bench/ep_flashinfer.py +++ b/experimental/CollectiveX/bench/ep_flashinfer.py @@ -120,6 +120,11 @@ 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: 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): 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 4c688ac556..c4e776d4c3 100644 --- a/experimental/CollectiveX/bench/ep_harness.py +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -213,6 +213,17 @@ 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 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 + 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. @@ -318,14 +329,20 @@ 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. + + `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() 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 +350,11 @@ 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 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" + return family def _reduce_vec(torch, dist, device, vals, op): @@ -693,6 +713,9 @@ 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()) + # 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 @@ -1027,8 +1050,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,8 +1187,9 @@ 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 - cuda_graph_output_applicable = cuda_graph and not backend.stage_excluded_from_roundtrip + chain_output_applicable = not backend.stage_excluded_from_roundtrip + # `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. # Per-iteration cross-rank MAX samples are pooled across trials. ---- @@ -1212,20 +1236,18 @@ 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: 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] 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 @@ -1235,7 +1257,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 +1346,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: @@ -1373,6 +1393,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 @@ -1385,23 +1406,37 @@ 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) + # 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")): + report = g[key] + oracle_verdicts[name] = bool( + _reduce_int(torch, dist, device, int(bool(report["passed"])), 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 + 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) + ) + # 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) @@ -1421,8 +1456,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 @@ -1477,18 +1512,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)), }, @@ -1514,15 +1549,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, ), }, @@ -1559,6 +1594,8 @@ 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, + "oracle_failed_checks": oracle_failed_checks, "passed": point_ok, }, "global_tokens": gt, @@ -1706,9 +1743,8 @@ 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 rows carry the chained family; graph mode captures it too. + "chained_period": True, "cuda_graph_replay": cuda_graph, "cuda_graph_supported": bool( getattr(backend, "cuda_graph_supported", False) @@ -1756,6 +1792,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 36cd3b889e..f34776526f 100644 --- a/experimental/CollectiveX/bench/ep_mori.py +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -42,6 +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") + # 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/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py index c0b08f59f0..26529c313d 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",) CUDA_GRAPH_MODES = ("normal", "low-latency") @@ -104,6 +106,15 @@ class NCCLEPBackend(EPBackend): receive_layout = "token-rank" combine_weight_semantics = "unweighted-rank-sum" zero_copy = True + _ll_expert_major = False + + @property + def cuda_graph_supported(self) -> bool: + # 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 def __init__(self, args, rank, world_size, local_rank, device): super().__init__(args, rank, world_size, local_rank, device) @@ -119,9 +130,23 @@ 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. + # "-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: self.kernel_generation = ( "nccl-ep-v02-ll-rm-zc" if self.zero_copy else "nccl-ep-v02-ll-rm" ) @@ -135,7 +160,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 +262,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 +370,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) @@ -372,18 +415,14 @@ 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. + """Read HT's received-token count and bind the combine input to the full receive plane. - 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. + 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()) - # 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). @@ -419,7 +458,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 +514,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 +564,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 +631,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] @@ -557,6 +659,11 @@ 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)) + # `_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() 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/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 76eefb4018..12ad1d5969 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 @@ -133,11 +134,31 @@ class UCCLEPBackend(EPBackend): kernel_generation = "uccl-legacy-buffer" SUPPORTED_MODES = ("normal", "low-latency") SUPPORTED_PRECISIONS = ("bf16", "fp8") + # 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 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: 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"): + 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/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 413b17ef8a..598dc3c674 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": { "arch": "gfx942", "product": "mi300x", @@ -171,7 +158,8 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-amds", - "backends": {}, + "runner_label": "cluster:mi325x-amds", + "backends": {"mori": [8], "uccl-ep": [8]}, "ll_backends": {}, "fabric": { "nic": "n/a (single-node scale-up)", diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index d98e78f076..3ce3da987d 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,40 @@ 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. -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. +Serving engines capture their decode step, so graph-compatible backend/mode pairs are measured +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. 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. +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 @@ -426,8 +442,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 @@ -647,7 +666,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 +679,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/summarize.py b/experimental/CollectiveX/summarize.py index 47e58826b3..ca498da423 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, diff --git a/experimental/CollectiveX/swap_matrix.py b/experimental/CollectiveX/swap_matrix.py index 040376c153..121083bfb5 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 919e88283f..e00bea9c54 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -28,6 +28,8 @@ 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"] + + SWEEP_BACKENDS = tuple(dict.fromkeys( backend for platform in PLATFORMS.values() for backend in platform["backends"] )) @@ -252,6 +254,8 @@ def resolve_matrix( shards_by_sku.setdefault(sku, []).append({ "id": f"{sku}-{target}{mode_segment}-{precision}-n{nodes}", "sku": 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 36d35cdbf3..64ddbe88b0 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 @@ -251,6 +253,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-em", "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) @@ -510,6 +545,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 @@ -540,15 +576,106 @@ 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) + + +@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 + 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) + + +def _deep_ep(*classes): + module = types.ModuleType("deep_ep") + for cls in classes: + setattr(module, cls, type(cls, (), {})) + return module + + +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 + + +class GraphReplayDefaults(unittest.TestCase): + """Graph replay is each adapter's default only where it was measured best and safe.""" + + 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 + + decode, prefill, ll = make("decode"), make("prefill"), make("decode", "low-latency") + + 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.assertEqual( + (prefill.cuda_graph_supported, prefill.kernel_generation), (False, "v2-elastic-buffer") + ) + self.assertTrue(ll.cuda_graph_supported) + + 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 ep_nccl + 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 965e75ba27..777b93716a 100644 --- a/experimental/CollectiveX/tests/test_chain.py +++ b/experimental/CollectiveX/tests/test_chain.py @@ -113,11 +113,13 @@ 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") ), ), 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( @@ -320,6 +322,45 @@ 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 _graph_backend(self): + backend = _ChainBackend(stage_device_work=False, fp8_consume="native", precision="bf16") + 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"] + 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 = 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) + 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(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) + def test_external_switch_restores_the_eager_component_pipeline(self): backend = _ChainBackend() backend.mode = "normal" @@ -337,6 +378,52 @@ 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 + 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): + backend = _ChainBackend(stage_device_work=True, fp8_consume="native", precision="fp8") + order, combined = [], _Combined(1.0) + handle = types.SimpleNamespace(recv_x="recv", recv_scales=None, combine_input=None) + 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)) 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) + 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). 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) + + 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.""" @@ -623,6 +710,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}" @@ -759,6 +850,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 @@ -860,8 +956,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 +970,20 @@ 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) + # 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) @@ -980,6 +1086,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. diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index 6f143c54cb..73d8382780 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -31,6 +31,11 @@ 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): + 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 = { "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 a9aa8480a8..d8ad34fe16 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"),