[v1] Move block management logic from KVCacheManager to SpecializedManager (#17474)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
This commit is contained in:
@@ -812,10 +812,11 @@ def _assert_right_kv_cache_manager(
|
||||
# Make sure the request stats are right.
|
||||
EXPECTED_TOTAL_BLOCKS = num_tokens // block_size
|
||||
for req_id in req_ids:
|
||||
blocks = scheduler.kv_cache_manager.req_to_blocks[req_id]
|
||||
blocks = (scheduler.kv_cache_manager.single_type_manager.
|
||||
req_to_blocks[req_id])
|
||||
hashes = scheduler.kv_cache_manager.req_to_block_hashes[req_id]
|
||||
assert (scheduler.kv_cache_manager.num_cached_block[req_id] ==
|
||||
EXPECTED_TOTAL_BLOCKS)
|
||||
assert (scheduler.kv_cache_manager.single_type_manager.
|
||||
num_cached_block[req_id] == EXPECTED_TOTAL_BLOCKS)
|
||||
assert len(blocks) == EXPECTED_TOTAL_BLOCKS
|
||||
assert len(hashes) == EXPECTED_TOTAL_BLOCKS
|
||||
|
||||
@@ -1195,9 +1196,11 @@ def assert_scheduler_empty(scheduler: Scheduler):
|
||||
assert len(scheduler.encoder_cache_manager.cached) == 0
|
||||
|
||||
# KVCache Manager.
|
||||
assert len(scheduler.kv_cache_manager.req_to_blocks) == 0
|
||||
assert len(
|
||||
scheduler.kv_cache_manager.single_type_manager.req_to_blocks) == 0
|
||||
assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
|
||||
assert len(scheduler.kv_cache_manager.num_cached_block) == 0
|
||||
assert len(
|
||||
scheduler.kv_cache_manager.single_type_manager.num_cached_block) == 0
|
||||
num_free_blocks = (
|
||||
scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
|
||||
assert num_free_blocks == (
|
||||
|
||||
Reference in New Issue
Block a user