[Perf] MiniMax H3 Qwen3-VL support fused RMSNorm on NPU (#5915)

Signed-off-by: Wenjia Li <wjialish@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
+
+
+def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
+    """Qwen keeps BF16 gamma while native fallback accumulates in FP32."""
+    eps = 1e-6
+    norm = MiniMaxH3Qwen3VLRMSNorm(hidden_size=4, eps=eps, dtype=torch.bfloat16)
+    norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
+    x = torch.tensor(
+        [[1.0, -2.0, 3.0, -4.0], [-5.0, 6.0, -7.0, 8.0]],
+        dtype=torch.bfloat16,
+    )
+
+    x_fp32 = x.float()
+    expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype)
+
+    assert isinstance(norm, RMSNorm)
+    assert norm.weight.dtype == torch.bfloat16
+    assert set(norm.state_dict()) == {"weight"}
+    torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)

vllm_omni/diffusion/models/minimax_h3/encoder.py

@@ -39,6 +39,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm
+
 MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50
 MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
 
@@ -291,18 +293,22 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module):
 # ---------------------------------------------------------------------------
 
 
-class MiniMaxH3Qwen3VLRMSNorm(nn.Module):
-    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
+class MiniMaxH3Qwen3VLRMSNorm(DiffusionRMSNorm):
+    """Qwen3-VL RMSNorm using the common implementation.
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
+    The model-specific name keeps checkpoint weight keys stable, while the
+    common RMSNorm dispatches to torch_npu.npu_rms_norm on Ascend. Gamma uses
+    the checkpoint/model dtype, but native fallbacks keep the RMS variance
+    reduction and scaling in float32.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        eps: float = 1e-6,
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__(hidden_size, eps=eps, dtype=dtype)
 
 
 def _rotate_half(x: torch.Tensor) -> torch.Tensor:
@@ -685,8 +691,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module):
             dtype=dtype,
             input_is_parallel=True,
         )
-        self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype)
+        self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype)
 
     def forward(
         self,
@@ -760,8 +766,10 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module):
         self.hidden_size = config.hidden_size
         self.self_attn = MiniMaxH3Qwen3VLTextAttention(group, config, dtype)
         self.mlp = MiniMaxH3Qwen3VLTextMLP(group, config, dtype)
-        self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype)
+        self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, dtype=dtype
+        )
 
     def forward(
         self,