[XPU] Route weightless RMSNorm to _C dispatch (#47121)
Signed-off-by: Yintong Lu <yintong.lu@intel.com> Co-authored-by: Yongqi Wang <yongqi.wang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Yintong Lu committed
Jul 29, 2026 at 18:59 UTC
72297d859ba34e019da4d7e5de58abbbd59961d4
1 file changed
-6
vllm/kernels/xpu_ops.py
-6
@@ -30,9 +30,6 @@ def rms_norm(
30
x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
31
) -> Tensor:
32
assert variance_size is None
33
- if weight is None:
34
- # Kernel requires weight tensor, pass ones
35
- weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
33
output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
34
torch.ops._C.rms_norm(output, x, weight, epsilon)
35
return output
@@ -58,8 +55,5 @@ def fused_add_rms_norm(
55
variance_size: int | None = None,
56
) -> tuple[Tensor, Tensor]:
57
assert variance_size is None
61
- if weight is None:
62
- # Kernel requires weight tensor, pass ones
63
- weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
58
torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
59
return x, x_residual