enable all2all overlap, and use rope overlap last v gemm

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)
 

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

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)