@triton.jit
def triton_rmsnorm_forward(Y_ptr, Y_row_stride, X_ptr, X_row_stride, W_ptr,
    RSTD_ptr, RSTD_row_stride, n_cols, eps, offset, BLOCK_SIZE: tl.constexpr):
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols
    Y_ptr += row_idx * Y_row_stride
    X_ptr += row_idx * X_row_stride
    RSTD_ptr += row_idx * RSTD_row_stride
    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
    X_row_dtype = X_row.dtype
    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
    X_row = X_row.to(tl.float32)
    mean_square = tl.sum(X_row * X_row, axis=0) / n_cols
    rstd = tl.libdevice.rsqrt(mean_square + eps)
    tl.store(RSTD_ptr, rstd)
    X_row = X_row * rstd
    X_row = X_row.to(X_row_dtype)
    Y_row = X_row * (offset + W_row)
    tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
