@triton.jit
def _topk_triton_kernel__parallel_bk(xs_ptr, values_out_ptr,
    indices_out_ptr, b: int, k: int, n: int, n_chunk: int, xs_stride: int,
    BLOCK_SIZE: tl.constexpr, PAD_VALUE: tl.constexpr, INTERLEAVED: tl.
    constexpr):
    pidx = tl.program_id(axis=0).to(tl.int64)
    bk_idx = BLOCK_SIZE * pidx + tl.arange(0, BLOCK_SIZE)
    b_idx, k_idx = bk_idx // k, bk_idx % k
    xs_ptr += b_idx * xs_stride
    if INTERLEAVED:
        k_stride, i_stride = 1, k
    else:
        k_stride, i_stride = n_chunk, 1
    mask = (b_idx < b) & (k_idx * k_stride < n)
    max_value = tl.load(xs_ptr + k_idx * k_stride, mask=mask, other=PAD_VALUE)
    max_i = tl.zeros((BLOCK_SIZE,), tl.int64)
    for i in tl.range(1, n_chunk):
        mask = (b_idx < b) & (k_idx * k_stride + i * i_stride < n)
        block = tl.load(xs_ptr + k_idx * k_stride + i * i_stride, mask=mask,
            other=PAD_VALUE)
        mask &= max_value < block
        max_value = tl.where(mask, block, max_value)
        max_i = tl.where(mask, i, max_i)
    max_index = k_idx * k_stride + max_i * i_stride
    tl.store(values_out_ptr + b_idx * k + k_idx, max_value, mask=b_idx < b)
    tl.store(indices_out_ptr + b_idx * k + k_idx, max_index, mask=b_idx < b)
