From 6dc77e490dd9428d1ec9bbf4ff45bc1abf060ac1 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 7 Aug 2026 12:36:23 +0800 Subject: [PATCH] [diffusion] chore: route zimage and hunyuanvideo attention through USPAttention (#33923) Co-authored-by: Claude Fable 5 --- .../runtime/layers/attention/layer.py | 88 ++++++++-- .../runtime/models/dits/hunyuanvideo.py | 29 ++-- .../runtime/models/dits/zimage.py | 57 ++----- .../multimodal_gen/test/server/gpu_cases.py | 3 + .../test_usp_replicated_parity_2_gpu.py | 155 ++++++++++++++++++ .../test_usp_attention_replicated_prefix.py | 20 +++ 6 files changed, 281 insertions(+), 71 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/single_test_file/test_usp_replicated_parity_2_gpu.py diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index aa23827b1..d88df2eee 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index 16fed4a6c..8bd36dcab 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index dbb801256..f558a7c15 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 534194e94..59e9be879 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -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, }, } diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_usp_replicated_parity_2_gpu.py b/python/sglang/multimodal_gen/test/single_test_file/test_usp_replicated_parity_2_gpu.py new file mode 100644 index 000000000..2da3bb6c0 --- /dev/null +++ b/python/sglang/multimodal_gen/test/single_test_file/test_usp_replicated_parity_2_gpu.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py b/python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py index bc0dcae36..40ce07404 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py +++ b/python/sglang/multimodal_gen/test/unit/test_usp_attention_replicated_prefix.py @@ -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)