From 40dcfdbb4343fca917b34287cba620eb436a856d Mon Sep 17 00:00:00 2001 From: XiongZhongwei <1228479724@qq.com> Date: Wed, 26 Aug 2026 09:14:05 +0800 Subject: [PATCH] [Perf][NPU] Enable native GQA & fused AddNorm for MiniMax-H3 encoder (#6040) Signed-off-by: MarkPoloChina <1228479724@qq.com> Signed-off-by: XiongZhongwei <1228479724@qq.com> --- tests/config/test_config_import_cycle.py | 2 + tests/diffusion/layers/test_norm.py | 91 ++++++++++++ tests/diffusion/layers/test_norm_cuda.py | 34 +++++ .../minimax_h3/test_minimax_h3_contract.py | 74 ++++++++++ .../minimax_h3/test_minimax_h3_qwen3vl_gqa.py | 132 ++++++++++++++++++ .../test_minimax_h3_qwen3vl_rope.py | 33 ++++- .../minimax_h3/test_minimax_h3_rmsnorm.py | 52 ++++++- vllm_omni/diffusion/layers/norm.py | 43 ++++-- .../diffusion/models/minimax_h3/encoder.py | 38 ++--- vllm_omni/platforms/npu/models/minimax_h3.py | 45 ++++++ vllm_omni/platforms/npu/platform.py | 4 +- 11 files changed, 520 insertions(+), 28 deletions(-) create mode 100644 tests/diffusion/layers/test_norm_cuda.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py diff --git a/tests/config/test_config_import_cycle.py b/tests/config/test_config_import_cycle.py index 8a28b7afb..e149b63ee 100644 --- a/tests/config/test_config_import_cycle.py +++ b/tests/config/test_config_import_cycle.py @@ -66,6 +66,8 @@ def test_npu_platform_defers_minimax_h3_encoder_patch() -> None: runtime_src = src[runtime_start:] assert "apply_minimax_h3_qwen3vl_patch" not in init_src assert "apply_minimax_h3_qwen3vl_patch" in runtime_src + assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src + assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index 71125684b..4794ca1c2 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/tests/diffusion/layers/test_norm.py @@ -340,6 +340,97 @@ def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.Mo assert captured["epsilon"].item() == pytest.approx(1e-5) +# ── RMSNorm residual tests ── + + +def test_rmsnorm_matches_native_residual_contract(): + """The non-NPU path preserves separate residual-add and RMSNorm math.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + eps = 1e-6 + norm = RMSNorm(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]], dtype=torch.bfloat16) + residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16) + + expected_residual = residual + x + expected_fp32 = expected_residual.float() + expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps) + expected_output = (expected_output * norm.weight.float()).to(x.dtype) + output, updated_residual = norm.forward_native(x, residual) + + torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "forward_name", + ["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"], +) +def test_rmsnorm_non_npu_residual_delegates_to_native( + monkeypatch: pytest.MonkeyPatch, + forward_name: str, +): + """Platform wrappers keep residual arithmetic in forward_native.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16) + x = torch.randn(2, 4, dtype=torch.bfloat16) + residual = torch.randn_like(x) + expected_output = torch.randn_like(x) + expected_residual = torch.randn_like(x) + captured: dict[str, object] = {} + + def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None): + captured.update(x=native_x, residual=native_residual) + return expected_output, expected_residual + + monkeypatch.setattr(norm, "forward_native", forward_native) + + output, updated_residual = getattr(norm, forward_name)(x, residual) + + assert output is expected_output + assert updated_residual is expected_residual + assert captured["x"] is x + assert captured["residual"] is residual + + +def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch): + """NPU keeps the direct fused AddRMSNorm path instead of IR fallback.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + captured: dict[str, object] = {} + fused_output = torch.randn(2, 4, dtype=torch.bfloat16) + fused_residual = torch.randn(2, 4, dtype=torch.bfloat16) + + def npu_add_rms_norm( + x: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + ): + captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon) + return fused_output, torch.empty(0), fused_residual + + monkeypatch.setitem( + sys.modules, + "torch_npu", + types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm), + ) + norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16) + x = torch.randn(2, 4, dtype=torch.bfloat16) + residual = torch.randn(2, 4, dtype=torch.bfloat16) + + output, updated_residual = norm.forward_npu(x, residual) + + assert output is fused_output + assert updated_residual is fused_residual + assert captured["x"] is x + assert captured["residual"] is residual + assert captured["gamma"] is norm.weight + assert captured["epsilon"] == 1e-5 + + # ── RMSNorm compile-path regression tests ── diff --git a/tests/diffusion/layers/test_norm_cuda.py b/tests/diffusion/layers/test_norm_cuda.py new file mode 100644 index 000000000..b791eab71 --- /dev/null +++ b/tests/diffusion/layers/test_norm_cuda.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.platforms import current_omni_platform + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.diffusion, + pytest.mark.cuda, + pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"), +] + + +def test_rmsnorm_cuda_residual_matches_native_contract() -> None: + eps = 1e-6 + device = torch.device(current_omni_platform.device_type) + torch.manual_seed(0) + norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device) + x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16) + residual = torch.randn_like(x) + x_before = x.clone() + residual_before = residual.clone() + + expected_output, expected_residual = norm.forward_native(x_before, residual_before) + output, updated_residual = norm.forward_cuda(x, residual) + + torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + torch.testing.assert_close(x, x_before, atol=0, rtol=0) + torch.testing.assert_close(residual, residual_before, atol=0, rtol=0) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 000e02f55..76adf7471 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -964,6 +964,80 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit(): assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9 +def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + + class FakeEncoderGroup: + rank_in_group = 0 + world_size = 2 + + def all_reduce(self, tensor): + del tensor + + config = SimpleNamespace( + hidden_size=8, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=2, + rms_norm_eps=1e-6, + ) + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + attention = encoder_module.MiniMaxH3Qwen3VLTextAttention(FakeEncoderGroup(), config, torch.float32) + # This CPU contract test only needs projection shapes and head reshaping. + # Avoid dispatching platform-specific norm and linear kernels (for example, + # rocm_unquantized_gemm when the test suite is collected on ROCm). + attention.q_norm = nn.Identity() + attention.k_norm = nn.Identity() + with torch.no_grad(): + attention.qkv_proj.weight.zero_() + attention.o_proj = nn.Linear( + attention.qkv_proj.local_num_heads * attention.head_dim, + config.hidden_size, + bias=False, + ) + attn_call = {} + + def fake_attention(query, key, value): + attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape) + return query + + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", fake_attention) + hidden_states = torch.randn(1, 3, config.hidden_size) + cos = torch.ones(1, 3, config.head_dim) + sin = torch.zeros_like(cos) + + output = attention(hidden_states, (cos, sin)) + + assert attn_call["query_shape"] == (1, config.num_attention_heads // 2, 3, config.head_dim) + assert attn_call["key_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim) + assert attn_call["value_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim) + assert output.shape == hidden_states.shape + + +def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + + sdpa_call = {} + + def fake_sdpa(query, key, value, **kwargs): + sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs) + return query + + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + output = encoder_module._scaled_dot_product_attention( + torch.randn(1, 4, 3, 8), + torch.randn(1, 2, 3, 8), + torch.randn(1, 2, 3, 8), + ) + + assert sdpa_call["query_shape"] == sdpa_call["key_shape"] == sdpa_call["value_shape"] == (1, 4, 3, 8) + assert sdpa_call["kwargs"] == {"dropout_p": 0.0, "is_causal": True} + assert output.shape == (1, 4, 3, 8) + + def test_encoder_forward_uses_hook_compatible_encode_entrypoint(): from vllm_omni.diffusion.models.minimax_h3.encoder import ( MiniMaxH3Qwen3VLEncoder, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py new file mode 100644 index 000000000..45a109b28 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py @@ -0,0 +1,132 @@ +# 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.npu] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_minimax_h3_npu_sdpa_patch_is_independent_and_idempotent(monkeypatch) -> None: + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + def sdpa_sentinel(query, key, value): + del key, value + return query + + def rope_sentinel(q, k, cos, sin): + del cos, sin + return q, k + + monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False) + monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False) + monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel) + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel) + + npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch() + + assert encoder_module._scaled_dot_product_attention is npu_minimax_h3._scaled_dot_product_attention_npu + assert encoder_module._apply_rotary_pos_emb is rope_sentinel + assert not npu_minimax_h3._ROPE_PATCHED + assert npu_minimax_h3._SDPA_PATCHED + + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel) + npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch() + + assert encoder_module._scaled_dot_product_attention is sdpa_sentinel + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize( + ("num_heads", "num_kv_heads", "expected_enable_gqa"), + [(2, 2, False), (4, 2, True)], +) +def test_minimax_h3_npu_sdpa_preserves_compressed_heads_and_selects_gqa( + monkeypatch, + num_heads: int, + num_kv_heads: int, + expected_enable_gqa: bool, +) -> None: + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + captured: dict[str, object] = {} + + def fake_sdpa(query, key, value, **kwargs): + captured.update(query=query, key=key, value=value, kwargs=kwargs) + return query + + monkeypatch.setattr(F, "scaled_dot_product_attention", fake_sdpa) + query = torch.randn(1, num_heads, 3, 8) + key = torch.randn(1, num_kv_heads, 3, 8) + value = torch.randn_like(key) + + output = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value) + + assert captured["query"] is query + assert captured["key"] is key + assert captured["value"] is value + assert captured["kwargs"] == { + "dropout_p": 0.0, + "is_causal": True, + "enable_gqa": expected_enable_gqa, + } + assert output is query + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize( + ("num_heads", "num_key_heads", "num_value_heads", "message"), + [ + (4, 2, 1, "key and value to have the same number of heads"), + (4, 0, 0, "at least one KV head"), + (3, 2, 2, "query heads to be a multiple of KV heads"), + ], +) +def test_minimax_h3_npu_sdpa_rejects_invalid_head_configurations( + num_heads: int, + num_key_heads: int, + num_value_heads: int, + message: str, +) -> None: + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + query = torch.randn(1, num_heads, 3, 8) + key = torch.randn(1, num_key_heads, 3, 8) + value = torch.randn(1, num_value_heads, 3, 8) + + with pytest.raises(ValueError, match=message): + npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_minimax_h3_npu_causal_gqa_matches_expanded_kv() -> None: + """Exercise the real torch-npu causal GQA operator.""" + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + torch.manual_seed(0) + batch_size, seq_len = 1, 16 + num_heads, num_kv_heads, head_size = 8, 2, 128 + device = torch.device("npu") + dtype = torch.bfloat16 + + query = torch.randn(batch_size, num_heads, seq_len, head_size, device=device, dtype=dtype) + key = torch.randn(batch_size, num_kv_heads, seq_len, head_size, device=device, dtype=dtype) + value = torch.randn_like(key) + + actual = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value) + + repeat_num = num_heads // num_kv_heads + expected = F.scaled_dot_product_attention( + query, + key.repeat_interleave(repeat_num, dim=1), + value.repeat_interleave(repeat_num, dim=1), + dropout_p=0.0, + is_causal=True, + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) 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 index 9ac3243bb..7422c0b5e 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py @@ -6,7 +6,38 @@ import torch from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_rope_patch_is_independent_and_idempotent(monkeypatch) -> None: + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + def rope_sentinel(q, k, cos, sin): + del cos, sin + return q, k + + def sdpa_sentinel(query, key, value): + del key, value + return query + + monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False) + monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False) + monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel) + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel) + + npu_minimax_h3.apply_minimax_h3_qwen3vl_patch() + + assert encoder_module._apply_rotary_pos_emb is npu_minimax_h3._apply_rotary_pos_emb_npu + assert encoder_module._scaled_dot_product_attention is sdpa_sentinel + assert npu_minimax_h3._ROPE_PATCHED + assert not npu_minimax_h3._SDPA_PATCHED + + monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel) + npu_minimax_h3.apply_minimax_h3_qwen3vl_patch() + + assert encoder_module._apply_rotary_pos_emb is rope_sentinel @pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py index 551fd68a5..65a0221ca 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py @@ -1,10 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import patch + +import pytest import torch +import torch.nn as nn from vllm_omni.diffusion.layers.norm import RMSNorm -from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm +from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLRMSNorm, + MiniMaxH3Qwen3VLTextDecoderLayer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None: @@ -24,3 +34,43 @@ def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None: 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) + + +def test_qwen3_vl_decoder_preserves_post_attention_residual_contract() -> None: + class ConstantAttention(nn.Module): + def forward(self, hidden_states, position_embeddings): + del position_embeddings + return torch.full_like(hidden_states, 2.0) + + config = SimpleNamespace( + hidden_size=8, + intermediate_size=16, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=2, + rms_norm_eps=1e-6, + ) + group = SimpleNamespace(rank_in_group=0, world_size=1) + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + layer = MiniMaxH3Qwen3VLTextDecoderLayer(group, config, torch.float32) + layer.input_layernorm = nn.Identity() + layer.self_attn = ConstantAttention() + layer.mlp = nn.Identity() + + hidden_states = torch.arange(24, dtype=torch.float32).reshape(1, 3, 8) + attention_output = torch.full_like(hidden_states, 2.0) + expected_residual = hidden_states + attention_output + expected_normalized = RMSNorm.forward_native(layer.post_attention_layernorm, expected_residual) + + with patch.object( + layer.post_attention_layernorm, + "forward", + wraps=layer.post_attention_layernorm.forward_native, + ): + output = layer(hidden_states, (torch.empty(0), torch.empty(0))) + + assert isinstance(layer.post_attention_layernorm, MiniMaxH3Qwen3VLRMSNorm) + torch.testing.assert_close(output, expected_residual + expected_normalized) diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 8d91df8f0..3618388e6 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -85,7 +85,10 @@ class RMSNorm(CustomOp): def forward_cuda( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + return self.forward_native(x, residual) # During torch.compile tracing, fused_rms_norm writes to `out` in-place # (returns None) and accesses self.weight.data, which is a DTensor under # HSDP. Both patterns confuse inductor's compute_ancestors scheduler. @@ -101,7 +104,10 @@ class RMSNorm(CustomOp): def forward_hip( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + return self.forward_native(x, residual) if torch.compiler.is_compiling(): return self.forward_native(x) try: @@ -112,7 +118,10 @@ class RMSNorm(CustomOp): def forward_musa( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + return self.forward_native(x, residual) # Preserve the aten::rms_norm graph so dynamic Inductor can fuse the # H3 Q/K norm with its inline RoPE path on MUSA. return F.rms_norm( @@ -125,9 +134,19 @@ class RMSNorm(CustomOp): def forward_npu( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: import torch_npu + if residual is not None: + output, _, updated_residual = torch_npu.npu_add_rms_norm( + x, + residual, + self.weight, + self.variance_epsilon, + ) + return output, updated_residual + output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0] return output @@ -135,19 +154,27 @@ class RMSNorm(CustomOp): def forward_xpu( self, x: torch.Tensor, - ) -> torch.Tensor: - return self.forward_native(x) + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is None: + return self.forward_native(x) + return self.forward_native(x, residual) def forward_native( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + residual = residual + x + x = residual input_dtype = x.dtype x = x.to(torch.float32) variance = x.pow(2).mean(-1, keepdim=True) out = x * torch.rsqrt(variance + self.variance_epsilon) out = self.weight.to(torch.float32) * out - return out.to(input_dtype) + out = out.to(input_dtype) + return (out, residual) if residual is not None else out class RMSNormVAE(CustomOp): diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 55393209a..96e488f75 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -402,6 +402,25 @@ def _apply_rotary_pos_emb( return q_embed, k_embed +def _scaled_dot_product_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> torch.Tensor: + """Run causal SDPA with the reference expanded-K/V GQA behavior.""" + num_key_value_groups = query.shape[1] // key.shape[1] + if num_key_value_groups != 1: + key = key.repeat_interleave(num_key_value_groups, dim=1) + value = value.repeat_interleave(num_key_value_groups, dim=1) + return F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.0, + is_causal=True, + ) + + def _apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor: """Reorganize chunked [TTT...HHH...WWW] into interleaved [THWTHW...] layout.""" freqs_t = freqs[0] @@ -810,19 +829,7 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): cos, sin = position_embeddings query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin) - # Expand GQA keys/values to the local query-head count so SDPA works - # uniformly across backends (mirrors the reference eager path). - num_key_value_groups = self.num_heads // self.num_kv_heads - key_states = key_states.repeat_interleave(num_key_value_groups, dim=1) - value_states = value_states.repeat_interleave(num_key_value_groups, dim=1) - - attn_output = F.scaled_dot_product_attention( - query_states, - key_states, - value_states, - dropout_p=0.0, - is_causal=True, - ) + attn_output = _scaled_dot_product_attention(query_states, key_states, value_states) attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output @@ -907,10 +914,7 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(hidden_states, position_embeddings=position_embeddings) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states diff --git a/vllm_omni/platforms/npu/models/minimax_h3.py b/vllm_omni/platforms/npu/models/minimax_h3.py index 78e0b3f14..b70f90e30 100644 --- a/vllm_omni/platforms/npu/models/minimax_h3.py +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -19,6 +19,7 @@ from vllm_omni.platforms.npu.layers.rotary_embedding import ( logger = init_logger(__name__) _ROPE_PATCHED = False +_SDPA_PATCHED = False _SWIGLU_PATCHED = False @@ -48,6 +49,50 @@ def apply_minimax_h3_qwen3vl_patch() -> None: logger.debug("Applied NPU fused RoPE patch for MiniMax H3 Qwen3-VL text encoder") +def _scaled_dot_product_attention_npu( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> torch.Tensor: + """Run causal SDPA with compressed K/V heads through NPU native GQA.""" + num_heads = query.shape[1] + num_key_value_heads = key.shape[1] + num_value_heads = value.shape[1] + if num_key_value_heads != num_value_heads: + raise ValueError( + "GQA requires key and value to have the same number of heads, " + f"got k_heads={num_key_value_heads} and v_heads={num_value_heads}." + ) + if num_key_value_heads == 0: + raise ValueError("GQA requires at least one KV head.") + if num_heads % num_key_value_heads != 0: + raise ValueError( + "GQA requires query heads to be a multiple of KV heads, " + f"got q_heads={num_heads} and kv_heads={num_key_value_heads}." + ) + return F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.0, + is_causal=True, + enable_gqa=num_heads != num_key_value_heads, + ) + + +def apply_minimax_h3_qwen3vl_sdpa_patch() -> None: + """Route MiniMax H3 Qwen3-VL text attention to NPU native GQA.""" + global _SDPA_PATCHED + if _SDPA_PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder._scaled_dot_product_attention = _scaled_dot_product_attention_npu + _SDPA_PATCHED = True + logger.debug("Applied NPU SDPA 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) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 8ad695fed..91415372a 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -107,16 +107,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): from vllm_omni.platforms.npu.models.minimax_h3 import ( apply_minimax_h3_qwen3vl_patch, + apply_minimax_h3_qwen3vl_sdpa_patch, apply_minimax_h3_qwen3vl_swiglu_patch, ) - # Both patches import the MiniMax encoder package, whose __init__ loads + # These patches import the MiniMax encoder package, whose __init__ loads # pipeline_minimax_h3 → diffusion.data. Doing that during platform # construction races vllm_omni/__init__.py (patch before config) and # closes a cycle through pipeline_registry → PI0_PIPELINE → # DiffusionOutput. Apply them only after the platform exists, before # the diffusion pipeline is loaded. apply_minimax_h3_qwen3vl_patch() + apply_minimax_h3_qwen3vl_sdpa_patch() apply_minimax_h3_qwen3vl_swiglu_patch() set_mc2_tokens_capacity(vllm_config, od_config.max_num_seqs, 1) set_mc2_mask(vllm_config, device)