[diffusion] model: Fix a performance bug in the Wan model about usp (#19340)

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
 )
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
 from sglang.multimodal_gen.runtime.layers.usp import (
+    _ulysses_input_split,
     _usp_input_all_to_all,
     _usp_output_all_to_all,
     ring_attn,
@@ -303,6 +304,7 @@ class USPAttention(nn.Module):
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
         prefix: str = "",
         dropout_rate: float = 0.0,
+        is_cross_attention: bool = False,
         **extra_impl_args,
     ) -> None:
         super().__init__()
@@ -335,6 +337,7 @@ class USPAttention(nn.Module):
         self.dtype = dtype
         self.causal = causal
         self.dropout_p = dropout_rate
+        self.is_cross_attention = is_cross_attention
 
     def forward(
         self,
@@ -366,11 +369,15 @@ class USPAttention(nn.Module):
         if get_ulysses_parallel_world_size() > 1:
             # -> [B, S, H_local, D]
             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 self.is_cross_attention:
+                k = _ulysses_input_split(k, dim=2)
+                v = _ulysses_input_split(v, dim=2)
+            else:
+                k = _usp_input_all_to_all(k, head_dim=2)
+                v = _usp_input_all_to_all(v, head_dim=2)
 
         # Ring Attention within subgroups or local attention
-        if get_ring_parallel_world_size() > 1:
+        if get_ring_parallel_world_size() > 1 and not self.is_cross_attention:
             out = ring_attn(
                 q,
                 k,

python/sglang/multimodal_gen/runtime/layers/usp.py

@@ -9,6 +9,7 @@ from torch.distributed.tensor.experimental._attention import _cp_options
 
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_sp_group,
+    get_ulysses_parallel_rank,
     get_ulysses_parallel_world_size,
 )
 from sglang.srt.utils.common import torch_release
@@ -158,6 +159,22 @@ def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor:
     return x
 
 
+def _ulysses_input_split(x: torch.Tensor, dim: int = 1) -> torch.Tensor:
+    world_size = get_ulysses_parallel_world_size()
+    if world_size <= 1:
+        return x
+    rank = get_ulysses_parallel_rank()
+    assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}"
+
+    dim_to_split_size = x.shape[dim]
+
+    assert (
+        dim_to_split_size % world_size == 0
+    ), f"The size of dimension {dim} ({dim_to_split_size}) must be divisible by world_size ({world_size})"
+
+    return torch.tensor_split(x, world_size, dim=dim)[rank].contiguous()
+
+
 def ring_attn(
     query: torch.Tensor,
     key: torch.Tensor,

python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py

@@ -220,6 +220,7 @@ class ConditionalCrossAttention(nn.Module):
             head_size=self.head_dim,
             causal=False,
             softmax_scale=None,
+            is_cross_attention=True,
         )
 
     def forward(

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -134,6 +134,7 @@ class WanSelfAttention(nn.Module):
         eps=1e-6,
         parallel_attention=False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        is_cross_attention: bool = False,
         quant_config: QuantizationConfig | None = None,
     ) -> None:
         assert dim % num_heads == 0
@@ -173,6 +174,7 @@ class WanSelfAttention(nn.Module):
             softmax_scale=None,
             causal=False,
             supported_attention_backends=supported_attention_backends,
+            is_cross_attention=is_cross_attention,
         )
 
     def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int):
@@ -187,6 +189,12 @@ class WanSelfAttention(nn.Module):
 
 
 class WanT2VCrossAttention(WanSelfAttention):
+    def __init__(self, *args, **kwargs):
+        super().__init__(
+            *args,
+            **kwargs,
+            is_cross_attention=True,
+        )
 
     def forward(self, x, context, context_lens):
         r"""
@@ -240,6 +248,7 @@ class WanI2VCrossAttention(WanSelfAttention):
             qk_norm,
             eps,
             supported_attention_backends=supported_attention_backends,
+            is_cross_attention=True,
             quant_config=quant_config,
         )
 
@@ -359,6 +368,7 @@ class WanTransformerBlock(nn.Module):
                 head_size=dim // num_heads,
                 causal=False,
                 supported_attention_backends=self_attn_backends,
+                is_cross_attention=False,
                 prefix=f"{prefix}.attn1",
             )