[Perf][MUSA] Restore MiniMax-H3 dynamic RoPE fusion (#5881)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -69,6 +69,7 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
 
     attention = object.__new__(MiniMaxH3Attention)
     nn.Module.__init__(attention)
+    attention.rot_dim = 96
     attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
     attention.rope._forward_method = attention.rope.forward_native
 

vllm_omni/diffusion/layers/norm.py

@@ -68,6 +68,7 @@ class RMSNorm(CustomOp):
     def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
+        self.hidden_size = hidden_size
         self.variance_epsilon = eps
 
     def _forward_fused(self, x: torch.Tensor) -> torch.Tensor:
@@ -107,6 +108,19 @@ class RMSNorm(CustomOp):
         except Exception:
             return self.forward_native(x)
 
+    def forward_musa(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        # Preserve the aten::rms_norm graph so dynamic Inductor can fuse the
+        # H3 Q/K norm with its inline RoPE path on MUSA.
+        return F.rms_norm(
+            x,
+            (self.hidden_size,),
+            self.weight,
+            self.variance_epsilon,
+        )
+
     def forward_npu(
         self,
         x: torch.Tensor,

vllm_omni/diffusion/layers/rope.py

@@ -193,7 +193,16 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        return self.forward_native(x, cos, sin)
+        # Keep the full-dimension NeoX form inline for MUSA compilation.
+        # Other layouts use the shared native implementation.
+        if self.half_head_dim or self.interleaved:
+            return self.forward_native(x, cos, sin)
+        if cos.dim() == 3:
+            cos, sin = cos[0], sin[0]
+        cos = cos.unsqueeze(-2)
+        sin = sin.unsqueeze(-2)
+        x1, x2 = x.chunk(2, dim=-1)
+        return x * cos + torch.cat((-x2, x1), dim=-1) * sin
 
     def forward_native(
         self,

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -333,6 +333,7 @@ class MiniMaxH3Attention(nn.Module):
         )
         self.num_heads = self.qkv_proj.num_heads
         self.num_kv_heads = self.qkv_proj.num_kv_heads
+        self.rot_dim = 6 * arch.rope_inv_freq_len
         self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
@@ -365,7 +366,7 @@ class MiniMaxH3Attention(nn.Module):
         x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin
         are cast to the activation dtype before the elementwise math.
         """
-        rot_dim = freqs.shape[-1]
+        rot_dim = self.rot_dim
         x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
         cos = torch.cos(freqs).to(x.dtype)  # [T, rot_dim]
         sin = torch.sin(freqs).to(x.dtype)