[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>
This commit is contained in:
Yintong Lu
2026-07-29 18:59:58 +08:00
committed by GitHub
co-authored by Yongqi Wang Kunshang Ji
parent f51193b9ae
commit 72297d859b
-6
View File
@@ -30,9 +30,6 @@ def rms_norm(
x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
) -> Tensor:
assert variance_size is None
if weight is None:
# Kernel requires weight tensor, pass ones
weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
torch.ops._C.rms_norm(output, x, weight, epsilon)
return output
@@ -58,8 +55,5 @@ def fused_add_rms_norm(
variance_size: int | None = None,
) -> tuple[Tensor, Tensor]:
assert variance_size is None
if weight is None:
# Kernel requires weight tensor, pass ones
weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
return x, x_residual