@triton.jit
def _topk_triton_kernel__parallel_bkn(xs_ptr, values_out_ptr,
    indices_out_ptr, xs_stride: int, n_stride: int, b: int, k: int, n: int,
    BLOCK_BK: tl.constexpr, BLOCK_N: tl.constexpr, PAD_VALUE: tl.constexpr,
    INTERLEAVED: tl.constexpr):
    idx = tl.program_id(axis=0) * BLOCK_BK + tl.arange(0, BLOCK_BK)
    b_idx, k_idx = idx // k, idx % k
    if INTERLEAVED:
        k_stride, i_stride = 1, k
    else:
        k_stride, i_stride = n_stride, 1
    ni = tl.arange(0, BLOCK_N)
    n_idx = k_idx[:, None] * k_stride + ni[None, :] * i_stride
    data = tl.load(xs_ptr + b_idx[:, None] * xs_stride + n_idx, mask=(b_idx
        [:, None] < b) & (n_idx < n) & (ni < n_stride), other=PAD_VALUE)
    max_value, max_i = tl.max(data, axis=1, return_indices=True)
    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)
