diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/__init__.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/__init__.py index 1060e304c..fdbaf595b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/__init__.py +++ b/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", diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py index 559703c71..c3e740e35 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py index 4f8772c74..616f9a447 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py +++ b/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) - k = rope_apply_head_dim(k, freqs, self.head_dim) + q = _apply_rotary_emb_complex(q, freqs) + 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