[XPU] Add fused GemmaRMSNorm path for eager execution (#53678)

Signed-off-by: Huanxing <huanxing.shen@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

vllm/_xpu_ops.py

@@ -113,6 +113,45 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
         return torch.empty((M, N), dtype=input.dtype, device=input.device)
 
 
+def _gemma_rms_norm_impl(
+    out: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    epsilon: float,
+) -> None:
+    # GemmaRMSNorm: computes out = (x_normed_fp32 * (1 + weight.float())
+    # ).to(dtype) with a raw (bf16/fp16) weight; the +1 offset and fp32
+    # multiply are done in-kernel. See vllm-xpu-kernels gemma_rms_norm.
+    torch.ops._C.gemma_rms_norm(out, input, weight, epsilon)
+
+
+def _gemma_rms_norm_fake(
+    out: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    epsilon: float,
+) -> None:
+    return None
+
+
+def _fused_add_gemma_rms_norm_impl(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    epsilon: float,
+) -> None:
+    torch.ops._C.fused_add_gemma_rms_norm(input, residual, weight, epsilon)
+
+
+def _fused_add_gemma_rms_norm_fake(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    epsilon: float,
+) -> None:
+    return None
+
+
 def _gdn_attention_core_xpu_impl(
     core_attn_out: torch.Tensor,
     z: torch.Tensor,
@@ -1218,6 +1257,21 @@ class xpu_ops:
         global _OPS_REGISTERED
         if not _OPS_REGISTERED:
             # register all the custom ops here
+            if hasattr(torch.ops._C, "gemma_rms_norm"):
+                direct_register_custom_op(
+                    op_name="xpu_gemma_rms_norm",
+                    op_func=_gemma_rms_norm_impl,
+                    mutates_args=["out"],
+                    fake_impl=_gemma_rms_norm_fake,
+                )
+
+                direct_register_custom_op(
+                    op_name="xpu_fused_add_gemma_rms_norm",
+                    op_func=_fused_add_gemma_rms_norm_impl,
+                    mutates_args=["input", "residual"],
+                    fake_impl=_fused_add_gemma_rms_norm_fake,
+                )
+
             direct_register_custom_op(
                 op_name="xpu_ops_deepseek_scaling_rope",
                 op_func=_xpu_ops_deepseek_scaling_rope_impl,

vllm/model_executor/layers/layernorm.py

@@ -166,6 +166,33 @@ class GemmaRMSNorm(CustomOp):
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         return self.forward_native(x, residual)
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        import vllm._xpu_ops  # noqa: F401 registers torch.ops.vllm.xpu_gemma_rms_norm
+
+        # Fall back to the native path if the fused gemma kernels are not
+        # available in the installed vllm-xpu-kernels package.
+        if not hasattr(torch.ops._C, "gemma_rms_norm"):
+            return self.forward_native(x, residual)
+
+        # Pass the raw (bf16/fp16) weight; the +1 offset and the fp32 multiply
+        # are folded into the kernel (matches forward_native numerics).
+        if residual is not None:
+            torch.ops.vllm.xpu_fused_add_gemma_rms_norm(
+                x, residual, self.weight.data, self.variance_epsilon
+            )
+            return x, residual
+        # empty_like preserves x's strides, but the kernel requires a
+        # contiguous out (unlike x, which it can handle non-contiguous).
+        out = torch.empty(x.shape, device=x.device, dtype=x.dtype)
+        torch.ops.vllm.xpu_gemma_rms_norm(
+            out, x, self.weight.data, self.variance_epsilon
+        )
+        return out
+
 
 # --8<-- [start:rms_norm_gated]
 @CustomOp.register("rms_norm_gated")