@wrap_jit_func(type_hint=dict(input=Tensor, weight=Tensor, output=Tensor,
    input_row_stride=int, eps=float, N_COLS=torch.int32, BLOCK_N=torch.int32))
@triton.jit
def rms_norm_kernel(input, weight, output, input_row_stride: tl.constexpr,
    eps: tl.constexpr, N_COLS: tl.constexpr, BLOCK_N: tl.constexpr):
    """rms norm kernel."""
    prog_id = tl.program_id(0)
    offsets = tl.arange(0, BLOCK_N)
    w = tl.load(weight + offsets, mask=offsets < N_COLS)
    x_ptr = input + prog_id * input_row_stride
    x = tl.load(x_ptr + offsets, mask=offsets < N_COLS)
    xf = x.to(tl.float32)
    var = tl.sum(xf * xf, 0) * float(1.0 / N_COLS)
    out = xf / tl.sqrt(var + eps)
    out = (w * out).to(x.dtype)
    out_ptr = output + prog_id * input_row_stride
    tl.store(out_ptr + offsets, out, mask=offsets < N_COLS)
