diff --git a/python/sglang/kernels/jit/csrc/distributed/ipc_a2a.cuh b/python/sglang/kernels/jit/csrc/distributed/ipc_a2a.cuh new file mode 100644 index 000000000..34247f610 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/distributed/ipc_a2a.cuh @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include + +#include + +#include + +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(); + TensorMatcher({1}).with_dtype().with_device(device).verify(flag).verify(target).verify(timed_out).verify( + peer_timed_out); + LaunchKernel(1, 1, device.unwrap())( + spin_wait_kernel, + static_cast(flag.data_ptr()), + static_cast(target.data_ptr()), + static_cast(timed_out.data_ptr()), + static_cast(peer_timed_out.data_ptr()), + static_cast(budget_ns)); +} + +inline void bump_signal(tvm::ffi::TensorView seq, tvm::ffi::TensorView peer_flag) { + using namespace host; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({1}).with_dtype().with_device(device).verify(seq).verify(peer_flag); + LaunchKernel(1, 1, device.unwrap())( + bump_signal_kernel, static_cast(seq.data_ptr()), static_cast(peer_flag.data_ptr())); +} + +} // namespace sglang::ipc_a2a diff --git a/python/sglang/kernels/ops/communication/ipc_a2a.py b/python/sglang/kernels/ops/communication/ipc_a2a.py new file mode 100644 index 000000000..8dd688fec --- /dev/null +++ b/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"), + ], + ) diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py index 90524471e..ee19f78e0 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py +++ b/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 -#include -__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(), target.data_ptr(), - timed_out.data_ptr(), peer_timed_out.data_ptr(), - (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(), (volatile int*)peer_flag.data_ptr()); -} -""" - 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") diff --git a/test/registered/e2e/diffusion/test_ipc_a2a_jit_recovery.py b/test/registered/e2e/diffusion/test_ipc_a2a_jit_recovery.py new file mode 100644 index 000000000..18ab04819 --- /dev/null +++ b/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()