[Perf][NPU] Fuse MiniMax H3 Qwen3-VL SwiGLU (#6167)

Signed-off-by: Wenjia Li <wjialish@gmail.com>

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,
+    )

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

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)