minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Mick
2026-07-31 19:35:48 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent fd28242b68
commit 754b692afc
9 changed files with 779 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/envs.py
+19
View File
@@ -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
+40
View File
@@ -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
+276
View File
@@ -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
+8 -2
View File
@@ -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,7 +771,11 @@ 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:
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)
@@ -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
+174
View File
@@ -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
+6
View File
@@ -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
+19
View File
@@ -762,6 +762,24 @@ 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).
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)
@@ -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
+3
View File
@@ -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
+228
View File
@@ -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()