From 8ecd1f6d5cc91aab8a475a861213720b336e2f65 Mon Sep 17 00:00:00 2001 From: Wenjia Li Date: Fri, 14 Aug 2026 19:47:17 +0800 Subject: [PATCH] [Perf][NPU] Fuse MiniMax H3 Qwen3-VL SwiGLU (#6167) Signed-off-by: Wenjia Li --- .../test_minimax_h3_qwen3vl_swiglu.py | 46 +++++++++++++++++++ vllm_omni/platforms/npu/models/minimax_h3.py | 36 +++++++++++++-- vllm_omni/platforms/npu/platform.py | 8 ++++ 3 files changed, 86 insertions(+), 4 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py new file mode 100644 index 000000000..2664990b4 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize("shape", [(1, 16, 256), (2, 1, 256)]) +def test_qwen3_vl_npu_swiglu_matches_split_reference( + shape: tuple[int, int, int], +) -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed + + packed = torch.randn(*shape, device="npu", dtype=torch.bfloat16) + gate, up = packed.chunk(2, dim=-1) + + torch.testing.assert_close( + npu_swiglu_from_packed(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed + + x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16) + weight = torch.randn(256, 128, device="npu", dtype=torch.bfloat16) + packed = F.linear(x, weight) + gate = F.linear(x, weight[:128]) + up = F.linear(x, weight[128:]) + + torch.testing.assert_close( + npu_swiglu_from_packed(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) diff --git a/vllm_omni/platforms/npu/models/minimax_h3.py b/vllm_omni/platforms/npu/models/minimax_h3.py index a05abd654..78e0b3f14 100644 --- a/vllm_omni/platforms/npu/models/minimax_h3.py +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -5,7 +5,11 @@ from __future__ import annotations +from typing import Any + import torch +import torch.nn.functional as F +import torch_npu from vllm.logger import init_logger from vllm_omni.platforms.npu.layers.rotary_embedding import ( @@ -14,7 +18,8 @@ from vllm_omni.platforms.npu.layers.rotary_embedding import ( logger = init_logger(__name__) -_PATCHED = False +_ROPE_PATCHED = False +_SWIGLU_PATCHED = False def _apply_rotary_pos_emb_npu( @@ -32,12 +37,35 @@ def _apply_rotary_pos_emb_npu( def apply_minimax_h3_qwen3vl_patch() -> None: """Route MiniMax H3 Qwen3-VL text RoPE to the Ascend fused operator.""" - global _PATCHED - if _PATCHED: + global _ROPE_PATCHED + if _ROPE_PATCHED: return from vllm_omni.diffusion.models.minimax_h3 import encoder encoder._apply_rotary_pos_emb = _apply_rotary_pos_emb_npu - _PATCHED = True + _ROPE_PATCHED = True logger.debug("Applied NPU fused RoPE patch for MiniMax H3 Qwen3-VL text encoder") + + +def npu_swiglu_from_packed(gate_up: torch.Tensor) -> torch.Tensor: + """Apply fused SwiGLU to a packed gate/up projection tensor.""" + return torch_npu.npu_swiglu(gate_up, dim=-1) + + +def _forward_minimax_h3_qwen3vl_text_mlp_npu(self: Any, x: torch.Tensor) -> torch.Tensor: + """Run the Qwen3-VL MLP with one packed GEMM and fused SwiGLU.""" + gate_up = F.linear(x, self.gate_up_proj.weight) + return self.down_proj(npu_swiglu_from_packed(gate_up)) + + +def apply_minimax_h3_qwen3vl_swiglu_patch() -> None: + """Route MiniMax H3 Qwen3-VL text MLP to the Ascend fused SwiGLU path.""" + global _SWIGLU_PATCHED + if _SWIGLU_PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder.MiniMaxH3Qwen3VLTextMLP.forward = _forward_minimax_h3_qwen3vl_text_mlp_npu + _SWIGLU_PATCHED = True diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 16f87b21d..b17f447c4 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -82,6 +82,14 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): def init_diffusion_model_runner_runtime(cls, vllm_config: Any, od_config: Any, device: torch.device) -> None: from vllm_ascend.ascend_forward_context import set_mc2_mask, set_mc2_tokens_capacity + from vllm_omni.platforms.npu.models.minimax_h3 import ( + apply_minimax_h3_qwen3vl_swiglu_patch, + ) + + # The patch imports the MiniMax encoder, which depends on + # current_omni_platform. Run it only after platform construction has + # completed, but before the diffusion pipeline is loaded. + apply_minimax_h3_qwen3vl_swiglu_patch() set_mc2_tokens_capacity(vllm_config, od_config.max_num_seqs, 1) set_mc2_mask(vllm_config, device)