[XPU] [RMSNorm] revert weightless change on xpu (#46987)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
zofia committed
Jun 30, 2026 at 03:03 UTC
72f639927f6caf3495d69dec63b9d4a87ed782ef
1 file changed
+4
-11
vllm/kernels/xpu_ops.py
+4
-11
@@ -31,10 +31,8 @@ def rms_norm(
31
) -> Tensor:
32
assert variance_size is None
33
if weight is None:
34
- # Weightless _C ops are CUDA-only; native skips the multiply on XPU.
35
- return ir.ops.rms_norm.impls["native"].impl_fn(
36
- x, weight, epsilon, variance_size
37
- )
34
+ # Kernel requires weight tensor, pass ones
35
+ weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
36
output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
37
torch.ops._C.rms_norm(output, x, weight, epsilon)
38
return output
@@ -61,12 +59,7 @@ def fused_add_rms_norm(
59
) -> tuple[Tensor, Tensor]:
60
assert variance_size is None
61
if weight is None:
64
- # Weightless _C ops are CUDA-only; native skips the multiply on XPU.
65
- output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
66
- x, x_residual, weight, epsilon, variance_size
67
- )
68
- x.copy_(output)
69
- x_residual.copy_(residual)
70
- return x, x_residual
62
+ # Kernel requires weight tensor, pass ones
63
+ weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
64
torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
65
return x, x_residual