From ffd11c181e5fa4e558ba2570f326fb41371c0421 Mon Sep 17 00:00:00 2001 From: Wenjia Li Date: Fri, 14 Aug 2026 17:49:26 +0800 Subject: [PATCH] [Perf][NPU] Fuse MiniMax H3 Qwen3-VL RoPE (#6061) Signed-off-by: Wenjia Li --- .../test_minimax_h3_qwen3vl_rope.py | 104 ++++++++++++++++++ vllm_omni/platforms/npu/models/minimax_h3.py | 43 ++++++++ vllm_omni/platforms/npu/platform.py | 2 + 3 files changed, 149 insertions(+) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py create mode 100644 vllm_omni/platforms/npu/models/minimax_h3.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py new file mode 100644 index 000000000..9ac3243bb --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +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("num_q_heads,num_kv_heads,seq_len", [(8, 2, 16), (64, 8, 1)]) +def test_qwen3_vl_npu_rope_matches_bnsd_reference( + num_q_heads: int, + num_kv_heads: int, + seq_len: int, +) -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu + + device = torch.device("npu") + dtype = torch.bfloat16 + q = torch.randn(1, num_q_heads, seq_len, 128, device=device, dtype=dtype) + k = torch.randn(1, num_kv_heads, seq_len, 128, device=device, dtype=dtype) + cos = torch.randn(1, seq_len, 128, device=device, dtype=dtype) + sin = torch.randn(1, seq_len, 128, device=device, dtype=dtype) + + q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin) + cos_bnsd = cos.unsqueeze(1) + sin_bnsd = sin.unsqueeze(1) + + def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + torch.testing.assert_close( + q_out, + q * cos_bnsd + rotate_half(q) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + torch.testing.assert_close( + k_out, + k * cos_bnsd + rotate_half(k) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_rope_preserves_batched_mrope_positions() -> None: + """Keep per-sample Qwen3-VL M-RoPE frequencies on the BNSD NPU path.""" + from types import SimpleNamespace + + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLTextRotaryEmbedding, + ) + from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu + + device = torch.device("npu") + dtype = torch.bfloat16 + batch_size, seq_len, head_dim = 2, 16, 128 + config = SimpleNamespace( + hidden_size=head_dim, + num_attention_heads=1, + rope_parameters={ + "rope_theta": 1_000_000.0, + "mrope_section": [24, 20, 20], + }, + ) + rope = MiniMaxH3Qwen3VLTextRotaryEmbedding(config).to(device) + positions = torch.arange(seq_len, device=device) + position_ids = torch.stack( + ( + torch.stack((positions, positions + 7)), + torch.stack((positions * 2, positions * 2 + 3)), + torch.stack((positions * 3, positions * 3 + 5)), + ) + ) + q = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn(batch_size, 2, seq_len, head_dim, device=device, dtype=dtype) + cos, sin = rope(q, position_ids) + + q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin) + cos_bnsd = cos.unsqueeze(1) + sin_bnsd = sin.unsqueeze(1) + + def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + assert not torch.equal(cos[0], cos[1]) + torch.testing.assert_close( + q_out, + q * cos_bnsd + rotate_half(q) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + torch.testing.assert_close( + k_out, + k * cos_bnsd + rotate_half(k) * sin_bnsd, + 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 new file mode 100644 index 000000000..a05abd654 --- /dev/null +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""NPU patches for the MiniMax H3 Qwen3-VL text encoder.""" + +from __future__ import annotations + +import torch +from vllm.logger import init_logger + +from vllm_omni.platforms.npu.layers.rotary_embedding import ( + npu_rotary_mul_with_bsnd_fallback, +) + +logger = init_logger(__name__) + +_PATCHED = False + + +def _apply_rotary_pos_emb_npu( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply Qwen3-VL text RoPE with fused BNSD rotary multiplication.""" + return ( + npu_rotary_mul_with_bsnd_fallback(q, cos, sin, unsqueeze_dim=1), + npu_rotary_mul_with_bsnd_fallback(k, cos, sin, unsqueeze_dim=1), + ) + + +def apply_minimax_h3_qwen3vl_patch() -> None: + """Route MiniMax H3 Qwen3-VL text RoPE to the Ascend fused operator.""" + global _PATCHED + if _PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder._apply_rotary_pos_emb = _apply_rotary_pos_emb_npu + _PATCHED = True + logger.debug("Applied NPU fused RoPE patch for MiniMax H3 Qwen3-VL text encoder") diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 525725fcf..16f87b21d 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -36,6 +36,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): from vllm_ascend.utils import adapt_patch from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches + from vllm_omni.platforms.npu.models.minimax_h3 import apply_minimax_h3_qwen3vl_patch from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import ( apply_qwen3_tts_code2wav_patch, ) @@ -44,6 +45,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): ) adapt_patch(is_global_patch=True) + apply_minimax_h3_qwen3vl_patch() apply_qwen3_tts_code2wav_patch() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches()