diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index e2450c662549..1b969bc1ed36 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -22,6 +22,10 @@ flowchart LR CPU <--> SN["..."] ``` +## Terminology: Chunks + +The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers. + ## Single-Tier Setup (CPU Only) ```bash diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index 58bbd4091560..b7d9d6c8efac 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -704,7 +704,7 @@ def test_prom_metrics_registers_tiering_metrics_from_spec(): ) metric = prom_metrics._offloading_metric_defs[ - TieringOffloadingMetrics.BLOCK_QUERIES + TieringOffloadingMetrics.CHUNK_QUERIES ] assert metric.kwargs["labelnames"] == ["model_name", "engine", "tier"] diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 14320c367cc6..47fb7103d742 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -96,7 +96,7 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): groups.append( KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1) ) - manager = CPUOffloadingManager(num_blocks=100) + manager = CPUOffloadingManager(num_chunks=100) spec = SimpleNamespace( tokens_per_block=(256, 64, 8), tokens_per_hash=8, @@ -2888,7 +2888,7 @@ def _make_req_status( req=req, req_context=ReqContext(req_id="test-req"), offloading_context=RequestOffloadingContext( - policy=OffloadPolicy.BLOCK_LEVEL + policy=OffloadPolicy.CHUNK_LEVEL ), num_locally_computed_tokens=num_computed_tokens, ) diff --git a/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py index 14ccf8b67e7b..338ec18d0e67 100644 --- a/tests/v1/kv_offload/cpu/policies/test_factory.py +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -7,7 +7,7 @@ from vllm.v1.kv_offload.base import OffloadKey, ReqContext from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy @@ -20,10 +20,10 @@ class _DummyCachePolicy(CachePolicy): def __init__(self, cache_capacity: int) -> None: self.cache_capacity = cache_capacity - def get(self, key: OffloadKey) -> BlockStatus | None: + def get(self, key: OffloadKey) -> ChunkStatus | None: return None - def insert(self, key: OffloadKey, block: BlockStatus) -> None: + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: pass def remove(self, key: OffloadKey) -> None: @@ -34,7 +34,7 @@ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: def evict( self, n: int, protected: set[OffloadKey] - ) -> list[tuple[OffloadKey, BlockStatus]] | None: + ) -> list[tuple[OffloadKey, ChunkStatus]] | None: return None def clear(self) -> None: @@ -72,7 +72,7 @@ def test_register_and_resolve_custom_policy(self): policy_cls = CachePolicyFactory.get_cache_policy_cls("dummy") assert policy_cls is _DummyCachePolicy - manager = CPUOffloadingManager(num_blocks=4, cache_policy="dummy") + manager = CPUOffloadingManager(num_chunks=4, cache_policy="dummy") assert isinstance(manager._policy, _DummyCachePolicy) def test_unregistered_policy_raises(self): @@ -98,7 +98,7 @@ def test_manager_resolves_policy_via_module_path(self): """End-to-end: CPUOffloadingManager resolves an unregistered policy purely from cache_policy_module_path.""" manager = CPUOffloadingManager( - num_blocks=4, + num_chunks=4, cache_policy="_DummyCachePolicy", cache_policy_module_path="tests.v1.kv_offload.cpu.policies.test_factory", ) diff --git a/tests/v1/kv_offload/cpu/test_canonical_layout.py b/tests/v1/kv_offload/cpu/test_canonical_layout.py index f10e534cc74b..4594e1eb0924 100644 --- a/tests/v1/kv_offload/cpu/test_canonical_layout.py +++ b/tests/v1/kv_offload/cpu/test_canonical_layout.py @@ -232,9 +232,9 @@ def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int): def canonical_view(rank: int, world_size: int) -> torch.Tensor: region = SharedOffloadRegion( engine_id=engine_id, - num_blocks=num_blocks, + num_chunks=num_blocks, rank=rank, - kv_bytes_per_block=row_stride, + kv_bytes_per_chunk=row_stride, cpu_page_size=row_stride // world_size, ) regions.append(region) diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 1b6e3a2586aa..5306d0557102 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -26,7 +26,7 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion NUM_GPU_BLOCKS = [64] -NUM_CPU_BLOCKS = [256] +NUM_CPU_CHUNKS = [256] GPU_PAGE_SIZES = [512, 1024] BLOCKS_PER_CHUNK_VALUES = [1, 3] NUM_TENSORS = [4] @@ -225,7 +225,7 @@ def record_region_cleanup() -> None: @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES) @pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES) @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS) -@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS) +@pytest.mark.parametrize("num_cpu_chunks", NUM_CPU_CHUNKS) @pytest.mark.parametrize("num_tensors", NUM_TENSORS) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", DEVICES) @@ -241,7 +241,7 @@ def test_transfer( gpu_page_size_bytes: int, blocks_per_chunk: int, num_gpu_blocks: int, - num_cpu_blocks: int, + num_cpu_chunks: int, num_tensors: int, seed: int, device: str, @@ -288,41 +288,41 @@ def test_transfer( SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT, ) simulated_world_size = 2 - kv_bytes_per_block = ( + kv_bytes_per_chunk = ( cpu_page_size if replicated_layout else cpu_page_size * simulated_world_size ) mmap_region = SharedOffloadRegion( engine_id=str(uuid.uuid4()), - num_blocks=num_cpu_blocks, + num_chunks=num_cpu_chunks, rank=0, - kv_bytes_per_block=kv_bytes_per_block, + kv_bytes_per_chunk=kv_bytes_per_chunk, cpu_page_size=cpu_page_size, ) worker = CPUOffloadingWorker( kv_caches=kv_caches, blocks_per_chunk=blocks_per_chunk, - num_cpu_blocks=num_cpu_blocks, + num_cpu_chunks=num_cpu_chunks, mmap_region=mmap_region, ) # select block mappings gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * blocks_per_chunk) - cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings) + cpu_chunks = random.sample(range(num_cpu_chunks), num_mappings) - # expand cpu blocks to gpu-page granularity for uniform comparison: - # each cpu block maps to blocks_per_chunk consecutive sub-blocks - cpu_blocks_expanded = [ - cpu_block * blocks_per_chunk + j - for cpu_block in cpu_blocks + # expand cpu chunks to gpu-page granularity for uniform comparison: + # each cpu chunk maps to blocks_per_chunk consecutive sub-blocks + cpu_chunks_expanded = [ + cpu_chunk * blocks_per_chunk + j + for cpu_chunk in cpu_chunks for j in range(blocks_per_chunk) ] - # maybe skip some GPU blocks to test reading/writing from the middle of a CPU block + # maybe skip some GPU blocks to test reading/writing from the middle of a CPU chunk blocks_to_skip = blocks_per_chunk - 1 if blocks_to_skip > 0: gpu_blocks = gpu_blocks[blocks_to_skip:] - cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:] + cpu_chunks_expanded = cpu_chunks_expanded[blocks_to_skip:] # set transfer direction if gpu_to_cpu: @@ -330,16 +330,16 @@ def test_transfer( src_spec = GPULoadStoreSpec( gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,) ) - dst_spec = CPULoadStoreSpec(cpu_blocks) - dst_to_src = dict(zip(cpu_blocks_expanded, gpu_blocks)) + dst_spec = CPULoadStoreSpec(cpu_chunks) + dst_to_src = dict(zip(cpu_chunks_expanded, gpu_blocks)) num_dst_sub_blocks = num_gpu_blocks else: handler = worker._load_handler - src_spec = CPULoadStoreSpec(cpu_blocks) + src_spec = CPULoadStoreSpec(cpu_chunks) dst_spec = GPULoadStoreSpec( gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,) ) - dst_to_src = dict(zip(gpu_blocks, cpu_blocks_expanded)) + dst_to_src = dict(zip(gpu_blocks, cpu_chunks_expanded)) num_dst_sub_blocks = num_gpu_blocks # randomize src and dst tensors before transfer @@ -410,7 +410,7 @@ def test_transfer( @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES) @pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES) @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS) -@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS) +@pytest.mark.parametrize("num_cpu_chunks", NUM_CPU_CHUNKS) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", DEVICES) @torch.inference_mode() @@ -421,7 +421,7 @@ def test_transfer_multi_group( gpu_page_size_bytes: int, blocks_per_chunk: int, num_gpu_blocks: int, - num_cpu_blocks: int, + num_cpu_chunks: int, seed: int, device: str, ) -> None: @@ -470,37 +470,37 @@ def test_transfer_multi_group( worker = CPUOffloadingWorker( kv_caches=canonical_kv_caches, blocks_per_chunk=blocks_per_chunk, - num_cpu_blocks=num_cpu_blocks, + num_cpu_chunks=num_cpu_chunks, ) # group 0: aligned, group 1: empty, group 2: unaligned on CPU->GPU - group_sizes_in_cpu_blocks = [num_mappings_per_group, 0, num_mappings_per_group] + group_sizes_in_cpu_chunks = [num_mappings_per_group, 0, num_mappings_per_group] - total_cpu_blocks = sum(group_sizes_in_cpu_blocks) - total_gpu_blocks_needed = total_cpu_blocks * blocks_per_chunk + total_cpu_chunks = sum(group_sizes_in_cpu_chunks) + total_gpu_blocks_needed = total_cpu_chunks * blocks_per_chunk gpu_blocks_all = random.sample(range(num_gpu_blocks), total_gpu_blocks_needed) - cpu_blocks_all = random.sample(range(num_cpu_blocks), total_cpu_blocks) + cpu_chunks_all = random.sample(range(num_cpu_chunks), total_cpu_chunks) - # split gpu/cpu blocks per group + # split gpu blocks / cpu chunks per group gpu_blocks_per_group: list[list[int]] = [] - cpu_blocks_per_group: list[list[int]] = [] + cpu_chunks_per_group: list[list[int]] = [] gpu_offset = 0 cpu_offset = 0 - for size in group_sizes_in_cpu_blocks: + for size in group_sizes_in_cpu_chunks: gpu_count = size * blocks_per_chunk gpu_blocks_per_group.append(gpu_blocks_all[gpu_offset : gpu_offset + gpu_count]) - cpu_blocks_per_group.append(cpu_blocks_all[cpu_offset : cpu_offset + size]) + cpu_chunks_per_group.append(cpu_chunks_all[cpu_offset : cpu_offset + size]) gpu_offset += gpu_count cpu_offset += size - # expand cpu blocks to gpu-page granularity - cpu_blocks_expanded_per_group = [ + # expand cpu chunks to gpu-page granularity + cpu_chunks_expanded_per_group = [ [ - cpu_block * blocks_per_chunk + j - for cpu_block in cpu_blocks + cpu_chunk * blocks_per_chunk + j + for cpu_chunk in cpu_chunks for j in range(blocks_per_chunk) ] - for cpu_blocks in cpu_blocks_per_group + for cpu_chunks in cpu_chunks_per_group ] # skip sub-blocks from group 2 to test unaligned transfers. @@ -509,7 +509,7 @@ def test_transfer_multi_group( gpu_blocks_per_group[2] = gpu_blocks_per_group[2][ sub_blocks_to_skip:-sub_blocks_to_skip ] - cpu_blocks_expanded_per_group[2] = cpu_blocks_expanded_per_group[2][ + cpu_chunks_expanded_per_group[2] = cpu_chunks_expanded_per_group[2][ sub_blocks_to_skip:-sub_blocks_to_skip ] @@ -520,10 +520,10 @@ def test_transfer_multi_group( gpu_blocks.extend(gpu_blks) group_sizes.append(len(gpu_blks)) - # build flat cpu_blocks list - cpu_blocks = [] - for cpu_blks in cpu_blocks_per_group: - cpu_blocks.extend(cpu_blks) + # build flat cpu_chunks list + cpu_chunks = [] + for cpu_chnks in cpu_chunks_per_group: + cpu_chunks.extend(cpu_chnks) # block_indices: only relevant for unaligned transfers block_indices: list[int] = [0, 0, sub_blocks_to_skip] @@ -533,18 +533,18 @@ def test_transfer_multi_group( src_spec = GPULoadStoreSpec( gpu_blocks, group_sizes=group_sizes, block_indices=block_indices ) - dst_spec = CPULoadStoreSpec(cpu_blocks) + dst_spec = CPULoadStoreSpec(cpu_chunks) # per-group mapping: cpu sub-block -> gpu sub-block dst_to_src_per_group = [ dict(zip(expanded, gpu_blks)) for expanded, gpu_blks in zip( - cpu_blocks_expanded_per_group, gpu_blocks_per_group + cpu_chunks_expanded_per_group, gpu_blocks_per_group ) ] - num_dst_sub_blocks = num_cpu_blocks * blocks_per_chunk + num_dst_sub_blocks = num_cpu_chunks * blocks_per_chunk else: handler = worker._load_handler - src_spec = CPULoadStoreSpec(cpu_blocks) + src_spec = CPULoadStoreSpec(cpu_chunks) dst_spec = GPULoadStoreSpec( gpu_blocks, group_sizes=group_sizes, block_indices=block_indices ) @@ -552,7 +552,7 @@ def test_transfer_multi_group( dst_to_src_per_group = [ dict(zip(gpu_blks, expanded)) for gpu_blks, expanded in zip( - gpu_blocks_per_group, cpu_blocks_expanded_per_group + gpu_blocks_per_group, cpu_chunks_expanded_per_group ) ] num_dst_sub_blocks = num_gpu_blocks @@ -644,7 +644,7 @@ def test_load_waits_for_pending_compute_stream_writes(default_vllm_config) -> No ], ), blocks_per_chunk=1, - num_cpu_blocks=num_blocks, + num_cpu_chunks=num_blocks, ) worker._load_handler.src_tensors[0].fill_(sentinel) expected = torch.full((page_size_bytes,), sentinel, dtype=torch.int8) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 22c29ecfb362..a79f480ff551 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -36,7 +36,7 @@ def make_req_context( def make_cpu_manager( - num_blocks: int = 4, + num_chunks: int = 4, cache_policy: str = "lru", cache_policy_module_path: str | None = None, enable_events: bool = False, @@ -44,7 +44,7 @@ def make_cpu_manager( max_tracker_size: int = 64_000, ) -> CPUOffloadingManager: return CPUOffloadingManager( - num_blocks=num_blocks, + num_chunks=num_chunks, cache_policy=cache_policy, cache_policy_module_path=cache_policy_module_path, enable_events=enable_events, @@ -56,7 +56,7 @@ def make_cpu_manager( @dataclass class ExpectedPrepareStoreOutput: keys_to_store: list[int] - store_block_ids: list[int] + store_chunk_ids: list[int] evicted_keys: list[int] @@ -93,9 +93,9 @@ def verify_store_output( store_spec = prepare_store_output.store_spec assert isinstance(store_spec, CPULoadStoreSpec) expected_array = np.array( - expected_prepare_store_output.store_block_ids, dtype=np.int64 + expected_prepare_store_output.store_chunk_ids, dtype=np.int64 ) - assert np.array_equal(expected_array, store_spec.block_ids) + assert np.array_equal(expected_array, store_spec.chunk_ids) def verify_load_output( @@ -103,7 +103,7 @@ def verify_load_output( ): assert isinstance(prepare_load_output, CPULoadStoreSpec) expected_array = np.array(expected_prepare_load_output, dtype=np.int64) - assert np.array_equal(expected_array, prepare_load_output.block_ids) + assert np.array_equal(expected_array, prepare_load_output.chunk_ids) def check_split_usage_stats( @@ -146,7 +146,7 @@ def to_key_sets( def test_cpu_eviction_removed_precedes_stored(): """An eviction is announced before the store that reuses its capacity.""" - manager = make_cpu_manager(num_blocks=2, enable_events=True) + manager = make_cpu_manager(num_chunks=2, enable_events=True) manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -164,23 +164,23 @@ def test_cpu_eviction_removed_precedes_stored(): @pytest.mark.parametrize("eviction_policy", ["lru", "arc"]) -def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): +def test_already_stored_chunk_not_evicted_during_prepare_store(eviction_policy): """ - Regression test: a block that is already stored must not be evicted - by prepare_store() when it needs to make room for new blocks. + Regression test: a chunk that is already stored must not be evicted + by prepare_store() when it needs to make room for new chunks. Applies to both lru and arc policies. Scenario: - - Store blocks [1, 2] and complete. - - touch([1]) makes block 2 the LRU candidate. + - Store chunks [1, 2] and complete. + - touch([1]) makes chunk 2 the LRU candidate. - prepare_store([2, 3, 4, 5]): - * block 2 is filtered out as "already stored" - * but without the fix, block 2 would be evicted as the LRU + * chunk 2 is filtered out as "already stored" + * but without the fix, chunk 2 would be evicted as the LRU candidate to make room for [3, 4, 5] - - After complete_store([2, 3, 4, 5]), block 2 must still be present. + - After complete_store([2, 3, 4, 5]), chunk 2 must still be present. """ manager = make_cpu_manager( - num_blocks=4, + num_chunks=4, cache_policy=eviction_policy, enable_events=True, ) @@ -189,33 +189,33 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) - # touch [1] to make block 2 the LRU candidate + # touch [1] to make chunk 2 the LRU candidate manager.touch(to_keys([1]), _EMPTY_REQ_CTX) # prepare_store([2, 3, 4, 5]): - # - block 2 is already stored -> filtered out of keys_to_store - # - block 2 must NOT be evicted even though it is the LRU candidate - # - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0 + # - chunk 2 is already stored -> filtered out of keys_to_store + # - chunk 2 must NOT be evicted even though it is the LRU candidate + # - chunk 1 (ID 0) is evicted instead; new chunks [3,4,5] get IDs 2,3,0 prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[3, 4, 5], - store_block_ids=[2, 3, 0], - evicted_keys=[1], # block 1 evicted, not block 2 + store_chunk_ids=[2, 3, 0], + evicted_keys=[1], # chunk 1 evicted, not chunk 2 ), ) - # complete_store must not silently drop block 2 + # complete_store must not silently drop chunk 2 manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) - # block 2 must still be present in the cache + # chunk 2 must still be present in the cache assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT def test_filter_reused_manager_reports_stores_skipped_counter(): manager = make_cpu_manager( - num_blocks=4, + num_chunks=4, cache_policy="lru", store_threshold=2, ) @@ -226,7 +226,7 @@ def test_filter_reused_manager_reports_stores_skipped_counter(): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[], - store_block_ids=[], + store_chunk_ids=[], evicted_keys=[], ), ) @@ -247,34 +247,34 @@ def check_usage_stats(manager: CPUOffloadingManager, value: float): ] == pytest.approx(value) # Zero-capacity manager always reports 0.0 - manager = make_cpu_manager(num_blocks=0) + manager = make_cpu_manager(num_chunks=0) check_usage_stats(manager, 0.0) - # Empty manager (4 blocks, none allocated): usage = 0.0 - manager = make_cpu_manager(num_blocks=4) + # Empty manager (4 chunks, none allocated): usage = 0.0 + manager = make_cpu_manager(num_chunks=4) check_usage_stats(manager, 0.0) - # After allocating 2 of 4 blocks: usage = 0.5 + # After allocating 2 of 4 chunks: usage = 0.5 manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) check_usage_stats(manager, 0.5) - # After filling all 4 blocks: usage = 1.0 + # After filling all 4 chunks: usage = 1.0 manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX) check_usage_stats(manager, 1.0) - # After completing store, the blocks becomes evictable as it is not actively used + # After completing store, the chunks become evictable as not actively used # and usage drops. manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) check_usage_stats(manager, 0.5) - # After completing store, the blocks becomes evictable as it is not actively used + # After completing store, the chunks become evictable as not actively used # and usage drops. manager.complete_store(to_keys([3, 4]), _EMPTY_REQ_CTX) check_usage_stats(manager, 0.0) def test_cpu_manager_reports_allocation_size_histogram(): - manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -297,12 +297,12 @@ def test_cpu_manager_reports_allocation_size_histogram(): def test_cpu_manager_reports_allocation_size_on_allocation_failure(monkeypatch): - manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") - def fail_allocate_blocks(keys): + def fail_allocate_chunks(keys): raise RuntimeError("allocation failed") - monkeypatch.setattr(manager, "_allocate_blocks", fail_allocate_blocks) + monkeypatch.setattr(manager, "_allocate_chunks", fail_allocate_chunks) with pytest.raises(RuntimeError, match="allocation failed"): manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) @@ -316,7 +316,7 @@ def fail_allocate_blocks(keys): def test_cpu_manager_reports_allocation_size_on_eviction_failure(): - manager = make_cpu_manager(num_blocks=1, cache_policy="lru") + manager = make_cpu_manager(num_chunks=1, cache_policy="lru") manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) manager.get_stats() @@ -332,7 +332,7 @@ def test_cpu_manager_reports_allocation_size_on_eviction_failure(): def test_cpu_manager_reports_cache_write_and_read_usage_gauges(): - manager = make_cpu_manager(num_blocks=4) + manager = make_cpu_manager(num_chunks=4) # Store path: pins write usage until complete_store. manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -356,7 +356,7 @@ def test_cpu_manager_reports_cache_write_and_read_usage_gauges(): def test_cpu_manager_clears_write_usage_after_failed_store(): - manager = make_cpu_manager(num_blocks=4) + manager = make_cpu_manager(num_chunks=4) manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5) @@ -369,8 +369,8 @@ def test_cpu_manager(): """ Tests CPUOffloadingManager with lru policy. """ - # initialize a CPU manager with a capacity of 4 blocks - cpu_manager = make_cpu_manager(num_blocks=4, cache_policy="lru", enable_events=True) + # initialize a CPU manager with a capacity of 4 chunks + cpu_manager = make_cpu_manager(num_chunks=4, cache_policy="lru", enable_events=True) # prepare store [1, 2] prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -378,7 +378,7 @@ def test_cpu_manager(): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[1, 2], - store_block_ids=[0, 1], + store_chunk_ids=[0, 1], evicted_keys=[], ), ) @@ -407,7 +407,7 @@ def test_cpu_manager(): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[3, 4, 5], - store_block_ids=[2, 3, 0], + store_chunk_ids=[2, 3, 0], evicted_keys=[1], ), ) @@ -445,7 +445,7 @@ def test_cpu_manager(): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[6, 7, 8], - store_block_ids=[1, 0, 3], + store_chunk_ids=[1, 0, 3], evicted_keys=[4, 5, 2], ), ) @@ -462,7 +462,7 @@ def test_cpu_manager(): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[9], - store_block_ids=[3], + store_chunk_ids=[3], evicted_keys=[8], ), ) @@ -482,38 +482,38 @@ def test_cpu_manager(): def test_prepare_load_preserves_key_order(): - """block_ids[i] must correspond to keys[i] (co-indexed invariant).""" - manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + """chunk_ids[i] must correspond to keys[i] (co-indexed invariant).""" + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") key_a, key_b, key_c = to_key(0), to_key(1), to_key(2) - # Store all three keys and learn their block ID assignments + # Store all three keys and learn their chunk ID assignments store_output = manager.prepare_store([key_a, key_b, key_c], _EMPTY_REQ_CTX) assert store_output is not None assert isinstance(store_output.store_spec, CPULoadStoreSpec) - key_to_block_id = { - k: int(bid) - for k, bid in zip(store_output.keys_to_store, store_output.store_spec.block_ids) + key_to_chunk_id = { + k: int(cid) + for k, cid in zip(store_output.keys_to_store, store_output.store_spec.chunk_ids) } manager.complete_store([key_a, key_b, key_c], _EMPTY_REQ_CTX) # Forward order: [a, b, c] spec_fwd = manager.prepare_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) assert isinstance(spec_fwd, CPULoadStoreSpec) - assert [int(x) for x in spec_fwd.block_ids] == [ - key_to_block_id[key_a], - key_to_block_id[key_b], - key_to_block_id[key_c], + assert [int(x) for x in spec_fwd.chunk_ids] == [ + key_to_chunk_id[key_a], + key_to_chunk_id[key_b], + key_to_chunk_id[key_c], ] manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant # Arbitrary permutation: [b, c, a] spec_perm = manager.prepare_load([key_b, key_c, key_a], _EMPTY_REQ_CTX) assert isinstance(spec_perm, CPULoadStoreSpec) - assert [int(x) for x in spec_perm.block_ids] == [ - key_to_block_id[key_b], - key_to_block_id[key_c], - key_to_block_id[key_a], + assert [int(x) for x in spec_perm.chunk_ids] == [ + key_to_chunk_id[key_b], + key_to_chunk_id[key_c], + key_to_chunk_id[key_a], ] manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant @@ -522,10 +522,10 @@ class TestARCPolicy: """Unit tests for CPUOffloadingManager with ARC eviction policy.""" def _make_manager( - self, num_blocks: int = 4, enable_events: bool = True + self, num_chunks: int = 4, enable_events: bool = True ) -> tuple[CPUOffloadingManager, ARCCachePolicy]: manager = make_cpu_manager( - num_blocks=num_blocks, + num_chunks=num_chunks, cache_policy="arc", enable_events=enable_events, ) @@ -548,7 +548,7 @@ def test_basic(self): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[1, 2], - store_block_ids=[0, 1], + store_chunk_ids=[0, 1], evicted_keys=[], ), ) @@ -569,36 +569,36 @@ def test_basic(self): assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS - # blocks should be in T1 (recent) + # chunks should be in T1 (recent) assert len(arc_policy.t1) == 2 assert len(arc_policy.t2) == 0 def test_t1_to_t2_promotion(self): """ - Tests that accessing a block in T1 promotes it to T2 (frequent). + Tests that accessing a chunk in T1 promotes it to T2 (frequent). This is a key feature of ARC's adaptive behavior. """ cpu_manager, arc_policy = self._make_manager(enable_events=False) - # store and complete block 1 + # store and complete chunk 1 cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX) - # block 1 starts in T1 (recent) + # chunk 1 starts in T1 (recent) assert to_keys([1])[0] in arc_policy.t1 assert to_keys([1])[0] not in arc_policy.t2 - # touch block 1 (simulate second access) + # touch chunk 1 (simulate second access) cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX) - # block 1 should now be in T2 (frequent) + # chunk 1 should now be in T2 (frequent) assert to_keys([1])[0] not in arc_policy.t1 assert to_keys([1])[0] in arc_policy.t2 def test_eviction_with_load(self): """ Tests ARC eviction behavior similar to LRU test. - Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted. + Verifies that chunks being loaded (ref_cnt > 0) cannot be evicted. """ cpu_manager, _ = self._make_manager() @@ -610,7 +610,7 @@ def test_eviction_with_load(self): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[1, 2, 3, 4], - store_block_ids=[0, 1, 2, 3], + store_chunk_ids=[0, 1, 2, 3], evicted_keys=[], ), ) @@ -628,36 +628,36 @@ def test_eviction_with_load(self): cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX) # now prepare store [5, 6, 7] should succeed - # ARC will evict blocks one at a time from T1 as needed + # ARC will evict chunks one at a time from T1 as needed prepare_store_output = cpu_manager.prepare_store( to_keys([5, 6, 7]), _EMPTY_REQ_CTX ) assert prepare_store_output is not None - # Should successfully evict enough blocks to make room (at least 1) + # Should successfully evict enough chunks to make room (at least 1) assert len(prepare_store_output.evicted_keys) >= 1 def test_adaptive_target(self): """ Tests ARC's adaptive target adjustment via ghost lists. - When a block in B1 (ghost list) is accessed, target_t1_size increases. - When a block in B2 is accessed, target_t1_size decreases. + When a chunk in B1 (ghost list) is accessed, target_t1_size increases. + When a chunk in B2 is accessed, target_t1_size decreases. """ - cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False) + cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False) - # store blocks 1, 2 (fills cache) + # store chunks 1, 2 (fills cache) cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) initial_target = arc_policy.target_t1_size - # store block 3, evicting block 1 (moves to B1 ghost list) + # store chunk 3, evicting chunk 1 (moves to B1 ghost list) cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX) - # block 1 should be in B1 (ghost list) + # chunk 1 should be in B1 (ghost list) assert to_keys([1])[0] in arc_policy.b1 - # touch block 1 (cache miss, but in B1) + # touch chunk 1 (cache miss, but in B1) # this should increase target_t1_size (favor recency) cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX) @@ -671,11 +671,11 @@ def test_t1_t2_eviction_policy(self): """ cpu_manager, arc_policy = self._make_manager(enable_events=False) - # store blocks 1, 2, 3, 4 + # store chunks 1, 2, 3, 4 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) - # promote blocks 3, 4 to T2 by touching them + # promote chunks 3, 4 to T2 by touching them cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX) # now: T1 = {1, 2}, T2 = {3, 4} @@ -686,22 +686,22 @@ def test_t1_t2_eviction_policy(self): # (when |T1| >= target, evict from T1) arc_policy.target_t1_size = 1 - # store block 5, should evict from T1 (block 1, LRU in T1) + # store chunk 5, should evict from T1 (chunk 1, LRU in T1) output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) assert output is not None assert to_keys([1]) == output.evicted_keys cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX) - # block 1 should be in B1 (ghost list) + # chunk 1 should be in B1 (ghost list) assert to_keys([1])[0] in arc_policy.b1 - # block 5 should be in T1 + # chunk 5 should be in T1 assert to_keys([5])[0] in arc_policy.t1 def test_batch_eviction_scans_t1_and_t2_once(self): """ARC batch eviction must preserve order without restarting scans.""" cpu_manager, arc_policy = self._make_manager( - num_blocks=256, enable_events=False + num_chunks=256, enable_events=False ) keys = to_keys(list(range(256))) cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX) @@ -721,13 +721,13 @@ def test_batch_eviction_scans_t1_and_t2_once(self): t2_order = list(arc_policy.t2) expected_t1 = [ key - for key, block in arc_policy.t1.items() - if block.ref_cnt == 0 and key not in protected + for key, chunk in arc_policy.t1.items() + if chunk.ref_cnt == 0 and key not in protected ][:num_t1_evictions] expected_t2 = [ key - for key, block in arc_policy.t2.items() - if block.ref_cnt == 0 and key not in protected + for key, chunk in arc_policy.t2.items() + if chunk.ref_cnt == 0 and key not in protected ][:num_t2_evictions] expected_t1_scans = t1_order.index(expected_t1[-1]) + 1 expected_t2_scans = t2_order.index(expected_t2[-1]) + 1 @@ -746,7 +746,7 @@ def test_batch_eviction_scans_t1_and_t2_once(self): def test_batch_eviction_falls_back_after_t1_iterator_exhausted(self): """An exhausted T1 scan must keep falling back to T2.""" - cpu_manager, arc_policy = self._make_manager(num_blocks=8, enable_events=False) + cpu_manager, arc_policy = self._make_manager(num_chunks=8, enable_events=False) keys = to_keys(list(range(8))) cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX) cpu_manager.complete_store(keys, _EMPTY_REQ_CTX) @@ -764,8 +764,8 @@ def test_batch_eviction_falls_back_after_t1_iterator_exhausted(self): # the target, so each remaining selection must retry T1 then use T2. eligible_t1 = [ key - for key, block in arc_policy.t1.items() - if block.ref_cnt == 0 and key not in protected + for key, chunk in arc_policy.t1.items() + if chunk.ref_cnt == 0 and key not in protected ] assert eligible_t1 == t1_order[:1] assert len(t1_order) - 1 >= int(arc_policy.target_t1_size) @@ -839,7 +839,7 @@ def test_store_falls_back_to_t1_when_t2_cannot_satisfy_eviction( def test_batch_eviction_failure_is_atomic(self): """Finding only some candidates must not partially evict the cache.""" - cpu_manager, arc_policy = self._make_manager(num_blocks=4, enable_events=False) + cpu_manager, arc_policy = self._make_manager(num_chunks=4, enable_events=False) keys = to_keys(list(range(4))) cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX) cpu_manager.complete_store(keys, _EMPTY_REQ_CTX) @@ -858,13 +858,13 @@ def test_ghost_list_bounds(self): Tests that ghost lists (B1, B2) don't grow unbounded. They should be capped at cache_capacity. """ - cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False) + cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False) - # fill cache with blocks 1, 2 + # fill cache with chunks 1, 2 cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) - # store many blocks to fill ghost lists + # store many chunks to fill ghost lists for i in range(3, 20): cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([i]), _EMPTY_REQ_CTX) @@ -880,7 +880,7 @@ def test_touch_ordering(self): """ cpu_manager, arc_policy = self._make_manager() - # store blocks 1, 2, 3, 4 + # store chunks 1, 2, 3, 4 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) @@ -895,13 +895,13 @@ def test_touch_ordering(self): assert len(arc_policy.t1) == 1 assert len(arc_policy.t2) == 3 - # store block 5, should evict from T1 (block 2, only one in T1) + # store chunk 5, should evict from T1 (chunk 2, only one in T1) prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[5], - store_block_ids=[1], # reuses block 2's storage + store_chunk_ids=[1], # reuses chunk 2's storage evicted_keys=[2], ), ) @@ -913,11 +913,11 @@ def test_failed_store(self): """ cpu_manager, arc_policy = self._make_manager() - # store blocks 1, 2, 3, 4 + # store chunks 1, 2, 3, 4 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) - # prepare store block 5 (will evict block 1) + # prepare store chunk 5 (will evict chunk 1) prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) assert prepare_store_output is not None assert len(prepare_store_output.evicted_keys) == 1 @@ -925,13 +925,13 @@ def test_failed_store(self): # complete store with failure cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False) - # block 5 should not be in cache + # chunk 5 should not be in cache assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.MISS - # block 5 should not be in T1 or T2 + # chunk 5 should not be in T1 or T2 assert to_keys([5])[0] not in arc_policy.t1 assert to_keys([5])[0] not in arc_policy.t2 - # evicted block should still be gone (in B1 ghost list) + # evicted chunk should still be gone (in B1 ghost list) evicted_hash = prepare_store_output.evicted_keys[0] assert evicted_hash in arc_policy.b1 @@ -954,7 +954,7 @@ def test_full_scenario(self): assert len(prepare_store_output.evicted_keys) == 1 cpu_manager.complete_store(to_keys([3, 4, 5]), _EMPTY_REQ_CTX) - # promote some blocks to T2 + # promote some chunks to T2 cpu_manager.touch(to_keys([2, 3]), _EMPTY_REQ_CTX) # T1 has {4, 5}, T2 has {2, 3} @@ -966,7 +966,7 @@ def test_full_scenario(self): assert prepare_store_output is not None cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX) - # verify blocks 2, 3 (in T2) are still present + # verify chunks 2, 3 (in T2) are still present assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT @@ -980,7 +980,7 @@ def test_filter_reused_manager(): Tests CPUOffloadingManager reuse filtering (store_threshold=2). """ manager = make_cpu_manager( - num_blocks=4, + num_chunks=4, cache_policy="lru", enable_events=True, store_threshold=2, @@ -1023,7 +1023,7 @@ def test_filter_reused_manager(): def test_filter_reused_manager_oversized_offer_makes_progress(): """An offer larger than the tracker capacity must not churn forever.""" manager = make_cpu_manager( - num_blocks=4, + num_chunks=4, cache_policy="lru", store_threshold=2, max_tracker_size=3, @@ -1042,69 +1042,69 @@ def test_filter_reused_manager_oversized_offer_makes_progress(): assert stored_keys == set(keys) -def test_evictable_cache_block_count(): +def test_evictable_cache_chunk_count(): """ - Verifies _num_evictable_cache_blocks is maintained correctly through the + Verifies _num_evictable_cache_chunks is maintained correctly through the full store/load lifecycle, eviction, failed stores, concurrent loads, reset_cache, and the early-exit fast path in prepare_store. """ - manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") - # Initially no blocks allocated. - assert manager._num_evictable_cache_blocks == 0 + # Initially no chunks allocated. + assert manager._num_evictable_cache_chunks == 0 # Initial cache state [x, x, x, x] - # We get 3 blocks from the cache. + # We get 3 chunks from the cache. manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) # cache state [1', 2', 3', x] <- 1', 2', 3' are actively being used. - assert manager._num_evictable_cache_blocks == 0 + assert manager._num_evictable_cache_chunks == 0 # Completing stores makes them idle. manager.complete_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) - # cache state [1, 2, 3, x] <- 1, 2, 3 blocks are idle. - assert manager._num_evictable_cache_blocks == 3 + # cache state [1, 2, 3, x] <- 1, 2, 3 chunks are idle. + assert manager._num_evictable_cache_chunks == 3 - # prepare_load pins a block: idle count decrements once even if the - # same block is loaded by two concurrent callers. + # prepare_load pins a chunk: idle count decrements once even if the + # same chunk is loaded by two concurrent callers. manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) - # cache state [1', 2, 3, x] <- 2, 3 blocks are idle. - assert manager._num_evictable_cache_blocks == 2 + # cache state [1', 2, 3, x] <- 2, 3 chunks are idle. + assert manager._num_evictable_cache_chunks == 2 manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) # 2nd concurrent load - # cache state [1', 2, 3, x] <- 2, 3 blocks are idle. - assert manager._num_evictable_cache_blocks == 2 # no double-decrement + # cache state [1', 2, 3, x] <- 2, 3 chunks are idle. + assert manager._num_evictable_cache_chunks == 2 # no double-decrement # First complete_load does not restore idle (ref_cnt still 1). manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX) - # cache state [1', 2, 3, x] <- 2, 3 blocks are idle. - assert manager._num_evictable_cache_blocks == 2 - # Second complete_load drops ref_cnt to 0 -> block becomes idle again. + # cache state [1', 2, 3, x] <- 2, 3 chunks are idle. + assert manager._num_evictable_cache_chunks == 2 + # Second complete_load drops ref_cnt to 0 -> chunk becomes idle again. manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX) - # cache state [1, 2, 3, x] <- 1, 2, 3 blocks are idle. - assert manager._num_evictable_cache_blocks == 3 + # cache state [1, 2, 3, x] <- 1, 2, 3 chunks are idle. + assert manager._num_evictable_cache_chunks == 3 # Eviction decrements idle count. - # Cache has 3 stored blocks and 1 free slot. Storing 3 new keys needs 2 eviction. + # Cache has 3 stored chunks and 1 free slot. Storing 3 new keys needs 2 eviction. manager.prepare_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX) - # cache state [1, 4', 5', 6'] <- block 1 is idle - assert manager._num_evictable_cache_blocks == 1 + # cache state [1, 4', 5', 6'] <- chunk 1 is idle + assert manager._num_evictable_cache_chunks == 1 - # Failed store does not increment idle count (block discarded from cache). + # Failed store does not increment idle count (chunk discarded from cache). manager.complete_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX, success=False) - # cache state [1, x, x, x] <- block 1 is idle. Other returned to cache. - assert manager._num_evictable_cache_blocks == 1 + # cache state [1, x, x, x] <- chunk 1 is idle. Other returned to cache. + assert manager._num_evictable_cache_chunks == 1 # reset_cache zeroes the count unconditionally. manager.reset_cache() # cache state [x, x, x, x] - assert manager._num_evictable_cache_blocks == 0 + assert manager._num_evictable_cache_chunks == 0 - # setup 3 blocks with loads so idle count drops to 0. + # setup 3 chunks with loads so idle count drops to 0. manager.prepare_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) manager.complete_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) manager.prepare_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) # cache state [10', 11', 12', x] - assert manager._num_evictable_cache_blocks == 0 + assert manager._num_evictable_cache_chunks == 0 # prepare_store requiring eviction must return None immediately (fast exit). # Spy on policy.evict to confirm the fast path short-circuits before calling it. @@ -1120,24 +1120,24 @@ def spy_evict(*args, **kwargs): # cache state [10', 11', 12', x] <- cannot evict anything assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is None assert not evict_called, ( - "_num_evictable_cache_blocks==0 should short-circuit before evict()" + "_num_evictable_cache_chunks==0 should short-circuit before evict()" ) # After releasing the loads, eviction becomes possible again. manager.complete_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) # cache state [10, 11, 12, x] <- 10, 11, 12 are idle - assert manager._num_evictable_cache_blocks == 3 + assert manager._num_evictable_cache_chunks == 3 assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is not None # cache state [10, 11, 14', 15'] <- 10, 11 are idle - assert manager._num_evictable_cache_blocks == 2 + assert manager._num_evictable_cache_chunks == 2 manager.complete_store(to_keys([14, 15]), _EMPTY_REQ_CTX) - # cache state [10, 11, 14, 15] <- all blocks idle - assert manager._num_evictable_cache_blocks == 4 + # cache state [10, 11, 14, 15] <- all chunks idle + assert manager._num_evictable_cache_chunks == 4 def test_touch_forwards_req_context_to_policy(monkeypatch): """Regression: CPUOffloadingManager.touch forwards ReqContext to policy.""" - manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") received = [] def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None: diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index f5f5c9720662..c7ddac4feffa 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -38,7 +38,7 @@ def _set_spawn_method(monkeypatch): def _make_region( engine_id: str, - num_blocks: int = 4, + num_chunks: int = 4, cpu_page_size: int = PAGE_SIZE, num_workers: int = 1, rank: int = 0, @@ -47,9 +47,9 @@ def _make_region( assert cpu_page_size % PAGE_SIZE == 0 return SharedOffloadRegion( engine_id=engine_id, - num_blocks=num_blocks, + num_chunks=num_chunks, rank=rank, - kv_bytes_per_block=num_workers * cpu_page_size, + kv_bytes_per_chunk=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, barrier=barrier, ) @@ -99,16 +99,16 @@ def _region(engine_id: str, **kwargs): def _multi_region( engine_id: str, num_workers: int, - num_blocks: int = 4, + num_chunks: int = 4, cpu_page_size: int = PAGE_SIZE, ): """Context manager: create one SharedOffloadRegion per rank, clean up on exit.""" regions = [ SharedOffloadRegion( engine_id=engine_id, - num_blocks=num_blocks, + num_chunks=num_chunks, rank=rank, - kv_bytes_per_block=num_workers * cpu_page_size, + kv_bytes_per_chunk=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) for rank in range(num_workers) @@ -124,7 +124,7 @@ def _multi_region( def _race_construct( engine_id: str, num_workers: int, - num_blocks: int = 4, + num_chunks: int = 4, cpu_page_size: int = PAGE_SIZE, ) -> tuple[list[SharedOffloadRegion], list[Exception]]: """Spawn num_workers threads that all race to construct SharedOffloadRegion.""" @@ -137,9 +137,9 @@ def worker(rank: int) -> None: try: regions[rank] = SharedOffloadRegion( engine_id=engine_id, - num_blocks=num_blocks, + num_chunks=num_chunks, rank=rank, - kv_bytes_per_block=num_workers * cpu_page_size, + kv_bytes_per_chunk=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) except Exception as e: @@ -156,7 +156,7 @@ def worker(rank: int) -> None: def _mp_race_construct_and_write( engine_id: str, - num_blocks: int, + num_chunks: int, rank: int, num_workers: int, cpu_page_size: int, @@ -170,9 +170,9 @@ def _mp_race_construct_and_write( try: region = SharedOffloadRegion( engine_id=engine_id, - num_blocks=num_blocks, + num_chunks=num_chunks, rank=rank, - kv_bytes_per_block=num_workers * cpu_page_size, + kv_bytes_per_chunk=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) t = region.create_next_worker_view(cpu_page_size) @@ -198,9 +198,9 @@ def _mp_barrier_construct_and_hold( try: region = SharedOffloadRegion( engine_id=engine_id, - num_blocks=2, + num_chunks=2, rank=rank, - kv_bytes_per_block=num_workers * PAGE_SIZE, + kv_bytes_per_chunk=num_workers * PAGE_SIZE, cpu_page_size=PAGE_SIZE, barrier=lambda: barrier.wait(30), ) @@ -233,9 +233,9 @@ def iid(): def test_create_next_worker_view_shape_and_stride(iid): - """Returned tensor must have shape (num_blocks, tensor_page_size) and + """Returned tensor must have shape (num_chunks, tensor_page_size) and stride (row_stride, 1) where row_stride = cpu_page_size * num_workers.""" - with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r: + with _region(iid, num_chunks=4, cpu_page_size=2 * PAGE_SIZE) as r: t = r.create_next_worker_view(PAGE_SIZE) assert t.shape == (4, PAGE_SIZE) # num_workers=1 → row_stride = cpu_page_size @@ -253,7 +253,7 @@ def test_create_next_worker_view_storage_offset_rank0(iid): def test_create_next_worker_view_storage_offset_rank1(iid): """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap.""" - with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1): + with _multi_region(iid, num_workers=2, num_chunks=4) as (r0, r1): t1 = r1.create_next_worker_view(PAGE_SIZE) assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE del t1 @@ -261,7 +261,7 @@ def test_create_next_worker_view_storage_offset_rank1(iid): def test_create_next_worker_view_row_stride_with_multiple_workers(iid): """With num_workers=4, row_stride must be 4 * cpu_page_size.""" - with _region(iid, num_blocks=2, num_workers=4) as r: + with _region(iid, num_chunks=2, num_workers=4) as r: t = r.create_next_worker_view(PAGE_SIZE) assert t.stride(0) == 4 * PAGE_SIZE del t @@ -332,12 +332,12 @@ def test_create_next_worker_view_overflow_does_not_mutate_cursor(iid): def test_create_next_worker_view_write_visible_in_raw_mmap(iid): """Writes into a create_next_worker_view view must appear at the correct raw mmap offset""" - with _region(iid, num_blocks=4) as r: + with _region(iid, num_chunks=4) as r: t = r.create_next_worker_view(PAGE_SIZE) - t[2, :] = 42 # write to block row 2 + t[2, :] = 42 # write to chunk row 2 raw = memoryview(r.mmap_obj) - # num_workers=1 → row_stride = PAGE_SIZE; block 2 starts at byte 2*PAGE_SIZE + # num_workers=1 → row_stride = PAGE_SIZE; chunk 2 starts at byte 2*PAGE_SIZE chunk = bytes(raw[2 * PAGE_SIZE : 3 * PAGE_SIZE]) assert all(b == 42 for b in chunk) del raw, t @@ -345,7 +345,7 @@ def test_create_next_worker_view_write_visible_in_raw_mmap(iid): def test_create_next_worker_view_multi_tensor_layout(iid): """Two tensors from the same worker land at consecutive byte offsets per row.""" - with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r: + with _region(iid, num_chunks=2, cpu_page_size=2 * PAGE_SIZE) as r: ta = r.create_next_worker_view(PAGE_SIZE) tb = r.create_next_worker_view(PAGE_SIZE) @@ -353,8 +353,8 @@ def test_create_next_worker_view_multi_tensor_layout(iid): tb[:, :] = 2 raw = memoryview(r.mmap_obj) - for blk in range(2): - row_offset = blk * 2 * PAGE_SIZE # num_workers=1 + for chunk in range(2): + row_offset = chunk * 2 * PAGE_SIZE # num_workers=1 assert all(b == 1 for b in raw[row_offset : row_offset + PAGE_SIZE]) assert all( b == 2 for b in raw[row_offset + PAGE_SIZE : row_offset + 2 * PAGE_SIZE] @@ -366,7 +366,7 @@ def test_create_next_worker_view_multiprocess_slots(iid): """Each worker process calls create_next_worker_view and writes distinct data; the parent verifies each slot lands at the correct interleaved offset.""" num_workers = 2 - num_blocks = 4 + num_chunks = 4 ctx = get_mp_context() done_queue = ctx.Queue() @@ -375,9 +375,9 @@ def test_create_next_worker_view_multiprocess_slots(iid): # Parent is rank 0 (creator); child is rank 1 (joiner). region = SharedOffloadRegion( engine_id=iid, - num_blocks=num_blocks, + num_chunks=num_chunks, rank=0, - kv_bytes_per_block=num_workers * PAGE_SIZE, + kv_bytes_per_chunk=num_workers * PAGE_SIZE, cpu_page_size=PAGE_SIZE, ) try: @@ -385,7 +385,7 @@ def test_create_next_worker_view_multiprocess_slots(iid): target=_mp_race_construct_and_write, args=( iid, - num_blocks, + num_chunks, 1, num_workers, PAGE_SIZE, @@ -403,12 +403,12 @@ def test_create_next_worker_view_multiprocess_slots(iid): assert result["error"] is None, result["error"] raw = memoryview(region.mmap_obj) - for blk in range(num_blocks): - row_start = blk * num_workers * PAGE_SIZE + for chunk in range(num_chunks): + row_start = chunk * num_workers * PAGE_SIZE w0 = bytes(raw[row_start : row_start + PAGE_SIZE]) w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE]) - assert all(b == 11 for b in w0), f"block {blk}: rank0 slot wrong" - assert all(b == 22 for b in w1), f"block {blk}: rank1 slot wrong" + assert all(b == 11 for b in w0), f"chunk {chunk}: rank0 slot wrong" + assert all(b == 22 for b in w1), f"chunk {chunk}: rank1 slot wrong" del raw, t0 # release before finally triggers cleanup cleanup_queue.put(True) @@ -422,8 +422,8 @@ def test_create_next_worker_view_multiprocess_slots(iid): def test_create_next_worker_view_worker_isolation(iid): """Writes by worker 0 must not affect worker 1's slot and vice versa.""" num_workers = 2 - num_blocks = 4 - with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions: + num_chunks = 4 + with _multi_region(iid, num_workers=num_workers, num_chunks=num_chunks) as regions: t0 = regions[0].create_next_worker_view(PAGE_SIZE) t1 = regions[1].create_next_worker_view(PAGE_SIZE) @@ -431,12 +431,12 @@ def test_create_next_worker_view_worker_isolation(iid): t1[:, :] = 22 raw = memoryview(regions[0].mmap_obj) - for blk in range(num_blocks): - row_start = blk * num_workers * PAGE_SIZE + for chunk in range(num_chunks): + row_start = chunk * num_workers * PAGE_SIZE w0 = bytes(raw[row_start : row_start + PAGE_SIZE]) w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE]) - assert all(b == 11 for b in w0), f"block {blk}: worker0 slot corrupted" - assert all(b == 22 for b in w1), f"block {blk}: worker1 slot corrupted" + assert all(b == 11 for b in w0), f"chunk {chunk}: worker0 slot corrupted" + assert all(b == 22 for b in w1), f"chunk {chunk}: worker1 slot corrupted" del raw, t0, t1 # release before finally triggers cleanup @@ -466,7 +466,7 @@ def test_file_exists_after_construction(iid): def test_file_has_correct_size(iid): """The mmap file size on disk must equal total_size_bytes.""" - with _region(iid, num_blocks=4) as r: + with _region(iid, num_chunks=4) as r: assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE @@ -491,11 +491,11 @@ def _spy_fallback(mm, off, ln): monkeypatch.setattr(sor, "_madvise_populate_write", _spy_madvise) monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback) - num_blocks = 3 + num_chunks = 3 num_workers = 2 - with _region(iid, num_blocks=num_blocks, num_workers=num_workers, rank=1): - # 1 probe (PAGESIZE) + N populate calls (one per block per worker column). - expected_populate = num_blocks # ranked path: one call per block + with _region(iid, num_chunks=num_chunks, num_workers=num_workers, rank=1): + # 1 probe (PAGESIZE) + N populate calls (one per chunk per worker column). + expected_populate = num_chunks # ranked path: one call per chunk mmap_id = madvise_calls[0][2] assert madvise_calls == [ (0, mmap.PAGESIZE, mmap_id), @@ -510,7 +510,7 @@ def _spy_fallback(mm, off, ln): def test_madvise_einval_selects_fallback_for_ranked_region(iid, monkeypatch): - """An EINVAL probe must select fallback for every ranked block.""" + """An EINVAL probe must select fallback for every ranked chunk.""" from vllm.v1.kv_offload.cpu import shared_offload_region as sor fallback_calls: list[tuple[int, int]] = [] @@ -526,7 +526,7 @@ def _spy_fallback(mm, off, ln): monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval) monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback) - with _region(iid, num_blocks=3, num_workers=2, rank=1): + with _region(iid, num_chunks=3, num_workers=2, rank=1): assert fallback_calls == [ (mmap.PAGESIZE, mmap.PAGESIZE), (3 * mmap.PAGESIZE, mmap.PAGESIZE), @@ -551,7 +551,7 @@ def _spy_fallback(mm, off, ln): monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval) monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback) - with _region(iid, num_blocks=3, num_workers=2, rank=None): + with _region(iid, num_chunks=3, num_workers=2, rank=None): assert fallback_calls == [(0, 6 * mmap.PAGESIZE)] @@ -651,7 +651,7 @@ def test_multiprocess_race_construct_and_write(iid): """N processes race to construct the same SharedOffloadRegion, each writes fill_value = rank+1 into their slot; parent verifies interleaved layout.""" num_workers = 4 - num_blocks = 3 + num_chunks = 3 ctx = get_mp_context() done_queue = ctx.Queue() @@ -662,7 +662,7 @@ def test_multiprocess_race_construct_and_write(iid): target=_mp_race_construct_and_write, args=( iid, - num_blocks, + num_chunks, rank, num_workers, PAGE_SIZE, @@ -689,13 +689,13 @@ def test_multiprocess_race_construct_and_write(iid): with open(mmap_path, "rb") as f: raw = f.read() - for blk in range(num_blocks): + for chunk in range(num_chunks): for w in range(num_workers): - slot_start = (blk * num_workers + w) * PAGE_SIZE + slot_start = (chunk * num_workers + w) * PAGE_SIZE slot = raw[slot_start : slot_start + PAGE_SIZE] expected = w + 1 # fill_value = rank + 1 assert all(b == expected for b in slot), ( - f"block {blk}, worker {w}: expected {expected} but got wrong bytes" + f"chunk {chunk}, worker {w}: expected {expected} but got wrong bytes" ) # Unblock all workers to clean up. @@ -839,9 +839,9 @@ def test_insufficient_space_raises_clear_error(monkeypatch): with pytest.raises(RuntimeError, match="Insufficient space"): SharedOffloadRegion( engine_id=engine_id, - num_blocks=4, + num_chunks=4, rank=0, - kv_bytes_per_block=PAGE_SIZE, + kv_bytes_per_chunk=PAGE_SIZE, cpu_page_size=PAGE_SIZE, ) @@ -870,9 +870,9 @@ def test_ftruncate_failure_cleans_up_creator(monkeypatch): with pytest.raises(OSError, match="ftruncate failed"): SharedOffloadRegion( engine_id=engine_id, - num_blocks=4, + num_chunks=4, rank=0, - kv_bytes_per_block=PAGE_SIZE, + kv_bytes_per_chunk=PAGE_SIZE, cpu_page_size=PAGE_SIZE, ) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 6d1c103fdf8c..251cef325143 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -136,7 +136,7 @@ def test_get_spec_cls_defaults_to_cpu(): def test_create_cpu_offloading_spec(): spec = _create_spec() assert isinstance(spec, CPUOffloadingSpec) - assert spec.num_blocks > 0 + assert spec.num_chunks > 0 def test_cpu_spec_sizes_normalized_worker_layout(): @@ -156,7 +156,7 @@ def test_cpu_spec_sizes_normalized_worker_layout(): assert isinstance(spec, CPUOffloadingSpec) assert spec.cpu_page_size_per_worker == 32 assert spec.kv_bytes_per_chunk == alignment - assert spec.num_blocks == 3 + assert spec.num_chunks == 3 def test_cpu_spec_zero_worker_bytes_produces_empty_cache(): @@ -165,7 +165,7 @@ def test_cpu_spec_zero_worker_bytes_produces_empty_cache(): assert isinstance(spec, CPUOffloadingSpec) assert spec.cpu_page_size_per_worker == 0 assert spec.kv_bytes_per_chunk == 0 - assert spec.num_blocks == 0 + assert spec.num_chunks == 0 def test_tiering_spec_aligns_row_size(): @@ -183,7 +183,7 @@ def test_tiering_spec_aligns_row_size(): assert isinstance(spec, TieringOffloadingSpec) assert spec.cpu_page_size_per_worker == 32 assert spec.kv_bytes_per_chunk == alignment - assert spec.num_blocks == 3 + assert spec.num_chunks == 3 @pytest.mark.parametrize("world_size", [2, 4, 8]) @@ -201,7 +201,7 @@ def test_tiering_spec_replicated_sizing_removes_world_factor(world_size: int): assert spec.replicated_layout is True assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block - assert spec.num_blocks == 8 + assert spec.num_chunks == 8 def test_tiering_spec_create_worker_uses_single_slot_for_replicated_layout(monkeypatch): @@ -239,7 +239,7 @@ def fake_worker_ctor(**kwargs): spec.create_worker(kv_caches) assert region_calls[0]["rank"] == 0 - assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block + assert region_calls[0]["kv_bytes_per_chunk"] == worker_kv_bytes_per_block assert worker_calls[0]["kv_caches"] is kv_caches assert worker_calls[0]["mmap_region"] is region @@ -292,7 +292,7 @@ def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: in assert spec.replicated_layout is True assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block - assert spec.num_blocks == 8 + assert spec.num_chunks == 8 @pytest.mark.parametrize("world_size", [2, 4, 8]) @@ -320,7 +320,7 @@ def test_cpu_spec_replicated_disabled_without_shared_region( assert spec.replicated_layout is False assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block * world_size - assert spec.num_blocks == 2 + assert spec.num_chunks == 2 @pytest.mark.parametrize("config_replicated", [True, False]) @@ -381,7 +381,7 @@ def fake_worker_ctor(**kwargs): spec.create_worker(kv_caches) assert region_calls[0]["engine_id"] == "test-engine" - assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block * 4 + assert region_calls[0]["kv_bytes_per_chunk"] == worker_kv_bytes_per_block * 4 assert worker_calls[0]["kv_caches"] is kv_caches assert worker_calls[0]["mmap_region"] is region @@ -419,11 +419,11 @@ def fake_worker_ctor(**kwargs): def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch): import vllm.v1.kv_offload.cpu.spec as cpu_spec_module - # worker_kv_bytes_per_block=0 yields num_blocks=0; a zero-byte region cannot + # worker_kv_bytes_per_block=0 yields num_chunks=0; a zero-byte region cannot # be mmap'd, so even on CUDA-alike this must fall back to the tensor path. spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4) assert isinstance(spec, CPUOffloadingSpec) - assert spec.num_blocks == 0 + assert spec.num_chunks == 0 region_calls: list[dict[str, Any]] = [] worker_calls: list[dict[str, Any]] = [] diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index e93eb77de3e4..dc0e384d7c2f 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -91,17 +91,17 @@ def _req_context(kv_params: dict | None = None) -> ReqContext: def _job_metadata( job_id: int, keys: list[bytes] | None = None, - block_ids: list[int] | None = None, + chunk_ids: list[int] | None = None, kv_params: dict | None = None, ) -> TransferJob: if keys is None: keys = [b"key1"] - if block_ids is None: - block_ids = list(range(len(keys))) + if chunk_ids is None: + chunk_ids = list(range(len(keys))) return TransferJob( job_id=job_id, keys=keys, - block_ids=np.array(block_ids), + chunk_ids=np.array(chunk_ids), is_promotion=False, req_context=_req_context(kv_params), ) @@ -254,9 +254,9 @@ def test_lookup_returns_miss_without_prefill_key(self): "kv_params,expected", [ (_remote_decoder_kv_params(), OffloadPolicy.REQUEST_LEVEL), - (None, OffloadPolicy.BLOCK_LEVEL), - (_remote_prefiller_kv_params(), OffloadPolicy.BLOCK_LEVEL), - ({"remote_decoder": {}}, OffloadPolicy.BLOCK_LEVEL), + (None, OffloadPolicy.CHUNK_LEVEL), + (_remote_prefiller_kv_params(), OffloadPolicy.CHUNK_LEVEL), + ({"remote_decoder": {}}, OffloadPolicy.CHUNK_LEVEL), ], ids=["producer", "plain", "consumer", "producer_no_id"], ) @@ -366,7 +366,7 @@ def test_no_binding_yet_parks_in_unbound_stores(self): job = _job_metadata( job_id=1, keys=[b"k1", b"k2"], - block_ids=[3, 4], + chunk_ids=[3, 4], kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -397,7 +397,7 @@ def test_routes_to_bound_session(self): job = _job_metadata( job_id=7, keys=[b"k1", b"k2"], - block_ids=[3, 4], + chunk_ids=[3, 4], kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -440,7 +440,7 @@ def test_empty_keys_succeeds_immediately(self): """Empty key list succeeds immediately.""" mgr = _make_manager() job = _job_metadata( - job_id=1, keys=[], block_ids=[], kv_params=_remote_prefiller_kv_params() + job_id=1, keys=[], chunk_ids=[], kv_params=_remote_prefiller_kv_params() ) mgr.submit_load(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] @@ -463,7 +463,7 @@ def test_happy_path_with_active_session(self): job = _job_metadata( job_id=42, keys=[b"k1", b"k2"], - block_ids=[5, 6], + chunk_ids=[5, 6], kv_params=_remote_prefiller_kv_params(kv_request_id="req-42"), ) mgr.submit_load(job) @@ -1155,7 +1155,7 @@ def test_both_loads_succeed(self): _job_metadata( job_id=100, keys=[b"a-block"], - block_ids=[0], + chunk_ids=[0], kv_params=a_prefiller_params, ) ) @@ -1163,7 +1163,7 @@ def test_both_loads_succeed(self): _job_metadata( job_id=200, keys=[b"b-block"], - block_ids=[0], + chunk_ids=[0], kv_params=b_prefiller_params, ) ) @@ -1173,7 +1173,7 @@ def test_both_loads_succeed(self): _job_metadata( job_id=101, keys=[b"b-block"], - block_ids=[0], + chunk_ids=[0], kv_params=a_decoder_params, ) ) @@ -1181,7 +1181,7 @@ def test_both_loads_succeed(self): _job_metadata( job_id=201, keys=[b"a-block"], - block_ids=[0], + chunk_ids=[0], kv_params=b_decoder_params, ) ) @@ -1573,7 +1573,7 @@ def test_dead_connection_with_pending_work_surfaces_failures(self): _job_metadata( job_id=900, keys=[b"a-block"], - block_ids=[0], + chunk_ids=[0], kv_params=a_prefiller_params, ) ) @@ -1581,7 +1581,7 @@ def test_dead_connection_with_pending_work_surfaces_failures(self): _job_metadata( job_id=901, keys=[b"b-block"], - block_ids=[0], + chunk_ids=[0], kv_params=a_decoder_params, ) ) diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 357f2b203bf1..2e4415fbb75d 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -259,13 +259,13 @@ def create_store_job( ) -> TransferJob: keys_list = list(keys) self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id)) - block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32) + chunk_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32) job_id = self._next_job_id self._next_job_id += 1 return TransferJob( job_id=job_id, keys=keys_list, - block_ids=block_ids, + chunk_ids=chunk_ids, is_promotion=False, req_context=ctx, ) diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index efeb59274fdb..4b495e1b315c 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -105,15 +105,15 @@ def key(n: int) -> OffloadKey: def make_job( job_id: int, keys: list[OffloadKey], - block_ids: list[int] | None = None, + chunk_ids: list[int] | None = None, is_promotion: bool = False, ) -> TransferJob: - if block_ids is None: - block_ids = list(range(len(keys))) + if chunk_ids is None: + chunk_ids = list(range(len(keys))) return TransferJob( job_id=job_id, keys=keys, - block_ids=np.array(block_ids, dtype=np.int64), + chunk_ids=np.array(chunk_ids, dtype=np.int64), is_promotion=is_promotion, req_context=_CTX, ) @@ -382,11 +382,11 @@ def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size): tensor[:] = _page_aligned_rand_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) keys = [key(i) for i in range(batch_size)] - store_block_ids = list(range(batch_size)) - load_block_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS)) + store_chunk_ids = list(range(batch_size)) + load_chunk_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS)) expected = tensor[:batch_size].clone() - tier.submit_store(make_job(1, keys, store_block_ids)) + tier.submit_store(make_job(1, keys, store_chunk_ids)) store_results = drain(tier) assert len(store_results) == 1 assert store_results[0].success @@ -396,15 +396,15 @@ def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size): tensor[:] = 0.0 # Load into a range disjoint by index from the store ids, to also - # exercise loading a block into a different id than it was stored from. - tier.submit_load(make_job(2, keys, load_block_ids, is_promotion=True)) + # exercise loading a chunk into a different id than it was stored from. + tier.submit_load(make_job(2, keys, load_chunk_ids, is_promotion=True)) load_results = drain(tier) assert len(load_results) == 1 assert load_results[0].success - for i, bid in enumerate(load_block_ids): - assert torch.allclose(tensor[bid], expected[i]), ( - f"Block {bid} data mismatch after store+load" + for i, cid in enumerate(load_chunk_ids): + assert torch.allclose(tensor[cid], expected[i]), ( + f"Chunk {cid} data mismatch after store+load" ) @@ -895,7 +895,7 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): view = memoryview(tensor.numpy()) mock_region = MagicMock() mock_region.create_kv_memoryview.return_value = view - primary = CPUPrimaryTierOffloadingManager(num_blocks=4, mmap_region=mock_region) + primary = CPUPrimaryTierOffloadingManager(num_chunks=4, mmap_region=mock_region) tier = FileSystemTierManager( offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), primary_kv_view=primary.get_kv_memoryview(), diff --git a/tests/v1/kv_offload/tiering/test_metrics.py b/tests/v1/kv_offload/tiering/test_metrics.py index 42b973943cae..1a155892f8d3 100644 --- a/tests/v1/kv_offload/tiering/test_metrics.py +++ b/tests/v1/kv_offload/tiering/test_metrics.py @@ -25,8 +25,8 @@ def to_keys(int_ids: Iterable[int]): def test_tiering_metrics_tracker_records_lookup_metrics(): tracker = TieringMetricsTracker( tier_types=["fs", "p2p"], - num_primary_blocks=5, - primary_block_size=16, + num_primary_chunks=5, + primary_chunk_size=16, ) tracker.on_new_request(_CTX) @@ -45,18 +45,18 @@ def test_tiering_metrics_tracker_records_lookup_metrics(): stats = tracker.take_stats() assert stats is not None values = stats.data["data"] - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("2:p2p",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_HITS][("2:p2p",)] == 1 - assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS] + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("0:primary",)] == 1 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("1:fs",)] == 1 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("2:p2p",)] == 1 + assert values[TieringOffloadingMetrics.CHUNK_HITS][("2:p2p",)] == 1 + assert ("1:fs",) not in values[TieringOffloadingMetrics.CHUNK_HITS] def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation(): tracker = TieringMetricsTracker( tier_types=["fs"], - num_primary_blocks=5, - primary_block_size=16, + num_primary_chunks=5, + primary_chunk_size=16, ) tracker.on_new_request(_CTX) key = to_keys([1])[0] @@ -73,9 +73,9 @@ def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation(): stats = tracker.take_stats() assert stats is not None values = stats.data["data"] - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:fs",)] == 1 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("0:primary",)] == 1 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("1:fs",)] == 1 + assert values[TieringOffloadingMetrics.CHUNK_HITS][("1:fs",)] == 1 tracker.on_request_allocated(_CTX) tracker.on_lookup( @@ -90,15 +90,15 @@ def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation(): stats = tracker.take_stats() assert stats is not None values = stats.data["data"] - assert TieringOffloadingMetrics.BLOCK_QUERIES not in values - assert TieringOffloadingMetrics.BLOCK_HITS not in values + assert TieringOffloadingMetrics.CHUNK_QUERIES not in values + assert TieringOffloadingMetrics.CHUNK_HITS not in values def test_tiering_metrics_tracker_records_finished_job_metrics(): tracker = TieringMetricsTracker( tier_types=["fs"], - num_primary_blocks=5, - primary_block_size=16, + num_primary_chunks=5, + primary_chunk_size=16, ) cascade_key, promotion_key, failed_cascade_key, failed_promotion_key = to_keys( range(4) @@ -149,8 +149,8 @@ def test_tiering_metrics_tracker_records_finished_job_metrics(): def test_tiering_metrics_tracker_records_partial_promotion_success_bytes(): tracker = TieringMetricsTracker( tier_types=["fs"], - num_primary_blocks=5, - primary_block_size=16, + num_primary_chunks=5, + primary_chunk_size=16, ) keys = to_keys(range(3)) job = JobMetadata( @@ -182,8 +182,8 @@ def test_tiering_metrics_tracker_records_partial_promotion_success_bytes(): def test_tiering_metrics_tracker_reports_active_job_and_primary_usage_gauges(): tracker = TieringMetricsTracker( tier_types=["fs", "p2p"], - num_primary_blocks=6, - primary_block_size=16, + num_primary_chunks=6, + primary_chunk_size=16, ) fs_job = JobMetadata( TransferJob(0, to_keys([0, 1]), np.array([0, 1]), False, _CTX), @@ -214,8 +214,8 @@ def test_tiering_metrics_tracker_reports_active_job_and_primary_usage_gauges(): def test_tiering_metrics_tracker_records_promotion_allocation_failures(): tracker = TieringMetricsTracker( tier_types=["fs"], - num_primary_blocks=1, - primary_block_size=16, + num_primary_chunks=1, + primary_chunk_size=16, ) tracker.on_promotion_allocation_failure() diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 1d631e778247..13840bac3599 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -106,14 +106,14 @@ def key(n: int) -> OffloadKey: def make_job( job_id: int, keys: list[OffloadKey], - block_ids: list[int] | None = None, + chunk_ids: list[int] | None = None, ) -> TransferJob: - if block_ids is None: - block_ids = list(range(len(keys))) + if chunk_ids is None: + chunk_ids = list(range(len(keys))) return TransferJob( job_id=job_id, keys=keys, - block_ids=np.array(block_ids, dtype=np.int32), + chunk_ids=np.array(chunk_ids, dtype=np.int32), is_promotion=False, req_context=_CTX, ) @@ -541,7 +541,7 @@ def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypat mmap_region = MagicMock() mmap_region.create_kv_memoryview.return_value = primary_kv_view primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=num_blocks, mmap_region=mmap_region + num_chunks=num_blocks, mmap_region=mmap_region ) obj_tier, agent = _make_tier( num_blocks=num_blocks, primary_kv_view=primary_kv_view diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 88d7c466cac5..f737685621b4 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -5,15 +5,16 @@ These tests verify: 1. Basic tiered offloading operations (store, load, lookup) -2. Cascade behavior (blocks stored to all secondary tiers) -3. Promotion behavior (blocks loaded from secondary to primary to GPU) -4. ref_cnt management (blocks protected during async transfers) +2. Cascade behavior (chunks stored to all secondary tiers) +3. Promotion behavior (chunks loaded from secondary to primary to GPU) +4. ref_cnt management (chunks protected during async transfers) 5. Eviction coordination between tiers """ from collections.abc import Iterable from unittest.mock import MagicMock +import numpy as np import pytest import torch @@ -57,10 +58,10 @@ _MOCK_OFFLOADING_SPEC = MagicMock() -def _mock_mmap_region(num_blocks: int, row_bytes: int = 16): +def _mock_mmap_region(num_chunks: int, row_bytes: int = 16): """Create a mock SharedOffloadRegion for testing.""" mock = MagicMock() - view = memoryview(torch.zeros((num_blocks, row_bytes), dtype=torch.int8).numpy()) + view = memoryview(torch.zeros((num_chunks, row_bytes), dtype=torch.int8).numpy()) mock.create_kv_memoryview.return_value = view return mock @@ -160,7 +161,7 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): metadata = metrics[MetricsSecondaryTierManager.MY_TIER_METRIC] assert metadata.documentation == "Number of bytes served by the test tier." assert metadata.labelnames == ("tier",) - assert metrics[TieringOffloadingMetrics.BLOCK_QUERIES].labelnames == ("tier",) + assert metrics[TieringOffloadingMetrics.CHUNK_QUERIES].labelnames == ("tier",) assert metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY].labelnames == ("tier",) assert metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY].labelnames == ("tier",) assert metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES].labelnames == ( @@ -183,7 +184,7 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): def test_tiering_manager_aggregates_secondary_stats(): mock_region = _mock_mmap_region(5) primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=5, mmap_region=mock_region + num_chunks=5, mmap_region=mock_region ) secondary_tier = MetricsSecondaryTierManager( offloading_spec=_MOCK_OFFLOADING_SPEC, @@ -232,19 +233,19 @@ def test_basic_store_and_lookup(self): ) # Initially empty - blocks = to_keys(range(3)) - assert tier.lookup(blocks[0], _CTX) is LookupResult.MISS + chunks = to_keys(range(3)) + assert tier.lookup(chunks[0], _CTX) is LookupResult.MISS - # Store blocks (simulate with direct insertion for testing) - tier.blocks[blocks[0]] = True - tier.blocks[blocks[1]] = True + # Store chunks (simulate with direct insertion for testing) + tier.chunks[chunks[0]] = True + tier.chunks[chunks[1]] = True - # Lookup should find first two blocks - assert tier.lookup(blocks[0], _CTX) is LookupResult.HIT - assert tier.lookup(blocks[1], _CTX) is LookupResult.HIT + # Lookup should find first two chunks + assert tier.lookup(chunks[0], _CTX) is LookupResult.HIT + assert tier.lookup(chunks[1], _CTX) is LookupResult.HIT - # Third block not present - assert tier.lookup(blocks[2], _CTX) is LookupResult.MISS + # Third chunk not present + assert tier.lookup(chunks[2], _CTX) is LookupResult.MISS # What a request-level cascade does with a key already present in the primary @@ -260,7 +261,7 @@ def test_basic_store_and_lookup(self): # The same test, before it lands. LookupResult.HIT_PENDING: "defer", # test_cascade_drops_deferred_keys_whose_primary_write_failed, where the - # failed write frees the block. + # failed write frees the chunk. LookupResult.MISS: "drop", # test_cascade_rejects_retry_from_primary. The primary tier resolves every # key it holds, so RETRY cannot reach the cascade, and parking on it would @@ -284,7 +285,7 @@ def manager_setup(self): # Create primary tier (CPU-based) mock_region = _mock_mmap_region(5) self.primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=5, mmap_region=mock_region + num_chunks=5, mmap_region=mock_region ) mock_view = mock_region.create_kv_memoryview() @@ -331,7 +332,7 @@ def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup): TransferJob( job_id=job_id, keys=to_keys([1, 2]), - block_ids=[0, 1], + chunk_ids=np.array([0, 1], dtype=np.int64), is_promotion=True, req_context=_CTX, ), @@ -358,7 +359,7 @@ def test_successful_promotion_finalizes_primary_with_success(self, manager_setup TransferJob( job_id=job_id, keys=to_keys([1]), - block_ids=[0], + chunk_ids=np.array([0], dtype=np.int64), is_promotion=True, req_context=_CTX, ), @@ -392,23 +393,23 @@ def test_take_events_aggregates_tier_owned_events(self, manager_setup): def test_basic_store_to_primary(self, manager_setup): """Test basic store operation to primary tier.""" - blocks = to_keys(range(3)) + chunks = to_keys(range(3)) # Prepare store self._start_request() - result = self.manager.prepare_store(blocks, _CTX) + result = self.manager.prepare_store(chunks, _CTX) assert result is not None assert len(result.keys_to_store) == 3 # Complete store - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.complete_store(chunks, _CTX, success=True) - # Blocks should be in primary tier - assert count_hits(self.primary_tier, blocks) == 3 + # Chunks should be in primary tier + assert count_hits(self.primary_tier, chunks) == 3 def test_cascade_to_all_secondary_tiers(self, manager_setup): - """Test that blocks are cascaded to ALL secondary tiers.""" - blocks = to_keys(range(3)) + """Test that chunks are cascaded to ALL secondary tiers.""" + chunks = to_keys(range(3)) self.secondary_tier1.submit_store = MagicMock( wraps=self.secondary_tier1.submit_store @@ -419,44 +420,44 @@ def test_cascade_to_all_secondary_tiers(self, manager_setup): # Store to primary self._start_request() - result = self.manager.prepare_store(blocks, _CTX) + result = self.manager.prepare_store(chunks, _CTX) assert result is not None # Complete store (triggers cascade via submit_store on each tier) - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.complete_store(chunks, _CTX, success=True) # submit_store was called once per secondary tier self.secondary_tier1.submit_store.assert_called_once() self.secondary_tier2.submit_store.assert_called_once() - # Blocks should be in both secondary tiers - assert self.secondary_tier1.get_num_blocks() == 3 - assert self.secondary_tier2.get_num_blocks() == 3 + # Chunks should be in both secondary tiers + assert self.secondary_tier1.get_num_chunks() == 3 + assert self.secondary_tier2.get_num_chunks() == 3 - # Verify blocks are present + # Verify chunks are present assert all( - self.secondary_tier1.lookup(b, _CTX) is LookupResult.HIT for b in blocks + self.secondary_tier1.lookup(c, _CTX) is LookupResult.HIT for c in chunks ) assert all( - self.secondary_tier2.lookup(b, _CTX) is LookupResult.HIT for b in blocks + self.secondary_tier2.lookup(c, _CTX) is LookupResult.HIT for c in chunks ) def test_ref_cnt_protection_during_cascade(self, manager_setup): - """Test that ref_cnt protects blocks during cascade.""" - blocks = to_keys(range(3)) + """Test that ref_cnt protects chunks during cascade.""" + chunks = to_keys(range(3)) # Store to primary self._start_request() - result = self.manager.prepare_store(blocks, _CTX) + result = self.manager.prepare_store(chunks, _CTX) assert result is not None - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.complete_store(chunks, _CTX, success=True) - # After complete_store, blocks should have ref_cnt > 0 + # After complete_store, chunks should have ref_cnt > 0 # (one for each secondary tier) - for block_hash in blocks: - block = self.primary_tier._policy.get(block_hash) + for key in chunks: + chunk = self.primary_tier._policy.get(key) # ref_cnt should be 2 (one for each secondary tier) - assert block.ref_cnt == 2 + assert chunk.ref_cnt == 2 # End of step 1: _maybe_process_finished_jobs() was already called by # prepare_store() above (setting the per-step flag), so on_schedule_end() @@ -466,9 +467,9 @@ def test_ref_cnt_protection_during_cascade(self, manager_setup): # ref_cnt still held: cascade jobs finished (sync tier) but haven't # been polled yet because the per-step guard skipped the second call. - for block_hash in blocks: - block = self.primary_tier._policy.get(block_hash) - assert block.ref_cnt == 2 + for key in chunks: + chunk = self.primary_tier._policy.get(key) + assert chunk.ref_cnt == 2 # Secondary tiers have completed jobs waiting to be drained assert len(self.secondary_tier1.completed_jobs) > 0 @@ -479,37 +480,37 @@ def test_ref_cnt_protection_during_cascade(self, manager_setup): self._simulate_on_schedule_end() # After cascade completes, ref_cnt should be 0 - for block_hash in blocks: - block = self.primary_tier._policy.get(block_hash) - assert block.ref_cnt == 0 + for key in chunks: + chunk = self.primary_tier._policy.get(key) + assert chunk.ref_cnt == 0 # All completed jobs have been drained assert len(self.secondary_tier1.completed_jobs) == 0 assert len(self.secondary_tier2.completed_jobs) == 0 def test_lookup_from_primary(self, manager_setup): - """Test lookup when blocks are in primary tier.""" - blocks = to_keys(range(3)) + """Test lookup when chunks are in primary tier.""" + chunks = to_keys(range(3)) - # Store blocks + # Store chunks self._start_request() - self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.prepare_store(chunks, _CTX) + self.manager.complete_store(chunks, _CTX, success=True) - # Lookup should find all blocks in primary - assert count_hits(self.manager, blocks) == 3 + # Lookup should find all chunks in primary + assert count_hits(self.manager, chunks) == 3 def test_promotion_from_secondary(self, manager_setup): - """Test promotion of blocks from secondary to primary tier.""" - blocks = to_keys(range(3)) + """Test promotion of chunks from secondary to primary tier.""" + chunks = to_keys(range(3)) - # Manually add blocks to secondary tier (simulate previous cascade) - for block in blocks: - self.secondary_tier1.blocks[block] = True + # Manually add chunks to secondary tier (simulate previous cascade) + for chunk in chunks: + self.secondary_tier1.chunks[chunk] = True - # Lookup each block to initiate promotion for all of them - for block in blocks: - result = self.manager.lookup(block, _CTX) + # Lookup each chunk to initiate promotion for all of them + for chunk in chunks: + result = self.manager.lookup(chunk, _CTX) assert result is LookupResult.HIT_PENDING # promotion initiated # End of step 1: flushes deferred submit_load() calls @@ -518,11 +519,11 @@ def test_promotion_from_secondary(self, manager_setup): # End of step 2: processes the completed promotion jobs self._simulate_on_schedule_end() - # Now blocks should be in primary tier - assert count_hits(self.primary_tier, blocks) == 3 + # Now chunks should be in primary tier + assert count_hits(self.primary_tier, chunks) == 3 # Next lookup should succeed - assert count_hits(self.manager, blocks) == 3 + assert count_hits(self.manager, chunks) == 3 @pytest.mark.parametrize( ("successful_indices", "expected_results"), @@ -538,18 +539,18 @@ def test_promotion_from_secondary(self, manager_setup): ], ids=["partial", "legacy-full-failure"], ) - def test_failed_promotion_keeps_only_successful_blocks( + def test_failed_promotion_keeps_only_successful_chunks( self, manager_setup, successful_indices, expected_results ): - blocks = to_keys(range(3)) - for block in blocks: - self.secondary_tier1.blocks[block] = True + chunks = to_keys(range(3)) + for chunk in chunks: + self.secondary_tier1.chunks[chunk] = True def submit_partial(job_metadata: TransferJob) -> None: successful_keys = ( None if successful_indices is None - else tuple(blocks[i] for i in successful_indices) + else tuple(chunks[i] for i in successful_indices) ) self.secondary_tier1.completed_jobs.append( JobResult( @@ -561,24 +562,24 @@ def submit_partial(job_metadata: TransferJob) -> None: self.secondary_tier1.submit_load = submit_partial - for block in blocks: - assert self.manager.lookup(block, _CTX) is LookupResult.HIT_PENDING + for chunk in chunks: + assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT_PENDING self._simulate_on_schedule_end() self._simulate_on_schedule_end() assert [ - self.primary_tier.lookup(block, _CTX) for block in blocks + self.primary_tier.lookup(chunk, _CTX) for chunk in chunks ] == expected_results def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): - """Resolved lookups report one sync delay sample per tier and block.""" + """Resolved lookups report one sync delay sample per tier and chunk.""" self._start_request() - blocks = to_keys(range(2)) + chunks = to_keys(range(2)) - # No tier has these blocks: they resolve immediately as misses. - for block in blocks: - assert self.manager.lookup(block, _CTX) is LookupResult.MISS + # No tier has these chunks: they resolve immediately as misses. + for chunk in chunks: + assert self.manager.lookup(chunk, _CTX) is LookupResult.MISS stats = self.manager.get_stats() assert stats is not None @@ -614,11 +615,11 @@ def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): def test_lookup_does_not_report_async_delay_for_promotion(self, manager_setup): """Promotion time is not included in lookup async delay.""" self._start_request() - block = to_keys(range(1))[0] - self.secondary_tier1.blocks[block] = True + chunk = to_keys(range(1))[0] + self.secondary_tier1.chunks[chunk] = True - # First lookup finds the block in a secondary tier and defers. - assert self.manager.lookup(block, _CTX) is LookupResult.HIT_PENDING + # First lookup finds the chunk in a secondary tier and defers. + assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT_PENDING # Promotion is not treated as unresolved secondary lookup time. self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id]) @@ -648,8 +649,8 @@ def test_lookup_does_not_report_async_delay_for_promotion(self, manager_setup): # Promotion completes on the next scheduler step. self._simulate_on_schedule_end() - # Next lookup resolves via the now-promoted primary-tier block. - assert self.manager.lookup(block, _CTX) is LookupResult.HIT + # Next lookup resolves via the now-promoted primary-tier chunk. + assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT stats = self.manager.get_stats() if stats is not None: @@ -663,13 +664,13 @@ def test_lookup_reports_async_delay_when_deferred_lookup_resolves( """Async delay is observed when a RETRY lookup later resolves.""" ctx = ReqContext(req_id="req_lookup_finish") self._start_request(ctx) - block = to_keys(range(1))[0] + chunk = to_keys(range(1))[0] self.secondary_tier1.lookup = MagicMock( side_effect=[LookupResult.RETRY, LookupResult.HIT] ) # First lookup is deferred by the secondary tier. - assert self.manager.lookup(block, ctx) is LookupResult.RETRY + assert self.manager.lookup(chunk, ctx) is LookupResult.RETRY stats = self.manager.get_stats() if stats is not None: @@ -677,7 +678,7 @@ def test_lookup_reports_async_delay_when_deferred_lookup_resolves( stats.reduce(), TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY ) - assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING + assert self.manager.lookup(chunk, ctx) is LookupResult.HIT_PENDING stats = self.manager.get_stats() assert stats is not None @@ -701,68 +702,68 @@ def test_lookup_reports_async_delay_when_deferred_lookup_resolves( def test_partial_lookup(self, manager_setup): """Test lookup with partial hits.""" - blocks = to_keys(range(5)) + chunks = to_keys(range(5)) - # Store first 3 blocks to primary + # Store first 3 chunks to primary self._start_request() - self.manager.prepare_store(blocks[:3], _CTX) - self.manager.complete_store(blocks[:3], _CTX, success=True) + self.manager.prepare_store(chunks[:3], _CTX) + self.manager.complete_store(chunks[:3], _CTX, success=True) - # Lookup all 5 blocks should return 3 (first 3 found) - assert count_hits(self.manager, blocks) == 3 + # Lookup all 5 chunks should return 3 (first 3 found) + assert count_hits(self.manager, chunks) == 3 def test_eviction_in_primary_tier(self, manager_setup): """Test eviction in primary tier when capacity is exceeded.""" - # Primary tier has capacity of 5 blocks + # Primary tier has capacity of 5 chunks # First, fill the primary tier - blocks = to_keys(range(5)) + chunks = to_keys(range(5)) self._start_request() - result = self.manager.prepare_store(blocks, _CTX) + result = self.manager.prepare_store(chunks, _CTX) assert result is not None assert len(result.keys_to_store) == 5 - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.complete_store(chunks, _CTX, success=True) # End of step: release ref_cnt from cascade self._simulate_on_schedule_end() - # Now try to store 2 more blocks (should trigger eviction) - more_blocks = to_keys(range(5, 7)) - result = self.manager.prepare_store(more_blocks, _CTX) + # Now try to store 2 more chunks (should trigger eviction) + more_chunks = to_keys(range(5, 7)) + result = self.manager.prepare_store(more_chunks, _CTX) - # Should evict 2 blocks from primary tier + # Should evict 2 chunks from primary tier assert result is not None assert len(result.evicted_keys) == 2 assert len(result.keys_to_store) == 2 def test_touch_propagates_to_all_tiers(self, manager_setup): """Test that touch() propagates to all tiers.""" - blocks = to_keys(range(3)) + chunks = to_keys(range(3)) - # Store blocks + # Store chunks self._start_request() - self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.prepare_store(chunks, _CTX) + self.manager.complete_store(chunks, _CTX, success=True) self._simulate_on_schedule_end() - # for secondary tiers to drain jobs, so primary tier's blocks are evictable. + # for secondary tiers to drain jobs, so primary tier's chunks are evictable. self._simulate_on_schedule_end() self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch) self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch) - # Touch blocks - self.manager.touch(blocks, _CTX) + # Touch chunks + self.manager.touch(chunks, _CTX) # Verify touch was called on primary tier (check LRU order) - primary_keys = list(self.primary_tier._policy.evictable_blocks.keys()) - assert primary_keys[-3:] == list(reversed(blocks)) + primary_keys = list(self.primary_tier._policy.evictable_chunks.keys()) + assert primary_keys[-3:] == list(reversed(chunks)) # Verify touch was propagated to all secondary tiers - self.secondary_tier1.touch.assert_called_once_with(blocks, _CTX) - self.secondary_tier2.touch.assert_called_once_with(blocks, _CTX) + self.secondary_tier1.touch.assert_called_once_with(chunks, _CTX) + self.secondary_tier2.touch.assert_called_once_with(chunks, _CTX) def test_failed_store_no_cascade(self, manager_setup): """Test that failed GPU→primary store doesn't cascade.""" - blocks = to_keys(range(3)) + chunks = to_keys(range(3)) self.secondary_tier1.submit_store = MagicMock( wraps=self.secondary_tier1.submit_store @@ -773,11 +774,11 @@ def test_failed_store_no_cascade(self, manager_setup): # Prepare store self._start_request() - result = self.manager.prepare_store(blocks, _CTX) + result = self.manager.prepare_store(chunks, _CTX) assert result is not None # Complete store with failure — cascade must not happen - self.manager.complete_store(blocks, _CTX, success=False) + self.manager.complete_store(chunks, _CTX, success=False) # submit_store was never called on either secondary tier self.secondary_tier1.submit_store.assert_not_called() @@ -786,12 +787,12 @@ def test_failed_store_no_cascade(self, manager_setup): def test_lookup_batches_submit_load_per_request(self, manager_setup): """lookup() defers submit_load until on_schedule_end(), one per request. - Blocks from different requests each get their own submit_load call, each + Chunks from different requests each get their own submit_load call, each carrying the correct req_context. """ - blocks = to_keys(range(4)) - for block in blocks: - self.secondary_tier1.blocks[block] = True + chunks = to_keys(range(4)) + for chunk in chunks: + self.secondary_tier1.chunks[chunk] = True self.secondary_tier1.submit_load = MagicMock( wraps=self.secondary_tier1.submit_load @@ -801,10 +802,10 @@ def test_lookup_batches_submit_load_per_request(self, manager_setup): ctx_b = ReqContext(req_id="req_b") # All lookups return HIT_PENDING: secondary hit triggers promotion - assert self.manager.lookup(blocks[0], ctx_a) is LookupResult.HIT_PENDING - assert self.manager.lookup(blocks[1], ctx_a) is LookupResult.HIT_PENDING - assert self.manager.lookup(blocks[2], ctx_b) is LookupResult.HIT_PENDING - assert self.manager.lookup(blocks[3], ctx_b) is LookupResult.HIT_PENDING + assert self.manager.lookup(chunks[0], ctx_a) is LookupResult.HIT_PENDING + assert self.manager.lookup(chunks[1], ctx_a) is LookupResult.HIT_PENDING + assert self.manager.lookup(chunks[2], ctx_b) is LookupResult.HIT_PENDING + assert self.manager.lookup(chunks[3], ctx_b) is LookupResult.HIT_PENDING # submit_load must not fire during lookup - only at end of step self.secondary_tier1.submit_load.assert_not_called() @@ -816,20 +817,20 @@ def test_lookup_batches_submit_load_per_request(self, manager_setup): calls = self.secondary_tier1.submit_load.call_args_list jm_a = calls[0].args[0] jm_b = calls[1].args[0] - assert set(jm_a.keys) == {blocks[0], blocks[1]} + assert set(jm_a.keys) == {chunks[0], chunks[1]} assert jm_a.req_context is ctx_a - assert set(jm_b.keys) == {blocks[2], blocks[3]} + assert set(jm_b.keys) == {chunks[2], chunks[3]} assert jm_b.req_context is ctx_b - def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup): - """A block looked up by two requests in the same step is promoted once. + def test_lookup_shared_chunk_no_duplicate_promotion(self, manager_setup): + """A chunk looked up by two requests in the same step is promoted once. The first lookup initiates promotion (returns None via secondary hit). The second lookup sees ref_cnt=-1 on the primary slot and returns None via the primary in-flight path — without triggering a second promotion. """ - shared_block = to_keys([0])[0] - self.secondary_tier1.blocks[shared_block] = True + shared_chunk = to_keys([0])[0] + self.secondary_tier1.chunks[shared_chunk] = True self.secondary_tier1.submit_load = MagicMock( wraps=self.secondary_tier1.submit_load @@ -838,10 +839,10 @@ def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup): ctx_a = ReqContext(req_id="req_a") ctx_b = ReqContext(req_id="req_b") - result_a = self.manager.lookup(shared_block, ctx_a) - result_b = self.manager.lookup(shared_block, ctx_b) + result_a = self.manager.lookup(shared_chunk, ctx_a) + result_b = self.manager.lookup(shared_chunk, ctx_b) - # Both lookups find the shared_block and trigger promotion + # Both lookups find the shared_chunk and trigger promotion # returning HIT_PENDING. assert result_a is LookupResult.HIT_PENDING assert result_b is LookupResult.HIT_PENDING @@ -851,12 +852,12 @@ def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup): # Only one submit_load call despite two lookups self.secondary_tier1.submit_load.assert_called_once() job_metadata = self.secondary_tier1.submit_load.call_args.args[0] - assert list(job_metadata.keys) == [shared_block] + assert list(job_metadata.keys) == [shared_chunk] assert job_metadata.req_context is ctx_a def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup): """complete_store cascades to secondary tiers with the correct req_context.""" - blocks = to_keys(range(2)) + chunks = to_keys(range(2)) self.secondary_tier1.submit_store = MagicMock( wraps=self.secondary_tier1.submit_store @@ -865,8 +866,8 @@ def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup ctx = ReqContext(req_id="req_ctx", kv_transfer_params={"key": "value"}) self._start_request(ctx) - self.manager.prepare_store(blocks, ctx) - self.manager.complete_store(blocks, ctx, success=True) + self.manager.prepare_store(chunks, ctx) + self.manager.complete_store(chunks, ctx, success=True) assert self.secondary_tier1.submit_store.call_count == 1 job_metadata = self.secondary_tier1.submit_store.call_args.args[0] @@ -876,7 +877,7 @@ def test_on_request_finished_delays_secondary_until_store_submitted( self, manager_setup ): """Manager hook is eager; secondary hooks wait for cascade submission.""" - blocks = to_keys(range(2)) + chunks = to_keys(range(2)) ctx = ReqContext(req_id="req_delayed_secondary") calls: list[tuple[str, str]] = [] @@ -911,14 +912,14 @@ def submit_store2(job_metadata): ) self._start_request(ctx) - self.manager.prepare_store(blocks, ctx) + self.manager.prepare_store(chunks, ctx) self.manager.on_request_finished(ctx) assert calls == [("primary_finish", ctx.req_id)] self.secondary_tier1.on_request_finished.assert_not_called() self.secondary_tier2.on_request_finished.assert_not_called() - self.manager.complete_store(blocks, ctx, success=True) + self.manager.complete_store(chunks, ctx, success=True) assert calls == [ ("primary_finish", ctx.req_id), @@ -930,7 +931,7 @@ def submit_store2(job_metadata): def test_failed_store_finalizes_finished_request(self, manager_setup): """Failed primary stores still unblock secondary finalization.""" - blocks = to_keys(range(2)) + chunks = to_keys(range(2)) ctx = ReqContext(req_id="req_failed_store_finalize") self.secondary_tier1.submit_store = MagicMock( @@ -947,13 +948,13 @@ def test_failed_store_finalizes_finished_request(self, manager_setup): ) self._start_request(ctx) - self.manager.prepare_store(blocks, ctx) + self.manager.prepare_store(chunks, ctx) self.manager.on_request_finished(ctx) self.secondary_tier1.on_request_finished.assert_not_called() self.secondary_tier2.on_request_finished.assert_not_called() - self.manager.complete_store(blocks, ctx, success=False) + self.manager.complete_store(chunks, ctx, success=False) self.secondary_tier1.submit_store.assert_not_called() self.secondary_tier2.submit_store.assert_not_called() @@ -981,7 +982,7 @@ def test_zero_store_request_finalizes_immediately(self, manager_setup): def test_reset_cache_finalizes_delayed_secondary_request(self, manager_setup): """reset_cache abandons pending primary stores and finalizes secondaries.""" - blocks = to_keys(range(2)) + chunks = to_keys(range(2)) ctx = ReqContext(req_id="req_reset_finalize_secondary") self.secondary_tier1.on_request_finished = MagicMock( @@ -992,7 +993,7 @@ def test_reset_cache_finalizes_delayed_secondary_request(self, manager_setup): ) self._start_request(ctx) - self.manager.prepare_store(blocks, ctx) + self.manager.prepare_store(chunks, ctx) self.manager.on_request_finished(ctx) self.secondary_tier1.on_request_finished.assert_not_called() @@ -1008,8 +1009,8 @@ def test_reset_cache_clears_pending_primary_stores_for_active_request( self, manager_setup ): """reset_cache drops active pending stores so resumed requests finalize.""" - initial_blocks = to_keys(range(2)) - resumed_blocks = to_keys(range(2, 4)) + initial_chunks = to_keys(range(2)) + resumed_chunks = to_keys(range(2, 4)) ctx = ReqContext(req_id="req_reset_resume") self.secondary_tier1.on_request_finished = MagicMock( @@ -1020,7 +1021,7 @@ def test_reset_cache_clears_pending_primary_stores_for_active_request( ) self._start_request(ctx) - self.manager.prepare_store(initial_blocks, ctx) + self.manager.prepare_store(initial_chunks, ctx) assert self.manager._req_state[ctx.req_id].pending_primary_stores == 1 self.manager.reset_cache() @@ -1030,8 +1031,8 @@ def test_reset_cache_clears_pending_primary_stores_for_active_request( self.secondary_tier1.on_request_finished.assert_not_called() self.secondary_tier2.on_request_finished.assert_not_called() - self.manager.prepare_store(resumed_blocks, ctx) - self.manager.complete_store(resumed_blocks, ctx, success=True) + self.manager.prepare_store(resumed_chunks, ctx) + self.manager.complete_store(resumed_chunks, ctx, success=True) self.manager.on_request_finished(ctx) self.secondary_tier1.on_request_finished.assert_called_once_with(ctx) @@ -1039,12 +1040,12 @@ def test_reset_cache_clears_pending_primary_stores_for_active_request( assert ctx.req_id not in self.manager._req_state def test_on_new_request_lifecycle(self, manager_setup): - """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it, + """Policy defaults to CHUNK_LEVEL, escalates when a tier requests it, and is cleaned up on on_request_finished.""" - # Default: all tiers return BLOCK_LEVEL + # Default: all tiers return CHUNK_LEVEL ctx = ReqContext(req_id="req_policy_lifecycle") result = self.manager.on_new_request(ctx) - assert result.policy == OffloadPolicy.BLOCK_LEVEL + assert result.policy == OffloadPolicy.CHUNK_LEVEL assert self.manager._req_state[ctx.req_id].request_level_tiers is None self.manager.on_request_finished(ctx) assert ctx.req_id not in self.manager._req_state @@ -1064,26 +1065,26 @@ def test_on_new_request_lifecycle(self, manager_setup): assert ctx.req_id not in self.manager._req_state @pytest.mark.parametrize("new_ids", [(), (3, 4)], ids=["fully_warm", "mixed"]) - def test_prepare_store_cascades_existing_blocks_to_request_level_tiers( + def test_prepare_store_cascades_existing_chunks_to_request_level_tiers( self, manager_setup, new_ids ): - """prepare_store cascades hit blocks to request-level tiers only. + """prepare_store cascades hit chunks to request-level tiers only. The fully-warm case is the producer shape from issue #52808: nothing is new, so primary returns an empty keys_to_store and no GPU->primary - transfer runs. A request-level tier must still be handed the blocks, + transfer runs. A request-level tier must still be handed the chunks, otherwise a peer waiting on them has nothing to fetch. """ - # Store some blocks to primary first - existing_blocks = to_keys(range(3)) + # Store some chunks to primary first + existing_chunks = to_keys(range(3)) self._start_request() - result = self.manager.prepare_store(existing_blocks, _CTX) + result = self.manager.prepare_store(existing_chunks, _CTX) assert result is not None - self.manager.complete_store(existing_blocks, _CTX, success=True) + self.manager.complete_store(existing_chunks, _CTX, success=True) # Drain cascade completions self._simulate_on_schedule_end() - # Make tier1 request-level, tier2 stays block-level + # Make tier1 request-level, tier2 stays chunk-level self.secondary_tier1.on_new_request = lambda req_context: ( RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL) ) @@ -1099,20 +1100,20 @@ def test_prepare_store_cascades_existing_blocks_to_request_level_tiers( wraps=self.secondary_tier2.submit_store ) - # Call prepare_store with existing + new blocks - new_blocks = to_keys(new_ids) - all_blocks = existing_blocks + new_blocks - result = self.manager.prepare_store(all_blocks, ctx) + # Call prepare_store with existing + new chunks + new_chunks = to_keys(new_ids) + all_chunks = existing_chunks + new_chunks + result = self.manager.prepare_store(all_chunks, ctx) assert result is not None - assert set(result.keys_to_store) == set(new_blocks) + assert set(result.keys_to_store) == set(new_chunks) - # Only tier1 (request-level) should get existing blocks cascaded now. - # New blocks are cascaded to ALL tiers later via complete_store(). + # Only tier1 (request-level) should get existing chunks cascaded now. + # New chunks are cascaded to ALL tiers later via complete_store(). self.secondary_tier1.submit_store.assert_called_once() job_metadata = self.secondary_tier1.submit_store.call_args.args[0] - assert set(job_metadata.keys) == set(existing_blocks) + assert set(job_metadata.keys) == set(existing_chunks) - # tier2 (block-level) does not get existing blocks here. + # tier2 (chunk-level) does not get existing chunks here. self.secondary_tier2.submit_store.assert_not_called() def _make_request_level_request(self, req_id: str) -> ReqContext: @@ -1150,7 +1151,7 @@ def test_cascade_rejects_retry_from_primary(self, manager_setup): self.manager.primary_tier.lookup = lambda key, req_context: (LookupResult.RETRY) with pytest.raises(AssertionError): - self.manager._cascade_existing_blocks_to_request_level_tiers(keys, ctx, {0}) + self.manager._cascade_existing_chunks_to_request_level_tiers(keys, ctx, {0}) def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext: """Leave a GPU->primary write for `keys` open, so they look present to @@ -1202,7 +1203,7 @@ def test_deferred_cascade_holds_request_from_finalization(self, manager_setup): def test_cascade_drops_deferred_keys_whose_primary_write_failed( self, manager_setup ): - """A failed write frees the block, so the deferred key resolves to MISS + """A failed write frees the chunk, so the deferred key resolves to MISS and the request finalizes instead of parking forever.""" keys = to_keys(range(3)) writer_ctx = self._start_in_flight_primary_write(keys) @@ -1224,21 +1225,21 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): """reset_cache wipes every kind of orchestrator state and resets primary tier; pending submissions are dropped without being sent to the secondary tier. Active request state is retained.""" - # Cascade — populates primary blocks and leaves cascade jobs + # Cascade — populates primary chunks and leaves cascade jobs # in _jobs (the synchronous example tier has already # queued completions); reset_cache's drain loop will pick them up. - blocks = to_keys(range(3)) + chunks = to_keys(range(3)) self._start_request() - self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.prepare_store(chunks, _CTX) + self.manager.complete_store(chunks, _CTX, success=True) assert self.manager._jobs # Pending promotion submission (deferred — no on_schedule_end after # the lookup that staged it). - promo_block = to_keys([99])[0] - self.secondary_tier1.blocks[promo_block] = True + promo_chunk = to_keys([99])[0] + self.secondary_tier1.chunks[promo_chunk] = True assert ( - self.manager.lookup(promo_block, ReqContext(req_id="pending")) + self.manager.lookup(promo_chunk, ReqContext(req_id="pending")) is LookupResult.HIT_PENDING ) assert self.manager._pending_load_submissions @@ -1268,10 +1269,10 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): assert self.manager._processed_jobs_this_step is False # Primary tier reset to a fresh state. - assert self.primary_tier._num_allocated_blocks == 0 + assert self.primary_tier._num_allocated_chunks == 0 assert self.primary_tier._free_list == [] - for block in blocks: - assert self.primary_tier.lookup(block, _CTX) is LookupResult.MISS + for chunk in chunks: + assert self.primary_tier.lookup(chunk, _CTX) is LookupResult.MISS # Pending submission was dropped, not submitted. self.secondary_tier1.submit_load.assert_not_called() @@ -1291,10 +1292,10 @@ def test_reset_cache_drains_all_tiers(self, manager_setup): ) # Drive a cascade so a job lands in _jobs. - blocks = to_keys(range(3)) + chunks = to_keys(range(3)) self._start_request() - self.manager.prepare_store(blocks, _CTX) - self.manager.complete_store(blocks, _CTX, success=True) + self.manager.prepare_store(chunks, _CTX) + self.manager.complete_store(chunks, _CTX, success=True) assert self.manager._jobs self.manager.reset_cache() @@ -1315,20 +1316,20 @@ def test_tier_filter_skips_filtered_secondary( self, manager_setup, load_tier_filter ): """Filter excluding secondary medium returns MISS from secondaries - even when they hold the block; primary is unaffected.""" - blocks = to_keys(range(2)) - # Put one block in primary, one only in secondary + even when they hold the chunk; primary is unaffected.""" + chunks = to_keys(range(2)) + # Put one chunk in primary, one only in secondary self._start_request() - self.manager.prepare_store(blocks[:1], _CTX) - self.manager.complete_store(blocks[:1], _CTX, success=True) - self.secondary_tier1.blocks[blocks[1]] = True + self.manager.prepare_store(chunks[:1], _CTX) + self.manager.complete_store(chunks[:1], _CTX, success=True) + self.secondary_tier1.chunks[chunks[1]] = True # Secondaries have medium=CPU, so load_tier_filter skips them. self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup) ctx = ReqContext(req_id="r1", load_tier_filter=load_tier_filter) - assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT - assert self.manager.lookup(blocks[1], ctx) is LookupResult.MISS + assert self.manager.lookup(chunks[0], ctx) is LookupResult.HIT + assert self.manager.lookup(chunks[1], ctx) is LookupResult.MISS self.secondary_tier1.lookup.assert_not_called() @pytest.mark.parametrize( @@ -1344,13 +1345,13 @@ def test_tier_filter_allows_matching_secondary( self, manager_setup, load_tier_filter ): """Filter that matches the secondary's medium allows lookup.""" - blocks = to_keys(range(1)) - self.secondary_tier1.blocks[blocks[0]] = True + chunks = to_keys(range(1)) + self.secondary_tier1.chunks[chunks[0]] = True self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup) ctx = ReqContext(req_id="r2", load_tier_filter=load_tier_filter) - assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT_PENDING + assert self.manager.lookup(chunks[0], ctx) is LookupResult.HIT_PENDING self.secondary_tier1.lookup.assert_called() @@ -1360,7 +1361,7 @@ class TestTieringOffloadingWithoutSecondaryTiers: def test_works_without_secondary_tiers(self): """Test that manager works with empty secondary_tiers list.""" primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=5, mmap_region=_mock_mmap_region(5) + num_chunks=5, mmap_region=_mock_mmap_region(5) ) # Create manager with no secondary tiers @@ -1368,15 +1369,15 @@ def test_works_without_secondary_tiers(self): primary_tier=primary_tier, secondary_tiers=[] ) - blocks = to_keys(range(3)) + chunks = to_keys(range(3)) # Should work like a regular OffloadingManager manager.on_new_request(_CTX) - result = manager.prepare_store(blocks, _CTX) + result = manager.prepare_store(chunks, _CTX) assert result is not None - manager.complete_store(blocks, _CTX, success=True) + manager.complete_store(chunks, _CTX, success=True) - assert count_hits(manager, blocks) == 3 + assert count_hits(manager, chunks) == 3 @pytest.mark.parametrize( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 972eeff2ee03..a3cf1a5a1e68 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1131,7 +1131,7 @@ def update_state_after_alloc( # Skip prefix-hit chunks for block-level policy; for # request-level, next_stored_chunk_idx stays at 0 so all # chunks (including hits) are offloaded. - if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL: + if req_status.offloading_context.policy == OffloadPolicy.CHUNK_LEVEL: group_state.next_stored_chunk_idx = num_chunks src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 5ad89eef5792..93f8e4bda4e9 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -114,17 +114,17 @@ class LookupResult(Enum): class OffloadPolicy(Enum): - # Offload only newly-computed blocks as they arrive; prefix-hit - # blocks (already offloaded by a prior request) are skipped. - BLOCK_LEVEL = "block_level" - # Offload all blocks for the request, including prefix hits. + # Offload only newly-computed chunks as they arrive; prefix-hit + # chunks (already offloaded by a prior request) are skipped. + CHUNK_LEVEL = "chunk_level" + # Offload all chunks for the request, including prefix hits. # Used by tiers that need the complete KV context for a request. REQUEST_LEVEL = "request_level" @dataclass class RequestOffloadingContext: - policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL + policy: OffloadPolicy = OffloadPolicy.CHUNK_LEVEL class ScheduleEndContext(NamedTuple): @@ -400,6 +400,9 @@ def shutdown(self) -> None: class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC): """ Spec for loading/storing KV blocks from given block numbers. + + Subclass semantics differ: GPULoadStoreSpec.block_ids are GPU block + indices; CPULoadStoreSpec.block_ids are CPU cache chunk indices. """ def __init__(self, block_ids: list[int]): diff --git a/vllm/v1/kv_offload/cpu/common.py b/vllm/v1/kv_offload/cpu/common.py index 7875d81501b2..61dce9365ea8 100644 --- a/vllm/v1/kv_offload/cpu/common.py +++ b/vllm/v1/kv_offload/cpu/common.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np + from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec @@ -12,6 +14,13 @@ class CPUOffloadingMetrics: class CPULoadStoreSpec(BlockIDsLoadStoreSpec): + """Spec for loading/storing KV chunks to/from CPU memory. + + The inherited block_ids field holds chunk indices into the + CPU cache (not GPU block IDs). The chunk_ids alias exposes + the same array under the name used by the tiering layer. """ - Spec for loading/storing a KV block to CPU memory. - """ + + @property + def chunk_ids(self) -> np.ndarray: + return self.block_ids diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index bcaf21ff7531..8db9e805ff1a 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -257,7 +257,7 @@ def __init__( gpu_tensors: list of GPU KV cache tensors. Each of shape (num_gpu_blocks, gpu_page_size_bytes) with dtype int8. cpu_tensors: list of CPU KV cache tensors. - Each of shape (num_cpu_blocks, cpu_page_size_bytes) with dtype int8. + Each of shape (num_cpu_chunks, cpu_page_size_bytes) with dtype int8. Order should match gpu_tensors. layer_refs_per_group: list of CanonicalKVCacheRef per group. gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU. @@ -527,20 +527,20 @@ def transfer_async( # 1. GPU -> CPU # 2. CPU -> GPU # - # transfers are also to CPU blocks, EXCEPT MAYBE for the first and last block. - # i.e. the first and last CPU blocks in src_blocks can match against + # transfers are also to CPU chunks, EXCEPT MAYBE for the first and last chunk. + # i.e. the first and last CPU chunks in src_blocks can match against # a smaller (byte-wise) set of GPU blocks in dst_blocks. # In such cases, we may need to skip some gpu-sized sub-blocks, - # and start reading/writing from the middle of the first CPU block. + # and start reading/writing from the middle of the first CPU chunk. # If we have multiple KV cache groups (when using HMA with hybrid models), - # we may have a partial first/last CPU block per each group. + # we may have a partial first/last CPU chunk per each group. # The group_sizes parameter encodes the size of each group of blocks # in the GPU dst_blocks. # If group_sizes is None, we assume all blocks belong to a single group. # The logical_offset parameter maps each group of blocks to its logical # offset inside the request, counting in GPU blocks. # This allows us to find the correct starting position - # in the matching first CPU block. + # in the matching first CPU chunk. # extract group_sizes from the GPU spec gpu_spec = src_spec if self.gpu_to_cpu else dst_spec @@ -751,7 +751,7 @@ def __init__( self, kv_caches: CanonicalKVCaches, blocks_per_chunk: int, - num_cpu_blocks: int, + num_cpu_chunks: int, mmap_region: SharedOffloadRegion | None = None, canonical_layout: bool = False, ): @@ -790,16 +790,16 @@ def __init__( else: t0 = time.monotonic() cpu_tensor = torch.zeros( - (num_cpu_blocks, cpu_page_size_bytes), + (num_cpu_chunks, cpu_page_size_bytes), dtype=torch.int8, device="cpu", pin_memory=pin_memory, ) logger.debug( "torch.zeros pinned tensor %d×%d (%.2f GB): %.3f s", - num_cpu_blocks, + num_cpu_chunks, cpu_page_size_bytes, - num_cpu_blocks * cpu_page_size_bytes / 1e9, + num_cpu_chunks * cpu_page_size_bytes / 1e9, time.monotonic() - t0, ) diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 9082c83de4f7..67b469b6a394 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -23,7 +23,7 @@ CPULoadStoreSpec, CPUOffloadingMetrics, ) -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory @@ -34,14 +34,14 @@ class CPUOffloadingManager(OffloadingManager): register their own or be loaded out-of-tree via cache_policy_module_path). The manager owns all shared logic: ref-counting, event emission, - block pool management, and the prepare_store/complete_store skeletons. - Policy-specific block organization and eviction decisions are delegated + chunk pool management, and the prepare_store/complete_store skeletons. + Policy-specific chunk organization and eviction decisions are delegated to the CachePolicy implementation. """ def __init__( self, - num_blocks: int, + num_chunks: int, cache_policy: str = "lru", cache_policy_module_path: str | None = None, enable_events: bool = False, @@ -49,59 +49,59 @@ def __init__( max_tracker_size: int = 64_000, ): self.medium: Medium = Medium.CPU - self._num_blocks: int = num_blocks - self._num_allocated_blocks: int = 0 + self._num_chunks: int = num_chunks + self._num_allocated_chunks: int = 0 self._free_list: list[int] = [] self.events: list[OffloadingEvent] | None = [] if enable_events else None policy_cls = CachePolicyFactory.get_cache_policy_cls( cache_policy, cache_policy_module_path ) - self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks) - # Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0. - self._num_evictable_cache_blocks: int = 0 - # Track blocks with an in-flight store (ref_cnt -1, not yet completed). - self._num_write_pending_blocks: int = 0 + self._policy: CachePolicy = policy_cls(cache_capacity=num_chunks) + # Track the number of chunks in the cache that are evictable. i.e. ref_cnt 0. + self._num_evictable_cache_chunks: int = 0 + # Track chunks with an in-flight store (ref_cnt -1, not yet completed). + self._num_write_pending_chunks: int = 0 self.store_threshold: int = store_threshold self.max_tracker_size: int = max_tracker_size self.stores_skipped_in_current_batch: int = 0 self.allocation_sizes_in_current_batch: list[int] = [] - # Number of block references. It is ordered so can evict the LRU entry in O(1). + # Number of chunk references. It is ordered so can evict the LRU entry in O(1). self.counts: OrderedDict[OffloadKey, int] | None = ( OrderedDict() if store_threshold >= 2 else None ) - # --- block pool --- + # --- chunk pool --- - def _get_num_free_blocks(self) -> int: - return len(self._free_list) + self._num_blocks - self._num_allocated_blocks + def _get_num_free_chunks(self) -> int: + return len(self._free_list) + self._num_chunks - self._num_allocated_chunks - def _allocate_blocks(self, keys: list[OffloadKey]) -> list[BlockStatus]: - num_fresh = min(len(keys), self._num_blocks - self._num_allocated_blocks) + def _allocate_chunks(self, keys: list[OffloadKey]) -> list[ChunkStatus]: + num_fresh = min(len(keys), self._num_chunks - self._num_allocated_chunks) num_reused = len(keys) - num_fresh assert len(self._free_list) >= num_reused - # allocate fresh blocks - blocks: list[BlockStatus] = [] + # allocate fresh chunks + chunks: list[ChunkStatus] = [] for _ in range(num_fresh): - blocks.append(BlockStatus(self._num_allocated_blocks)) - self._num_allocated_blocks += 1 + chunks.append(ChunkStatus(self._num_allocated_chunks)) + self._num_allocated_chunks += 1 - # allocate reused blocks + # allocate reused chunks for _ in range(num_reused): - blocks.append(BlockStatus(self._free_list.pop())) - return blocks + chunks.append(ChunkStatus(self._free_list.pop())) + return chunks - def _free_block(self, block: BlockStatus) -> None: - self._free_list.append(block.block_id) + def _free_chunk(self, chunk: ChunkStatus) -> None: + self._free_list.append(chunk.chunk_id) def _get_load_store_spec( self, keys: Iterable[OffloadKey], - blocks: Iterable[BlockStatus], + chunks: Iterable[ChunkStatus], ) -> CPULoadStoreSpec: - return CPULoadStoreSpec([block.block_id for block in blocks]) + return CPULoadStoreSpec([chunk.chunk_id for chunk in chunks]) def _record_accesses(self, keys: Collection[OffloadKey]) -> None: """Record an offer without evicting its tracked candidates.""" @@ -132,10 +132,10 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - block = self._policy.get(key) - if block is None: + chunk = self._policy.get(key) + if chunk is None: return LookupResult.MISS - if not block.is_ready: + if not chunk.is_ready: return LookupResult.HIT_PENDING return LookupResult.HIT @@ -145,18 +145,18 @@ def prepare_load( keys: Collection[OffloadKey], req_context: ReqContext, ) -> LoadStoreSpec: - blocks = [] + chunks = [] for key in keys: - block = self._policy.get(key) - assert block is not None, f"Block {key!r} not found in cache" - assert block.is_ready, f"Block {key!r} is not ready for reading" - if block.ref_cnt == 0: + chunk = self._policy.get(key) + assert chunk is not None, f"Chunk {key!r} not found in cache" + assert chunk.is_ready, f"Chunk {key!r} is not ready for reading" + if chunk.ref_cnt == 0: self._policy.mark_non_evictable(key) - self._num_evictable_cache_blocks -= 1 # ref_cnt 0 -> 1 - assert self._num_evictable_cache_blocks >= 0 - block.ref_cnt += 1 - blocks.append(block) - return self._get_load_store_spec(keys, blocks) + self._num_evictable_cache_chunks -= 1 # ref_cnt 0 -> 1 + assert self._num_evictable_cache_chunks >= 0 + chunk.ref_cnt += 1 + chunks.append(chunk) + return self._get_load_store_spec(keys, chunks) @override def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None: @@ -167,12 +167,12 @@ def complete_load( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> None: for key in keys: - block = self._policy.get(key) - assert block is not None, f"Block {key!r} not found" - assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0" - block.ref_cnt -= 1 - if block.ref_cnt == 0: - self._num_evictable_cache_blocks += 1 # ref_cnt 1 -> 0 + chunk = self._policy.get(key) + assert chunk is not None, f"Chunk {key!r} not found" + assert chunk.ref_cnt > 0, f"Chunk {key!r} ref_cnt is already 0" + chunk.ref_cnt -= 1 + if chunk.ref_cnt == 0: + self._num_evictable_cache_chunks += 1 # ref_cnt 1 -> 0 self._policy.mark_evictable(key) @override @@ -186,7 +186,7 @@ def prepare_store( self._record_accesses(keys) keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold] self.stores_skipped_in_current_batch += num_keys - len(keys) - # filter out blocks that are already stored + # filter out chunks that are already stored keys_to_store = [k for k in keys if self._policy.get(k) is None] if not keys_to_store: @@ -197,29 +197,29 @@ def prepare_store( ) self.allocation_sizes_in_current_batch.append(len(keys_to_store)) - num_blocks_to_evict = len(keys_to_store) - self._get_num_free_blocks() + num_chunks_to_evict = len(keys_to_store) - self._get_num_free_chunks() to_evict: list[OffloadKey] = [] - if num_blocks_to_evict > 0: - if num_blocks_to_evict > self._num_evictable_cache_blocks: + if num_chunks_to_evict > 0: + if num_chunks_to_evict > self._num_evictable_cache_chunks: # Eviction will fail. return None # There is a still a chance for eviction failure as some of the - # idle blocks might be in the protected list. + # idle chunks might be in the protected list. - # Blocks from the original input are excluded from eviction candidates: - # a block that was already stored must remain in the cache after this call. + # Chunks from the original input are excluded from eviction candidates: + # a chunk that was already stored must remain in the cache after this call. protected = set(keys) - evicted = self._policy.evict(num_blocks_to_evict, protected) + evicted = self._policy.evict(num_chunks_to_evict, protected) if evicted is None: return None - # cache-policy removes only idle blocks. - self._num_evictable_cache_blocks -= len(evicted) - assert self._num_evictable_cache_blocks >= 0 + # cache-policy removes only idle chunks. + self._num_evictable_cache_chunks -= len(evicted) + assert self._num_evictable_cache_chunks >= 0 - for key, block in evicted: - self._free_block(block) + for key, chunk in evicted: + self._free_chunk(chunk) to_evict.append(key) if to_evict and self.events is not None: @@ -231,17 +231,17 @@ def prepare_store( ) ) - blocks = self._allocate_blocks(keys_to_store) - assert len(blocks) == len(keys_to_store), ( - "Block pool did not allocate the expected number of blocks" + chunks = self._allocate_chunks(keys_to_store) + assert len(chunks) == len(keys_to_store), ( + "Chunk pool did not allocate the expected number of chunks" ) - for key, block in zip(keys_to_store, blocks): - self._policy.insert(key, block) - self._num_write_pending_blocks += len(keys_to_store) + for key, chunk in zip(keys_to_store, chunks): + self._policy.insert(key, chunk) + self._num_write_pending_chunks += len(keys_to_store) - # build store specs for allocated blocks - store_spec = self._get_load_store_spec(keys_to_store, blocks) + # build store specs for allocated chunks + store_spec = self._get_load_store_spec(keys_to_store, chunks) return PrepareStoreOutput( keys_to_store=keys_to_store, @@ -260,20 +260,20 @@ def complete_store( if success: for key in keys: - block = self._policy.get(key) - if block is not None and not block.is_ready: - block.ref_cnt = 0 - self._num_write_pending_blocks -= 1 - self._num_evictable_cache_blocks += 1 + chunk = self._policy.get(key) + if chunk is not None and not chunk.is_ready: + chunk.ref_cnt = 0 + self._num_write_pending_chunks -= 1 + self._num_evictable_cache_chunks += 1 self._policy.mark_evictable(key) stored_keys.append(key) else: for key in keys: - block = self._policy.get(key) - if block is not None and not block.is_ready: - self._num_write_pending_blocks -= 1 + chunk = self._policy.get(key) + if chunk is not None and not chunk.is_ready: + self._num_write_pending_chunks -= 1 self._policy.remove(key) - self._free_block(block) + self._free_chunk(chunk) if stored_keys and self.events is not None: self.events.append( @@ -286,17 +286,17 @@ def complete_store( @override def reset_cache(self) -> None: - # Clear ALL blocks unconditionally. The scheduler's _stale_job_threshold + # Clear ALL chunks unconditionally. The scheduler's _stale_job_threshold # guarantees that complete_load / complete_store are never called for # pre-reset jobs, so no lazy cleanup is needed. The scheduler also # flushes in-flight load job IDs to the workers before any new stores - # can begin, preventing a cross-direction data race on reused offload block IDs. + # can begin, preventing a cross-direction data race on reused offload chunk IDs. self._policy.clear() - self._num_evictable_cache_blocks = 0 - self._num_write_pending_blocks = 0 + self._num_evictable_cache_chunks = 0 + self._num_write_pending_chunks = 0 self._free_list.clear() - self._num_allocated_blocks = 0 + self._num_allocated_chunks = 0 @override def take_events(self) -> Iterable[OffloadingEvent]: @@ -309,11 +309,11 @@ def get_stats(self) -> OffloadingConnectorStats | None: # Compute cache usage. num_used = ( - self._num_allocated_blocks + self._num_allocated_chunks - len(self._free_list) - - self._num_evictable_cache_blocks + - self._num_evictable_cache_chunks ) - usage = num_used / self._num_blocks if self._num_blocks > 0 else 0.0 + usage = num_used / self._num_chunks if self._num_chunks > 0 else 0.0 stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC, usage) for allocation_size in self.allocation_sizes_in_current_batch: @@ -323,8 +323,8 @@ def get_stats(self) -> OffloadingConnectorStats | None: self.allocation_sizes_in_current_batch.clear() write_usage = ( - self._num_write_pending_blocks / self._num_blocks - if self._num_blocks > 0 + self._num_write_pending_chunks / self._num_chunks + if self._num_chunks > 0 else 0.0 ) read_usage = max(usage - write_usage, 0.0) diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index e8d65dafceaa..e413bd603eb4 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -6,7 +6,7 @@ from typing_extensions import override from vllm.v1.kv_offload.base import OffloadKey, ReqContext -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus class ARCCachePolicy(CachePolicy): @@ -14,15 +14,15 @@ class ARCCachePolicy(CachePolicy): ARC (Adaptive Replacement Cache) cache policy. Data Structures: - T1: Recent cache containing blocks accessed once. - T2: Frequent cache containing blocks accessed multiple times. - B1/B2: Ghost lists tracking recently evicted blocks from T1/T2. + T1: Recent cache containing chunks accessed once. + T2: Frequent cache containing chunks accessed multiple times. + B1/B2: Ghost lists tracking recently evicted chunks from T1/T2. target_t1_size: Adaptive target size for the T1 partition. Algorithm Flow: 1. Cache lookup (lookup): - Searches T1 and T2 for block hashes and counts consecutive hits - until a miss or non-ready block is encountered. + Searches T1 and T2 for chunk hashes and counts consecutive hits + until a miss or non-ready chunk is encountered. 2. Cache touch (touch) - Adaptive Learning: For each key (in reverse order): @@ -31,15 +31,15 @@ class ARCCachePolicy(CachePolicy): - If in B1 ghost list: Increase target_t1_size. - If in B2 ghost list: Decrease target_t1_size. - 3. Block eviction (evict) - Adaptive Replacement: + 3. Chunk eviction (evict) - Adaptive Replacement: Determines eviction source based on adaptive target: - If T1 size >= target_t1_size: Evict from T1, add to B1. - Otherwise: Evict from T2, add to B2. Finally, bound each ghost list size. - 4. Block insertion (insert): - New blocks are always inserted into T1 and removed from B1/B2 if - present. Blocks may later be promoted to T2 during touch operations. + 4. Chunk insertion (insert): + New chunks are always inserted into T1 and removed from B1/B2 if + present. Chunks may later be promoted to T2 during touch operations. Adaptive Behavior: The algorithm self-tunes the recency vs. frequency trade-off: @@ -50,19 +50,19 @@ class ARCCachePolicy(CachePolicy): def __init__(self, cache_capacity: int): super().__init__(cache_capacity) self.target_t1_size: float = 0.0 - self.t1: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() - self.t2: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() + self.t1: OrderedDict[OffloadKey, ChunkStatus] = OrderedDict() + self.t2: OrderedDict[OffloadKey, ChunkStatus] = OrderedDict() # key -> None (only care about presence) self.b1: OrderedDict[OffloadKey, None] = OrderedDict() self.b2: OrderedDict[OffloadKey, None] = OrderedDict() @override - def get(self, key: OffloadKey) -> BlockStatus | None: + def get(self, key: OffloadKey) -> ChunkStatus | None: return self.t1.get(key) or self.t2.get(key) @override - def insert(self, key: OffloadKey, block: BlockStatus) -> None: - self.t1[key] = block + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: + self.t1[key] = chunk self.b1.pop(key, None) self.b2.pop(key, None) @@ -75,13 +75,13 @@ def remove(self, key: OffloadKey) -> None: def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): if key in self.t1: - block = self.t1.pop(key) - if not block.is_ready: - # block was just prepared to be stored, not really touched + chunk = self.t1.pop(key) + if not chunk.is_ready: + # chunk was just prepared to be stored, not really touched # twice — keep it in T1 and mark as most recently used - self.t1[key] = block + self.t1[key] = chunk else: - self.t2[key] = block + self.t2[key] = chunk elif key in self.t2: self.t2.move_to_end(key) @@ -111,31 +111,31 @@ def clear(self) -> None: @override def evict( self, n: int, protected: set[OffloadKey] - ) -> list[tuple[OffloadKey, BlockStatus]] | None: + ) -> list[tuple[OffloadKey, ChunkStatus]] | None: if n == 0: return [] # Collect candidates atomically: simulate T1 size changes as we select, # but do not modify actual data structures until all n are found. candidates: list[ - tuple[OffloadKey, BlockStatus, bool] - ] = [] # (key, block, from_t1) + tuple[OffloadKey, ChunkStatus, bool] + ] = [] # (key, chunk, from_t1) virtual_t1_size = len(self.t1) # Keep the scans monotonic: restarting from the LRU end after every - # selection makes a batch eviction quadratic in the number of blocks. + # selection makes a batch eviction quadratic in the number of chunks. t1_iter = iter(self.t1.items()) t2_iter = iter(self.t2.items()) def next_candidate( - entries: Iterator[tuple[OffloadKey, BlockStatus]], - ) -> tuple[OffloadKey, BlockStatus] | None: - for key, block in entries: - if block.ref_cnt == 0 and key not in protected: - return key, block + entries: Iterator[tuple[OffloadKey, ChunkStatus]], + ) -> tuple[OffloadKey, ChunkStatus] | None: + for key, chunk in entries: + if chunk.ref_cnt == 0 and key not in protected: + return key, chunk return None for _ in range(n): - candidate: tuple[OffloadKey, BlockStatus, bool] | None = None + candidate: tuple[OffloadKey, ChunkStatus, bool] | None = None if virtual_t1_size >= int(self.target_t1_size): entry = next_candidate(t1_iter) @@ -158,15 +158,15 @@ def next_candidate( candidates.append(candidate) # Apply all evictions now that we know n candidates exist. - result: list[tuple[OffloadKey, BlockStatus]] = [] - for key, block, from_t1 in candidates: + result: list[tuple[OffloadKey, ChunkStatus]] = [] + for key, chunk, from_t1 in candidates: if from_t1: del self.t1[key] self.b1[key] = None else: del self.t2[key] self.b2[key] = None - result.append((key, block)) + result.append((key, chunk)) # Trim ghost lists to cache_capacity. for ghost in (self.b1, self.b2): diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index 908907d326f0..4059dd742f7c 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -7,36 +7,36 @@ from vllm.v1.kv_offload.base import OffloadKey, ReqContext -class BlockStatus(ctypes.Structure): +class ChunkStatus(ctypes.Structure): """ - Offloading status for a single block of KV data. + Offloading status for a single chunk of KV data. Holds the following information: - ref_cnt - the current number of transfers using this block as a source. - A value of -1 indicates the block is not yet ready to be read. - block_id - index of the physical CPU buffer slot. + ref_cnt - the current number of transfers using this chunk as a source. + A value of -1 indicates the chunk is not yet ready to be read. + chunk_id - index of the physical CPU buffer slot. """ - _fields_ = [("ref_cnt", ctypes.c_int32), ("block_id", ctypes.c_int64)] + _fields_ = [("ref_cnt", ctypes.c_int32), ("chunk_id", ctypes.c_int64)] - def __init__(self, block_id: int): + def __init__(self, chunk_id: int): super().__init__() - # initialize block as "not ready" (ref_cnt = -1) + # initialize chunk as "not ready" (ref_cnt = -1) self.ref_cnt = -1 - self.block_id = block_id + self.chunk_id = chunk_id @property def is_ready(self) -> bool: """ - Returns whether the block is ready to be read. + Returns whether the chunk is ready to be read. """ return self.ref_cnt >= 0 class CachePolicy(ABC): """ - Encapsulates both block organization (data structures) and replacement - decisions (which block to evict). LRU and ARC differ in both dimensions — + Encapsulates both chunk organization (data structures) and replacement + decisions (which chunk to evict). LRU and ARC differ in both dimensions — ARC's ghost lists and target_t1_size live at the intersection of storage and eviction, so they cannot be separated cleanly. """ @@ -45,35 +45,35 @@ def __init__(self, cache_capacity: int) -> None: self.cache_capacity = cache_capacity @abstractmethod - def get(self, key: OffloadKey) -> BlockStatus | None: - """Find block in data structures. Returns None if not present.""" + def get(self, key: OffloadKey) -> ChunkStatus | None: + """Find chunk in data structures. Returns None if not present.""" @abstractmethod - def insert(self, key: OffloadKey, block: BlockStatus) -> None: - """Add a newly allocated block. For ARC: also removes from ghost lists.""" + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: + """Add a newly allocated chunk. For ARC: also removes from ghost lists.""" @abstractmethod def remove(self, key: OffloadKey) -> None: - """Remove a block (used to clean up after a failed store).""" + """Remove a chunk (used to clean up after a failed store).""" @abstractmethod def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: """ - Mark blocks as recently used. + Mark chunks as recently used. Args: - keys: Blocks to mark as recently used. - req_context: Per-request context for the request touching these blocks. + keys: Chunks to mark as recently used. + req_context: Per-request context for the request touching these chunks. """ @abstractmethod def evict( self, n: int, protected: set[OffloadKey] - ) -> list[tuple[OffloadKey, BlockStatus]] | None: + ) -> list[tuple[OffloadKey, ChunkStatus]] | None: """ - Evict exactly n blocks, skipping any in protected. + Evict exactly n chunks, skipping any in protected. - Returns a list of (key, block) for the evicted blocks, + Returns a list of (key, chunk) for the evicted chunks, or None if n evictions cannot be satisfied. The operation is atomic: if None is returned, no state changes are made. @@ -84,15 +84,15 @@ def evict( @abstractmethod def clear(self) -> None: """ - Remove ALL blocks regardless of ref_cnt. + Remove ALL chunks regardless of ref_cnt. Ghost lists and adaptive state are also reset. """ def mark_evictable(self, key: OffloadKey) -> None: - """Called when a block's ref_cnt transitions to 0.""" + """Called when a chunk's ref_cnt transitions to 0.""" return def mark_non_evictable(self, key: OffloadKey) -> None: - """Called when a block's ref_cnt transitions from 0.""" + """Called when a chunk's ref_cnt transitions from 0.""" return diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index e8ccf0bdef5e..c1dc631d4fe9 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -6,7 +6,7 @@ from typing_extensions import override from vllm.v1.kv_offload.base import OffloadKey, ReqContext -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus class LRUCachePolicy(CachePolicy): @@ -20,71 +20,71 @@ class LRUCachePolicy(CachePolicy): def __init__(self, cache_capacity: int): super().__init__(cache_capacity) - # Blocks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU - self.evictable_blocks: OrderedDict[OffloadKey, None] = OrderedDict() - self.blocks: dict[OffloadKey, BlockStatus] = {} + # Chunks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU + self.evictable_chunks: OrderedDict[OffloadKey, None] = OrderedDict() + self.chunks: dict[OffloadKey, ChunkStatus] = {} @override - def get(self, key: OffloadKey) -> BlockStatus | None: - return self.blocks.get(key) + def get(self, key: OffloadKey) -> ChunkStatus | None: + return self.chunks.get(key) @override - def insert(self, key: OffloadKey, block: BlockStatus) -> None: - self.blocks[key] = block - if block.ref_cnt == 0: - self.evictable_blocks[key] = None + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: + self.chunks[key] = chunk + if chunk.ref_cnt == 0: + self.evictable_chunks[key] = None @override def remove(self, key: OffloadKey) -> None: - del self.blocks[key] - self.evictable_blocks.pop(key, None) + del self.chunks[key] + self.evictable_chunks.pop(key, None) @override def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): - if key in self.evictable_blocks: - self.evictable_blocks.move_to_end(key) - # active blocks are untouched as they are non-evictable now. They - # will eventually reach the end of evictable_blocks when they finish. + if key in self.evictable_chunks: + self.evictable_chunks.move_to_end(key) + # active chunks are untouched as they are non-evictable now. They + # will eventually reach the end of evictable_chunks when they finish. @override def clear(self) -> None: - self.evictable_blocks.clear() - self.blocks.clear() + self.evictable_chunks.clear() + self.chunks.clear() @override def evict( self, n: int, protected: set[OffloadKey] - ) -> list[tuple[OffloadKey, BlockStatus]] | None: + ) -> list[tuple[OffloadKey, ChunkStatus]] | None: if n == 0: return [] - candidates: list[tuple[OffloadKey, BlockStatus]] = [] - for key, _ in self.evictable_blocks.items(): + candidates: list[tuple[OffloadKey, ChunkStatus]] = [] + for key, _ in self.evictable_chunks.items(): if key in protected: continue - block = self.blocks[key] - assert block.ref_cnt == 0 - candidates.append((key, block)) + chunk = self.chunks[key] + assert chunk.ref_cnt == 0 + candidates.append((key, chunk)) if len(candidates) == n: break if len(candidates) < n: return None for key, _ in candidates: - del self.evictable_blocks[key] - del self.blocks[key] + del self.evictable_chunks[key] + del self.chunks[key] return candidates @override def mark_evictable(self, key: OffloadKey) -> None: - # blocks can become evictable when, + # chunks can become evictable when, # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list # all loads complete - i.e ref_cnt 1 -> 0 # not in evictable list - self.evictable_blocks[key] = None + self.evictable_chunks[key] = None @override def mark_non_evictable(self, key: OffloadKey) -> None: # key must have been in the evictable list. - del self.evictable_blocks[key] + del self.evictable_chunks[key] diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 33639aab4594..1a5dde21e62d 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -82,24 +82,24 @@ class SharedOffloadRegion: def __init__( self, engine_id: str, - num_blocks: int, + num_chunks: int, rank: int | None, - kv_bytes_per_block: int, + kv_bytes_per_chunk: int, cpu_page_size: int, barrier: Callable[[], None] | None = None, ) -> None: self.page_size = mmap.PAGESIZE - assert kv_bytes_per_block % self.page_size == 0 + assert kv_bytes_per_chunk % self.page_size == 0 - self.num_blocks = num_blocks - self._row_stride = kv_bytes_per_block - self.total_size_bytes = self.num_blocks * self._row_stride + self.num_chunks = num_chunks + self._row_stride = kv_bytes_per_chunk + self.total_size_bytes = self.num_chunks * self._row_stride self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap" self._creator = False # set True only if this worker creates the file self.rank = rank if rank is not None: - # byte offset to this worker's first slot within each block row + # byte offset to this worker's first slot within each chunk row self._worker_offset = rank * cpu_page_size # exclusive upper bound for this worker's area within each row self._worker_area_end = (rank + 1) * cpu_page_size @@ -182,19 +182,19 @@ def __init__( populate_write_fn = _get_populate_write_fn(self.mmap_obj) if rank is not None: - # Populate only this worker's pages (one slot per block row). + # Populate only this worker's pages (one slot per chunk row). worker_offset = rank * cpu_page_size _t0 = time.perf_counter() page_size = self.page_size - for block in range(num_blocks): - raw_offset = block * self._row_stride + worker_offset + for chunk in range(num_chunks): + raw_offset = chunk * self._row_stride + worker_offset aligned_offset = (raw_offset // page_size) * page_size end = raw_offset + cpu_page_size aligned_length = end - aligned_offset populate_write_fn(self.mmap_obj, aligned_offset, aligned_length) logger.debug( - "MADV_POPULATE_WRITE loop: %d blocks in %.3f s", - num_blocks, + "MADV_POPULATE_WRITE loop: %d chunks in %.3f s", + num_chunks, time.perf_counter() - _t0, ) else: @@ -215,22 +215,22 @@ def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor: Must be called once per canonical tensor. The full mmap layout is: - worker0_block0 | worker1_block0 | ... | worker{M-1}_block0 - worker0_block1 | worker1_block1 | ... | worker{M-1}_block1 + worker0_chunk0 | worker1_chunk0 | ... | worker{M-1}_chunk0 + worker0_chunk1 | worker1_chunk1 | ... | worker{M-1}_chunk1 ... - Each worker_block cell is cpu_page_size bytes and holds all canonical - tensors for that worker and block concatenated: + Each worker_chunk cell is cpu_page_size bytes and holds all canonical + tensors for that worker and chunk concatenated: [ tensor0_data | tensor1_data | ... | tensor{L-1}_data ] Consecutive rows are separated by row_stride = cpu_page_size * M. - Returns an int8 tensor of shape (num_blocks, tensor_page_size) with stride + Returns an int8 tensor of shape (num_chunks, tensor_page_size) with stride (row_stride, 1). Using int8 keeps stride == bytes, so swap_blocks address arithmetic works without any dtype conversion. Args: - tensor_page_size: Bytes per block for this tensor. + tensor_page_size: Bytes per chunk for this tensor. """ assert self.rank is not None new_offset = self._worker_offset + tensor_page_size @@ -241,7 +241,7 @@ def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor: ) worker_layer_view = torch.as_strided( self._base, - size=(self.num_blocks, tensor_page_size), + size=(self.num_chunks, tensor_page_size), stride=(self._row_stride, 1), storage_offset=self._worker_offset, ) @@ -266,8 +266,8 @@ def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor: _canonical_offset=0, then advances by each tensor's size Each canonical_t{i} cell is that tensor's canonical page for the - block. Canonical areas are carved consecutively from the start of - each block row; consecutive rows are separated by row_stride. Every + chunk. Canonical areas are carved consecutively from the start of + each chunk row; consecutive rows are separated by row_stride. Every worker gets the identical byte ranges and writes only its disjoint bytes within them, as described by its canonical mappings — unlike create_next_worker_view, which gives each worker a private @@ -279,13 +279,13 @@ def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor: where one canonical copy replaces world_size worker copies. Args: - tensor_page_size: Canonical bytes per block for this tensor. + tensor_page_size: Canonical bytes per chunk for this tensor. """ new_offset = self._canonical_offset + tensor_page_size assert new_offset <= self._row_stride view = torch.as_strided( self._base, - size=(self.num_blocks, tensor_page_size), + size=(self.num_chunks, tensor_page_size), stride=(self._row_stride, 1), storage_offset=self._canonical_offset, ) @@ -296,10 +296,10 @@ def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor: def create_kv_memoryview(self) -> memoryview: """Return a zero-copy memoryview over the entire KV buffer. - Shape: (num_blocks, row_stride_bytes). Secondary tiers address - block *b* as ``view[b]``. + Shape: (num_chunks, row_stride_bytes). Secondary tiers address + chunk *b* as ``view[b]``. """ - kv_tensor = self._base.view(self.num_blocks, self._row_stride) + kv_tensor = self._base.view(self.num_chunks, self._row_stride) np_arr = kv_tensor.numpy() assert np_arr.ctypes.data == self._base.data_ptr(), ( "view()/numpy() created a copy instead of sharing the mmap buffer; " diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 0a1ac6d6a05d..00b36e863a07 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -73,7 +73,7 @@ def build_metric_definitions( ), CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata( documentation=( - "Histogram of the number of CPU blocks requested by each " + "Histogram of the number of CPU chunks requested by each " "KV offload prepare_store call." ), buckets=(1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144), @@ -101,7 +101,7 @@ def __init__(self, config: OffloadingConfig): ) world_size = config.parallel.world_size - self.num_blocks = 0 + self.num_chunks = 0 self.kv_bytes_per_chunk = 0 self.cpu_page_size_per_worker = 0 self.replicated_layout = config.replicated_layout and self._uses_shared_region() @@ -113,17 +113,17 @@ def __init__(self, config: OffloadingConfig): # calculate cpu_page_size_per_worker self.cpu_page_size_per_worker = kv_bytes_per_chunk // num_copies - # calculate num_blocks + # calculate num_chunks aligned_kv_bytes_per_chunk = round_up( kv_bytes_per_chunk, self.BLOCK_SIZE_ALIGNMENT ) - self.num_blocks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk + self.num_chunks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk # Expose aligned_kv_bytes_per_chunk as # kv_bytes_per_chunk. Note that this might contain - # some padding. i.e. each offloaded block is of the form, - # |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** | - # or |--- B0 (single copy) ---| *** maybe-pad *** | + # some padding. i.e. each offloaded chunk is of the form, + # |--- W0-C0---|---- W1-C0---| ... |---- Wn-C0---| *** maybe-pad *** | + # or |--- C0 (single copy) ---| *** maybe-pad *** | self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk # scheduler-side @@ -140,7 +140,7 @@ def __init__(self, config: OffloadingConfig): @override def get_manager(self) -> OffloadingManager: if not self._manager: - # store_threshold: how many times a block must be offered for + # store_threshold: how many times a chunk must be offered for # storage before it is eligible for CPU offloading. Values < 2 # disable filtering (a threshold of 1 equals no filter; 0 is the # default). @@ -150,7 +150,7 @@ def get_manager(self) -> OffloadingManager: max_tracker_size = int(self.extra_config.get("max_tracker_size", 64_000)) self._manager = CPUOffloadingManager( - num_blocks=self.num_blocks, + num_chunks=self.num_chunks, cache_policy=self.eviction_policy, cache_policy_module_path=self.cache_policy_module_path, enable_events=self.kv_events_config.enable_kv_cache_events, @@ -166,9 +166,9 @@ def _uses_shared_region(self) -> bool: def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: mmap_region: SharedOffloadRegion | None = None - # num_blocks == 0 would size the region to zero bytes, which cannot be + # num_chunks == 0 would size the region to zero bytes, which cannot be # mmap'd; fall back to the tensor path (empty tensors) as before. - if self._uses_shared_region() and self.num_blocks > 0: + if self._uses_shared_region() and self.num_chunks > 0: # Replicated layout puts all ranks on slot 0 (single MLA copy); # otherwise each rank takes its own slot by physical device index. if self.replicated_layout: @@ -178,9 +178,9 @@ def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: rank = torch.accelerator.current_device_index() % world_size mmap_region = SharedOffloadRegion( engine_id=self.config.engine_id, - num_blocks=self.num_blocks, + num_chunks=self.num_chunks, rank=rank, - kv_bytes_per_block=self.kv_bytes_per_chunk, + kv_bytes_per_chunk=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, barrier=_all_workers_barrier, ) @@ -188,7 +188,7 @@ def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: return CPUOffloadingWorker( kv_caches=kv_caches, blocks_per_chunk=self.blocks_per_chunk, - num_cpu_blocks=self.num_blocks, + num_cpu_chunks=self.num_chunks, mmap_region=mmap_region, ) except Exception: diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index c1e39811e341..34aa6ddf78ad 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -45,8 +45,8 @@ class TieringOffloadingMetrics: WRITE_TIME = "vllm:kv_offload_tiering_write_time" PROMOTION_JOB_FAILURES = "vllm:kv_offload_tiering_promotion_job_failures" CASCADE_JOB_FAILURES = "vllm:kv_offload_tiering_cascade_job_failures" - BLOCK_QUERIES = "vllm:kv_offload_tiering_chunk_queries" - BLOCK_HITS = "vllm:kv_offload_tiering_chunk_hits" + CHUNK_QUERIES = "vllm:kv_offload_tiering_chunk_queries" + CHUNK_HITS = "vllm:kv_offload_tiering_chunk_hits" PRIMARY_WRITE_USAGE_PERC = "vllm:kv_offload_tiering_primary_write_usage_perc" PRIMARY_READ_USAGE_PERC = "vllm:kv_offload_tiering_primary_read_usage_perc" PROMOTION_ALLOCATION_FAILURES = ( @@ -62,7 +62,7 @@ class TransferJob: job_id: JobId keys: Collection[OffloadKey] - block_ids: np.ndarray + chunk_ids: np.ndarray is_promotion: bool req_context: ReqContext @@ -90,14 +90,14 @@ class ParentManager(ABC): Required call sequence for each remote request: 1. on_new_request(req_context) — set up per-request state - 2. lookup(key, req_context) — check block availability - (repeat per block) - 3. create_store_job(keys, req_context) — pin blocks and get a + 2. lookup(key, req_context) — check chunk availability + (repeat per chunk) + 3. create_store_job(keys, req_context) — pin chunks and get a job handle 4. on_request_finished(req_context) — clean up per-request state Steps 2-3 may be interleaved. Step 4 must be called even if no - blocks were found, to avoid leaking async lookup state (e.g. in + chunks were found, to avoid leaking async lookup state (e.g. in the fs tier's AsyncLookupManager). """ @@ -155,23 +155,23 @@ def __init__( @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ - Check whether a block exists in this secondary tier. + Check whether a chunk exists in this secondary tier. Args: key: Offload key to look up. req_context: per-request context (e.g. kv_transfer_params). Returns: - HIT if the block is present and ready, + HIT if the chunk is present and ready, MISS if not found, - or RETRY if the block is being transferred (retry later). + or RETRY if the chunk is being transferred (retry later). """ pass @abstractmethod def submit_store(self, job_metadata: TransferJob) -> None: """ - Submit an async job to store blocks from the primary tier to this + Submit an async job to store chunks from the primary tier to this secondary tier. This method must be lightweight and non-blocking: allocate metadata @@ -179,19 +179,19 @@ def submit_store(self, job_metadata: TransferJob) -> None: calling thread. Preconditions (guaranteed by the framework): - - ``job_metadata.block_ids`` are valid primary-tier slots, pinned + - ``job_metadata.chunk_ids`` are valid primary-tier slots, pinned (ref-counted) for the duration of the transfer. The implementation is responsible for: - 1. Filtering out blocks already present in this tier - 2. Evicting blocks if capacity is needed + 1. Filtering out chunks already present in this tier + 2. Evicting chunks if capacity is needed 3. Allocating space in this tier - 4. Submitting the async transfer (read from primary via block_ids) + 4. Submitting the async transfer (read from primary via chunk_ids) Report completion via ``get_finished_jobs()``. Args: - job_metadata: Job metadata including job_id, keys, and block_ids + job_metadata: Job metadata including job_id, keys, and chunk_ids identifying the primary-tier slots to read from. """ pass @@ -199,24 +199,24 @@ def submit_store(self, job_metadata: TransferJob) -> None: @abstractmethod def submit_load(self, job_metadata: TransferJob) -> None: """ - Submit an async job to load blocks from this secondary tier to the + Submit an async job to load chunks from this secondary tier to the primary tier. - This method must be lightweight and non-blocking: mark blocks as + This method must be lightweight and non-blocking: mark chunks as in-flight and submit the transfer, but do NOT perform the data copy on the calling thread. Preconditions (guaranteed by the framework): - - ``job_metadata.block_ids`` are allocated primary-tier slots + - ``job_metadata.chunk_ids`` are allocated primary-tier slots ready to receive data. The implementation must copy data from this tier into the - primary-tier slots identified by ``block_ids``. + primary-tier slots identified by ``chunk_ids``. Report completion via ``get_finished_jobs()``. Args: - job_metadata: Job metadata including job_id, keys, and block_ids + job_metadata: Job metadata including job_id, keys, and chunk_ids identifying the primary-tier slots to write into. """ pass @@ -249,7 +249,7 @@ def take_events(self) -> Iterable[OffloadingEvent]: def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ - Mark blocks as recently used for eviction policy. + Mark chunks as recently used for eviction policy. Args: keys: Offload keys to mark as recently used. @@ -263,7 +263,7 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: Called when a new request is first seen by the scheduler. Returns a RequestOffloadingContext expressing this tier's preference - for how blocks should be offloaded for this request. + for how chunks should be offloaded for this request. Args: req_context: Per-request context. diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index b2652a6d555e..967f7cc49826 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -3,7 +3,7 @@ """ ExampleSecondaryTierManager: A simple in-memory secondary tier. -This implementation provides a minimal secondary tier that stores blocks +This implementation provides a minimal secondary tier that stores chunks in memory (using a dictionary) with immediate completion. It serves as a reference for writing new tiers and is useful for testing the TieringOffloadingManager without requiring actual storage or network backends. @@ -39,7 +39,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): A simple in-memory secondary tier. This implementation: - - Stores blocks in a dictionary (key -> True) + - Stores chunks in a dictionary (key -> True) - Completes transfers immediately (synchronous) """ @@ -69,45 +69,45 @@ def __init__( ) # key -> True (only care about presence) - self.blocks: dict[OffloadKey, bool] = {} + self.chunks: dict[OffloadKey, bool] = {} # Completed jobs waiting to be retrieved by get_finished_jobs() self.completed_jobs: list[JobResult] = [] assert primary_kv_view.strides is not None - self._block_size = primary_kv_view.strides[0] + self._chunk_size = primary_kv_view.strides[0] @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ - Check whether a block exists in this secondary tier. + Check whether a chunk exists in this secondary tier. Args: key: Offload key to look up. req_context: Per-request context. Returns: - HIT if the block is present, MISS if not found. + HIT if the chunk is present, MISS if not found. """ - return LookupResult.HIT if key in self.blocks else LookupResult.MISS + return LookupResult.HIT if key in self.chunks else LookupResult.MISS @override def submit_store(self, job_metadata: TransferJob) -> None: """ - Submit a job to store blocks from primary tier to this tier. + Submit a job to store chunks from primary tier to this tier. Args: job_metadata: Job metadata including job_id, keys, and - spec for reading blocks from the primary tier. + spec for reading chunks from the primary tier. """ keys = job_metadata.keys - block_ids = job_metadata.block_ids + chunk_ids = job_metadata.chunk_ids - assert len(keys) == len(block_ids), ( - f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" + assert len(keys) == len(chunk_ids), ( + f"Length mismatch: {len(keys)} keys but {len(chunk_ids)} chunk_ids" ) for key in keys: - self.blocks[key] = True + self.chunks[key] = True self.completed_jobs.append( JobResult( job_id=job_metadata.job_id, @@ -119,21 +119,21 @@ def submit_store(self, job_metadata: TransferJob) -> None: @override def submit_load(self, job_metadata: TransferJob) -> None: """ - Submit a job to load blocks from this tier to primary tier. + Submit a job to load chunks from this tier to primary tier. Args: job_metadata: Job metadata including job_id, keys, and - spec for writing blocks into the primary tier. + spec for writing chunks into the primary tier. """ keys = job_metadata.keys - block_ids = job_metadata.block_ids + chunk_ids = job_metadata.chunk_ids - assert len(keys) == len(block_ids), ( - f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids" + assert len(keys) == len(chunk_ids), ( + f"Length mismatch: {len(keys)} keys but {len(chunk_ids)} chunk_ids" ) for key in keys: - if key not in self.blocks: + if key not in self.chunks: self.completed_jobs.append( JobResult(job_id=job_metadata.job_id, success=False) ) @@ -170,6 +170,6 @@ def drain_jobs(self) -> None: completes, so there is nothing to wait for.""" return - def get_num_blocks(self) -> int: - """Get the number of blocks currently stored in this tier.""" - return len(self.blocks) + def get_num_chunks(self) -> int: + """Get the number of chunks currently stored in this tier.""" + return len(self.chunks) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 6de2d074605a..f73f70e77ffb 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -223,7 +223,7 @@ def submit_store(self, job_metadata: TransferJob) -> None: batch_store_block, [self.file_mapper.get_file_name(key) for key in keys], self._primary_kv_view, - [int(bid) * self._block_size for bid in job_metadata.block_ids], + [int(cid) * self._block_size for cid in job_metadata.chunk_ids], self._block_size, self._use_o_direct, ) @@ -237,7 +237,7 @@ def submit_load(self, job_metadata: TransferJob) -> None: keys = list(job_metadata.keys) self._load_job_keys[job_id] = keys paths = [self.file_mapper.get_file_name(key) for key in keys] - offsets = [int(bid) * self._block_size for bid in job_metadata.block_ids] + offsets = [int(cid) * self._block_size for cid in job_metadata.chunk_ids] def load_task() -> None: try: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index d2499859f689..0940aaeae07b 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -8,16 +8,16 @@ hierarchical KV cache offloading. Key Design Principles: -1. Always offload to all tiers — When a block is stored to the primary tier, +1. Always offload to all tiers — When a chunk is stored to the primary tier, it is cascaded to ALL secondary tiers 2. Primary tier is the gateway — Secondary tiers cannot access GPU memory directly; all data flows through the CPU primary tier -3. Staged promotion — Blocks in secondary tiers must be promoted to the +3. Staged promotion — Chunks in secondary tiers must be promoted to the primary tier before GPU can access them 4. Transparent retry mechanism — Return None from lookup() to signal "data is being promoted, try later" 5. ref_cnt as eviction protection — primary.prepare_read() increments ref_cnt, - protecting blocks from eviction until complete_read() is called + protecting chunks from eviction until complete_read() is called """ import time @@ -61,11 +61,11 @@ @dataclass class PendingPromotion: - """Accumulator for blocks awaiting submit_load() for one (tier, request).""" + """Accumulator for chunks awaiting submit_load() for one (tier, request).""" req_context: ReqContext keys: list[OffloadKey] = field(default_factory=list) - block_ids: list[int] = field(default_factory=list) + chunk_ids: list[int] = field(default_factory=list) @dataclass(slots=True) @@ -94,14 +94,14 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): def __init__( self, - num_blocks: int, + num_chunks: int, mmap_region: SharedOffloadRegion, cache_policy: str = "lru", cache_policy_module_path: str | None = None, enable_events: bool = False, ): super().__init__( - num_blocks=num_blocks, + num_chunks=num_chunks, cache_policy=cache_policy, cache_policy_module_path=cache_policy_module_path, enable_events=enable_events, @@ -120,9 +120,9 @@ def __init__( def get_kv_memoryview(self) -> memoryview: """Return the memoryview over the primary tier's KV cache buffer. - The view has shape (num_blocks, row_stride_bytes) and is backed by the - SharedOffloadRegion mmap. Secondary tiers address block *b* as - ``view[b]``. + The view has shape (num_chunks, row_stride_bytes) and is backed by the + SharedOffloadRegion mmap. Secondary tiers address chunk *c* as + ``view[c]``. """ return self._kv_memoryview @@ -206,14 +206,14 @@ def __init__( assert primary_view.strides is not None self._metrics = TieringMetricsTracker( tier_types=[tier.tier_type for tier in self.secondary_tiers], - num_primary_blocks=self.primary_tier._num_blocks, - primary_block_size=primary_view.strides[0], + num_primary_chunks=self.primary_tier._num_chunks, + primary_chunk_size=primary_view.strides[0], ) # Pending promotion requests accumulated during lookup() calls; flushed # as one batched submit_load() per (tier, request) in on_schedule_end(). # Outer key: tier index. Inner key: req_context.req_id — the same ReqContext - # object is reused for all block lookups of a given request per engine step. + # object is reused for all chunk lookups of a given request per engine step. self._pending_load_submissions: dict[int, dict[str, PendingPromotion]] = {} # Gate for once-per-step execution of _maybe_process_finished_jobs(). @@ -305,7 +305,7 @@ def _process_finished_jobs(self): 2. For completed stores (primary→secondary): calls primary.complete_read() to decrement ref_cnt 3. For completed loads (secondary→primary): calls primary.complete_write() - to make blocks available + to make chunks available """ for i, tier in enumerate(self.secondary_tiers): for completed_job in tier.get_finished_jobs(): @@ -324,11 +324,11 @@ def _process_finished_jobs(self): if transfer_job.is_promotion: # secondary→primary transfer (promotion) completed. - # Make blocks available in primary tier. + # Make chunks available in primary tier. self._complete_promotion(job_metadata, completed_job) else: # primary→secondary transfer completed. - # Decrement ref_cnt on primary blocks. + # Decrement ref_cnt on primary chunks. self.primary_tier.complete_read( transfer_job.keys, transfer_job.req_context ) @@ -342,7 +342,7 @@ def lookup( exclude_tier_idx: int | None = None, ) -> LookupResult: """ - Check whether a single block is offloaded and ready. + Check whether a single chunk is offloaded and ready. Algorithm: 1. Process any completed async jobs first. @@ -351,20 +351,20 @@ def lookup( hit and initiate promotion. Args: - key: Block hash to look up. + key: Chunk hash to look up. req_context: Per-request context. Returns: - HIT — block is ready in the primary tier. - HIT_PENDING — block found but not yet readable (write + HIT — chunk is ready in the primary tier. + HIT_PENDING — chunk found but not yet readable (write in-flight on the primary tier). RETRY — promotion started or a secondary tier is busy. - MISS — block not found in any tier, or primary is full + MISS — chunk not found in any tier, or primary is full and cannot accept a promotion. """ # Poll first so a promotion that finished since the last call is # already reflected as HIT (not stale HIT_PENDING/MISS) below, and - # so blocks freed by cascade or promotion completions are evictable + # so chunks freed by cascade or promotion completions are evictable # in time for a promotion this lookup may initiate. self._maybe_process_finished_jobs() @@ -424,30 +424,30 @@ def _initiate_promotion( req_context: ReqContext, ) -> bool: """ - Queue a block for promotion from a secondary tier to the primary tier. + Queue a chunk for promotion from a secondary tier to the primary tier. Allocates space in the primary tier immediately (sets ref_cnt=-1 so subsequent lookups within the same step see the slot as in-flight), then defers the actual submit_load() call to _flush_pending_promotions() - so all blocks queued during one engine step are submitted as a single + so all chunks queued during one engine step are submitted as a single batched job. Args: tier_idx: The secondary tier index to promote from - key: Block to promote + key: Chunk to promote req_context: Per-request context forwarded to primary.prepare_write(). Returns: True if promotion was initiated, False if primary tier is full. """ - # Allocate space in primary tier for promoted block. + # Allocate space in primary tier for promoted chunk. # Must happen immediately so primary.lookup() returns None (in-flight) # for this key on any subsequent lookup() call within the same step, # preventing duplicate promotion attempts. primary_write_result = self.primary_tier.prepare_write([key], req_context) if primary_write_result is None: - # Primary tier is full; caller should treat the block as unavailable + # Primary tier is full; caller should treat the chunk as unavailable # rather than retrying indefinitely. self._metrics.on_promotion_allocation_failure() return False @@ -455,16 +455,16 @@ def _initiate_promotion( store_spec = primary_write_result.store_spec assert isinstance(store_spec, CPULoadStoreSpec) # Defer submit_load to on_schedule_end(). Group by (tier, request) so - # each request's blocks are submitted as one batched job per tier. + # each request's chunks are submitted as one batched job per tier. tier_pending = self._pending_load_submissions.setdefault(tier_idx, {}) ctx_id = req_context.req_id if ctx_id not in tier_pending: tier_pending[ctx_id] = PendingPromotion( - keys=[], block_ids=[], req_context=req_context + keys=[], chunk_ids=[], req_context=req_context ) entry = tier_pending[ctx_id] entry.keys.extend(primary_write_result.keys_to_store) - entry.block_ids.extend(store_spec.block_ids) + entry.chunk_ids.extend(store_spec.chunk_ids) return True def _flush_pending_promotions(self) -> None: @@ -483,7 +483,7 @@ def _flush_pending_promotions(self) -> None: job_metadata = TransferJob( job_id=job_id, keys=entry.keys, - block_ids=np.array(entry.block_ids, dtype=np.int32), + chunk_ids=np.array(entry.chunk_ids, dtype=np.int32), is_promotion=True, req_context=entry.req_context, ) @@ -497,16 +497,16 @@ def prepare_load( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> LoadStoreSpec: """ - Prepare blocks to be loaded from primary tier to GPU. + Prepare chunks to be loaded from primary tier to GPU. Callers only pass keys already confirmed HIT by lookup() earlier this step. - This increments ref_cnt on the blocks in the primary tier, protecting + This increments ref_cnt on the chunks in the primary tier, protecting them from eviction during the transfer. Args: - keys: Blocks to prepare for loading. + keys: Chunks to prepare for loading. req_context: Per-request context. Returns: @@ -517,10 +517,10 @@ def prepare_load( @override def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ - Mark blocks as recently used in all tiers. + Mark chunks as recently used in all tiers. Args: - keys: Blocks to mark as recently used. + keys: Chunks to mark as recently used. req_context: Per-request context. """ self.primary_tier.touch(keys, req_context) @@ -530,13 +530,13 @@ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): @override def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): """ - Mark blocks as done loading from primary tier to GPU. + Mark chunks as done loading from primary tier to GPU. - This decrements ref_cnt on the blocks in the primary tier, allowing + This decrements ref_cnt on the chunks in the primary tier, allowing them to be evicted again. Args: - keys: Blocks that finished loading. + keys: Chunks that finished loading. req_context: Per-request context. """ self.primary_tier.complete_load(keys, req_context) @@ -546,37 +546,37 @@ def prepare_store( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> PrepareStoreOutput | None: """ - Prepare blocks to be stored from GPU to primary tier. + Prepare chunks to be stored from GPU to primary tier. CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure that any completed async transfers have their ref_cnt decremented before the primary tier makes eviction decisions. - For request-level tiers, blocks already present in the primary tier + For request-level tiers, chunks already present in the primary tier are immediately cascaded via submit_store(). Args: - keys: Blocks to prepare for storing. + keys: Chunks to prepare for storing. req_context: Per-request context. Returns: - PrepareStoreOutput describing where to store blocks and what was + PrepareStoreOutput describing where to store chunks and what was evicted, or None if store cannot proceed. """ # Step 1: Poll for completed async jobs FIRST # _process_finished_jobs() handles two kinds of completions here: # - Cascade completions (store to a secondary tier, either a local # cascade or a store job created for a remote requester via - # create_store_job()): decrements ref_cnt on the primary blocks + # create_store_job()): decrements ref_cnt on the primary chunks # that were read, making them evictable again once ref_cnt hits 0. # - Promotion completions (secondary->primary loads): sets a - # not-yet-ready block's ref_cnt from -1 to 0 via complete_write(), + # not-yet-ready chunk's ref_cnt from -1 to 0 via complete_write(), # making it evictable for the first time. # Both must be accounted for before the eviction decision below. self._maybe_process_finished_jobs() - # Step 2: Store to primary tier (new blocks only). - # Cascading of these newly-stored blocks to ALL secondary tiers + # Step 2: Store to primary tier (new chunks only). + # Cascading of these newly-stored chunks to ALL secondary tiers # happens later in complete_store(), after the GPU→Primary transfer # completes. primary_result = self.primary_tier.prepare_store(keys, req_context) @@ -588,7 +588,7 @@ def prepare_store( state = self._req_state[req_context.req_id] state.pending_primary_stores += 1 - # Step 3: For request-level tiers, cascade blocks already in primary + # Step 3: For request-level tiers, cascade chunks already in primary request_level_tiers = self._req_state[req_context.req_id].request_level_tiers if request_level_tiers: keys_to_store_set = set(primary_result.keys_to_store) @@ -596,13 +596,13 @@ def prepare_store( k for k in keys if k not in keys_to_store_set ) if keys_already_in_primary: - self._cascade_existing_blocks_to_request_level_tiers( + self._cascade_existing_chunks_to_request_level_tiers( keys_already_in_primary, req_context, request_level_tiers ) return primary_result - def _cascade_existing_blocks_to_request_level_tiers( + def _cascade_existing_chunks_to_request_level_tiers( self, keys: Sequence[OffloadKey], req_context: ReqContext, @@ -610,7 +610,7 @@ def _cascade_existing_blocks_to_request_level_tiers( ) -> None: """ For tiers that requested request-level policy, submit_store() for - blocks that are already present in the primary tier. + chunks that are already present in the primary tier. A key whose primary write is still in flight (HIT_PENDING) cannot be dropped: prepare_store already excluded it as present, and the @@ -653,7 +653,7 @@ def _flush_pending_cascades(self) -> None: continue assert state.request_level_tiers keys, state.pending_cascade_keys = state.pending_cascade_keys, [] - self._cascade_existing_blocks_to_request_level_tiers( + self._cascade_existing_chunks_to_request_level_tiers( keys, state.req_context, state.request_level_tiers ) self._maybe_finalize_request(req_id) @@ -666,30 +666,30 @@ def complete_store( success: bool = True, ) -> None: """ - Mark blocks as done storing from GPU to primary tier. + Mark chunks as done storing from GPU to primary tier. - This is where secondary tier cascading happens — after blocks are + This is where secondary tier cascading happens — after chunks are confirmed to be in the primary tier, they are cascaded to ALL secondary tiers. For each secondary tier: 1. Call primary.prepare_read() to get LoadStoreSpec AND increment - ref_cnt (protecting blocks during async transfer) + ref_cnt (protecting chunks during async transfer) 2. Call tier.submit_store() to start async transfer: primary→secondary 3. Track the job in _store_jobs dictionary Args: - keys: Blocks that finished storing. + keys: Chunks that finished storing. success: Whether the GPU→primary transfer succeeded. req_context: Per-request context forwarded to primary.prepare_read(). """ - # Step 1: Complete store in primary tier (makes blocks loadable) + # Step 1: Complete store in primary tier (makes chunks loadable) self.primary_tier.complete_store(keys, req_context, success) if success: # Step 2: Cascade to ALL secondary tiers # For each secondary tier, call primary.prepare_read() to get the - # LoadStoreSpec AND to increment ref_cnt (protecting blocks from + # LoadStoreSpec AND to increment ref_cnt (protecting chunks from # eviction during the async transfer). One prepare_read() call per # secondary tier. for tier_idx, tier in enumerate(self.secondary_tiers): @@ -710,22 +710,22 @@ def create_store_job( req_context: ReqContext, tier_idx: int = 0, ) -> TransferJob: - """Pin blocks in the primary tier and create a tracked store job. + """Pin chunks in the primary tier and create a tracked store job. - Calls prepare_read() to increment ref_cnt (protecting blocks + Calls prepare_read() to increment ref_cnt (protecting chunks from eviction during the async transfer), allocates a job ID, and registers the job in _jobs. The caller is responsible for the actual data transfer and reporting completion via get_finished_jobs(). """ - primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context) - assert isinstance(primary_blocks_spec, CPULoadStoreSpec) + primary_chunks_spec = self.primary_tier.prepare_read(keys, req_context) + assert isinstance(primary_chunks_spec, CPULoadStoreSpec) job_id = self._next_job_id() job_metadata = TransferJob( job_id=job_id, keys=keys, - block_ids=primary_blocks_spec.block_ids, + chunk_ids=primary_chunks_spec.chunk_ids, is_promotion=False, req_context=req_context, ) @@ -760,7 +760,7 @@ def on_new_request( policy = ( OffloadPolicy.REQUEST_LEVEL if state.request_level_tiers - else OffloadPolicy.BLOCK_LEVEL + else OffloadPolicy.CHUNK_LEVEL ) return RequestOffloadingContext(policy=policy) diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 60907142ac38..55d56999398f 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -39,20 +39,20 @@ class _RequestMetricsState: class _TierState: active_promotion_count: int = 0 active_cascade_count: int = 0 - primary_write_block_count: int = 0 - primary_read_block_count: int = 0 + primary_write_chunk_count: int = 0 + primary_read_chunk_count: int = 0 class TieringMetricsTracker: def __init__( self, tier_types: list[str], - num_primary_blocks: int, - primary_block_size: int, + num_primary_chunks: int, + primary_chunk_size: int, ) -> None: self._tier_types = tier_types - self._num_primary_blocks = num_primary_blocks - self._primary_block_size = primary_block_size + self._num_primary_chunks = num_primary_chunks + self._primary_chunk_size = primary_chunk_size self._request_states: dict[str, _RequestMetricsState] = {} self._tier_states = [_TierState() for _ in tier_types] self._stats = OffloadingConnectorStats() @@ -111,13 +111,13 @@ def on_lookup( def on_job_registered(self, job_metadata: _JobMetadataLike) -> None: transfer_job = job_metadata.transfer_job state = self._tier_states[job_metadata.tier_idx] - block_count = len(transfer_job.block_ids) + chunk_count = len(transfer_job.chunk_ids) if transfer_job.is_promotion: state.active_promotion_count += 1 - state.primary_write_block_count += block_count + state.primary_write_chunk_count += chunk_count else: state.active_cascade_count += 1 - state.primary_read_block_count += block_count + state.primary_read_chunk_count += chunk_count def on_job_finished( self, job_metadata: _JobMetadataLike, result: JobResult @@ -149,25 +149,25 @@ def assert_idle(self) -> None: assert all( state.active_promotion_count == 0 and state.active_cascade_count == 0 - and state.primary_write_block_count == 0 - and state.primary_read_block_count == 0 + and state.primary_write_chunk_count == 0 + and state.primary_read_chunk_count == 0 for state in self._tier_states ) def _decrement_tier_state(self, job_metadata: _JobMetadataLike) -> None: transfer_job = job_metadata.transfer_job state = self._tier_states[job_metadata.tier_idx] - block_count = len(transfer_job.block_ids) + chunk_count = len(transfer_job.chunk_ids) if transfer_job.is_promotion: assert state.active_promotion_count > 0 state.active_promotion_count -= 1 - state.primary_write_block_count -= block_count - assert state.primary_write_block_count >= 0 + state.primary_write_chunk_count -= chunk_count + assert state.primary_write_chunk_count >= 0 else: assert state.active_cascade_count > 0 state.active_cascade_count -= 1 - state.primary_read_block_count -= block_count - assert state.primary_read_block_count >= 0 + state.primary_read_chunk_count -= chunk_count + assert state.primary_read_chunk_count >= 0 def _observe_finished_job_stats( self, @@ -199,7 +199,7 @@ def _observe_finished_job_stats( if transfer_job.is_promotion else TieringOffloadingMetrics.WRITE_TIME ) - transfer_size = completed_key_count * self._primary_block_size + transfer_size = completed_key_count * self._primary_chunk_size self._stats.increase_counter(bytes_metric, transfer_size, labelvalues) if completed_job.transfer_time is not None: self._stats.increase_counter( @@ -210,13 +210,13 @@ def _observe_active_transfer_stats(self, stats: OffloadingConnectorStats) -> Non for tier_idx, state in enumerate(self._tier_states): labelvalues = self.tier_label(tier_idx) write_usage = ( - state.primary_write_block_count / self._num_primary_blocks - if self._num_primary_blocks > 0 + state.primary_write_chunk_count / self._num_primary_chunks + if self._num_primary_chunks > 0 else 0.0 ) read_usage = ( - state.primary_read_block_count / self._num_primary_blocks - if self._num_primary_blocks > 0 + state.primary_read_chunk_count / self._num_primary_chunks + if self._num_primary_chunks > 0 else 0.0 ) stats.set_gauge( @@ -248,12 +248,12 @@ def _observe_resolved_lookup( async_start_time: float | None, ) -> None: self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_QUERIES, + TieringOffloadingMetrics.CHUNK_QUERIES, labelvalues=tier_label, ) if result is LookupResult.HIT: self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_HITS, + TieringOffloadingMetrics.CHUNK_HITS, labelvalues=tier_label, ) self._stats.observe_histogram( diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index ed3c5524593c..50a7e611f674 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -287,14 +287,14 @@ def submit_store(self, job_metadata: TransferJob) -> None: self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys) obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) self._submit_transfer( - job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE + job_metadata.job_id, job_metadata.chunk_ids, obj_keys, NIXL_WRITE ) def submit_load(self, job_metadata: TransferJob) -> None: self._load_job_keys[job_metadata.job_id] = list(job_metadata.keys) obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) self._submit_transfer( - job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_READ + job_metadata.job_id, job_metadata.chunk_ids, obj_keys, NIXL_READ ) def on_request_finished(self, req_context: ReqContext) -> None: diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 70ed3da4567f..2cd490171006 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -417,7 +417,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: def submit_store(self, job_metadata: TransferJob) -> None: job_id = job_metadata.job_id keys = list(job_metadata.keys) - block_ids = job_metadata.block_ids.tolist() + block_ids = job_metadata.chunk_ids.tolist() assert len(keys) == len(block_ids) @@ -476,7 +476,7 @@ def submit_store(self, job_metadata: TransferJob) -> None: def submit_load(self, job_metadata: TransferJob) -> None: job_id = job_metadata.job_id keys = list(job_metadata.keys) - block_ids = job_metadata.block_ids + block_ids = job_metadata.chunk_ids source = job_metadata.req_context.get_state(P2PSourceInfo) logger.debug( diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 7d24f3773fb3..c78918323266 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -593,7 +593,7 @@ def _pin_and_register_hits( self.add_stored_blocks( lookup.kv_request_id, list(meta.keys), - meta.block_ids.tolist(), + meta.chunk_ids.tolist(), meta.job_id, round_seq=lookup.round_seq, from_lookup=True, diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index bb0325ca2842..9922a4095c4e 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -8,7 +8,7 @@ Configuration via kv_connector_extra_config: - cpu_bytes_to_use: (required) Bytes to allocate for CPU primary tier - - block_size: (optional) Block size for offloaded blocks (default: GPU block size) + - block_size: (optional) Tokens per offloaded chunk (default: GPU block size) - eviction_policy: (optional) Primary tier eviction policy: built-in "lru"/ "arc", or the name of a policy registered via CachePolicyFactory, or an out-of-tree CachePolicy class name paired with cache_policy_module_path @@ -105,7 +105,7 @@ def build_metric_definitions( metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY] = ( OffloadingHistogramMetadata( documentation=( - "Histogram of blocking time spent in a per-block tier lookup " + "Histogram of blocking time spent in a per-chunk tier lookup " "that resolved as a hit or miss, labeled by tier, in seconds." ), labelnames=("tier",), @@ -127,7 +127,7 @@ def build_metric_definitions( metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY] = ( OffloadingHistogramMetadata( documentation=( - "Histogram of wall-clock time from a per-block tier lookup " + "Histogram of wall-clock time from a per-chunk tier lookup " "first returning retry until that same tier lookup resolves " "as a hit or miss, labeled by tier, in seconds." ), @@ -191,14 +191,14 @@ def build_metric_definitions( labelnames=("tier",), ) ) - metrics[TieringOffloadingMetrics.BLOCK_QUERIES] = OffloadingCounterMetadata( + metrics[TieringOffloadingMetrics.CHUNK_QUERIES] = OffloadingCounterMetadata( documentation=( - "Number of block lookup queries sent to a tier, labeled by tier." + "Number of chunk lookup queries sent to a tier, labeled by tier." ), labelnames=("tier",), ) - metrics[TieringOffloadingMetrics.BLOCK_HITS] = OffloadingCounterMetadata( - documentation="Number of block lookup hits in a tier, labeled by tier.", + metrics[TieringOffloadingMetrics.CHUNK_HITS] = OffloadingCounterMetadata( + documentation="Number of chunk lookup hits in a tier, labeled by tier.", labelnames=("tier",), ) metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES] = ( @@ -300,16 +300,16 @@ def get_manager(self) -> OffloadingManager: # primary tier can eagerly create a memoryview over _base. scheduler_mmap = SharedOffloadRegion( engine_id=self._engine_id, - num_blocks=self.num_blocks, + num_chunks=self.num_chunks, rank=None, - kv_bytes_per_block=self.kv_bytes_per_chunk, + kv_bytes_per_chunk=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, ) self._scheduler_mmap = scheduler_mmap # Create primary tier (CPU-based) primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=self.num_blocks, + num_chunks=self.num_chunks, cache_policy=self.eviction_policy, cache_policy_module_path=self.cache_policy_module_path, enable_events=self.kv_events_config.enable_kv_cache_events, @@ -367,9 +367,9 @@ def get_manager(self) -> OffloadingManager: logger.info( "Created TieringOffloadingManager with primary tier " - "(%s, %s blocks) and %s secondary tier(s)", + "(%s, %s chunks) and %s secondary tier(s)", self.eviction_policy, - self.num_blocks, + self.num_chunks, len(secondary_tiers), ) @@ -393,9 +393,9 @@ def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: rank = torch.accelerator.current_device_index() % world_size worker_mmap = SharedOffloadRegion( engine_id=self._engine_id, - num_blocks=self.num_blocks, + num_chunks=self.num_chunks, rank=rank, - kv_bytes_per_block=self.kv_bytes_per_chunk, + kv_bytes_per_chunk=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, ) try: @@ -404,7 +404,7 @@ def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: return CPUOffloadingWorker( kv_caches=kv_caches, blocks_per_chunk=self.blocks_per_chunk, - num_cpu_blocks=self.num_blocks, + num_cpu_chunks=self.num_chunks, mmap_region=worker_mmap, canonical_layout=self.config.canonical_layout, )