[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(
-            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.ops = load_ipc_a2a_sync()
         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")