Unverified Commit 3f331398 authored by Ning Xie's avatar Ning Xie Committed by GitHub
Browse files

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


Signed-off-by: default avatarAndy Xie <andy.xning@gmail.com>
parent 78818dd1
...@@ -243,6 +243,18 @@ def test_free_kv_cache_block_queue_append_n(): ...@@ -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 blocks[3].next_free_block is queue.fake_free_list_tail
assert queue.fake_free_list_tail.prev_free_block is blocks[3] 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(): def test_free_kv_cache_block_queue_popleft_n():
blocks = [KVCacheBlock(block_id=i) for i in range(6)] blocks = [KVCacheBlock(block_id=i) for i in range(6)]
......
...@@ -370,7 +370,6 @@ class FreeKVCacheBlockQueue: ...@@ -370,7 +370,6 @@ class FreeKVCacheBlockQueue:
""" """
if len(blocks) == 0: if len(blocks) == 0:
return return
self.num_free_blocks += len(blocks)
last_block = self.fake_free_list_tail.prev_free_block last_block = self.fake_free_list_tail.prev_free_block
assert last_block is not None, ( assert last_block is not None, (
...@@ -385,6 +384,8 @@ class FreeKVCacheBlockQueue: ...@@ -385,6 +384,8 @@ class FreeKVCacheBlockQueue:
last_block.next_free_block = self.fake_free_list_tail last_block.next_free_block = self.fake_free_list_tail
self.fake_free_list_tail.prev_free_block = last_block self.fake_free_list_tail.prev_free_block = last_block
self.num_free_blocks += len(blocks)
def get_all_free_blocks(self) -> list[KVCacheBlock]: def get_all_free_blocks(self) -> list[KVCacheBlock]:
"""Get all free blocks in the free list. Mainly used for testing. """Get all free blocks in the free list. Mainly used for testing.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment