-
-
Notifications
You must be signed in to change notification settings - Fork 5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
1 changed file
with
88 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,88 @@ | ||
import pytest | ||
|
||
from vllm.v1.core.kv_cache_utils import FreeKVCacheBlockQueue, KVCacheBlock | ||
|
||
|
||
def test_free_kv_cache_block_queue_initialization(): | ||
# Test with a single block | ||
block = KVCacheBlock(block_id=0) | ||
queue = FreeKVCacheBlockQueue([block]) | ||
assert queue.num_free_blocks == 1 | ||
assert queue.free_list_head == block | ||
assert queue.free_list_tail == block | ||
|
||
# Test with an empty list | ||
with pytest.raises(IndexError): | ||
queue = FreeKVCacheBlockQueue([]) | ||
|
||
|
||
def test_free_kv_cache_block_queue_operations(): | ||
# Create a list of KVCacheBlock objects | ||
blocks = [KVCacheBlock(block_id=i) for i in range(5)] | ||
|
||
# Create a FreeKVCacheBlockQueue with these blocks | ||
queue = FreeKVCacheBlockQueue(blocks) | ||
|
||
# Check initial state | ||
assert queue.num_free_blocks == 5 | ||
assert queue.free_list_head == blocks[0] | ||
assert queue.free_list_tail == blocks[4] | ||
|
||
# Pop the first block | ||
block1 = queue.popleft() | ||
assert block1 == blocks[0] | ||
assert queue.num_free_blocks == 4 | ||
assert queue.free_list_head == blocks[1] | ||
assert queue.free_list_tail == blocks[4] | ||
|
||
# Remove a block from the middle | ||
block_to_remove = blocks[2] | ||
queue.remove(block_to_remove) | ||
assert queue.num_free_blocks == 3 | ||
assert blocks[1].next_free_block == blocks[3] | ||
assert blocks[3].prev_free_block == blocks[1] | ||
|
||
# Append a block back | ||
queue.append(block_to_remove) | ||
assert queue.num_free_blocks == 4 | ||
assert queue.free_list_tail == block_to_remove | ||
assert block_to_remove.prev_free_block == blocks[4] | ||
assert block_to_remove.next_free_block is None | ||
|
||
# Pop blocks until empty | ||
for _ in range(4): | ||
queue.popleft() | ||
assert queue.num_free_blocks == 0 | ||
assert queue.free_list_head is None | ||
assert queue.free_list_tail is None | ||
|
||
# Attempt to pop from an empty queue | ||
with pytest.raises(ValueError) as e: | ||
queue.popleft() | ||
assert str(e.value) == "No free blocks available" | ||
|
||
|
||
def test_free_kv_cache_block_queue_get_all_free_blocks(): | ||
# Create a list of KVCacheBlock objects | ||
blocks = [KVCacheBlock(block_id=i) for i in range(5)] | ||
|
||
# Create a FreeKVCacheBlockQueue with these blocks | ||
queue = FreeKVCacheBlockQueue(blocks) | ||
|
||
# Check all blocks are correctly retrieved | ||
assert queue.get_all_free_blocks() == blocks | ||
|
||
# Pop a block and check again | ||
queue.popleft() | ||
assert queue.get_all_free_blocks() == blocks[1:] | ||
|
||
# Remove a block and check again | ||
block_to_remove = blocks[2] | ||
queue.remove(block_to_remove) | ||
assert queue.get_all_free_blocks() == blocks[1:2] + blocks[3:] | ||
|
||
# Append a block back and check again | ||
queue.append(block_to_remove) | ||
assert queue.get_all_free_blocks() == blocks[1:2] + blocks[3:] + [ | ||
block_to_remove | ||
] |