diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py index bcd074483..2b168499e 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/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 diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 14ab677a0..83897eb15 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/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, diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 633dbd275..fae1a3414 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/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, diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 610c20d19..d41395bb5 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/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)