[Bugfix] Block manager v2 with preemption and lookahead slots

[sroy745]

In this pr we address an issue with the logic in BlockManager V2 when computing whether a sequence group can be swapped in.

In the current implementation we don't consider the unseen tokens when deciding whether a sequence group can be swapped in or not. Later when we try to allocate blocks for the sequence group we might hit a NoFreeBlocksAvailable exception depending on how may unseen tokens are there.

In this pr we add the logic for considering these unseen tokens when determining if a sequence group is eligible for swapping in. We have added unit-tests to verify these cases both in unit-tests for block_manager_v2, naive_block and prefix_caching_block.

[sroy745]

Thanks for the review. Addressed your comments. PTAL

[comaniac]

Should not be here?

[comaniac]

Rewrite the above "NOTE" to explain the new logic (considering unseen tokens).

I feel we could simplify the logic, because we only need the number of touched blocks instead of their IDs. The idea is:

1. We first count the number of full blocks, because neither lookahead slots or unseen tokens won't affect them.
2. Then use the existing block_table.get_num_blocks_touched_by_append_slots to calculate unfilled blocks. This function already considers unseens tokens so it should be usable (if not, then could we add support to this function)?

# In _can_swap
block_table_and_unseen_ids: List[Tuple(BlockTable, int)] = []
for seq in seq_group.get_seqs(status=status):
    block_table = self.block_tables[seq.seq_id]
    if block_table.blocks is not None:
        block_table_and_unseen_ids.append(
            block_table, 
            len(block_table.get_unseen_token_ids(seq.get_token_ids()))
        )
num_blocks_touched = self.block_allocator.get_num_blocks_touched(
    block_table_and_unseen_ids, device,
    num_lookahead_slots=num_lookahead_slots
)

# In get_num_blocks_touched
old_block_set = set()
new_block_count = 0

for block_table, unseen_tokens in block_table_and_unseen_ids:
    blocks = block_table.blocks
    num_full_blocks = len(blocks)
    # 1. I suppose only the last block can be unfilled, but please double check.
    # 2. I suppose this can also handle prefix caching as the same block IDs will be de-duplicated by the set.
    if not blocks[-1].is_full:
        num_full_blocks -= 1
    old_block_set.update(blocks[:num_full_blocks])

    # This includes full blocks.
    num_touched_blocks = block_table.get_num_blocks_touched_by_append_slots(...)
    # New block count should be the number of touched blocks - full blocks
    new_block_count += num_touched_blocks - num_full_blocks
    return new_block_count + len(old_block_set)

[sroy745]

Thanks for the review. Addressed comments. PTAL

[comaniac]

LGTM! Pretty clean now.

[sroy745]

@comaniac
Thanks for the review. Synced to head, all tests passing and ready to merge.