[kv cache] update num_free_blocks in the end (#24228)

Signed-off-by: Andy Xie <andy.xning@gmail.com>
This commit is contained in:
Ning Xie 2025-09-15 13:15:12 +08:00 committed by GitHub
parent 78818dd1b0
commit 3f3313981c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 14 additions and 1 deletions

View File

@ -243,6 +243,18 @@ def test_free_kv_cache_block_queue_append_n():
assert blocks[3].next_free_block is queue.fake_free_list_tail
assert queue.fake_free_list_tail.prev_free_block is blocks[3]
# Create an empty FreeKVCacheBlockQueue
invalid_queue = FreeKVCacheBlockQueue([])
# set prev_free_block to None and this will cause assertation in append_n
invalid_queue.fake_free_list_tail.prev_free_block = None
with pytest.raises(AssertionError):
# Append 1 block
# fake_head->fake_tail
invalid_queue.append_n(blocks[0:1])
assert invalid_queue.num_free_blocks == 0
assert (invalid_queue.fake_free_list_head.next_free_block ==
invalid_queue.fake_free_list_tail)
def test_free_kv_cache_block_queue_popleft_n():
blocks = [KVCacheBlock(block_id=i) for i in range(6)]

View File

@ -370,7 +370,6 @@ class FreeKVCacheBlockQueue:
"""
if len(blocks) == 0:
return
self.num_free_blocks += len(blocks)
last_block = self.fake_free_list_tail.prev_free_block
assert last_block is not None, (
@ -385,6 +384,8 @@ class FreeKVCacheBlockQueue:
last_block.next_free_block = self.fake_free_list_tail
self.fake_free_list_tail.prev_free_block = last_block
self.num_free_blocks += len(blocks)
def get_all_free_blocks(self) -> list[KVCacheBlock]:
"""Get all free blocks in the free list. Mainly used for testing.