Function rmsnorm_kernel

kernels/rmsnorm.py:14–42

Row-parallel RMS normalization.

(
    X_ptr, W_ptr, OUT_ptr,
    M, N,
    stride_xm, stride_xn,
    stride_om, stride_on,
    eps,
    BLOCK_SIZE: tl.constexpr,
)

@triton.jit
def rmsnorm_kernel(
    X_ptr, W_ptr, OUT_ptr,
    M, N,
    stride_xm, stride_xn,
    stride_om, stride_on,
    eps,
    BLOCK_SIZE: tl.constexpr,
):
    """Row-parallel RMS normalization."""
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_SIZE)
    mask = offs < N

    # Load row into float32 for numerical stability
    x = tl.load(X_ptr + row * stride_xm + offs * stride_xn, mask=mask, other=0.0).to(tl.float32)

    # Compute RMS
    sq_mean = tl.sum(x * x, axis=0) / N
    rms = tl.sqrt(sq_mean + eps)

    # Normalize
    x_norm = x / rms

    # Scale by weight
    w = tl.load(W_ptr + offs, mask=mask, other=0.0).to(tl.float32)
    out = x_norm * w

    # Store (cast back to input dtype via the output tensor's dtype)
    tl.store(OUT_ptr + row * stride_om + offs * stride_on, out, mask=mask)


def kernel_fn(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: