@triton.jit
def triton_rmsnorm_backward(dY_ptr, dY_row_stride, X_ptr, X_row_stride,
    X_dtype: tl.constexpr, W_ptr, RSTD_ptr, RSTD_row_stride, dW_ptr,
    dW_row_stride, n_rows, n_cols, offset, rows_per_program: tl.constexpr,
    BLOCK_SIZE: tl.constexpr):
    row_block_id = tl.program_id(0)
    row_start = row_block_id * rows_per_program
    row_end = min((row_block_id + 1) * rows_per_program, n_rows)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols
    dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
    dY_ptr += row_start * dY_row_stride
    X_ptr += row_start * X_row_stride
    RSTD_ptr += row_start
    W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
    W_row = W_row + offset
    for _ in range(row_start, row_end):
        dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0)
        X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
        rstd_row = tl.load(RSTD_ptr)
        X_row = X_row.to(tl.float32)
        m = (dY_row * W_row).to(tl.float32)
        dX_row = rstd_row * m
        dX_row += rstd_row * (-(1 / n_cols) * rstd_row * rstd_row * tl.sum(
            m * X_row, axis=0) * X_row)
        dW_row += dY_row * (X_row * rstd_row).to(X_dtype)
        tl.store(dY_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
        dY_ptr += dY_row_stride
        X_ptr += X_row_stride
        RSTD_ptr += RSTD_row_stride
    tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row,
        mask=mask)
