# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for SimpleCPUOffloadScheduler.""" from __future__ import annotations from dataclasses import dataclass import torch from vllm import SamplingParams from vllm.config import ( CacheConfig, DeviceConfig, KVTransferConfig, ModelConfig, SchedulerConfig, VllmConfig, ) from vllm.utils.hashing import sha256 from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import ( get_request_block_hasher, init_none_hash, make_block_hash_with_group_id, ) from vllm.v1.core.sched.output import ( CachedRequestData, NewRequestData, SchedulerOutput, ) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, ) from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadWorkerMetadata # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- BLOCK_SIZE = 16 HEAD_SIZE = 16 NUM_KV_HEADS = 1 DTYPE = torch.float16 # bytes per block per tensor: # block_size * num_kv_heads * head_size * 2 (K+V) * element_size _BYTES_PER_BLOCK = BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * 2 * DTYPE.itemsize # Ensure none_hash is initialized once init_none_hash(sha256) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_kv_cache_config( num_blocks: int, num_groups: int = 1, ) -> KVCacheConfig: """Build a KVCacheConfig with non-empty kv_cache_tensors.""" groups = [] tensors = [] for g in range(num_groups): layer_names = [f"layer_{g}"] groups.append( KVCacheGroupSpec( layer_names, FullAttentionSpec( block_size=BLOCK_SIZE, num_kv_heads=NUM_KV_HEADS, head_size=HEAD_SIZE, dtype=DTYPE, ), ) ) tensors.append( KVCacheTensor( size=_BYTES_PER_BLOCK * num_blocks, shared_by=layer_names, ) ) return KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=tensors, kv_cache_groups=groups, ) def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig: """Minimal VllmConfig for scheduler tests (no GPU).""" model_config = ModelConfig( model="facebook/opt-125m", trust_remote_code=True, dtype="float16", seed=42, ) scheduler_config = SchedulerConfig( max_num_seqs=16, max_num_batched_tokens=64, max_model_len=10000, enable_chunked_prefill=True, is_encoder_decoder=False, ) cache_config = CacheConfig( block_size=block_size, gpu_memory_utilization=0.9, enable_prefix_caching=True, ) kv_transfer_config = KVTransferConfig( kv_connector="SimpleCPUOffloadConnector", kv_role="kv_both", ) return VllmConfig( scheduler_config=scheduler_config, model_config=model_config, cache_config=cache_config, kv_transfer_config=kv_transfer_config, device_config=DeviceConfig("cpu"), ) @dataclass class SchedulerFixture: """Bundle returned by make_scheduler for convenient access.""" scheduler: SimpleCPUOffloadScheduler gpu_block_pool: BlockPool vllm_config: VllmConfig kv_cache_config: KVCacheConfig num_groups: int = 1 def make_scheduler( num_cpu_blocks: int = 8, num_gpu_blocks: int = 16, num_groups: int = 1, lazy: bool = False, ) -> SchedulerFixture: """Build a SimpleCPUOffloadScheduler with small block pools.""" kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) vllm_config = _make_vllm_config() cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups sched = SimpleCPUOffloadScheduler( vllm_config=vllm_config, kv_cache_config=kv_cache_config, cpu_capacity_bytes=cpu_capacity_bytes, lazy_offload=lazy, ) # Build a real GPU block pool and bind it gpu_block_pool = BlockPool( num_gpu_blocks=num_gpu_blocks, enable_caching=True, hash_block_size=BLOCK_SIZE, ) sched.bind_gpu_block_pool(gpu_block_pool) return SchedulerFixture( scheduler=sched, gpu_block_pool=gpu_block_pool, vllm_config=vllm_config, kv_cache_config=kv_cache_config, num_groups=num_groups, ) _req_counter = 0 def make_request( num_blocks: int = 2, request_id: str | None = None, ) -> Request: """Create a Request with deterministic block hashes.""" global _req_counter _req_counter += 1 if request_id is None: request_id = f"req-{_req_counter}" num_tokens = num_blocks * BLOCK_SIZE start = _req_counter * 10000 prompt_token_ids = list(range(start, start + num_tokens)) sampling_params = SamplingParams(max_tokens=1) req = Request( request_id=request_id, prompt_token_ids=prompt_token_ids, sampling_params=sampling_params, pooling_params=None, mm_features=None, block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256), ) return req def make_scheduler_output( req_id_to_num_tokens: dict[str, int], *, new_reqs: dict[str, tuple[list[int], ...]] | None = None, cached_req_new_blocks: dict[str, tuple[list[int], ...] | None] | None = None, ) -> SchedulerOutput: """Build a minimal SchedulerOutput with num_scheduled_tokens. Args: new_reqs: For first-time requests, maps req_id -> block_ids tuple. These are placed into ``scheduled_new_reqs`` as ``NewRequestData``. cached_req_new_blocks: For returning (cached) requests, maps req_id -> new_block_ids (incremental) or None. These are placed into ``scheduled_cached_reqs``. """ scheduled_new_reqs: list[NewRequestData] = [] if new_reqs: for req_id, block_ids in new_reqs.items(): scheduled_new_reqs.append( NewRequestData( req_id=req_id, prompt_token_ids=None, mm_features=[], sampling_params=None, pooling_params=None, block_ids=block_ids, num_computed_tokens=0, lora_request=None, ) ) if cached_req_new_blocks: cached_req_ids = list(cached_req_new_blocks.keys()) cached_new_block_ids = [cached_req_new_blocks[rid] for rid in cached_req_ids] cached_reqs = CachedRequestData( req_ids=cached_req_ids, resumed_req_ids=set(), new_token_ids=[[] for _ in cached_req_ids], all_token_ids={}, new_block_ids=cached_new_block_ids, num_computed_tokens=[0] * len(cached_req_ids), num_output_tokens=[0] * len(cached_req_ids), ) else: cached_reqs = CachedRequestData.make_empty() return SchedulerOutput( scheduled_new_reqs=scheduled_new_reqs, scheduled_cached_reqs=cached_reqs, num_scheduled_tokens=req_id_to_num_tokens, total_num_scheduled_tokens=sum(req_id_to_num_tokens.values()), scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, num_common_prefix_blocks=[], preempted_req_ids=set(), finished_req_ids=set(), free_encoder_mm_hashes=[], ) def simulate_store_completion( scheduler: SimpleCPUOffloadScheduler, event_idx: int, ) -> None: """Simulate worker reporting a store event completion.""" output = KVConnectorOutput( finished_recving=set(), kv_connector_worker_meta=SimpleCPUOffloadWorkerMetadata( completed_store_events={event_idx: scheduler._expected_worker_count}, ), ) scheduler.update_connector_output(output) def simulate_load_completion( scheduler: SimpleCPUOffloadScheduler, req_ids: set[str], ) -> None: """Simulate worker reporting load completions for requests.""" output = KVConnectorOutput( finished_sending=set(), finished_recving=req_ids, ) scheduler.update_connector_output(output) def get_cpu_free_blocks(scheduler: SimpleCPUOffloadScheduler) -> int: """Return number of free CPU blocks.""" return scheduler.cpu_block_pool.get_num_free_blocks() def _allocate_gpu_blocks( gpu_block_pool: BlockPool, request: Request, num_blocks: int, group_id: int = 0, ) -> list: """Allocate GPU blocks, cache them with hashes, return block list. Mimics what KVCacheManager does: allocate blocks from pool, then register them in the prefix cache via cache_full_blocks so that re-allocation properly evicts stale hashes. """ blocks = gpu_block_pool.get_new_blocks(num_blocks) num_full = min(num_blocks, len(request.block_hashes)) if num_full > 0: gpu_block_pool.cache_full_blocks( request=request, blocks=blocks, num_cached_blocks=0, num_full_blocks=num_full, block_size=BLOCK_SIZE, kv_cache_group_id=group_id, ) return blocks def _alloc_and_register( fix: SchedulerFixture, request: Request, num_blocks: int, *, confirmed: bool = True, group_id: int = 0, ) -> KVCacheBlocks: """Allocate GPU blocks and return KVCacheBlocks. Block IDs are no longer registered in a mock KVCacheManager; instead tests pass them through ``make_scheduler_output`` so that ``yield_req_data`` can pick them up. If ``confirmed`` is True, advance ``request.num_computed_tokens`` to simulate the scheduler's ``_update_after_schedule`` from a prior step. """ gpu_blocks = _allocate_gpu_blocks( fix.gpu_block_pool, request, num_blocks, group_id=group_id ) kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,)) if confirmed: request.num_computed_tokens = num_blocks * BLOCK_SIZE return kv_blocks # --------------------------------------------------------------------------- # Test 1a: Eager store-and-load roundtrip # --------------------------------------------------------------------------- def test_eager_store_and_load_roundtrip() -> None: """Eager mode: store blocks on compute, complete store, verify cache hit.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler num_blocks = 2 req = make_request(num_blocks=num_blocks) kv_blocks = _alloc_and_register(fix, req, num_blocks) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) block_ids = kv_blocks.get_block_ids() sched_out = make_scheduler_output( {req.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req.request_id: block_ids}, ) meta = sched.build_connector_meta(sched_out) assert meta.store_event >= 0, "Expected a store event to be scheduled" assert len(meta.store_gpu_blocks) > 0 assert len(meta.store_cpu_blocks) == len(meta.store_gpu_blocks) simulate_store_completion(sched, meta.store_event) # New request with same tokens should get CPU cache hit req2 = Request( request_id="req-eager-load", prompt_token_ids=req.prompt_token_ids, sampling_params=req.sampling_params, pooling_params=None, mm_features=None, block_hasher=req._block_hasher, ) hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens == num_blocks * BLOCK_SIZE assert is_async is True gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) block_ids2 = kv_blocks2.get_block_ids() sched_out2 = make_scheduler_output( {req2.request_id: 1}, new_reqs={req2.request_id: block_ids2}, ) meta2 = sched.build_connector_meta(sched_out2) assert meta2.load_event >= 0, "Expected a load event to be assigned" assert len(meta2.load_gpu_blocks) > 0 assert len(meta2.load_cpu_blocks) == len(meta2.load_gpu_blocks) # --------------------------------------------------------------------------- # Test 1b: Lazy store-and-load roundtrip # --------------------------------------------------------------------------- def _flush_old_blocks_to_lru_head( gpu_pool: BlockPool, num_filler_blocks: int, ) -> list: """Allocate filler blocks so that previously-freed (hashed) blocks migrate to the LRU head of the free queue. Returns the filler blocks (caller must free them later to restore pool capacity). In a real engine the same thing happens naturally: after one request finishes and frees its blocks, subsequent requests allocate from the LRU head, consuming the unhashed blocks and leaving the old hashed blocks at the front of the queue. """ fillers = gpu_pool.get_new_blocks(num_filler_blocks) return fillers def test_lazy_store_and_load_roundtrip() -> None: """Lazy mode: schedule a request, finish it so its hashed blocks are freed, then schedule new requests so the old blocks migrate to the LRU head. The lazy scanner offloads them to CPU. Re-scheduling the old request triggers a CPU cache hit + load. GPU pool: 8 blocks (7 usable). _target_free = ceil(64/16) = 4. """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True) sched = fix.scheduler gpu_pool = fix.gpu_block_pool num_blocks = 2 # --- Step 1: Schedule req_old, compute, and finish --- req_old = make_request(num_blocks=num_blocks) gpu_blocks_old = _allocate_gpu_blocks(gpu_pool, req_old, num_blocks, group_id=0) gpu_pool.free_blocks(gpu_blocks_old) # Allocate filler blocks so req_old's hashed blocks move to LRU head. # 7 usable - 2 (req_old freed) = 5 other free blocks to consume. fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5) # --- Step 2: Lazy scanner should offload req_old's blocks --- sched_out = make_scheduler_output({}) meta = sched.build_connector_meta(sched_out) assert meta.store_event >= 0, "Expected lazy store to offload old blocks" assert len(meta.store_gpu_blocks) == num_blocks simulate_store_completion(sched, meta.store_event) # Free fillers to restore pool capacity. gpu_pool.free_blocks(fillers) # --- Step 3: Re-schedule req_old — should get CPU cache hit --- req_old2 = Request( request_id="req-old-reload", prompt_token_ids=req_old.prompt_token_ids, sampling_params=req_old.sampling_params, pooling_params=None, mm_features=None, block_hasher=req_old._block_hasher, ) hit_tokens, is_async = sched.get_num_new_matched_tokens( req_old2, num_computed_tokens=0 ) assert hit_tokens == num_blocks * BLOCK_SIZE, ( f"Expected {num_blocks * BLOCK_SIZE} hit tokens, got {hit_tokens}" ) assert is_async is True # Allocate fresh GPU blocks for the load. gpu_blocks_load = gpu_pool.get_new_blocks(num_blocks) kv_blocks_load = KVCacheBlocks(blocks=(gpu_blocks_load,)) sched.update_state_after_alloc( req_old2, kv_blocks_load, num_external_tokens=hit_tokens ) sched_out2 = make_scheduler_output({req_old2.request_id: 1}) meta2 = sched.build_connector_meta(sched_out2) assert meta2.load_event >= 0, "Expected a load event to be assigned" assert len(meta2.load_gpu_blocks) > 0 # --------------------------------------------------------------------------- # Test 2a: Eager duplicate store is skipped # --------------------------------------------------------------------------- def test_eager_duplicate_store_skipped() -> None: """Eager: storing the same block hashes twice should not allocate new CPU blocks.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler num_blocks = 2 req = make_request(num_blocks=num_blocks) kv_blocks = _alloc_and_register(fix, req, num_blocks) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) block_ids = kv_blocks.get_block_ids() sched_out = make_scheduler_output( {req.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req.request_id: block_ids}, ) meta1 = sched.build_connector_meta(sched_out) assert meta1.store_event >= 0 simulate_store_completion(sched, meta1.store_event) cpu_free_after_first = get_cpu_free_blocks(sched) # Second request with identical hashes — should skip store req2 = Request( request_id="req-dup-eager", prompt_token_ids=req.prompt_token_ids, sampling_params=req.sampling_params, pooling_params=None, mm_features=None, block_hasher=req._block_hasher, ) kv_blocks2 = _alloc_and_register(fix, req2, num_blocks) sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=0) block_ids2 = kv_blocks2.get_block_ids() sched_out2 = make_scheduler_output( {req2.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req2.request_id: block_ids2}, ) meta2 = sched.build_connector_meta(sched_out2) if meta2.store_event >= 0: assert len(meta2.store_cpu_blocks) == 0, ( "Expected no new CPU blocks for duplicate hashes" ) assert get_cpu_free_blocks(sched) == cpu_free_after_first # --------------------------------------------------------------------------- # Test 2b: Lazy duplicate store is skipped # --------------------------------------------------------------------------- def test_lazy_duplicate_store_skipped() -> None: """Lazy: blocks already offloaded to CPU should not be offloaded again. Same pattern as the lazy roundtrip: flush old blocks to LRU head, offload, then repeat with the same hashes and verify no new CPU allocation. """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True) sched = fix.scheduler gpu_pool = fix.gpu_block_pool num_blocks = 2 req = make_request(num_blocks=num_blocks) # Schedule + finish → hashed blocks in free queue gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) gpu_pool.free_blocks(gpu_blocks) # Flush old blocks to LRU head, then trigger lazy offload. fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5) meta1 = sched.build_connector_meta(make_scheduler_output({})) assert meta1.store_event >= 0 simulate_store_completion(sched, meta1.store_event) gpu_pool.free_blocks(fillers) cpu_free_after_first = get_cpu_free_blocks(sched) # Allocate blocks with the same hashes and free them again. # The scanner should see they are already in CPU cache and skip them. req2 = Request( request_id="req-dup-lazy", prompt_token_ids=req.prompt_token_ids, sampling_params=req.sampling_params, pooling_params=None, mm_features=None, block_hasher=req._block_hasher, ) gpu_blocks2 = _allocate_gpu_blocks(gpu_pool, req2, num_blocks, group_id=0) gpu_pool.free_blocks(gpu_blocks2) # Flush again so the hashed blocks are at LRU head for the scanner. fillers2 = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5) meta2 = sched.build_connector_meta(make_scheduler_output({})) gpu_pool.free_blocks(fillers2) # Either no store event, or zero new CPU blocks (already cached). if meta2.store_event >= 0: assert len(meta2.store_cpu_blocks) == 0, ( "Expected no new CPU blocks for duplicate hashes" ) assert get_cpu_free_blocks(sched) == cpu_free_after_first # --------------------------------------------------------------------------- # Test 3: LRU eviction order # --------------------------------------------------------------------------- def test_lru_eviction_order() -> None: """With limited CPU space, oldest blocks should be evicted first. CPU block pool: num_cpu_blocks=5 -> 4 free usable blocks (1 taken by null_block). After storing 4 blocks (2 req_a + 2 req_b), all free slots are occupied by cached blocks (ref_cnt=0, in hash map). When 2 more are stored (req_c), 2 LRU blocks from req_a get evicted from the cache to make room. """ # 5 total = 4 usable (null_block takes 1), filling exactly with 4 blocks fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False) sched = fix.scheduler # Fill CPU with 4 blocks: 2 requests x 2 blocks (in LRU insertion order) req_a = make_request(num_blocks=2) req_b = make_request(num_blocks=2) kv_a = _alloc_and_register(fix, req_a, 2) kv_b = _alloc_and_register(fix, req_b, 2) sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0) sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0) ids_a = kv_a.get_block_ids() ids_b = kv_b.get_block_ids() sched_out = make_scheduler_output( { req_a.request_id: 2 * BLOCK_SIZE, req_b.request_id: 2 * BLOCK_SIZE, }, new_reqs={ req_a.request_id: ids_a, req_b.request_id: ids_b, }, ) meta = sched.build_connector_meta(sched_out) assert meta.store_event >= 0 simulate_store_completion(sched, meta.store_event) # Verify all 4 blocks are cached in CPU hash map for i, bhash in enumerate(req_a.block_hashes[:2]): bhash_with_group = make_block_hash_with_group_id(bhash, 0) assert ( sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( bhash_with_group ) is not None ), f"req_a block {i} should be cached after store" for i, bhash in enumerate(req_b.block_hashes[:2]): bhash_with_group = make_block_hash_with_group_id(bhash, 0) assert ( sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( bhash_with_group ) is not None ), f"req_b block {i} should be cached after store" # Store 2 more blocks from a new request - must evict 2 LRU blocks (req_a) req_c = make_request(num_blocks=2) kv_c = _alloc_and_register(fix, req_c, 2) sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0) ids_c = kv_c.get_block_ids() sched_out2 = make_scheduler_output( {req_c.request_id: 2 * BLOCK_SIZE}, new_reqs={req_c.request_id: ids_c}, ) meta2 = sched.build_connector_meta(sched_out2) assert meta2.store_event >= 0 simulate_store_completion(sched, meta2.store_event) # req_a hashes should be evicted from CPU (they were LRU) for i, bhash in enumerate(req_a.block_hashes[:2]): bhash_with_group = make_block_hash_with_group_id(bhash, 0) cache_map = sched.cpu_block_pool.cached_block_hash_to_block cached = cache_map.get_one_block(bhash_with_group) assert cached is None, f"req_a block {i} should have been evicted" # req_b and req_c hashes should be present for i, bhash in enumerate(req_b.block_hashes[:2]): bhash_with_group = make_block_hash_with_group_id(bhash, 0) cache_map = sched.cpu_block_pool.cached_block_hash_to_block cached = cache_map.get_one_block(bhash_with_group) assert cached is not None, f"req_b block {i} should still be cached" for i, bhash in enumerate(req_c.block_hashes[:2]): bhash_with_group = make_block_hash_with_group_id(bhash, 0) cache_map = sched.cpu_block_pool.cached_block_hash_to_block cached = cache_map.get_one_block(bhash_with_group) assert cached is not None, f"req_c block {i} should still be cached" # --------------------------------------------------------------------------- # Test 4: Touched blocks survive eviction # --------------------------------------------------------------------------- def test_touched_blocks_survive_eviction() -> None: """Touching CPU blocks updates their LRU position, protecting them from eviction.""" # 5 total = 4 usable (null_block takes 1) fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False) sched = fix.scheduler # Fill CPU with 4 blocks (req_a: 2, req_b: 2) in LRU order req_a = make_request(num_blocks=2) req_b = make_request(num_blocks=2) kv_a = _alloc_and_register(fix, req_a, 2) kv_b = _alloc_and_register(fix, req_b, 2) sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0) sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0) ids_a = kv_a.get_block_ids() ids_b = kv_b.get_block_ids() sched_out = make_scheduler_output( { req_a.request_id: 2 * BLOCK_SIZE, req_b.request_id: 2 * BLOCK_SIZE, }, new_reqs={ req_a.request_id: ids_a, req_b.request_id: ids_b, }, ) meta = sched.build_connector_meta(sched_out) simulate_store_completion(sched, meta.store_event) # Touch req_a's CPU blocks to make them most-recently-used cpu_pool = sched.cpu_block_pool for bhash in req_a.block_hashes[:2]: bhash_with_group = make_block_hash_with_group_id(bhash, 0) cached_blk = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group) assert cached_blk is not None cpu_pool.touch([cached_blk]) # Undo touch to return ref_cnt to 0 # (so it's a free candidate but at MRU position) cpu_pool.free_blocks([cached_blk]) # Now store 2 more blocks; req_b (LRU front) should be evicted, not req_a req_c = make_request(num_blocks=2) kv_c = _alloc_and_register(fix, req_c, 2) sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0) ids_c = kv_c.get_block_ids() sched_out2 = make_scheduler_output( {req_c.request_id: 2 * BLOCK_SIZE}, new_reqs={req_c.request_id: ids_c}, ) meta2 = sched.build_connector_meta(sched_out2) simulate_store_completion(sched, meta2.store_event) # req_b should be evicted (LRU), req_a and req_c should survive for i, bhash in enumerate(req_b.block_hashes[:2]): bhash_with_group = make_block_hash_with_group_id(bhash, 0) cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group) assert cached is None, f"req_b block {i} should have been evicted (it was LRU)" for i, bhash in enumerate(req_a.block_hashes[:2]): bhash_with_group = make_block_hash_with_group_id(bhash, 0) cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group) assert cached is not None, f"req_a block {i} should survive (was touched/MRU)" # --------------------------------------------------------------------------- # Test 5: Preemption no CPU block leak # --------------------------------------------------------------------------- def test_preemption_no_cpu_block_leak() -> None: """request_finished during in-flight load defers cleanup; completes after load done.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler num_blocks = 2 # First: store blocks to CPU req = make_request(num_blocks=num_blocks) kv_blocks = _alloc_and_register(fix, req, num_blocks) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) block_ids = kv_blocks.get_block_ids() sched_out = make_scheduler_output( {req.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req.request_id: block_ids}, ) meta = sched.build_connector_meta(sched_out) simulate_store_completion(sched, meta.store_event) # Create new request with same tokens, check hit req2 = Request( request_id="req-preempt-load", prompt_token_ids=req.prompt_token_ids, sampling_params=req.sampling_params, pooling_params=None, mm_features=None, block_hasher=req._block_hasher, ) hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens > 0 gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) # Assign load_event via build_connector_meta block_ids2 = kv_blocks2.get_block_ids() sched_out2 = make_scheduler_output( {req2.request_id: 1}, new_reqs={req2.request_id: block_ids2}, ) meta2 = sched.build_connector_meta(sched_out2) assert meta2.load_event >= 0 # Request finishes BEFORE load completes -> deferred sched.request_finished(req2, block_ids=[]) assert req2.request_id in sched._reqs_to_load assert sched._reqs_to_load[req2.request_id].finished is True # Now simulate load completion -> cleanup fires simulate_load_completion(sched, {req2.request_id}) assert req2.request_id not in sched._reqs_to_load # --------------------------------------------------------------------------- # Test 6: Eager store preemption cleanup # --------------------------------------------------------------------------- def test_eager_store_preemption_cleanup() -> None: """In eager mode, finishing a request during in-flight store defers cleanup.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler num_blocks = 2 req = make_request(num_blocks=num_blocks) kv_blocks = _alloc_and_register(fix, req, num_blocks) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) block_ids = kv_blocks.get_block_ids() sched_out = make_scheduler_output( {req.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req.request_id: block_ids}, ) meta = sched.build_connector_meta(sched_out) store_event = meta.store_event assert store_event >= 0 # The request gets store_events populated assert req.request_id in sched._reqs_to_store store_state = sched._reqs_to_store[req.request_id] assert store_event in store_state.store_events # Finish request while store still in-flight -> deferred sched.request_finished(req, block_ids=[]) assert req.request_id in sched._reqs_to_store assert sched._reqs_to_store[req.request_id].finished is True # Simulate store completion -> deferred cleanup fires simulate_store_completion(sched, store_event) assert req.request_id not in sched._reqs_to_store # --------------------------------------------------------------------------- # Test 7: In-flight finish deferred cleanup (load variant) # --------------------------------------------------------------------------- def test_inflight_finish_deferred_cleanup() -> None: """Store, then start a load, request_finished defers, load completion fires cleanup.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler num_blocks = 2 # Store req = make_request(num_blocks=num_blocks) kv_blocks = _alloc_and_register(fix, req, num_blocks) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) block_ids = kv_blocks.get_block_ids() sched_out = make_scheduler_output( {req.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req.request_id: block_ids}, ) meta = sched.build_connector_meta(sched_out) simulate_store_completion(sched, meta.store_event) # Load req2 = Request( request_id="req-inflight-load", prompt_token_ids=req.prompt_token_ids, sampling_params=req.sampling_params, pooling_params=None, mm_features=None, block_hasher=req._block_hasher, ) hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens > 0 gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) block_ids2 = kv_blocks2.get_block_ids() sched_out2 = make_scheduler_output( {req2.request_id: 1}, new_reqs={req2.request_id: block_ids2}, ) meta2 = sched.build_connector_meta(sched_out2) assert meta2.load_event >= 0 # Finish before load completes sched.request_finished(req2, block_ids=[]) assert req2.request_id in sched._reqs_to_load # Simulate load completion -> request removed simulate_load_completion(sched, {req2.request_id}) assert req2.request_id not in sched._reqs_to_load # --------------------------------------------------------------------------- # Test 8: Null GPU blocks are skipped in store and load transfer pairs # --------------------------------------------------------------------------- def test_multi_group_null_blocks_skipped() -> None: """Null GPU blocks (no block_hash) must not appear in store or load pairs. In eager store mode, _prepare_eager_store_specs skips blocks whose block_hash is None (null blocks have no hash). We verify this by mixing real hashed blocks with unhashed (null-like) blocks in a single group and checking that only real blocks appear in the store list. """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False) sched = fix.scheduler gpu_pool = fix.gpu_block_pool num_blocks = 2 req = make_request(num_blocks=num_blocks) # Allocate real blocks (with hashes) and use the null_block as a placeholder gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) null_block = gpu_pool.null_block # Mix: [real_block, null_block] — null_block has no hash, should be skipped mixed_blocks = [gpu_blocks[0], null_block] kv_blocks = KVCacheBlocks(blocks=(mixed_blocks,)) req.num_computed_tokens = num_blocks * BLOCK_SIZE sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) block_ids = kv_blocks.get_block_ids() sched_out = make_scheduler_output( {req.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req.request_id: block_ids}, ) meta = sched.build_connector_meta(sched_out) # Null block's ID should NOT appear in store_gpu_blocks null_block_id = null_block.block_id assert null_block_id not in meta.store_gpu_blocks, ( f"Null block id {null_block_id} should not appear in store transfer pairs" ) # Only real block should be scheduled for store assert len(meta.store_gpu_blocks) == 1 assert gpu_blocks[0].block_id in meta.store_gpu_blocks # Complete the store assert meta.store_event >= 0 simulate_store_completion(sched, meta.store_event) # Create matching request and get load hit req2 = Request( request_id="req-null-load", prompt_token_ids=req.prompt_token_ids, sampling_params=req.sampling_params, pooling_params=None, mm_features=None, block_hasher=req._block_hasher, ) hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) # Only 1 block was stored (the real one) assert hit_tokens == BLOCK_SIZE assert is_async is True # Allocate new GPU blocks for the load gpu_blocks2 = gpu_pool.get_new_blocks(1) kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],)) sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) sched_out2 = make_scheduler_output({req2.request_id: 1}) meta2 = sched.build_connector_meta(sched_out2) # Null block's ID should NOT appear in load_gpu_blocks assert null_block_id not in meta2.load_gpu_blocks, ( f"Null block id {null_block_id} should not appear in load transfer pairs" ) # --------------------------------------------------------------------------- # Test 9: Chunked prefill accumulates block_ids across steps # --------------------------------------------------------------------------- def test_chunked_prefill_reads_live_block_ids() -> None: """With chunked prefill, block IDs accumulate across scheduler steps. _prepare_eager_store_specs reads block IDs from scheduler_output via yield_req_data, so the store should reflect the updated (larger) block list, not a stale snapshot.""" fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler num_blocks = 4 req = make_request(num_blocks=num_blocks) # First chunk: allocate 2 blocks kv_blocks_first = _alloc_and_register(fix, req, 2) sched.update_state_after_alloc(req, kv_blocks_first, num_external_tokens=0) assert req.request_id in sched._reqs_to_store # Should still be exactly 1 entry in _reqs_to_store assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1 # Build connector meta with 2 blocks — stores the first 2 ids_first = kv_blocks_first.get_block_ids() sched_out1 = make_scheduler_output( {req.request_id: 2 * BLOCK_SIZE}, new_reqs={req.request_id: ids_first}, ) meta1 = sched.build_connector_meta(sched_out1) assert meta1.store_event >= 0 assert len(meta1.store_gpu_blocks) == 2 simulate_store_completion(sched, meta1.store_event) # Second chunk: allocate 4 blocks total (2 new ones) kv_blocks_second = _alloc_and_register(fix, req, num_blocks) # update_state_after_alloc is idempotent for store registration sched.update_state_after_alloc(req, kv_blocks_second, num_external_tokens=0) # Still exactly 1 entry assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1 # The second chunk's NEW block IDs (positions 2,3) are passed as # cached_req_new_blocks. The full block_ids include both old and new, # but yield_req_data only appends the new_block_ids for cached reqs. ids_second_full = kv_blocks_second.get_block_ids() # New blocks are those beyond the first chunk new_block_ids = tuple(ids_second_full[g][2:] for g in range(len(ids_second_full))) sched_out2 = make_scheduler_output( {req.request_id: 2 * BLOCK_SIZE}, cached_req_new_blocks={req.request_id: new_block_ids}, ) meta2 = sched.build_connector_meta(sched_out2) assert meta2.store_event >= 0 # Only the 2 NEW blocks should be stored (first 2 already done) assert len(meta2.store_gpu_blocks) == 2 # --------------------------------------------------------------------------- # Test 10: Partial GPU prefix hit + CPU load + new compute blocks # --------------------------------------------------------------------------- def test_partial_gpu_prefix_plus_cpu_load() -> None: """When GPU has a prefix cache hit for the first N blocks, CPU has a hit for the next M blocks, and there are P new blocks needing fresh compute, the block layout is: | comp (N) | ext_comp (M) | new (P) | External blocks sit in the middle — not at the beginning or end. The load path must target hashes at positions [N, N+M). Request: 6 blocks (0..5). - Store all 6 to CPU. - New request: GPU prefix cache hits blocks 0,1 (hashed). CPU hits blocks 2,3. Blocks 4,5 are new (need compute). - update_state_after_alloc receives 6 GPU blocks: [0,1] hashed (comp), [2,3] unhashed (ext_comp), [4,5] unhashed (new). - Load must target hash positions 2,3. """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) sched = fix.scheduler gpu_pool = fix.gpu_block_pool num_blocks = 6 req = make_request(num_blocks=num_blocks) # Store all 6 blocks to CPU via eager store. kv_blocks = _alloc_and_register(fix, req, num_blocks) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) block_ids = kv_blocks.get_block_ids() sched_out = make_scheduler_output( {req.request_id: num_blocks * BLOCK_SIZE}, new_reqs={req.request_id: block_ids}, ) meta = sched.build_connector_meta(sched_out) assert meta.store_event >= 0 simulate_store_completion(sched, meta.store_event) # New request with same tokens — but only partial GPU prefix hit. req2 = Request( request_id="req-partial-gpu", prompt_token_ids=req.prompt_token_ids, sampling_params=req.sampling_params, pooling_params=None, mm_features=None, block_hasher=req._block_hasher, ) # GPU prefix cache hits the first 2 blocks. gpu_local_computed = 2 * BLOCK_SIZE hit_tokens, is_async = sched.get_num_new_matched_tokens( req2, num_computed_tokens=gpu_local_computed ) # CPU should hit blocks 2,3 (not 4,5 — those are beyond the CPU range). num_cpu_hit_blocks = 2 # Actually CPU has all 6 stored; it returns hits starting from position 2. # The number of CPU hit blocks = min(remaining request blocks, CPU cached). # Here remaining = 6 - 2 = 4 blocks are in CPU, so hit = 4 * BLOCK_SIZE. num_cpu_hit_blocks = 4 assert hit_tokens == num_cpu_hit_blocks * BLOCK_SIZE, ( f"Expected {num_cpu_hit_blocks * BLOCK_SIZE} CPU hit tokens, got {hit_tokens}" ) assert is_async is True # Simulate what the real scheduler does: only accept 2 of the 4 CPU hit # blocks as external (e.g. due to budget constraints), leaving 2 new # blocks for fresh compute. num_ext_blocks = 2 num_new_blocks = 2 external_tokens = num_ext_blocks * BLOCK_SIZE # Build block list matching real layout: | comp(2) | ext_comp(2) | new(2) | # comp: GPU prefix cache hit — blocks with hashes gpu_comp = _allocate_gpu_blocks(gpu_pool, req2, 2, group_id=0) # ext_comp + new: freshly allocated, no hashes gpu_ext_and_new = gpu_pool.get_new_blocks(num_ext_blocks + num_new_blocks) all_gpu_blocks = gpu_comp + gpu_ext_and_new kv_blocks2 = KVCacheBlocks(blocks=(all_gpu_blocks,)) # Critical call: with 2 hashed comp blocks and 2 external tokens worth # of blocks, the manager must derive skipped=2 and load hashes [2,3]. sched.update_state_after_alloc( req2, kv_blocks2, num_external_tokens=external_tokens ) block_ids2 = kv_blocks2.get_block_ids() sched_out2 = make_scheduler_output( {req2.request_id: num_new_blocks * BLOCK_SIZE}, new_reqs={req2.request_id: block_ids2}, ) meta2 = sched.build_connector_meta(sched_out2) assert meta2.load_event >= 0, "Expected a load event for partial GPU + CPU hit" assert len(meta2.load_gpu_blocks) == num_ext_blocks assert len(meta2.load_cpu_blocks) == num_ext_blocks # Verify the load targets the ext_comp GPU blocks (positions 2,3), # not the comp blocks (0,1) or new blocks (4,5). ext_block_ids = [b.block_id for b in gpu_ext_and_new[:num_ext_blocks]] for bid in meta2.load_gpu_blocks: assert bid in ext_block_ids, ( f"Load GPU block {bid} should be an ext_comp block, not a comp or new block" )