@triton.jit
def rms_norm(output_ptr, input_ptr, weights_ptr, stride, N, eps, DTYPE: tl.
    constexpr, BLOCK_SIZE: tl.constexpr):
    """
    RMS Norm Triton Kernel

    Params:
        - input_ptr (tensor): Pointer to Input
        - output_ptr (tensor): Pointer to Output
        - weights_ptr (tensor): Pointer to Scale applied to the normalized input
        - stride (int): Stride to be applied when accessing elements in the input and output tensors
        - N (int): Number of elements to be reduced == input_ptr.shape[-1]
        - eps (half/float): Epsilon value added to the variance to prevent division by zero
        - BLOCK_SIZE (constexpr): Size of the block for computation, provided as a compile-time constant

    Usage:
        _rms_norm[grid, block](x, y, self.w, input_stride , N, eps, BLOCK_SIZE)
    """
    row = tl.program_id(0)
    output_ptr += row * stride
    input_ptr += row * stride
    tmp = 0
    tmp = tl.zeros([BLOCK_SIZE], dtype=DTYPE)
    for offset in range(0, N, BLOCK_SIZE):
        cols = offset + tl.arange(0, BLOCK_SIZE)
        mask = cols < N
        a = tl.load(input_ptr + cols, mask=mask, other=0.0).to(DTYPE)
        tmp += a * a
    rms = tl.sqrt(tl.sum(tmp) / N + eps)
    for offset in range(0, N, BLOCK_SIZE):
        cols = offset + tl.arange(0, BLOCK_SIZE)
        mask = cols < N
        x = tl.load(input_ptr + cols, mask=mask, other=0.0, eviction_policy
            ='evict_first').to(DTYPE)
        w = tl.load(weights_ptr + cols, mask=mask)
        x_hat = x / rms
        y = x_hat * w
        tl.store(output_ptr + cols, y, mask=mask)
