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

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