A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 23:45 UTC
[Core] Add basic unit test for maybe_evict_cached_block (#21400)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Jul 23, 2025 at 00:02 UTC
a1f3610fc650cf1d9e8761b17d23cd25bb8f8563
1 file changed
+67
tests/v1/core/test_prefix_caching.py
+67
@@ -1097,6 +1097,73 @@ def test_prefix_cache_stats_disabled():
1097
assert manager.prefix_cache_stats is None
1098
1099
1100
+def test_maybe_evict_cached_block():
1101
+ pool = BlockPool(num_gpu_blocks=4, enable_caching=True)
1102
+ block_hash0 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=10,
1103
+ token_ids=(100, )),
1104
+ group_id=1000)
1105
+ block_hash1 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=20,
1106
+ token_ids=(200, )),
1107
+ group_id=2000)
1108
+ block_hash2 = BlockHashWithGroupId(block_hash=BlockHash(hash_value=30,
1109
+ token_ids=(300, )),
1110
+ group_id=3000)
1111
+ block_hashes = [
1112
+ block_hash0,
1113
+ block_hash1,
1114
+ block_hash2,
1115
+ # block3 had the exact same block_hash as the first block
1116
+ block_hash0,
1117
+ ]
1118
+ assert len(pool.blocks) == len(block_hashes)
1119
+ # Manually add all blocks to cached_blocks
1120
+ for block, block_hash in zip(pool.blocks, block_hashes):
1121
+ block.block_hash = block_hash
1122
+ pool.cached_block_hash_to_block[block_hash][block.block_id] = block
1123
+
1124
+ block0, block1, block2, block3 = pool.blocks
1125
+ assert pool.cached_block_hash_to_block == {
1126
+ block_hash0: {
1127
+ block0.block_id: block0,
1128
+ block3.block_id: block3
1129
+ },
1130
+ block_hash1: {
1131
+ block1.block_id: block1
1132
+ },
1133
+ block_hash2: {
1134
+ block2.block_id: block2
1135
+ }
1136
+ }
1137
+ # Evict block1
1138
+ pool._maybe_evict_cached_block(block1)
1139
+ assert pool.cached_block_hash_to_block == {
1140
+ block_hash0: {
1141
+ block0.block_id: block0,
1142
+ block3.block_id: block3
1143
+ },
1144
+ block_hash2: {
1145
+ block2.block_id: block2
1146
+ }
1147
+ }
1148
+ # Evict block0: block_hash0 entry should NOT be removed, as block3
1149
+ # also use the same hash
1150
+ pool._maybe_evict_cached_block(block0)
1151
+ assert pool.cached_block_hash_to_block == {
1152
+ block_hash0: {
1153
+ block3.block_id: block3
1154
+ },
1155
+ block_hash2: {
1156
+ block2.block_id: block2
1157
+ }
1158
+ }
1159
+ # Evict block2
1160
+ pool._maybe_evict_cached_block(block2)
1161
+ assert pool.cached_block_hash_to_block == {block_hash0: {3: block3}}
1162
+ # Evict block3
1163
+ pool._maybe_evict_cached_block(block3)
1164
+ assert pool.cached_block_hash_to_block == {}
1165
+
1166
+
1167
@pytest.mark.parametrize("blocks_to_cache", [2, 3, 10])
1168
def test_kv_cache_events(blocks_to_cache: int):
1169
block_size = 16