CollectiveX: graph-replay every capturable EP backend with rank alignment and the chained period

.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

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):

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"))

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(

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"