[Diffusion][Refactor] Refactor and extract complex RoPE implementation to layers/rotary_embedding for MOVA DiT (#31453)

Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/__init__.py

@@ -35,12 +35,14 @@ from .mrope import (
 )
 from .utils import (
     _apply_rotary_emb,
+    _apply_rotary_emb_complex,
     apply_flashinfer_rope_qk_inplace,
 )
 
 __all__ = [
     # _utils
     "_apply_rotary_emb",
+    "_apply_rotary_emb_complex",
     "apply_flashinfer_rope_qk_inplace",
     # _base
     "RotaryEmbedding",

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py

@@ -65,6 +65,33 @@ def _apply_rotary_emb(
         return apply_rotary_embedding(x, cos, sin, interleaved)
 
 
+def _apply_rotary_emb_complex(
+    x: torch.Tensor,  # [b, s, h, d]
+    freqs: torch.Tensor,  # [s, 1, d // 2]
+) -> torch.Tensor:  # [b, s, h, d]
+    """
+    Apply complex rotary positional embeddings designed for interleaved=True, neox_style=False.
+    Works by mathematically mapping the complex multiplication
+    (a + ib) * (cos + isin) to the interleaved layout.
+
+    Args:
+        x: Input activation tensor in bf16/fp16.
+            Shape: [batch, num_tokens, num_heads, head_size]
+        freqs: Complex-valued frequency tensor in complex64 format.
+            Shape: [num_tokens, 1, head_size // 2]
+
+    Returns:
+        torch.Tensor: The same shape and dtype as x.
+    """
+    b, s, h, d = x.shape
+    dtype_c = torch.float64
+
+    x_complex = torch.view_as_complex(x.to(dtype_c).reshape(b, s, h, d // 2, 2))
+    x_out = torch.view_as_real(x_complex * freqs)
+    x_out = x_out.view(b, s, h, d)
+    return x_out.to(x.dtype)
+
+
 @debug_kernel_api
 def apply_flashinfer_rope_qk_inplace(
     q: torch.Tensor,

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

@@ -34,6 +34,9 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
+    _apply_rotary_emb_complex,
+)
 from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
     LayerwiseOffloadableModuleMixin,
 )
@@ -83,16 +86,6 @@ def precompute_freqs_cis(
     return freqs_cis
 
 
-def rope_apply_head_dim(x, freqs, head_dim):
-    x = rearrange(x, "b s (n d) -> b s n d", d=head_dim)
-    x_out = torch.view_as_complex(
-        x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)
-    )
-    # print(f"{x_out.shape = }, {freqs.shape = }")
-    x_out = torch.view_as_real(x_out * freqs).flatten(2)
-    return x_out.to(x.dtype)
-
-
 class SelfAttention(nn.Module):
     """
     Self-Attention module for MOVA DiT with Sequence Parallelism support.
@@ -173,19 +166,20 @@ class SelfAttention(nn.Module):
             q = self.norm_q(q)
             k = self.norm_k(k)
 
+        b, s, _ = q.shape
+        q = q.view(b, s, self.num_heads_per_rank, self.head_dim)
+        k = k.view(b, s, self.num_heads_per_rank, self.head_dim)
+        v = v.view(b, s, self.num_heads_per_rank, self.head_dim)
+
         # Apply RoPE
-        q = rope_apply_head_dim(q, freqs, self.head_dim)
+        q = _apply_rotary_emb_complex(q, freqs)
-        k = rope_apply_head_dim(k, freqs, self.head_dim)
+        k = _apply_rotary_emb_complex(k, freqs)
 
         # USPAttention expects [B, S_local, H, D] format
-        q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
-        k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
-        v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
-
         # USPAttention handles SP communication internally; the tail meta keeps
         # SP padding out of the softmax.
         out = self.attn(q, k, v, attn_mask_meta=attn_mask_meta)
-        out = rearrange(out, "b s n d -> b s (n d)")
+        out = out.view(b, s, -1)
 
         out, _ = self.o(out)
         return out