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