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

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

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.
+        q_shard = _usp_input_all_to_all(q_shard, head_dim=2)
-        # 1. rotate the replicated suffix to the front
+        k_shard = _usp_input_all_to_all(k_shard, head_dim=2)
-        # 2. reuse the validated replicated-prefix path, then
+        v_shard = _usp_input_all_to_all(v_shard, head_dim=2)
-        # 3. rotate the output back
+
-        out = self._forward_with_replicated_prefix(
+        h_local = q_shard.shape[2]
-            torch.cat([q_rep, q_shard], dim=1),
+        kv_h_local = k_shard.shape[2]
-            torch.cat([k_rep, k_shard], dim=1),
+        h_start = sp_rank * h_local
-            torch.cat([v_rep, v_shard], dim=1),
+        kv_h_start = sp_rank * kv_h_local
-            ctx_attn_metadata,
+        q_rep = q_rep[:, :, h_start : h_start + h_local, :].contiguous()
-            num_rep,
+        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

@@ -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 = USPAttention(
-        self.attn = UlyssesAttention(
             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 = USPAttention(
-        self.attn = UlyssesAttention(
             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(
+            attn_output = self.attn(
-                img_q, img_k, img_v, txt_q, txt_k, txt_v
+                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

@@ -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 (
+        hidden_states = self.attn(
-            num_replicated_suffix > 0
+            q,
-            and get_sp_world_size() > 1
+            k,
-            and get_ring_parallel_world_size() == 1
+            v,
-        ):
+            attn_mask=attn_mask,
-            # the cap (last num_replicated_suffix tokens), as condition, should be replicated
+            attn_mask_meta=attn_mask_meta,
-            q_shard, q_rep = (
+            num_replicated_prefix=num_replicated_prefix,
-                q[:, :-num_replicated_suffix],
+            num_replicated_suffix=num_replicated_suffix,
-                q[:, -num_replicated_suffix:],
+            skip_sequence_parallel_override=skip_sequence_parallel_override,
-            )
+        )
-            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 = hidden_states.flatten(2)
 
         hidden_states, _ = self.to_out[0](hidden_states)

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

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

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)