[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/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 ──
 
 

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)

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,

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)

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")

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)