def rms_norm(
x: torch.Tensor,
weight: torch.Tensor,
eps: float,
*,
out_dtype: torch.dtype | None = None,
weight_bias: float = 0.0,
residual: torch.Tensor | None = None,
out: torch.Tensor | None = None,
) -> torch.Tensor:
"""Compute in FP32 and round at the output store.
When supplied, the FP32 residual is updated in place with x + residual.
"""
if not supports_rms_norm(x, weight, residual):
raise ValueError("Unsupported RMSNorm inputs")
out_dtype = out_dtype or x.dtype
if out_dtype not in _DTYPES:
raise ValueError(f"Unsupported output dtype: {out_dtype}")
if out is None:
out = torch.empty_like(x, dtype=out_dtype)
elif (
out.shape != x.shape
or out.dtype != out_dtype
or out.device != x.device
or not out.is_contiguous()
or out.data_ptr() % 16
):
raise ValueError("Output must match shape/device/dtype and be contiguous")
rows = x.numel() // x.shape[-1]
if rows == 0:
return out
compiled = _compile(
x.dtype, weight.dtype, out_dtype, x.shape[-1], residual is not None
)
r_ptr = residual.data_ptr() if residual is not None else x.data_ptr()
compiled(
make_ptr(_DTYPES[x.dtype], x.data_ptr(), assumed_align=16),
make_ptr(_DTYPES[weight.dtype], weight.data_ptr(), assumed_align=16),
make_ptr(Float32, r_ptr, assumed_align=16),
make_ptr(_DTYPES[out_dtype], out.data_ptr(), assumed_align=16),
Int32(rows),
Float32(eps),
Float32(weight_bias),
CUstream(current_stream().cuda_stream),
)
return out