@triton.jit
def k_mean(X, Mean, Var, stride, N, BLOCK_SIZE_N: tl.constexpr):
    """
        Fused layernorm kernel over a 3d tensor.
        The layer norm is applied over the last dimension.

        Compute
            y = (x - E(x))/(sqrt(var(x) + epsilon)) * gamma + beta
        """
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE_N)
    x_ptrs = X + row * stride + cols
    x = tl.load(x_ptrs, mask=cols < N, other=0.0).to(tl.float32)
    x = tl.where(cols < N, x, 0.0)
    x_mean = tl.sum(x, axis=0) / N
    x_zm = x - x_mean
    x_zm = tl.where(cols < N, x_zm, 0.0)
    x_var = tl.sum(x_zm * x_zm, axis=0) / N
    tl.store(Mean + row, x_mean)
    tl.store(Var + row, x_var)
