[diffusion] fix: recover ipc jit initialization after interrupted builds (#39034)

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>

python/sglang/kernels/jit/csrc/distributed/ipc_a2a.cuh

@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <cstdint>
+
+namespace sglang::ipc_a2a {
+
+__device__ __forceinline__ unsigned long long now_ns() {
+  unsigned long long t;
+  asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
+  return t;
+}
+
+__global__ void spin_wait_kernel(
+    volatile int* flag, const int* target, int* timed_out, int* peer_timed_out, unsigned long long budget_ns) {
+  int t = *target;
+  unsigned long long start = now_ns();
+  while (*flag < t) {
+    if (now_ns() - start > budget_ns) {
+      // both ranks must retire the transport at the same request boundary
+      *timed_out = 1;
+      *peer_timed_out = 1;
+      __threadfence_system();
+      return;
+    }
+  }
+  __threadfence_system();
+}
+
+__global__ void bump_signal_kernel(int* seq, volatile int* peer_flag) {
+  int v = *seq + 1;
+  *seq = v;
+  __threadfence_system();
+  *peer_flag = v;
+}
+
+inline void spin_wait(
+    tvm::ffi::TensorView flag,
+    tvm::ffi::TensorView target,
+    tvm::ffi::TensorView timed_out,
+    tvm::ffi::TensorView peer_timed_out,
+    int64_t budget_ns) {
+  using namespace host;
+  auto device = SymbolicDevice{};
+  device.set_options<kDLCUDA>();
+  TensorMatcher({1}).with_dtype<int32_t>().with_device(device).verify(flag).verify(target).verify(timed_out).verify(
+      peer_timed_out);
+  LaunchKernel(1, 1, device.unwrap())(
+      spin_wait_kernel,
+      static_cast<int*>(flag.data_ptr()),
+      static_cast<const int*>(target.data_ptr()),
+      static_cast<int*>(timed_out.data_ptr()),
+      static_cast<int*>(peer_timed_out.data_ptr()),
+      static_cast<unsigned long long>(budget_ns));
+}
+
+inline void bump_signal(tvm::ffi::TensorView seq, tvm::ffi::TensorView peer_flag) {
+  using namespace host;
+  auto device = SymbolicDevice{};
+  device.set_options<kDLCUDA>();
+  TensorMatcher({1}).with_dtype<int32_t>().with_device(device).verify(seq).verify(peer_flag);
+  LaunchKernel(1, 1, device.unwrap())(
+      bump_signal_kernel, static_cast<int*>(seq.data_ptr()), static_cast<int*>(peer_flag.data_ptr()));
+}
+
+}  // namespace sglang::ipc_a2a

python/sglang/kernels/ops/communication/ipc_a2a.py

@@ -0,0 +1,17 @@
+"""CUDA-IPC sequence-counter kernels for two-rank Ulysses."""
+
+from sglang.kernels.jit.utils import cache_once, load_jit
+
+
+@cache_once
+def load_ipc_a2a_sync():
+    # the shared jit loader releases its lock on process death and publishes
+    # complete builds atomically, unlike load_inline's persistent lock file
+    return load_jit(
+        "ipc_a2a_sync",
+        cuda_files=["distributed/ipc_a2a.cuh"],
+        cuda_wrappers=[
+            ("spin_wait", "ipc_a2a::spin_wait"),
+            ("bump_signal", "ipc_a2a::bump_signal"),
+        ],
+    )

python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py

@@ -14,72 +14,17 @@ read of that slot.
 """
 
 import logging
-import os
 import socket
 from collections import OrderedDict
 
 import torch
 import torch.distributed as dist
 
+from sglang.kernels.ops.communication.ipc_a2a import load_ipc_a2a_sync
 from sglang.multimodal_gen import envs
 
 logger = logging.getLogger(__name__)
 
-_SYNC_DECL = (
-    "void spin_wait(torch::Tensor flag, torch::Tensor target, torch::Tensor timed_out,"
-    " torch::Tensor peer_timed_out, int64_t budget_ns);\n"
-    "void bump_signal(torch::Tensor seq, torch::Tensor peer_flag);"
-)
-_SYNC_SRC = """
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-__device__ __forceinline__ unsigned long long now_ns() {
-    // %globaltimer is a nanosecond wall clock, so the budget needs no SM-clock
-    // conversion -- cudaDevAttrClockRate is not dependable across architectures
-    // (B200 reports 120 MHz, which would shrink the timeout ~16x).
-    unsigned long long t;
-    asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t));
-    return t;
-}
-__global__ void spin_wait_kernel(volatile int* flag, const int* target,
-                                 int* timed_out, int* peer_timed_out,
-                                 unsigned long long budget_ns) {
-    int t = *target;
-    unsigned long long start = now_ns();
-    while (*flag < t) {
-        if (now_ns() - start > budget_ns) {
-            // Give up rather than hang the stream forever. The peer never
-            // published, so this exchange's data is incomplete. Flag it on the
-            // peer as well as here: both ranks must retire the transport at the
-            // same request boundary, or the one that switched to NCCL would post
-            // a collective the other never posts.
-            *timed_out = 1;
-            *peer_timed_out = 1;
-            __threadfence_system();
-            return;
-        }
-    }
-    __threadfence_system();
-}
-__global__ void bump_signal_kernel(int* seq, volatile int* peer_flag) {
-    int v = *seq + 1;
-    *seq = v;
-    __threadfence_system();
-    *peer_flag = v;
-}
-void spin_wait(torch::Tensor flag, torch::Tensor target, torch::Tensor timed_out,
-               torch::Tensor peer_timed_out, int64_t budget_ns) {
-    spin_wait_kernel<<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(
-        (volatile int*)flag.data_ptr<int>(), target.data_ptr<int>(),
-        timed_out.data_ptr<int>(), peer_timed_out.data_ptr<int>(),
-        (unsigned long long)budget_ns);
-}
-void bump_signal(torch::Tensor seq, torch::Tensor peer_flag) {
-    bump_signal_kernel<<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(
-        seq.data_ptr<int>(), (volatile int*)peer_flag.data_ptr<int>());
-}
-"""
-
 
 class _Unsupported(RuntimeError):
     """This topology cannot run the transport -- an expected outcome, not a bug."""
@@ -183,8 +128,6 @@ class IpcA2AState:
     def init(self, group):
         import ctypes
 
-        from torch.utils.cpp_extension import load_inline
-
         self.rank = dist.get_rank(group=group)
         self.group = group
         dev = torch.cuda.current_device()
@@ -201,19 +144,7 @@ class IpcA2AState:
             )
         # kernel-level dereference of peer mappings needs explicit peer access
         ctypes.CDLL("libcudart.so").cudaDeviceEnablePeerAccess(peer_dev, 0)
-        build_dir = os.path.join(
+        self.ops = load_ipc_a2a_sync()
-            envs.SGLANG_DIFFUSION_CACHE_ROOT, f"ipc_a2a_sync_r{dev}"
-        )
-        os.makedirs(build_dir, exist_ok=True)
-        self.ops = load_inline(
-            name="ipc_a2a_sync",
-            cpp_sources=_SYNC_DECL,
-            cuda_sources=_SYNC_SRC,
-            functions=["spin_wait", "bump_signal"],
-            extra_cuda_cflags=["-O3"],
-            build_directory=build_dir,
-            verbose=False,
-        )
         self.flag = torch.zeros(1, dtype=torch.int32, device="cuda")
         self.my_seq = torch.zeros(1, dtype=torch.int32, device="cuda")
         self.timed_out = torch.zeros(1, dtype=torch.int32, device="cuda")

test/registered/e2e/diffusion/test_ipc_a2a_jit_recovery.py

@@ -0,0 +1,136 @@
+"""IPC initialization must survive a lock left by an interrupted JIT build.
+
+Unlike the parity test, this test needs isolated cold caches and a bounded
+subprocess lifetime so a blocked compiler cannot consume the whole CI job.
+"""
+
+import os
+import signal
+import socket
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch.utils.file_baton import FileBaton
+
+from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import (
+    IPC_A2A,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_sp_group,
+    maybe_init_distributed_environment_and_model_parallel,
+)
+from sglang.multimodal_gen.runtime.layers.usp import (
+    _usp_input_all_to_all,
+    _usp_output_all_to_all,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=60, stage="base-b", runner_config="diffusion-2-gpu-h100")
+
+
+def _worker():
+    rank = int(os.environ["RANK"])
+    torch.cuda.set_device(rank)
+    maybe_init_distributed_environment_and_model_parallel(
+        tp_size=1, sp_size=2, ulysses_degree=2
+    )
+    IPC_A2A.init(get_sp_group().ulysses_group)
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        value = torch.full((1, 32, 4, 64), rank, dtype=torch.bfloat16, device="cuda")
+        result = _usp_input_all_to_all(value, head_dim=2)
+        expected = torch.cat(
+            [torch.zeros_like(result[:, :32]), torch.ones_like(result[:, :32])],
+            dim=1,
+        )
+        torch.testing.assert_close(result, expected, rtol=0, atol=0)
+        restored = _usp_output_all_to_all(result, head_dim=2)
+        torch.testing.assert_close(restored, value, rtol=0, atol=0)
+    stream.synchronize()
+    # The replacement binding must preserve current-stream and capture semantics.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        captured = _usp_input_all_to_all(value, head_dim=2)
+        restored = _usp_output_all_to_all(captured, head_dim=2)
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            graph.replay()
+            torch.testing.assert_close(captured, expected, rtol=0, atol=0)
+            torch.testing.assert_close(restored, value, rtol=0, atol=0)
+    stream.synchronize()
+    assert IPC_A2A.inited and IPC_A2A.calls > 0
+    IPC_A2A.check_timeout()
+    print(f"IPC_RECOVERY PASS rank={rank}", flush=True)
+    dist.barrier()
+    IPC_A2A.reset()
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+class TestIpcA2AJitRecovery(CustomTestCase):
+    def test_initialization_with_orphaned_legacy_lock(self):
+        # Every diffusion partition invokes registered files; run this once.
+        if os.environ.get("DIFFUSION_PARTITION_ID", "0") != "0":
+            self.skipTest("cold-cache recovery runs on diffusion partition 0")
+        if not current_platform.is_cuda() or torch.cuda.device_count() < 2:
+            self.skipTest("requires two CUDA GPUs")
+        if not torch.cuda.can_device_access_peer(0, 1):
+            self.skipTest("requires CUDA peer access")
+        with tempfile.TemporaryDirectory(prefix="ipc-jit-recovery-") as cache_dir:
+            for rank in range(2):
+                build_dir = Path(cache_dir) / f"ipc_a2a_sync_r{rank}"
+                build_dir.mkdir()
+                baton = FileBaton(str(build_dir / "lock"))
+                self.assertTrue(baton.try_acquire())
+                # model a dead owner: its descriptor is closed, its marker remains
+                os.close(baton.fd)
+            with socket.socket() as listener:
+                listener.bind(("127.0.0.1", 0))
+                port = listener.getsockname()[1]
+            processes = []
+            try:
+                for rank in range(2):
+                    env = os.environ.copy()
+                    env.update(
+                        RANK=str(rank),
+                        LOCAL_RANK=str(rank),
+                        WORLD_SIZE="2",
+                        MASTER_ADDR="127.0.0.1",
+                        MASTER_PORT=str(port),
+                        SGLANG_DIFFUSION_CACHE_ROOT=cache_dir,
+                        SGLANG_JIT_CACHE_DIR=str(Path(cache_dir) / "jit"),
+                        SGLANG_DIFFUSION_IPC_A2A="1",
+                    )
+                    processes.append(
+                        subprocess.Popen(
+                            [sys.executable, __file__, "--worker"],
+                            env=env,
+                            stdout=subprocess.PIPE,
+                            stderr=subprocess.STDOUT,
+                            text=True,
+                            start_new_session=True,
+                        )
+                    )
+                for rank, process in enumerate(processes):
+                    output, _ = process.communicate(timeout=120)
+                    self.assertEqual(process.returncode, 0, output)
+                    self.assertIn(f"IPC_RECOVERY PASS rank={rank}", output)
+            finally:
+                for process in processes:
+                    if process.poll() is None:
+                        os.killpg(process.pid, signal.SIGKILL)
+                    process.communicate()
+
+
+if __name__ == "__main__":
+    if "--worker" in sys.argv:
+        _worker()
+    else:
+        unittest.main()