[diffusion] optimization: support cuda-ipc zero-staging all-to-all for 2-rank Ulysses (#31854)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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),

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, (

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 <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."""
+
+
+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

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:
+            if sp_size > 1 and not qkv_pre_all_to_all:
-                q = _usp_input_all_to_all(q, head_dim=2)
+                qkv_fast = _ipc_input_a2a_qkv(q, k, v)
-                k = _usp_input_all_to_all(k, head_dim=2)
+                if qkv_fast is not None:
-                v = _usp_input_all_to_all(v, head_dim=2)
+                    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(

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 (

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(

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)
+        seg_qkv = None
-        joint_key = join_seqs(txt_key, img_key, sp_txt_pad)
+        if sp_text_sharded:
-        joint_value = join_seqs(txt_value, img_value, sp_txt_pad)
+            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

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,
     },
 }
 

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()