@triton.jit
def padding_index_kernel(seqlen_ptr, new_seqlen_ptr, new_max_seqlen,
    idx_ptr, window_size, BLK_N: tl.constexpr):
    pid_b = tl.program_id(0)
    i1 = tl.load(seqlen_ptr + pid_b).to(tl.int32)
    j1 = tl.load(seqlen_ptr + pid_b + 1).to(tl.int32)
    i2 = tl.load(new_seqlen_ptr + pid_b).to(tl.int32)
    j2 = tl.load(new_seqlen_ptr + pid_b + 1).to(tl.int32)
    for pid_n in range(tl.cdiv(new_max_seqlen, BLK_N)):
        offs_idx = pid_n * BLK_N + tl.arange(0, BLK_N)
        mask_idx = offs_idx < j2 - i2
        idx_ptrs = idx_ptr + i2 + offs_idx
        idx = i1 + offs_idx.to(tl.int32)
        tmp = clamp(idx - window_size, i1, j1 - 1)
        idx_out = tl.where(idx < j1, idx, tmp)
        tl.store(idx_ptrs, idx_out, mask=mask_idx)
