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

[diffusion] chore: route zimage and hunyuanvideo attention through USPAttention (#33923)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Mick
2026-08-07 12:36:23 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 698f019a5a
commit 6dc77e490d
6 changed files with 281 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+76 -12
View File
@@ -744,6 +744,7 @@ class USPAttention(nn.Module):
skip_sequence_parallel_override: bool = False,
attn_mask_meta: dict | None = None,
qkv_pre_all_to_all: bool = False,
seq_lens: list[int] | None = None,
) -> torch.Tensor:
"""
Forward pass for USPAttention.
@@ -778,6 +779,24 @@ class USPAttention(nn.Module):
effective_skip_sp = (
self.skip_sequence_parallel or skip_sequence_parallel_override
)
if seq_lens is not None:
assert (
attn_mask is None
and attn_mask_meta is None
and not num_replicated_prefix
and not num_replicated_suffix
and not num_replicated_kv_prefix
), "Varlen USPAttention does not support masks or replicated tokens"
if effective_skip_sp or get_sequence_parallel_world_size() == 1:
return self.attn_impl.forward(q, k, v, ctx_attn_metadata)
qkv = torch.cat([q, k, v], dim=0)
qkv = _usp_input_all_to_all_varlen(qkv, seq_lens, head_dim=2)
qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata)
q, k, v = qkv.chunk(3, dim=0)
out = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
out = self.attn_impl.postprocess_output(out, ctx_attn_metadata)
return _usp_output_all_to_all_varlen(out, seq_lens, head_dim=2)
if isinstance(attn_mask_meta, DynamicVarlenMaskMeta):
attn_mask_meta = attn_mask_meta.resolve(attn_mask)
@@ -830,6 +849,19 @@ class USPAttention(nn.Module):
raise NotImplementedError(unsupported)
if attn_mask is not None or meta_only_pad:
if (
num_replicated_prefix
or num_replicated_suffix
or num_replicated_kv_prefix
):
# This path shards every row through the all-to-all; a
# replicated prefix/suffix would be duplicated across ranks and
# silently corrupt the output, so refuse loudly instead.
raise NotImplementedError(
"USPAttention's masked path does not support replicated "
"prefix/suffix tokens; drop attn_mask/attn_mask_meta or "
"the replicated segment."
)
def _prepare_sdpa_mask(
mask: torch.Tensor, *, dtype: torch.dtype, device: torch.device
@@ -1516,26 +1548,58 @@ class USPAttention(nn.Module):
num_rep: int,
) -> torch.Tensor:
"""Ulysses attention where the last num_rep tokens are replicated
across SP ranks and should not be duplicated by the all-to-all."""
across SP ranks and should not be duplicated by the all-to-all.
The suffix stays at the sequence tail so every query scans K/V in the
same order as a single rank (bitwise-stable across SP degrees);
rotating it to the front reorders the reduction, and few-step models
amplify that into visible drift.
"""
if num_rep <= 0:
raise ValueError("num_rep must be positive for replicated suffix.")
if get_ring_parallel_world_size() > 1:
raise NotImplementedError(
"USPAttention replicated-prefix/suffix path does not support "
"ring parallelism yet."
)
sp_rank = get_sp_parallel_rank()
q_shard, q_rep = q[:, :-num_rep], q[:, -num_rep:]
k_shard, k_rep = k[:, :-num_rep], k[:, -num_rep:]
v_shard, v_rep = v[:, :-num_rep], v[:, -num_rep:]
# dense self-attention is permutation equivariant for non-causal use.
# 1. rotate the replicated suffix to the front
# 2. reuse the validated replicated-prefix path, then
# 3. rotate the output back
out = self._forward_with_replicated_prefix(
torch.cat([q_rep, q_shard], dim=1),
torch.cat([k_rep, k_shard], dim=1),
torch.cat([v_rep, v_shard], dim=1),
ctx_attn_metadata,
num_rep,
q_shard = _usp_input_all_to_all(q_shard, head_dim=2)
k_shard = _usp_input_all_to_all(k_shard, head_dim=2)
v_shard = _usp_input_all_to_all(v_shard, head_dim=2)
h_local = q_shard.shape[2]
kv_h_local = k_shard.shape[2]
h_start = sp_rank * h_local
kv_h_start = sp_rank * kv_h_local
q_rep = q_rep[:, :, h_start : h_start + h_local, :].contiguous()
k_rep = k_rep[:, :, kv_h_start : kv_h_start + kv_h_local, :].contiguous()
v_rep = v_rep[:, :, kv_h_start : kv_h_start + kv_h_local, :].contiguous()
q = torch.cat([q_shard, q_rep], dim=1)
k = torch.cat([k_shard, k_rep], dim=1)
v = torch.cat([v_shard, v_rep], dim=1)
out = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
out_shard = out[:, :-num_rep]
out_rep = out[:, -num_rep:]
out_shard = _usp_output_all_to_all(out_shard, head_dim=2)
sp_size = get_ulysses_parallel_world_size()
gathered = [torch.empty_like(out_rep) for _ in range(sp_size)]
torch.distributed.all_gather(
gathered,
out_rep.contiguous(),
group=get_sp_group().ulysses_group,
)
out_rep, out_shard = out[:, :num_rep], out[:, num_rep:]
out_rep = torch.cat(gathered, dim=2)
return torch.cat([out_shard, out_rep], dim=1)
python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
+17 -12
View File
@@ -18,7 +18,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
)
from sglang.multimodal_gen.runtime.layers.attention import (
LocalAttention,
UlyssesAttention,
USPAttention,
)
from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
from sglang.multimodal_gen.runtime.layers.layernorm import (
@@ -212,8 +212,7 @@ class MMDoubleStreamBlock(nn.Module):
quant_config=quant_config,
)
# Use UlyssesAttention to replace Distributed attention
self.attn = UlyssesAttention(
self.attn = USPAttention(
num_heads=self.local_num_attention_heads,
head_size=head_dim,
causal=False,
@@ -292,15 +291,20 @@ class MMDoubleStreamBlock(nn.Module):
# Run distributed attention
if txt_is_sharded:
attn, _ = self.attn(
attn = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
seq_lens=seq_lens,
)
img_attn, txt_attn = attn.split([image_seq_len, text_seq_len], dim=1)
else:
img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v)
attn = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
num_replicated_suffix=text_seq_len,
)
img_attn, txt_attn = attn.split([image_seq_len, text_seq_len], dim=1)
img_attn_out, _ = self.img_attn_proj(
img_attn.reshape(batch_size, image_seq_len, -1)
)
@@ -406,8 +410,7 @@ class MMSingleStreamBlock(nn.Module):
prefix=f"{prefix}.modulation",
)
# Use UlyssesAttention to replace Distributed attention
self.attn = UlyssesAttention(
self.attn = USPAttention(
num_heads=self.local_num_attention_heads,
head_size=head_dim,
causal=False,
@@ -463,17 +466,19 @@ class MMSingleStreamBlock(nn.Module):
# Run distributed attention
if txt_is_sharded:
attn_output, _ = self.attn(
attn_output = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
seq_lens=seq_lens,
)
else:
img_attn_output, txt_attn_output = self.attn(
img_q, img_k, img_v, txt_q, txt_k, txt_v
attn_output = self.attn(
torch.cat((img_q, txt_q), dim=1),
torch.cat((img_k, txt_k), dim=1),
torch.cat((img_v, txt_v), dim=1),
num_replicated_suffix=txt_len,
)
attn_output = torch.cat((img_attn_output, txt_attn_output), dim=1)
attn_output = attn_output.view(batch_size, seq_len, -1)
# Process MLP activation
mlp_output = self.mlp_act(mlp)
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+10 -47
View File
@@ -16,7 +16,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
)
from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
from sglang.multimodal_gen.runtime.layers.attention import (
UlyssesAttention,
USPAttention,
build_varlen_mask_meta_from_lengths,
build_varlen_mask_meta_from_ranges,
@@ -304,13 +303,6 @@ class ZImageAttention(nn.Module):
softmax_scale=None,
causal=False,
)
self.ulysses_attn = UlyssesAttention(
num_heads=self.local_num_heads,
head_size=self.head_dim,
num_kv_heads=self.local_num_kv_heads,
softmax_scale=None,
causal=False,
)
def forward(
self,
@@ -441,45 +433,16 @@ class ZImageAttention(nn.Module):
allow_inplace=self.enable_zimage_qk_fusion,
)
if (
num_replicated_suffix > 0
and get_sp_world_size() > 1
and get_ring_parallel_world_size() == 1
):
# the cap (last num_replicated_suffix tokens), as condition, should be replicated
q_shard, q_rep = (
q[:, :-num_replicated_suffix],
q[:, -num_replicated_suffix:],
)
k_shard, k_rep = (
k[:, :-num_replicated_suffix],
k[:, -num_replicated_suffix:],
)
v_shard, v_rep = (
v[:, :-num_replicated_suffix],
v[:, -num_replicated_suffix:],
)
hidden_states, hidden_states_rep = self.ulysses_attn(
q_shard,
k_shard,
v_shard,
replicated_q=q_rep,
replicated_k=k_rep,
replicated_v=v_rep,
)
assert hidden_states_rep is not None
hidden_states = torch.cat([hidden_states, hidden_states_rep], dim=1)
else:
hidden_states = self.attn(
q,
k,
v,
attn_mask=attn_mask,
attn_mask_meta=attn_mask_meta,
num_replicated_prefix=num_replicated_prefix,
num_replicated_suffix=num_replicated_suffix,
skip_sequence_parallel_override=skip_sequence_parallel_override,
)
hidden_states = self.attn(
q,
k,
v,
attn_mask=attn_mask,
attn_mask_meta=attn_mask_meta,
num_replicated_prefix=num_replicated_prefix,
num_replicated_suffix=num_replicated_suffix,
skip_sequence_parallel_override=skip_sequence_parallel_override,
)
hidden_states = hidden_states.flatten(2)
hidden_states, _ = self.to_out[0](hidden_states)
python/sglang/multimodal_gen/test/server/gpu_cases.py
+3
View File
@@ -1145,6 +1145,7 @@ STANDALONE_FILES = {
"../single_test_file/test_ipc_a2a_2_gpu.py",
"../single_test_file/test_dp_serving_2_gpu.py",
"../single_test_file/test_pynccl_a2a_capture_2_gpu.py",
"../single_test_file/test_usp_replicated_parity_2_gpu.py",
],
}
@@ -1183,6 +1184,8 @@ STANDALONE_FILE_EST_TIMES = {
"../single_test_file/test_dp_serving_2_gpu.py": 900.0,
# one capture plus three replays on a 32K-element exchange
"../single_test_file/test_pynccl_a2a_capture_2_gpu.py": 180.0,
# two SDPA parity checks on 128+6 rows
"../single_test_file/test_usp_replicated_parity_2_gpu.py": 180.0,
},
}
python/sglang/multimodal_gen/test/single_test_file/test_usp_replicated_parity_2_gpu.py
+155
View File
@@ -0,0 +1,155 @@
"""USPAttention replicated prefix/suffix must match a single-rank reference.
The suffix path keeps the replicated tokens at the sequence tail so every
query scans K/V in the same order as a single rank — bitwise-stable across SP
degrees. A rotate-to-front implementation is numerically valid but reorders
the reduction, and few-step (turbo) models amplify that into visible drift;
this test pins the order-preserving behavior.
pytest -v python/sglang/multimodal_gen/test/single_test_file/test_usp_replicated_parity_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:
from types import SimpleNamespace
import torch.nn.functional as F
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
rank = int(os.environ["RANK"])
world = int(os.environ["WORLD_SIZE"])
torch.cuda.set_device(rank)
init_distributed_environment(world_size=world, rank=rank, local_rank=rank)
initialize_model_parallel(
sequence_parallel_degree=world, ulysses_degree=world, ring_degree=1
)
import sglang.multimodal_gen.runtime.layers.attention.layer as L
L.get_forward_context = lambda: SimpleNamespace(attn_metadata=None)
class Sdpa:
def __init__(self, scale):
self.scale = scale
def forward(self, q, k, v, _ctx):
return F.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
dropout_p=0.0,
is_causal=False,
scale=self.scale,
).transpose(1, 2)
B, SHARD, H, D, REP = 1, 64, 30, 128, 6
S = SHARD * world
scale = D**-0.5
torch.manual_seed(0)
dev = torch.device(f"cuda:{rank}")
qf = torch.randn(B, S + REP, H, D, device=dev, dtype=torch.bfloat16)
kf = torch.randn(B, S + REP, H, D, device=dev, dtype=torch.bfloat16)
vf = torch.randn(B, S + REP, H, D, device=dev, dtype=torch.bfloat16)
ref = F.scaled_dot_product_attention(
qf.transpose(1, 2), kf.transpose(1, 2), vf.transpose(1, 2), scale=scale
).transpose(1, 2)
attn = L.USPAttention.__new__(L.USPAttention)
attn.causal = False
attn.softmax_scale = scale
attn.attn_impl = Sdpa(scale)
attn.skip_sequence_parallel = False
attn.enable_packed_qkv_input_a2a = False
attn.allow_cudnn_sdp = False
attn.backend = L.AttentionBackendEnum.TORCH_SDPA
attn.dtype = torch.bfloat16
attn.dropout_p = 0.0
attn.sp_attention_mode = "ulysses"
attn.sp_attention_mode_is_auto = False
failures = []
sl = slice(rank * SHARD, (rank + 1) * SHARD)
out = attn.forward(
torch.cat([qf[:, sl], qf[:, S:]], dim=1),
torch.cat([kf[:, sl], kf[:, S:]], dim=1),
torch.cat([vf[:, sl], vf[:, S:]], dim=1),
num_replicated_suffix=REP,
)
exp = torch.cat([ref[:, sl], ref[:, S:]], dim=1)
if not torch.equal(out, exp):
d = (out.float() - exp.float()).abs()
failures.append(f"suffix not bitwise: mae={d.mean():.3e} max={d.max():.3e}")
out_p = attn.forward(
torch.cat([qf[:, S:], qf[:, sl]], dim=1),
torch.cat([kf[:, S:], kf[:, sl]], dim=1),
torch.cat([vf[:, S:], vf[:, sl]], dim=1),
num_replicated_prefix=REP,
)
exp_p = torch.cat([ref[:, S:], ref[:, sl]], dim=1)
dp = (out_p.float() - exp_p.float()).abs()
if dp.max().item() > 1e-2:
failures.append(f"prefix drift: mae={dp.mean():.3e} max={dp.max():.3e}")
for f in failures:
print(f"FAILURE rank{rank}: {f}", flush=True)
return 1 if failures else 0
class TestUSPReplicatedParity(CustomTestCase):
def test_replicated_parity_two_ranks(self):
if not current_platform.is_cuda():
self.skipTest("CUDA-only test")
if torch.cuda.device_count() < _WORLD:
self.skipTest(f"needs {_WORLD} GPUs")
procs = []
for rank in range(_WORLD):
env = os.environ.copy()
env.update(
{
"RANK": str(rank),
"LOCAL_RANK": str(rank),
"WORLD_SIZE": str(_WORLD),
"MASTER_ADDR": "127.0.0.1",
"MASTER_PORT": "29751",
}
)
procs.append(
subprocess.Popen(
[sys.executable, __file__],
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
)
)
outputs = [p.communicate(timeout=300)[0] for p in procs]
codes = [p.returncode for p in procs]
if any(codes):
self.fail("worker failed:\n" + "\n".join(outputs))
if __name__ == "__main__":
if "RANK" in os.environ:
sys.exit(_worker())
unittest.main()
python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py
+20
View File
@@ -107,3 +107,23 @@ class TestUSPAttentionReplicatedPrefix(unittest.TestCase):
if __name__ == "__main__":
unittest.main()
class TestUSPAttentionMaskedReplicatedGuard(unittest.TestCase):
def test_masked_path_rejects_replicated_tokens(self):
obj = USPAttention.__new__(USPAttention)
obj.attn_impl = _CaptureAttn()
obj.skip_sequence_parallel = False
obj.sp_attention_mode = "ulysses"
obj.sp_attention_mode_is_auto = False
q = torch.randn(1, 6, 2, 4)
mask = torch.ones(1, 6, dtype=torch.bool)
with (
patch(
f"{_LAYER}.get_forward_context",
return_value=MagicMock(attn_metadata=None),
),
patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2),
):
with self.assertRaisesRegex(NotImplementedError, "replicated"):
obj.forward(q, q, q, attn_mask=mask, num_replicated_suffix=2)