diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 16e195e56618..b600256841f7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -32,6 +32,10 @@ ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.utils import get_compute_dtype +from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( + NDRotaryEmbedding, + _apply_rotary_emb, +) class UlyssesAttention(nn.Module): @@ -346,6 +350,8 @@ def forward( replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, + freqs_cis: torch.Tensor | None = None, + is_neox_style: bool = False, ) -> torch.Tensor: """ Forward pass for USPAttention. @@ -385,6 +391,11 @@ def forward( dropout_p=self.dropout_p, ) else: + if freqs_cis is not None: + cos, sin = freqs_cis + q, k = _apply_rotary_emb( + q, cos, sin, is_neox_style=is_neox_style + ), _apply_rotary_emb(k, cos, sin, is_neox_style=is_neox_style) # -> [B, S, H_local, D] out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py index e13ea05ad7ff..56d3d8bbc99c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py @@ -524,7 +524,7 @@ def _forward_cached_from_grid( # Apply sequence parallel sharding to the sizes and compute shard offset shard_sizes = list(sizes) shard_offsets = [0] * self.ndim - if sp_world_size > 1 and shard_dim >= 0: + if sp_world_size > 1 and shard_dim >= 0 and False: assert sizes[shard_dim] % sp_world_size == 0, ( f"Dimension {shard_dim} with size {sizes[shard_dim]} is not divisible " f"by sequence parallel world size {sp_world_size}" diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 711477192f12..2901861818e4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -400,11 +400,11 @@ def forward( value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) # Apply rotary embeddings - cos, sin = freqs_cis - query, key = _apply_rotary_emb( - query, cos, sin, is_neox_style=False - ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) - attn_output = self.attn1(query, key, value) + # cos, sin = freqs_cis + # query, key = _apply_rotary_emb( + # query, cos, sin, is_neox_style=False + # ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) + attn_output = self.attn1(query, key, value, freqs_cis=freqs_cis, is_neox_style=False) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1)