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