[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>
Huanxing committed
Sep 3, 2026 at 06:11 UTC
a0d3e5c16e96b62ec5c7f741952cfddc99534faf
2 files changed
+81
vllm/_xpu_ops.py
+54
@@ -113,6 +113,45 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
113
return torch.empty((M, N), dtype=input.dtype, device=input.device)
114
115
116
+def _gemma_rms_norm_impl(
117
+ out: torch.Tensor,
118
+ input: torch.Tensor,
119
+ weight: torch.Tensor,
120
+ epsilon: float,
121
+) -> None:
122
+ # GemmaRMSNorm: computes out = (x_normed_fp32 * (1 + weight.float())
123
+ # ).to(dtype) with a raw (bf16/fp16) weight; the +1 offset and fp32
124
+ # multiply are done in-kernel. See vllm-xpu-kernels gemma_rms_norm.
125
+ torch.ops._C.gemma_rms_norm(out, input, weight, epsilon)
126
+
127
+
128
+def _gemma_rms_norm_fake(
129
+ out: torch.Tensor,
130
+ input: torch.Tensor,
131
+ weight: torch.Tensor,
132
+ epsilon: float,
133
+) -> None:
134
+ return None
135
+
136
+
137
+def _fused_add_gemma_rms_norm_impl(
138
+ input: torch.Tensor,
139
+ residual: torch.Tensor,
140
+ weight: torch.Tensor,
141
+ epsilon: float,
142
+) -> None:
143
+ torch.ops._C.fused_add_gemma_rms_norm(input, residual, weight, epsilon)
144
+
145
+
146
+def _fused_add_gemma_rms_norm_fake(
147
+ input: torch.Tensor,
148
+ residual: torch.Tensor,
149
+ weight: torch.Tensor,
150
+ epsilon: float,
151
+) -> None:
152
+ return None
153
+
154
+
155
def _gdn_attention_core_xpu_impl(
156
core_attn_out: torch.Tensor,
157
z: torch.Tensor,
@@ -1218,6 +1257,21 @@ class xpu_ops:
1257
global _OPS_REGISTERED
1258
if not _OPS_REGISTERED:
1259
# register all the custom ops here
1260
+ if hasattr(torch.ops._C, "gemma_rms_norm"):
1261
+ direct_register_custom_op(
1262
+ op_name="xpu_gemma_rms_norm",
1263
+ op_func=_gemma_rms_norm_impl,
1264
+ mutates_args=["out"],
1265
+ fake_impl=_gemma_rms_norm_fake,
1266
+ )
1267
+
1268
+ direct_register_custom_op(
1269
+ op_name="xpu_fused_add_gemma_rms_norm",
1270
+ op_func=_fused_add_gemma_rms_norm_impl,
1271
+ mutates_args=["input", "residual"],
1272
+ fake_impl=_fused_add_gemma_rms_norm_fake,
1273
+ )
1274
+
1275
direct_register_custom_op(
1276
op_name="xpu_ops_deepseek_scaling_rope",
1277
op_func=_xpu_ops_deepseek_scaling_rope_impl,
vllm/model_executor/layers/layernorm.py
+27
@@ -166,6 +166,33 @@ class GemmaRMSNorm(CustomOp):
166
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
167
return self.forward_native(x, residual)
168
169
+ def forward_xpu(
170
+ self,
171
+ x: torch.Tensor,
172
+ residual: torch.Tensor | None = None,
173
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
174
+ import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_gemma_rms_norm
175
+
176
+ # Fall back to the native path if the fused gemma kernels are not
177
+ # available in the installed vllm-xpu-kernels package.
178
+ if not hasattr(torch.ops._C, "gemma_rms_norm"):
179
+ return self.forward_native(x, residual)
180
+
181
+ # Pass the raw (bf16/fp16) weight; the +1 offset and the fp32 multiply
182
+ # are folded into the kernel (matches forward_native numerics).
183
+ if residual is not None:
184
+ torch.ops.vllm.xpu_fused_add_gemma_rms_norm(
185
+ x, residual, self.weight.data, self.variance_epsilon
186
+ )
187
+ return x, residual
188
+ # empty_like preserves x's strides, but the kernel requires a
189
+ # contiguous out (unlike x, which it can handle non-contiguous).
190
+ out = torch.empty(x.shape, device=x.device, dtype=x.dtype)
191
+ torch.ops.vllm.xpu_gemma_rms_norm(
192
+ out, x, self.weight.data, self.variance_epsilon
193
+ )
194
+ return out
195
+
196
197
# --8<-- [start:rms_norm_gated]
198
@CustomOp.register("rms_norm_gated")