From 754b692afc2948137c4315989481e406acae0159 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 31 Jul 2026 19:35:48 +0800 Subject: [PATCH] [diffusion] optimization: support cuda-ipc zero-staging all-to-all for 2-rank Ulysses (#31854) Co-authored-by: Claude Fable 5 --- python/sglang/multimodal_gen/envs.py | 19 ++ .../base_device_communicator.py | 40 +++ .../device_communicators/ipc_a2a.py | 276 ++++++++++++++++++ .../runtime/layers/attention/layer.py | 16 +- .../multimodal_gen/runtime/layers/usp.py | 174 +++++++++++ .../runtime/managers/gpu_worker.py | 6 + .../runtime/models/dits/qwen_image.py | 25 +- .../multimodal_gen/test/server/gpu_cases.py | 3 + .../single_test_file/test_ipc_a2a_2_gpu.py | 228 +++++++++++++++ 9 files changed, 779 insertions(+), 8 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py create mode 100644 python/sglang/multimodal_gen/test/single_test_file/test_ipc_a2a_2_gpu.py diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 64d349735..3943f0164 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -32,6 +32,17 @@ if TYPE_CHECKING: SGLANG_DIFFUSION_STAGE_LOGGING: bool = False SGLANG_DIFFUSION_CFG_GATE_STEP: float = 1.0 # cache-dit env vars (primary transformer) + # on by default; engages only on 2 ranks with peer-to-peer access and falls + # back to NCCL when unavailable. Set 0 to force NCCL. Keep this in step with + # the resolver below -- that is the value the runtime reads. + SGLANG_DIFFUSION_IPC_A2A: bool = True + # a deadlock backstop, not a per-step budget: a rank can legitimately stall + # for seconds (layerwise offload, wan2.2 expert-tower swaps), and expiry now + # retires the transport on every rank and fails the request + SGLANG_DIFFUSION_IPC_A2A_TIMEOUT_MS: float = 10000.0 + # distinct (n_local, n_peer, dtype) staging pairs kept; each is two + # slots and is never freed, so multi-resolution serving needs a cap + SGLANG_DIFFUSION_IPC_A2A_MAX_BUFFERS: int = 16 SGLANG_CACHE_DIT_ENABLED: bool = False SGLANG_CACHE_DIT_FN: int = 1 SGLANG_CACHE_DIT_BN: int = 0 @@ -219,6 +230,14 @@ environment_variables: dict[str, Callable[[], Any]] = { ), # ================== cache-dit Env Vars ================== # Enable cache-dit acceleration for DiT inference + # CUDA-IPC transport for 2-rank Ulysses all-to-all (NVLink same-node) + "SGLANG_DIFFUSION_IPC_A2A": _lazy_bool("SGLANG_DIFFUSION_IPC_A2A", "true"), + "SGLANG_DIFFUSION_IPC_A2A_TIMEOUT_MS": _lazy_float( + "SGLANG_DIFFUSION_IPC_A2A_TIMEOUT_MS", 10000.0 + ), + "SGLANG_DIFFUSION_IPC_A2A_MAX_BUFFERS": _lazy_int( + "SGLANG_DIFFUSION_IPC_A2A_MAX_BUFFERS", 16 + ), "SGLANG_CACHE_DIT_ENABLED": _lazy_bool("SGLANG_CACHE_DIT_ENABLED"), # Number of first blocks to always compute (DBCache F parameter) "SGLANG_CACHE_DIT_FN": _lazy_int("SGLANG_CACHE_DIT_FN", 1), diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py index 228c5c1cf..71da9c977 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py @@ -12,6 +12,41 @@ from torch import Tensor from torch.distributed import ProcessGroup, ReduceOp +def _ipc_all_to_all_4d(group, input_, scatter_dim): + """2-rank IPC path for AllToAll4D; None when the transport is unavailable.""" + from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IPC_A2A, + ipc_a2a_ready, + ) + + if not ipc_a2a_ready(group): + return None + r = IPC_A2A.rank + if scatter_dim == 2: + # [bs, s_local, h, d] -> [bs, 2*s_local, h/2, d] + bs, shard_seqlen, hn, hd = input_.shape + half = hn // 2 + send = input_[:, :, (1 - r) * half : (2 - r) * half].contiguous() + out = input_.new_empty(bs, 2 * shard_seqlen, half, hd) + out.narrow(1, r * shard_seqlen, shard_seqlen).copy_( + input_[:, :, r * half : (r + 1) * half] + ) + theirs = IPC_A2A.exchange(group, send, (bs, shard_seqlen, half, hd)) + out.narrow(1, (1 - r) * shard_seqlen, shard_seqlen).copy_(theirs) + return out + # scatter_dim == 1: [bs, s, h_local, d] -> [bs, s/2, 2*h_local, d] + bs, seqlen, shard_hn, hd = input_.shape + shard_seqlen = seqlen // 2 + send = input_.narrow(1, (1 - r) * shard_seqlen, shard_seqlen).contiguous() + out = input_.new_empty(bs, shard_seqlen, 2 * shard_hn, hd) + out[:, :, r * shard_hn : (r + 1) * shard_hn].copy_( + input_.narrow(1, r * shard_seqlen, shard_seqlen) + ) + theirs = IPC_A2A.exchange(group, send, (bs, shard_seqlen, shard_hn, hd)) + out[:, :, (1 - r) * shard_hn : (2 - r) * shard_hn].copy_(theirs) + return out + + class DistributedAutograd: """Collection of autograd functions for distributed operations. @@ -132,6 +167,11 @@ class DistributedAutograd: input_.dim() == 4 ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + if world_size == 2 and scatter_dim in (1, 2): + fast = _ipc_all_to_all_4d(group, input_, scatter_dim) + if fast is not None: + return fast + if scatter_dim == 2 and gather_dim == 1: bs, shard_seqlen, hn, hd = input_.shape assert hn % world_size == 0, ( 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 new file mode 100644 index 000000000..1c06560e8 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA-IPC transport for 2-rank Ulysses all-to-all. + +Each rank maps the peer's staging buffers into its own device context +(handles are re-opened locally, so every access is same-device semantics over +NVLink) and writes its half of the exchange directly into them. Rank +synchronization is a GPU-side sequence counter: bump_signal publishes my new +sequence number into the peer's flag after my writes, spin_wait blocks my +stream until the peer has published the same number. Both primitives are +plain kernels on local memory, so the whole exchange is CUDA-graph +capturable. Double-buffered slots alternate per call; a slot is only rewritten +after an intervening spin_wait, which orders the rewrite after the peer's +read of that slot. +""" + +import logging +import os +from collections import OrderedDict + +import torch +import torch.distributed as dist + +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.""" + + +class IpcA2AState: + def __init__(self): + self.ops = None + # insertion-ordered so eviction is identical on every rank (all ranks + # create and touch the same keys in the same order) + self.staging = OrderedDict() + self.flag = None + self.peer_flag = None + self.my_seq = None + self.timed_out = None + self.peer_timed_out = None + self.budget_ns = 0 + self.max_buffers = 0 + self.calls = 0 + self.rank = None + self.failed = False + self.inited = False + + def _share(self, t, group): + """Exchange `t` with the peer via torch IPC, re-opening the handle in + the LOCAL device context (the mapping is only dereferenceable from the + context that opened it).""" + from torch.multiprocessing.reductions import reduce_tensor + + fn, args = reduce_tensor(t) + mine = [(fn, args)] + theirs = [None] + r0 = dist.get_global_rank(group, 0) + r1 = dist.get_global_rank(group, 1) + if self.rank == 0: + dist.broadcast_object_list(mine, src=r0, group=group) + dist.broadcast_object_list(theirs, src=r1, group=group) + else: + dist.broadcast_object_list(theirs, src=r0, group=group) + dist.broadcast_object_list(mine, src=r1, group=group) + pf, pa = theirs[0] + pa = list(pa) + dev = torch.cuda.current_device() + for i, v in enumerate(pa): + if isinstance(v, torch.device): + pa[i] = torch.device(f"cuda:{dev}") + elif isinstance(v, int) and i == 6: + # rebuild_cuda_tensor positional device index + pa[i] = dev + return pf(*pa) + + def init(self, group): + import ctypes + + from torch.utils.cpp_extension import load_inline + + self.rank = dist.get_rank(group=group) + dev = torch.cuda.current_device() + if not torch.cuda.can_device_access_peer(dev, 1 - dev): + raise _Unsupported("no peer-to-peer access between the two devices") + # kernel-level dereference of peer mappings needs explicit peer access + ctypes.CDLL("libcudart.so").cudaDeviceEnablePeerAccess(1 - 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.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") + self.budget_ns = int(envs.SGLANG_DIFFUSION_IPC_A2A_TIMEOUT_MS * 1e6) + self.max_buffers = envs.SGLANG_DIFFUSION_IPC_A2A_MAX_BUFFERS + self.peer_flag = self._share(self.flag, group) + self.peer_timed_out = self._share(self.timed_out, group) + self.inited = True + + def get_staging(self, n_local, n_peer, dtype, group): + """Local buffer of `n_local` elements (the peer writes into it) paired + with the peer's mapped buffer of `n_peer` elements (we write into it). + Creation is a paired collective, so both ranks must reach a new key at + the same call site. Returns None on a miss during CUDA graph capture: + the IPC handle exchange allocates and broadcasts, which is illegal + inside capture, so callers fall back to NCCL and the graph bakes that + path; pre-warmed keys keep the IPC fast path (its copies and + spin/bump kernels are capture-safe).""" + key = (n_local, n_peer, dtype) + pair = self.staging.get(key) + if pair is None: + if torch.cuda.is_current_stream_capturing(): + return None + local = torch.zeros(2, n_local, dtype=dtype, device="cuda") + peer = self._share(local, group) + pair = (local, peer) + if len(self.staging) >= self.max_buffers: + # Evicting drops my buffer while the peer still maps it, so both + # ranks must drop the same key: they share the insertion order. + self.staging.popitem(last=False) + self.staging[key] = pair + return pair + + def exchange(self, group, send, recv_shape): + """Symmetric exchange: write my contiguous `send` into the peer's + staging slot, return my staging slot viewed as `recv_shape`.""" + n_send = send.numel() + n_recv = 1 + for v in recv_shape: + n_recv *= v + pair = self.get_staging(n_recv, n_send, send.dtype, group) + if pair is None: + return None + local, peer = pair + slot = self.next_slot() + peer[slot].narrow(0, 0, n_send).copy_(send.view(-1), non_blocking=True) + self.signal_and_wait() + return local[slot].narrow(0, 0, n_recv).view(recv_shape) + + def next_slot(self): + slot = self.calls % 2 + self.calls += 1 + return slot + + def signal_and_wait(self): + self.signal() + self.wait() + + def signal(self): + self.ops.bump_signal(self.my_seq, self.peer_flag) + + def wait(self): + self.ops.spin_wait( + self.flag, + self.my_seq, + self.timed_out, + self.peer_timed_out, + self.budget_ns, + ) + + def check_timeout(self) -> None: + """Retire the transport on every rank if any rank's spin expired. + + The flag is a device read, so this belongs at a request boundary, never + inside a capture. It reads as local, but the spin kernel sets the flag on + both sides, so a timeout retires the transport on both: a rank that + retired alone would post an all_to_all its peer -- still on IPC -- never + posts, and the NCCL watchdog would take the process down. + + Expiry means the peer went silent for the whole budget, which is a hang, + not a slow step, and the exchange that gave up returned incomplete data. + So this raises rather than quietly serving a corrupted result. + """ + if self.failed or not self.inited or self.timed_out.item() == 0: + return + self.failed = True + raise RuntimeError( + "IPC all-to-all gave up waiting for its peer after " + f"{envs.SGLANG_DIFFUSION_IPC_A2A_TIMEOUT_MS:g} ms, so that exchange " + "returned incomplete data. The transport is now disabled on every " + "rank; retry the request over NCCL. Raise " + "SGLANG_DIFFUSION_IPC_A2A_TIMEOUT_MS if a rank can legitimately " + "stall this long (layerwise offload, expert-tower swaps), or set " + "SGLANG_DIFFUSION_IPC_A2A=0." + ) + + +IPC_A2A = IpcA2AState() + + +def ipc_a2a_ready(group) -> bool: + """True when the IPC transport is enabled and initialized (initializes + lazily on the first eager call; never inside a graph capture).""" + if not envs.SGLANG_DIFFUSION_IPC_A2A or IPC_A2A.failed: + return False + if IPC_A2A.inited: + return True + if torch.cuda.is_current_stream_capturing(): + return False + try: + IPC_A2A.init(group) + return True + except _Unsupported as e: + logger.info("IPC all-to-all unavailable (%s); using NCCL", e) + IPC_A2A.failed = True + return False + except Exception: + logger.exception("IPC all-to-all init failed; falling back to NCCL") + IPC_A2A.failed = True + return False diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 1cc1d21f1..7d3ed8872 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -44,6 +44,7 @@ from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import ( async_a2a_communicate, ) from sglang.multimodal_gen.runtime.layers.usp import ( + _ipc_input_a2a_qkv, _usp_input_all_to_all, _usp_input_all_to_all_varlen, _usp_output_all_to_all, @@ -619,6 +620,7 @@ class USPAttention(nn.Module): num_replicated_kv_prefix: int = 0, skip_sequence_parallel_override: bool = False, attn_mask_meta: dict | None = None, + qkv_pre_all_to_all: bool = False, ) -> torch.Tensor: """ Forward pass for USPAttention. @@ -769,10 +771,14 @@ class USPAttention(nn.Module): ) sp_size = get_ulysses_parallel_world_size() - if sp_size > 1: - q = _usp_input_all_to_all(q, head_dim=2) - k = _usp_input_all_to_all(k, head_dim=2) - v = _usp_input_all_to_all(v, head_dim=2) + if sp_size > 1 and not qkv_pre_all_to_all: + qkv_fast = _ipc_input_a2a_qkv(q, k, v) + if qkv_fast is not None: + q, k, v = qkv_fast + else: + q = _usp_input_all_to_all(q, head_dim=2) + k = _usp_input_all_to_all(k, head_dim=2) + v = _usp_input_all_to_all(v, head_dim=2) if ( _VARLEN_FA_ENABLED @@ -953,7 +959,7 @@ class USPAttention(nn.Module): ) # Ulysses-style All-to-All for sequence/head sharding - if sp_size > 1: + if sp_size > 1 and not qkv_pre_all_to_all: # -> [B, S, H_local, D] if self.enable_packed_qkv_input_a2a and q.device.type == "cuda": q, k, v = async_a2a_communicate( diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index aaabb3303..bd81eb066 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/python/sglang/multimodal_gen/runtime/layers/usp.py @@ -66,6 +66,159 @@ def _usp_all_to_all_single_varlen( return output +def _ipc_ready_group(): + """The ulysses group when the 2-rank IPC transport is usable, else None.""" + from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + ipc_a2a_ready, + ) + + group = get_sp_group().ulysses_group + return group if ipc_a2a_ready(group) else None + + +def _ipc_varlen_fast(x, seq_lens, head_dim, direction): + """2-rank IPC path for the varlen A2A pair; None when unavailable.""" + if head_dim != 2: + return None + group = _ipc_ready_group() + if group is None: + return None + from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IPC_A2A, + ) + + r = IPC_A2A.rank + off = [0, seq_lens[0]] + if direction == "input": + # [b, s_local, h_global, d] -> [b, sum(seq_lens), h_global/2, d] + b, s_local, h_global, d = x.shape + half = h_global // 2 + peer_len = seq_lens[1 - r] + send = x[:, :, (1 - r) * half : (2 - r) * half].contiguous() + out = x.new_empty(b, seq_lens[0] + seq_lens[1], half, d) + out.narrow(1, off[r], s_local).copy_(x[:, :, r * half : (r + 1) * half]) + theirs = IPC_A2A.exchange(group, send, (b, peer_len, half, d)) + if theirs is None: + return None + out.narrow(1, off[1 - r], peer_len).copy_(theirs) + return out + # output: [b, s_global, h_local, d] -> [b, seq_lens[r], 2*h_local, d]. + # The staging buffer IS the gathered result: each rank writes its head + # half of the peer's slot directly; no intermediate copy. + b, s_global, h_local, d = x.shape + my_len = seq_lens[r] + peer_len = seq_lens[1 - r] + n_out = b * my_len * 2 * h_local * d + n_peer_out = b * peer_len * 2 * h_local * d + pair = IPC_A2A.get_staging(n_out, n_peer_out, x.dtype, group) + if pair is None: + return None + local, peer = pair + slot = IPC_A2A.next_slot() + pst = peer[slot].narrow(0, 0, n_peer_out).view(b, peer_len, 2 * h_local, d) + pst[:, :, r * h_local : (r + 1) * h_local].copy_( + x.narrow(1, off[1 - r], peer_len), non_blocking=True + ) + IPC_A2A.signal() + out = local[slot].narrow(0, 0, n_out).view(b, my_len, 2 * h_local, d) + out[:, :, r * h_local : (r + 1) * h_local].copy_(x.narrow(1, off[r], my_len)) + IPC_A2A.wait() + return out + + +def _ipc_input_a2a_qkv(q, k, v): + """The three input A2As of one attention as a single IPC exchange; + None when unavailable.""" + if get_ulysses_parallel_world_size() != 2: + return None + group = _ipc_ready_group() + if group is None: + return None + from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IPC_A2A, + ) + + b, s_local, h_global, d = q.shape + half = h_global // 2 + r = IPC_A2A.rank + n = b * s_local * half * d + pair = IPC_A2A.get_staging(3 * n, 3 * n, q.dtype, group) + if pair is None: + return None + local, peer = pair + slot = IPC_A2A.next_slot() + outs = [] + for i, t in enumerate((q, k, v)): + send = t[:, :, (1 - r) * half : (2 - r) * half].contiguous() + peer[slot].narrow(0, i * n, n).copy_(send.view(-1), non_blocking=True) + out = t.new_empty(b, 2 * s_local, half, d) + out.narrow(1, r * s_local, s_local).copy_(t[:, :, r * half : (r + 1) * half]) + outs.append(out) + IPC_A2A.signal_and_wait() + for i, out in enumerate(outs): + theirs = local[slot].narrow(0, i * n, n).view(b, s_local, half, d) + out.narrow(1, (1 - r) * s_local, s_local).copy_(theirs) + return tuple(outs) + + +def _ipc_input_a2a_qkv_segmented(txt_q, img_q, txt_k, img_k, txt_v, img_v, local_pad): + """Joint-attention input A2A that reads the (txt, img) pair directly in + join_seqs layout [txt_real | img | txt_pad], skipping the per-projection + joint cats. The staging buffer IS the gathered q/k/v: each rank writes its + own sequence span of the peer's slot; returns (q, k, v) staging views or + None when unavailable.""" + if get_ulysses_parallel_world_size() != 2: + return None + group = _ipc_ready_group() + if group is None: + return None + from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IPC_A2A, + ) + + b, txt_len, h_global, d = txt_q.shape + img_len = img_q.shape[1] + half = h_global // 2 + r = IPC_A2A.rank + L = txt_len + img_len + real = txt_len - local_pad + n = b * 2 * L * half * d + pair = IPC_A2A.get_staging(3 * n, 3 * n, txt_q.dtype, group) + if pair is None: + return None + local, peer = pair + slot = IPC_A2A.next_slot() + ph = slice((1 - r) * half, (2 - r) * half) + lh = slice(r * half, (r + 1) * half) + base = r * L + outs = [] + peer_dsts, peer_srcs = [], [] + loc_dsts, loc_srcs = [], [] + for i, (txt, img) in enumerate(((txt_q, img_q), (txt_k, img_k), (txt_v, img_v))): + pst = peer[slot].narrow(0, i * n, n).view(b, 2 * L, half, d) + pspan = pst[:, base : base + L] + peer_dsts += [pspan[:, 0:real], pspan[:, real : real + img_len]] + peer_srcs += [txt[:, 0:real, ph], img[:, :, ph]] + if local_pad: + peer_dsts.append(pspan[:, real + img_len :]) + peer_srcs.append(txt[:, real:, ph]) + out = local[slot].narrow(0, i * n, n).view(b, 2 * L, half, d) + lspan = out[:, base : base + L] + loc_dsts += [lspan[:, 0:real], lspan[:, real : real + img_len]] + loc_srcs += [txt[:, 0:real, lh], img[:, :, lh]] + if local_pad: + loc_dsts.append(lspan[:, real + img_len :]) + loc_srcs.append(txt[:, real:, lh]) + outs.append(out) + torch._foreach_copy_(peer_dsts, peer_srcs) + # signal as soon as the peer's data is in flight; our local-half writes + # overlap with the peer's wait + IPC_A2A.signal() + torch._foreach_copy_(loc_dsts, loc_srcs) + IPC_A2A.wait() + return tuple(outs) + + def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: """ Perform Ulysses-style input all-to-all over the head dimension. @@ -88,6 +241,11 @@ def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: if world_size <= 1: return x + if world_size == 2 and head_dim == 2: + fast = _ipc_varlen_fast(x, [x.shape[1], x.shape[1]], 2, "input") + if fast is not None: + return fast + assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" @@ -148,6 +306,11 @@ def _usp_input_all_to_all_varlen( if world_size <= 1: return x + if world_size == 2: + fast = _ipc_varlen_fast(x, seq_lens, head_dim, "input") + if fast is not None: + return fast + assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" assert ( @@ -221,6 +384,12 @@ def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: if world_size <= 1: return x + if world_size == 2 and head_dim == 2 and x.shape[1] % 2 == 0: + half_len = x.shape[1] // 2 + fast = _ipc_varlen_fast(x, [half_len, half_len], 2, "output") + if fast is not None: + return fast + assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" @@ -281,6 +450,11 @@ def _usp_output_all_to_all_varlen( if world_size <= 1: return x + if world_size == 2: + fast = _ipc_varlen_fast(x, seq_lens, head_dim, "output") + if fast is not None: + return fast + assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" assert ( diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index cc3168d2f..db2a0a761 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -23,6 +23,9 @@ from sglang.multimodal_gen.runtime.distributed import ( maybe_init_distributed_environment_and_model_parallel, model_parallel_is_initialized, ) +from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IPC_A2A, +) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, @@ -326,6 +329,9 @@ class GPUWorker(GPUWorkerPostTrainingMixin): Used by disaggregated pipelines to access intermediate tensors. """ assert self.pipeline is not None + # request boundary: the IPC watchdog flag is a device read, illegal + # inside a graph capture and too costly per exchange + IPC_A2A.check_timeout() if len(batch) > 1: if return_req: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 68396b3e8..21714c3d4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -762,9 +762,27 @@ class QwenImageCrossAttention(nn.Module): # Joint order [text, image]; join_seqs relocates any SP text tail-pad # behind the image (see sp_shard.join_seqs for why). - joint_query = join_seqs(txt_query, img_query, sp_txt_pad) - joint_key = join_seqs(txt_key, img_key, sp_txt_pad) - joint_value = join_seqs(txt_value, img_value, sp_txt_pad) + seg_qkv = None + if sp_text_sharded: + from sglang.multimodal_gen.runtime.layers.usp import ( + _ipc_input_a2a_qkv_segmented, + ) + + seg_qkv = _ipc_input_a2a_qkv_segmented( + txt_query, + img_query, + txt_key, + img_key, + txt_value, + img_value, + sp_txt_pad, + ) + if seg_qkv is not None: + joint_query, joint_key, joint_value = seg_qkv + else: + joint_query = join_seqs(txt_query, img_query, sp_txt_pad) + joint_key = join_seqs(txt_key, img_key, sp_txt_pad) + joint_value = join_seqs(txt_value, img_value, sp_txt_pad) if attn_mask is None and encoder_hidden_states_mask is not None: image_mask = torch.ones( (hidden_states.shape[0], img_query.shape[1]), @@ -784,6 +802,7 @@ class QwenImageCrossAttention(nn.Module): attn_mask=attn_mask, attn_mask_meta=attn_mask_meta, num_replicated_prefix=0 if sp_text_sharded else seq_len_txt, + qkv_pre_all_to_all=seg_qkv is not None, ) # Reshape back diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index a78aef40d..5a0ae5b20 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -994,6 +994,7 @@ STANDALONE_FILES = { "2-gpu": [ "../single_test_file/test_disagg_server.py", "../single_test_file/test_ar_models.py", + "../single_test_file/test_ipc_a2a_2_gpu.py", ], } @@ -1026,6 +1027,8 @@ STANDALONE_FILE_EST_TIMES = { # Raise if CI reports a higher measured time. "../single_test_file/test_disagg_server.py": 600.0, "../single_test_file/test_ar_models.py": 600.0, + # no model load; the cost is the one-time JIT build of the sync kernels + "../single_test_file/test_ipc_a2a_2_gpu.py": 240.0, }, } diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_ipc_a2a_2_gpu.py b/python/sglang/multimodal_gen/test/single_test_file/test_ipc_a2a_2_gpu.py new file mode 100644 index 000000000..37da6bf5d --- /dev/null +++ b/python/sglang/multimodal_gen/test/single_test_file/test_ipc_a2a_2_gpu.py @@ -0,0 +1,228 @@ +"""The CUDA-IPC all-to-all must match NCCL bit for bit. + +This transport only activates on 2 ranks with peer-to-peer access, so nothing in +the single-GPU suite reaches it. Without this test a regression in the IPC path +lands silently. + +Runs the whole USPAttention A2A family twice per shape -- once over NCCL, once +over IPC -- and requires identical bytes: + + pytest -v python/sglang/multimodal_gen/test/single_test_file/test_ipc_a2a_2_gpu.py +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import unittest + +import torch + +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.test.test_utils import CustomTestCase + +_WORLD = 2 + + +def _worker() -> int: + """One rank: compare every A2A path against its NCCL result.""" + import torch.distributed as dist + + from sglang.multimodal_gen import envs + from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + ) + + rank = int(os.environ["RANK"]) + torch.cuda.set_device(rank) + maybe_init_distributed_environment_and_model_parallel( + tp_size=1, sp_size=_WORLD, ulysses_degree=_WORLD + ) + + from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IPC_A2A, + ) + from sglang.multimodal_gen.runtime.layers.usp import ( + _usp_input_all_to_all, + _usp_output_all_to_all, + ) + + if not torch.cuda.can_device_access_peer(rank, 1 - rank): + print("SKIP no peer-to-peer access between the two devices", flush=True) + return 0 + + # b, s_local, h_global, d -- h_global must split across the two ranks + shapes = [ + (1, 256, 8, 64), + (1, 1152, 24, 128), # qwen-image 1024x1024 class + (2, 64, 4, 64), # batched, and a short sequence like a warmup prompt + ] + failures = [] + for b, s_local, h_global, d in shapes: + torch.manual_seed(1234) # same input on both ranks' RNG, sliced per rank + full = torch.randn( + b, s_local * _WORLD, h_global, d, dtype=torch.bfloat16, device="cuda" + ) + x_in = full.narrow(1, rank * s_local, s_local).contiguous() + + def both_paths(fn, arg, **kw): + envs.SGLANG_DIFFUSION_IPC_A2A = False + nccl = fn(arg, **kw) + envs.SGLANG_DIFFUSION_IPC_A2A = True + ipc = fn(arg, **kw) + envs.SGLANG_DIFFUSION_IPC_A2A = False + return nccl, ipc + + nccl_out, ipc_out = both_paths(_usp_input_all_to_all, x_in, head_dim=2) + if not torch.equal(nccl_out, ipc_out): + failures.append(f"input a2a {(b, s_local, h_global, d)}") + + # the output A2A consumes [b, s_global, h_local, d] + x_out = torch.randn( + b, + s_local * _WORLD, + h_global // _WORLD, + d, + dtype=torch.bfloat16, + device="cuda", + ) + nccl_out, ipc_out = both_paths(_usp_output_all_to_all, x_out, head_dim=2) + if not torch.equal(nccl_out, ipc_out): + failures.append(f"output a2a {(b, s_local, h_global, d)}") + + # AllToAll4D is a second entry point (stacked-qkv UlyssesAttention: zimage's + # secondary attention, wan-VSA, hunyuanvideo) and routes through the same + # exchange, so it needs its own parity check. + from sglang.multimodal_gen.runtime.distributed.communication_op import ( + sequence_model_parallel_all_to_all_4D as all_to_all_4D, + ) + + # count exchanges, not staging keys: the key is (n_local, n_peer, dtype), + # so these shapes reuse the buffers the earlier arms already allocated + calls_before_a2a4d = IPC_A2A.calls + for b, s_local, h_global, d in shapes: + torch.manual_seed(4321) + for scatter_dim in (1, 2): + if scatter_dim == 2: + x = torch.randn( + b, s_local, h_global, d, dtype=torch.bfloat16, device="cuda" + ) + else: + x = torch.randn( + b, + s_local * _WORLD, + h_global // _WORLD, + d, + dtype=torch.bfloat16, + device="cuda", + ) + envs.SGLANG_DIFFUSION_IPC_A2A = False + nccl = all_to_all_4D(x, scatter_dim=scatter_dim, gather_dim=3 - scatter_dim) + envs.SGLANG_DIFFUSION_IPC_A2A = True + ipc = all_to_all_4D(x, scatter_dim=scatter_dim, gather_dim=3 - scatter_dim) + envs.SGLANG_DIFFUSION_IPC_A2A = False + if nccl.shape != ipc.shape: + failures.append( + f"a2a4d scatter={scatter_dim} shape {tuple(nccl.shape)} != " + f"{tuple(ipc.shape)} for {(b, s_local, h_global, d)}" + ) + elif not torch.equal(nccl, ipc): + failures.append( + f"a2a4d scatter={scatter_dim} {(b, s_local, h_global, d)}" + ) + + # A parity check that never reached the IPC branch would pass while proving + # nothing, so require that exchanges actually happened. + if IPC_A2A.calls == calls_before_a2a4d: + failures.append( + f"AllToAll4D never took the IPC path (exchange count stayed at " + f"{calls_before_a2a4d})" + ) + + # the transport must survive a shape it has never staged before, and the + # eviction that a capped cache performs + envs.SGLANG_DIFFUSION_IPC_A2A = True + IPC_A2A.max_buffers = 1 + for s_local in (32, 48, 64): + x = torch.randn(1, s_local, 8, 64, dtype=torch.bfloat16, device="cuda") + if _usp_input_all_to_all(x, head_dim=2) is None: + failures.append(f"eviction path returned None at s_local={s_local}") + envs.SGLANG_DIFFUSION_IPC_A2A = False + + # A comparison where both arms quietly fell back to NCCL would pass while + # testing nothing, so require evidence the transport actually ran. + if not IPC_A2A.inited or IPC_A2A.failed or not IPC_A2A.staging: + failures.append( + f"IPC never engaged (inited={IPC_A2A.inited} failed={IPC_A2A.failed} " + f"staged={len(IPC_A2A.staging)})" + ) + + # A timeout must retire the transport on BOTH ranks. When only the rank that + # timed out switched to NCCL, it posted an all_to_all its peer -- still on + # IPC -- never posted, and the NCCL watchdog took the process down (seen on + # wan2_2_t2v_a14b_2gpu in CI). Rank 0 waits for a signal rank 1 never sends, + # so rank 1 may only learn of it through the peer-side flag write. + IPC_A2A.budget_ns = 5 * 1000 * 1000 + if rank == 0: + IPC_A2A.signal_and_wait() + torch.cuda.synchronize() + dist.barrier() + if IPC_A2A.timed_out.item() == 0: + failures.append(f"rank{rank} never saw the timeout flag") + try: + IPC_A2A.check_timeout() + failures.append(f"rank{rank} check_timeout did not raise after a timeout") + except RuntimeError: + pass + + verdict = torch.tensor([len(failures)], device="cuda") + dist.all_reduce(verdict) + if failures: + print(f"rank{rank} MISMATCH: {failures}", flush=True) + if rank == 0: + print( + f"IPC_A2A_PARITY {'FAIL' if verdict.item() else 'PASS'} " + f"(staged keys={len(IPC_A2A.staging)})", + flush=True, + ) + dist.barrier() + dist.destroy_process_group() + return 1 if verdict.item() else 0 + + +class TestIpcA2ATwoGpu(CustomTestCase): + def test_ipc_matches_nccl_bitwise(self): + # CUDA only: the transport opens torch IPC handles through libcudart and + # cudaDeviceEnablePeerAccess. On ROCm/NPU it correctly refuses and both + # arms run over NCCL, which the evidence assertion below reads -- rightly + # -- as "the transport never engaged". That is unsupported, not broken. + if not current_platform.is_cuda(): + self.skipTest("CUDA-IPC transport is unavailable on this platform") + if torch.cuda.device_count() < _WORLD: + self.skipTest(f"needs {_WORLD} GPUs") + proc = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc-per-node={_WORLD}", + "--master-port=29517", + __file__, + "--worker", + ], + capture_output=True, + text=True, + timeout=1200, + ) + print(proc.stdout[-4000:]) + if proc.returncode != 0: + print(proc.stderr[-4000:], file=sys.stderr) + self.assertEqual(proc.returncode, 0, "IPC all-to-all diverged from NCCL") + self.assertIn("IPC_A2A_PARITY PASS", proc.stdout) + + +if __name__ == "__main__": + if "--worker" in sys.argv: + raise SystemExit(_worker()) + unittest.main()