CollectiveX: graph-replay MoRI (ROCm event capture via hipEventRecordWithFlags)

.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,185 @@ 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)
+
+    @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).
+        """
+        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` 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)]
+
+        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])
+
+        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 +673,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 +728,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 +799,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"))