xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-09 14:35:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Core] Add basic unit test for maybe_evict_cached_block (#21400)

Browse Source 
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
This commit is contained in:
Jialin Ouyang 2025-07-23 00:02:02 -07:00 committed by GitHub
parent 4ecedd1806
commit a1f3610fc6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 67 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

67
tests/v1/core/test_prefix_caching.py
View File

@ -1097,6 +1097,73 @@ def test_prefix_cache_stats_disabled():
assert manager.prefix_cache_stats is None
def test_maybe_evict_cached_block():
pool = BlockPool(num_gpu_blocks=4, enable_caching=True)
block_hash0 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=10,
token_ids=(100, )),
group_id=1000)
block_hash1 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=20,
token_ids=(200, )),
group_id=2000)
block_hash2 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=30,
token_ids=(300, )),
group_id=3000)
block_hashes = [
block_hash0,
block_hash1,
block_hash2,
# block3 had the exact same block_hash as the first block
block_hash0,
]
assert len(pool.blocks) == len(block_hashes)
# Manually add all blocks to cached_blocks
for block, block_hash in zip(pool.blocks, block_hashes):
block.block_hash = block_hash
pool.cached_block_hash_to_block[block_hash][block.block_id] = block
block0, block1, block2, block3 = pool.blocks
assert pool.cached_block_hash_to_block == {
block_hash0: {
block0.block_id: block0,
block3.block_id: block3
},
block_hash1: {
block1.block_id: block1
},
block_hash2: {
block2.block_id: block2
}
}
# Evict block1
pool._maybe_evict_cached_block(block1)
assert pool.cached_block_hash_to_block == {
block_hash0: {
block0.block_id: block0,
block3.block_id: block3
},
block_hash2: {
block2.block_id: block2
}
}
# Evict block0: block_hash0 entry should NOT be removed, as block3
# also use the same hash
pool._maybe_evict_cached_block(block0)
assert pool.cached_block_hash_to_block == {
block_hash0: {
block3.block_id: block3
},
block_hash2: {
block2.block_id: block2
}
}
# Evict block2
pool._maybe_evict_cached_block(block2)
assert pool.cached_block_hash_to_block == {block_hash0: {3: block3}}
# Evict block3
pool._maybe_evict_cached_block(block3)
assert pool.cached_block_hash_to_block == {}
@pytest.mark.parametrize("blocks_to_cache", [2, 3, 10])
def test_kv_cache_events(blocks_to_cache: int):
block_size = 16