From 38a2500bdf887cd0d23c1c1616c59faec6ffd317 Mon Sep 17 00:00:00 2001 From: Wenjia Li Date: Wed, 12 Aug 2026 19:12:54 +0800 Subject: [PATCH] [Perf] MiniMax H3 Qwen3-VL support fused RMSNorm on NPU (#5915) Signed-off-by: Wenjia Li Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- .../minimax_h3/test_minimax_h3_rmsnorm.py | 26 +++++++++++++ .../diffusion/models/minimax_h3/encoder.py | 38 +++++++++++-------- 2 files changed, 49 insertions(+), 15 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py new file mode 100644 index 000000000..551fd68a5 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 15098b098..2e07ace96 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/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,