@triton.jit
def seqlen_to_batch_index_kernel(seqlen_ptr, idx_ptr, BLK: tl.constexpr):
    pid = tl.program_id(0)
    i = tl.load(seqlen_ptr + pid)
    j = tl.load(seqlen_ptr + pid + 1)
    idx = tl.arange(0, BLK)
    tl.store(idx_ptr + i + idx, pid, mask=idx < j - i)
