@triton.jit
def rz_linear_idx_kernel(bh_ptr, b_ptr, K, N, H, R3, R2, R1, R0, stride_bk,
    stride_bn, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):
    pid = tl.program_id(axis=0)
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    pid_k = pid // grid_n
    pid_n = pid % grid_n
    bh_offset = bh_ptr + tl.arange(0, BLOCK_SIZE_K)[:, None
        ] * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)[None, :]
    bh_ptrs = bh_offset + (pid_k * R3 + pid_n * R2 + R1) % R0 % (H - 
        BLOCK_SIZE_K * BLOCK_SIZE_N)
    b_ptrs = (b_ptr + pid_k * BLOCK_SIZE_K * stride_bk + pid_n *
        BLOCK_SIZE_N * stride_bn + tl.arange(0, BLOCK_SIZE_K)[:, None] *
        stride_bk + tl.arange(0, BLOCK_SIZE_N)[None, :])
    bh = tl.load(bh_ptrs)
    tl.store(b_ptrs, bh)
