diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 8a10b18a9..3228757bc 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index e82235009..f4de11797 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py index cbc73d6e8..35a797ede 100644 --- a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 059d5b582..39b98554c 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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", )