[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

@@ -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
 
 

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)

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

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

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)

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)