@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 138d137b5b

[Refactor][kv_offload]: rename `block`→`chunk` (#52615)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>

Ronen Schaffer committed Sep 9, 2026 at 17:45 UTC 138d137b5b955a5ebee98dfc946ecdd65d7b87ce
34 files changed +933 -916
docs/features/kv_offloading_usage.md
+4
@@ -22,6 +22,10 @@ flowchart LR
22 CPU <--> SN["..."]
23 ```
24
25 +## Terminology: Chunks
26 +
27 +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.
28 +
29 ## Single-Tier Setup (CPU Only)
30
31 ```bash
tests/v1/kv_connector/unit/offloading_connector/test_metrics.py
+1 -1
@@ -704,7 +704,7 @@ def test_prom_metrics_registers_tiering_metrics_from_spec():
704 )
705
706 metric = prom_metrics._offloading_metric_defs[
707 - TieringOffloadingMetrics.BLOCK_QUERIES
707 + TieringOffloadingMetrics.CHUNK_QUERIES
708 ]
709 assert metric.kwargs["labelnames"] == ["model_name", "engine", "tier"]
710
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+2 -2
@@ -96,7 +96,7 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
96 groups.append(
97 KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1)
98 )
99 - manager = CPUOffloadingManager(num_blocks=100)
99 + manager = CPUOffloadingManager(num_chunks=100)
100 spec = SimpleNamespace(
101 tokens_per_block=(256, 64, 8),
102 tokens_per_hash=8,
@@ -2888,7 +2888,7 @@ class TestEagle:
2888 req=req,
2889 req_context=ReqContext(req_id="test-req"),
2890 offloading_context=RequestOffloadingContext(
2891 - policy=OffloadPolicy.BLOCK_LEVEL
2891 + policy=OffloadPolicy.CHUNK_LEVEL
2892 ),
2893 num_locally_computed_tokens=num_computed_tokens,
2894 )
tests/v1/kv_offload/cpu/policies/test_factory.py
+6 -6
@@ -7,7 +7,7 @@ import pytest
7 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
8 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
9 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
10 -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
10 +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
11 from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
12 from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
13
@@ -20,10 +20,10 @@ class _DummyCachePolicy(CachePolicy):
20 def __init__(self, cache_capacity: int) -> None:
21 self.cache_capacity = cache_capacity
22
23 - def get(self, key: OffloadKey) -> BlockStatus | None:
23 + def get(self, key: OffloadKey) -> ChunkStatus | None:
24 return None
25
26 - def insert(self, key: OffloadKey, block: BlockStatus) -> None:
26 + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None:
27 pass
28
29 def remove(self, key: OffloadKey) -> None:
@@ -34,7 +34,7 @@ class _DummyCachePolicy(CachePolicy):
34
35 def evict(
36 self, n: int, protected: set[OffloadKey]
37 - ) -> list[tuple[OffloadKey, BlockStatus]] | None:
37 + ) -> list[tuple[OffloadKey, ChunkStatus]] | None:
38 return None
39
40 def clear(self) -> None:
@@ -72,7 +72,7 @@ class TestCachePolicyFactory:
72 policy_cls = CachePolicyFactory.get_cache_policy_cls("dummy")
73 assert policy_cls is _DummyCachePolicy
74
75 - manager = CPUOffloadingManager(num_blocks=4, cache_policy="dummy")
75 + manager = CPUOffloadingManager(num_chunks=4, cache_policy="dummy")
76 assert isinstance(manager._policy, _DummyCachePolicy)
77
78 def test_unregistered_policy_raises(self):
@@ -98,7 +98,7 @@ class TestCachePolicyFactory:
98 """End-to-end: CPUOffloadingManager resolves an unregistered policy
99 purely from cache_policy_module_path."""
100 manager = CPUOffloadingManager(
101 - num_blocks=4,
101 + num_chunks=4,
102 cache_policy="_DummyCachePolicy",
103 cache_policy_module_path="tests.v1.kv_offload.cpu.policies.test_factory",
104 )
tests/v1/kv_offload/cpu/test_canonical_layout.py
+2 -2
@@ -232,9 +232,9 @@ def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int):
232 def canonical_view(rank: int, world_size: int) -> torch.Tensor:
233 region = SharedOffloadRegion(
234 engine_id=engine_id,
235 - num_blocks=num_blocks,
235 + num_chunks=num_blocks,
236 rank=rank,
237 - kv_bytes_per_block=row_stride,
237 + kv_bytes_per_chunk=row_stride,
238 cpu_page_size=row_stride // world_size,
239 )
240 regions.append(region)
tests/v1/kv_offload/cpu/test_gpu_worker.py
+46 -46
@@ -26,7 +26,7 @@ from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker
26 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
27
28 NUM_GPU_BLOCKS = [64]
29 -NUM_CPU_BLOCKS = [256]
29 +NUM_CPU_CHUNKS = [256]
30 GPU_PAGE_SIZES = [512, 1024]
31 BLOCKS_PER_CHUNK_VALUES = [1, 3]
32 NUM_TENSORS = [4]
@@ -225,7 +225,7 @@ def test_worker_syncs_before_cleanup_after_handler_failure(
225 @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
226 @pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES)
227 @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
228 -@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
228 +@pytest.mark.parametrize("num_cpu_chunks", NUM_CPU_CHUNKS)
229 @pytest.mark.parametrize("num_tensors", NUM_TENSORS)
230 @pytest.mark.parametrize("seed", SEEDS)
231 @pytest.mark.parametrize("device", DEVICES)
@@ -241,7 +241,7 @@ def test_transfer(
241 gpu_page_size_bytes: int,
242 blocks_per_chunk: int,
243 num_gpu_blocks: int,
244 - num_cpu_blocks: int,
244 + num_cpu_chunks: int,
245 num_tensors: int,
246 seed: int,
247 device: str,
@@ -288,41 +288,41 @@ def test_transfer(
288 SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT,
289 )
290 simulated_world_size = 2
291 - kv_bytes_per_block = (
291 + kv_bytes_per_chunk = (
292 cpu_page_size if replicated_layout else cpu_page_size * simulated_world_size
293 )
294 mmap_region = SharedOffloadRegion(
295 engine_id=str(uuid.uuid4()),
296 - num_blocks=num_cpu_blocks,
296 + num_chunks=num_cpu_chunks,
297 rank=0,
298 - kv_bytes_per_block=kv_bytes_per_block,
298 + kv_bytes_per_chunk=kv_bytes_per_chunk,
299 cpu_page_size=cpu_page_size,
300 )
301
302 worker = CPUOffloadingWorker(
303 kv_caches=kv_caches,
304 blocks_per_chunk=blocks_per_chunk,
305 - num_cpu_blocks=num_cpu_blocks,
305 + num_cpu_chunks=num_cpu_chunks,
306 mmap_region=mmap_region,
307 )
308
309 # select block mappings
310 gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * blocks_per_chunk)
311 - cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
311 + cpu_chunks = random.sample(range(num_cpu_chunks), num_mappings)
312
313 - # expand cpu blocks to gpu-page granularity for uniform comparison:
314 - # each cpu block maps to blocks_per_chunk consecutive sub-blocks
315 - cpu_blocks_expanded = [
316 - cpu_block * blocks_per_chunk + j
317 - for cpu_block in cpu_blocks
313 + # expand cpu chunks to gpu-page granularity for uniform comparison:
314 + # each cpu chunk maps to blocks_per_chunk consecutive sub-blocks
315 + cpu_chunks_expanded = [
316 + cpu_chunk * blocks_per_chunk + j
317 + for cpu_chunk in cpu_chunks
318 for j in range(blocks_per_chunk)
319 ]
320
321 - # maybe skip some GPU blocks to test reading/writing from the middle of a CPU block
321 + # maybe skip some GPU blocks to test reading/writing from the middle of a CPU chunk
322 blocks_to_skip = blocks_per_chunk - 1
323 if blocks_to_skip > 0:
324 gpu_blocks = gpu_blocks[blocks_to_skip:]
325 - cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:]
325 + cpu_chunks_expanded = cpu_chunks_expanded[blocks_to_skip:]
326
327 # set transfer direction
328 if gpu_to_cpu:
@@ -330,16 +330,16 @@ def test_transfer(
330 src_spec = GPULoadStoreSpec(
331 gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,)
332 )
333 - dst_spec = CPULoadStoreSpec(cpu_blocks)
334 - dst_to_src = dict(zip(cpu_blocks_expanded, gpu_blocks))
333 + dst_spec = CPULoadStoreSpec(cpu_chunks)
334 + dst_to_src = dict(zip(cpu_chunks_expanded, gpu_blocks))
335 num_dst_sub_blocks = num_gpu_blocks
336 else:
337 handler = worker._load_handler
338 - src_spec = CPULoadStoreSpec(cpu_blocks)
338 + src_spec = CPULoadStoreSpec(cpu_chunks)
339 dst_spec = GPULoadStoreSpec(
340 gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,)
341 )
342 - dst_to_src = dict(zip(gpu_blocks, cpu_blocks_expanded))
342 + dst_to_src = dict(zip(gpu_blocks, cpu_chunks_expanded))
343 num_dst_sub_blocks = num_gpu_blocks
344
345 # randomize src and dst tensors before transfer
@@ -410,7 +410,7 @@ def test_transfer(
410 @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
411 @pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES)
412 @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
413 -@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
413 +@pytest.mark.parametrize("num_cpu_chunks", NUM_CPU_CHUNKS)
414 @pytest.mark.parametrize("seed", SEEDS)
415 @pytest.mark.parametrize("device", DEVICES)
416 @torch.inference_mode()
@@ -421,7 +421,7 @@ def test_transfer_multi_group(
421 gpu_page_size_bytes: int,
422 blocks_per_chunk: int,
423 num_gpu_blocks: int,
424 - num_cpu_blocks: int,
424 + num_cpu_chunks: int,
425 seed: int,
426 device: str,
427 ) -> None:
@@ -470,37 +470,37 @@ def test_transfer_multi_group(
470 worker = CPUOffloadingWorker(
471 kv_caches=canonical_kv_caches,
472 blocks_per_chunk=blocks_per_chunk,
473 - num_cpu_blocks=num_cpu_blocks,
473 + num_cpu_chunks=num_cpu_chunks,
474 )
475
476 # group 0: aligned, group 1: empty, group 2: unaligned on CPU->GPU
477 - group_sizes_in_cpu_blocks = [num_mappings_per_group, 0, num_mappings_per_group]
477 + group_sizes_in_cpu_chunks = [num_mappings_per_group, 0, num_mappings_per_group]
478
479 - total_cpu_blocks = sum(group_sizes_in_cpu_blocks)
480 - total_gpu_blocks_needed = total_cpu_blocks * blocks_per_chunk
479 + total_cpu_chunks = sum(group_sizes_in_cpu_chunks)
480 + total_gpu_blocks_needed = total_cpu_chunks * blocks_per_chunk
481 gpu_blocks_all = random.sample(range(num_gpu_blocks), total_gpu_blocks_needed)
482 - cpu_blocks_all = random.sample(range(num_cpu_blocks), total_cpu_blocks)
482 + cpu_chunks_all = random.sample(range(num_cpu_chunks), total_cpu_chunks)
483
484 - # split gpu/cpu blocks per group
484 + # split gpu blocks / cpu chunks per group
485 gpu_blocks_per_group: list[list[int]] = []
486 - cpu_blocks_per_group: list[list[int]] = []
486 + cpu_chunks_per_group: list[list[int]] = []
487 gpu_offset = 0
488 cpu_offset = 0
489 - for size in group_sizes_in_cpu_blocks:
489 + for size in group_sizes_in_cpu_chunks:
490 gpu_count = size * blocks_per_chunk
491 gpu_blocks_per_group.append(gpu_blocks_all[gpu_offset : gpu_offset + gpu_count])
492 - cpu_blocks_per_group.append(cpu_blocks_all[cpu_offset : cpu_offset + size])
492 + cpu_chunks_per_group.append(cpu_chunks_all[cpu_offset : cpu_offset + size])
493 gpu_offset += gpu_count
494 cpu_offset += size
495
496 - # expand cpu blocks to gpu-page granularity
497 - cpu_blocks_expanded_per_group = [
496 + # expand cpu chunks to gpu-page granularity
497 + cpu_chunks_expanded_per_group = [
498 [
499 - cpu_block * blocks_per_chunk + j
500 - for cpu_block in cpu_blocks
499 + cpu_chunk * blocks_per_chunk + j
500 + for cpu_chunk in cpu_chunks
501 for j in range(blocks_per_chunk)
502 ]
503 - for cpu_blocks in cpu_blocks_per_group
503 + for cpu_chunks in cpu_chunks_per_group
504 ]
505
506 # skip sub-blocks from group 2 to test unaligned transfers.
@@ -509,7 +509,7 @@ def test_transfer_multi_group(
509 gpu_blocks_per_group[2] = gpu_blocks_per_group[2][
510 sub_blocks_to_skip:-sub_blocks_to_skip
511 ]
512 - cpu_blocks_expanded_per_group[2] = cpu_blocks_expanded_per_group[2][
512 + cpu_chunks_expanded_per_group[2] = cpu_chunks_expanded_per_group[2][
513 sub_blocks_to_skip:-sub_blocks_to_skip
514 ]
515
@@ -520,10 +520,10 @@ def test_transfer_multi_group(
520 gpu_blocks.extend(gpu_blks)
521 group_sizes.append(len(gpu_blks))
522
523 - # build flat cpu_blocks list
524 - cpu_blocks = []
525 - for cpu_blks in cpu_blocks_per_group:
526 - cpu_blocks.extend(cpu_blks)
523 + # build flat cpu_chunks list
524 + cpu_chunks = []
525 + for cpu_chnks in cpu_chunks_per_group:
526 + cpu_chunks.extend(cpu_chnks)
527
528 # block_indices: only relevant for unaligned transfers
529 block_indices: list[int] = [0, 0, sub_blocks_to_skip]
@@ -533,18 +533,18 @@ def test_transfer_multi_group(
533 src_spec = GPULoadStoreSpec(
534 gpu_blocks, group_sizes=group_sizes, block_indices=block_indices
535 )
536 - dst_spec = CPULoadStoreSpec(cpu_blocks)
536 + dst_spec = CPULoadStoreSpec(cpu_chunks)
537 # per-group mapping: cpu sub-block -> gpu sub-block
538 dst_to_src_per_group = [
539 dict(zip(expanded, gpu_blks))
540 for expanded, gpu_blks in zip(
541 - cpu_blocks_expanded_per_group, gpu_blocks_per_group
541 + cpu_chunks_expanded_per_group, gpu_blocks_per_group
542 )
543 ]
544 - num_dst_sub_blocks = num_cpu_blocks * blocks_per_chunk
544 + num_dst_sub_blocks = num_cpu_chunks * blocks_per_chunk
545 else:
546 handler = worker._load_handler
547 - src_spec = CPULoadStoreSpec(cpu_blocks)
547 + src_spec = CPULoadStoreSpec(cpu_chunks)
548 dst_spec = GPULoadStoreSpec(
549 gpu_blocks, group_sizes=group_sizes, block_indices=block_indices
550 )
@@ -552,7 +552,7 @@ def test_transfer_multi_group(
552 dst_to_src_per_group = [
553 dict(zip(gpu_blks, expanded))
554 for gpu_blks, expanded in zip(
555 - gpu_blocks_per_group, cpu_blocks_expanded_per_group
555 + gpu_blocks_per_group, cpu_chunks_expanded_per_group
556 )
557 ]
558 num_dst_sub_blocks = num_gpu_blocks
@@ -644,7 +644,7 @@ def test_load_waits_for_pending_compute_stream_writes(default_vllm_config) -> No
644 ],
645 ),
646 blocks_per_chunk=1,
647 - num_cpu_blocks=num_blocks,
647 + num_cpu_chunks=num_blocks,
648 )
649 worker._load_handler.src_tensors[0].fill_(sentinel)
650 expected = torch.full((page_size_bytes,), sentinel, dtype=torch.int8)
tests/v1/kv_offload/cpu/test_manager.py
+144 -144
@@ -36,7 +36,7 @@ _EMPTY_REQ_CTX = make_req_context()
36
37
38 def make_cpu_manager(
39 - num_blocks: int = 4,
39 + num_chunks: int = 4,
40 cache_policy: str = "lru",
41 cache_policy_module_path: str | None = None,
42 enable_events: bool = False,
@@ -44,7 +44,7 @@ def make_cpu_manager(
44 max_tracker_size: int = 64_000,
45 ) -> CPUOffloadingManager:
46 return CPUOffloadingManager(
47 - num_blocks=num_blocks,
47 + num_chunks=num_chunks,
48 cache_policy=cache_policy,
49 cache_policy_module_path=cache_policy_module_path,
50 enable_events=enable_events,
@@ -56,7 +56,7 @@ def make_cpu_manager(
56 @dataclass
57 class ExpectedPrepareStoreOutput:
58 keys_to_store: list[int]
59 - store_block_ids: list[int]
59 + store_chunk_ids: list[int]
60 evicted_keys: list[int]
61
62
@@ -93,9 +93,9 @@ def verify_store_output(
93 store_spec = prepare_store_output.store_spec
94 assert isinstance(store_spec, CPULoadStoreSpec)
95 expected_array = np.array(
96 - expected_prepare_store_output.store_block_ids, dtype=np.int64
96 + expected_prepare_store_output.store_chunk_ids, dtype=np.int64
97 )
98 - assert np.array_equal(expected_array, store_spec.block_ids)
98 + assert np.array_equal(expected_array, store_spec.chunk_ids)
99
100
101 def verify_load_output(
@@ -103,7 +103,7 @@ def verify_load_output(
103 ):
104 assert isinstance(prepare_load_output, CPULoadStoreSpec)
105 expected_array = np.array(expected_prepare_load_output, dtype=np.int64)
106 - assert np.array_equal(expected_array, prepare_load_output.block_ids)
106 + assert np.array_equal(expected_array, prepare_load_output.chunk_ids)
107
108
109 def check_split_usage_stats(
@@ -146,7 +146,7 @@ def verify_events(
146
147 def test_cpu_eviction_removed_precedes_stored():
148 """An eviction is announced before the store that reuses its capacity."""
149 - manager = make_cpu_manager(num_blocks=2, enable_events=True)
149 + manager = make_cpu_manager(num_chunks=2, enable_events=True)
150
151 manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
152 manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
@@ -164,23 +164,23 @@ def test_cpu_eviction_removed_precedes_stored():
164
165
166 @pytest.mark.parametrize("eviction_policy", ["lru", "arc"])
167 -def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
167 +def test_already_stored_chunk_not_evicted_during_prepare_store(eviction_policy):
168 """
169 - Regression test: a block that is already stored must not be evicted
170 - by prepare_store() when it needs to make room for new blocks.
169 + Regression test: a chunk that is already stored must not be evicted
170 + by prepare_store() when it needs to make room for new chunks.
171 Applies to both lru and arc policies.
172
173 Scenario:
174 - - Store blocks [1, 2] and complete.
175 - - touch([1]) makes block 2 the LRU candidate.
174 + - Store chunks [1, 2] and complete.
175 + - touch([1]) makes chunk 2 the LRU candidate.
176 - prepare_store([2, 3, 4, 5]):
177 - * block 2 is filtered out as "already stored"
178 - * but without the fix, block 2 would be evicted as the LRU
177 + * chunk 2 is filtered out as "already stored"
178 + * but without the fix, chunk 2 would be evicted as the LRU
179 candidate to make room for [3, 4, 5]
180 - - After complete_store([2, 3, 4, 5]), block 2 must still be present.
180 + - After complete_store([2, 3, 4, 5]), chunk 2 must still be present.
181 """
182 manager = make_cpu_manager(
183 - num_blocks=4,
183 + num_chunks=4,
184 cache_policy=eviction_policy,
185 enable_events=True,
186 )
@@ -189,33 +189,33 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
189 manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
190 manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
191
192 - # touch [1] to make block 2 the LRU candidate
192 + # touch [1] to make chunk 2 the LRU candidate
193 manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
194
195 # prepare_store([2, 3, 4, 5]):
196 - # - block 2 is already stored -> filtered out of keys_to_store
197 - # - block 2 must NOT be evicted even though it is the LRU candidate
198 - # - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0
196 + # - chunk 2 is already stored -> filtered out of keys_to_store
197 + # - chunk 2 must NOT be evicted even though it is the LRU candidate
198 + # - chunk 1 (ID 0) is evicted instead; new chunks [3,4,5] get IDs 2,3,0
199 prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
200 verify_store_output(
201 prepare_store_output,
202 ExpectedPrepareStoreOutput(
203 keys_to_store=[3, 4, 5],
204 - store_block_ids=[2, 3, 0],
205 - evicted_keys=[1], # block 1 evicted, not block 2
204 + store_chunk_ids=[2, 3, 0],
205 + evicted_keys=[1], # chunk 1 evicted, not chunk 2
206 ),
207 )
208
209 - # complete_store must not silently drop block 2
209 + # complete_store must not silently drop chunk 2
210 manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
211
212 - # block 2 must still be present in the cache
212 + # chunk 2 must still be present in the cache
213 assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT
214
215
216 def test_filter_reused_manager_reports_stores_skipped_counter():
217 manager = make_cpu_manager(
218 - num_blocks=4,
218 + num_chunks=4,
219 cache_policy="lru",
220 store_threshold=2,
221 )
@@ -226,7 +226,7 @@ def test_filter_reused_manager_reports_stores_skipped_counter():
226 prepare_store_output,
227 ExpectedPrepareStoreOutput(
228 keys_to_store=[],
229 - store_block_ids=[],
229 + store_chunk_ids=[],
230 evicted_keys=[],
231 ),
232 )
@@ -247,34 +247,34 @@ def test_cpu_manager_reports_cache_usage_gauge():
247 ] == pytest.approx(value)
248
249 # Zero-capacity manager always reports 0.0
250 - manager = make_cpu_manager(num_blocks=0)
250 + manager = make_cpu_manager(num_chunks=0)
251 check_usage_stats(manager, 0.0)
252
253 - # Empty manager (4 blocks, none allocated): usage = 0.0
254 - manager = make_cpu_manager(num_blocks=4)
253 + # Empty manager (4 chunks, none allocated): usage = 0.0
254 + manager = make_cpu_manager(num_chunks=4)
255 check_usage_stats(manager, 0.0)
256
257 - # After allocating 2 of 4 blocks: usage = 0.5
257 + # After allocating 2 of 4 chunks: usage = 0.5
258 manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
259 check_usage_stats(manager, 0.5)
260
261 - # After filling all 4 blocks: usage = 1.0
261 + # After filling all 4 chunks: usage = 1.0
262 manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
263 check_usage_stats(manager, 1.0)
264
265 - # After completing store, the blocks becomes evictable as it is not actively used
265 + # After completing store, the chunks become evictable as not actively used
266 # and usage drops.
267 manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
268 check_usage_stats(manager, 0.5)
269
270 - # After completing store, the blocks becomes evictable as it is not actively used
270 + # After completing store, the chunks become evictable as not actively used
271 # and usage drops.
272 manager.complete_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
273 check_usage_stats(manager, 0.0)
274
275
276 def test_cpu_manager_reports_allocation_size_histogram():
277 - manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
277 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
278
279 manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
280 manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
@@ -297,12 +297,12 @@ def test_cpu_manager_reports_allocation_size_histogram():
297
298
299 def test_cpu_manager_reports_allocation_size_on_allocation_failure(monkeypatch):
300 - manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
300 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
301
302 - def fail_allocate_blocks(keys):
302 + def fail_allocate_chunks(keys):
303 raise RuntimeError("allocation failed")
304
305 - monkeypatch.setattr(manager, "_allocate_blocks", fail_allocate_blocks)
305 + monkeypatch.setattr(manager, "_allocate_chunks", fail_allocate_chunks)
306
307 with pytest.raises(RuntimeError, match="allocation failed"):
308 manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
@@ -316,7 +316,7 @@ def test_cpu_manager_reports_allocation_size_on_allocation_failure(monkeypatch):
316
317
318 def test_cpu_manager_reports_allocation_size_on_eviction_failure():
319 - manager = make_cpu_manager(num_blocks=1, cache_policy="lru")
319 + manager = make_cpu_manager(num_chunks=1, cache_policy="lru")
320
321 manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
322 manager.get_stats()
@@ -332,7 +332,7 @@ def test_cpu_manager_reports_allocation_size_on_eviction_failure():
332
333
334 def test_cpu_manager_reports_cache_write_and_read_usage_gauges():
335 - manager = make_cpu_manager(num_blocks=4)
335 + manager = make_cpu_manager(num_chunks=4)
336
337 # Store path: pins write usage until complete_store.
338 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():
356
357
358 def test_cpu_manager_clears_write_usage_after_failed_store():
359 - manager = make_cpu_manager(num_blocks=4)
359 + manager = make_cpu_manager(num_chunks=4)
360
361 manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
362 check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5)
@@ -369,8 +369,8 @@ def test_cpu_manager():
369 """
370 Tests CPUOffloadingManager with lru policy.
371 """
372 - # initialize a CPU manager with a capacity of 4 blocks
373 - cpu_manager = make_cpu_manager(num_blocks=4, cache_policy="lru", enable_events=True)
372 + # initialize a CPU manager with a capacity of 4 chunks
373 + cpu_manager = make_cpu_manager(num_chunks=4, cache_policy="lru", enable_events=True)
374
375 # prepare store [1, 2]
376 prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
@@ -378,7 +378,7 @@ def test_cpu_manager():
378 prepare_store_output,
379 ExpectedPrepareStoreOutput(
380 keys_to_store=[1, 2],
381 - store_block_ids=[0, 1],
381 + store_chunk_ids=[0, 1],
382 evicted_keys=[],
383 ),
384 )
@@ -407,7 +407,7 @@ def test_cpu_manager():
407 prepare_store_output,
408 ExpectedPrepareStoreOutput(
409 keys_to_store=[3, 4, 5],
410 - store_block_ids=[2, 3, 0],
410 + store_chunk_ids=[2, 3, 0],
411 evicted_keys=[1],
412 ),
413 )
@@ -445,7 +445,7 @@ def test_cpu_manager():
445 prepare_store_output,
446 ExpectedPrepareStoreOutput(
447 keys_to_store=[6, 7, 8],
448 - store_block_ids=[1, 0, 3],
448 + store_chunk_ids=[1, 0, 3],
449 evicted_keys=[4, 5, 2],
450 ),
451 )
@@ -462,7 +462,7 @@ def test_cpu_manager():
462 prepare_store_output,
463 ExpectedPrepareStoreOutput(
464 keys_to_store=[9],
465 - store_block_ids=[3],
465 + store_chunk_ids=[3],
466 evicted_keys=[8],
467 ),
468 )
@@ -482,38 +482,38 @@ def test_cpu_manager():
482
483
484 def test_prepare_load_preserves_key_order():
485 - """block_ids[i] must correspond to keys[i] (co-indexed invariant)."""
486 - manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
485 + """chunk_ids[i] must correspond to keys[i] (co-indexed invariant)."""
486 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
487
488 key_a, key_b, key_c = to_key(0), to_key(1), to_key(2)
489
490 - # Store all three keys and learn their block ID assignments
490 + # Store all three keys and learn their chunk ID assignments
491 store_output = manager.prepare_store([key_a, key_b, key_c], _EMPTY_REQ_CTX)
492 assert store_output is not None
493 assert isinstance(store_output.store_spec, CPULoadStoreSpec)
494 - key_to_block_id = {
495 - k: int(bid)
496 - for k, bid in zip(store_output.keys_to_store, store_output.store_spec.block_ids)
494 + key_to_chunk_id = {
495 + k: int(cid)
496 + for k, cid in zip(store_output.keys_to_store, store_output.store_spec.chunk_ids)
497 }
498 manager.complete_store([key_a, key_b, key_c], _EMPTY_REQ_CTX)
499
500 # Forward order: [a, b, c]
501 spec_fwd = manager.prepare_load([key_a, key_b, key_c], _EMPTY_REQ_CTX)
502 assert isinstance(spec_fwd, CPULoadStoreSpec)
503 - assert [int(x) for x in spec_fwd.block_ids] == [
504 - key_to_block_id[key_a],
505 - key_to_block_id[key_b],
506 - key_to_block_id[key_c],
503 + assert [int(x) for x in spec_fwd.chunk_ids] == [
504 + key_to_chunk_id[key_a],
505 + key_to_chunk_id[key_b],
506 + key_to_chunk_id[key_c],
507 ]
508 manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant
509
510 # Arbitrary permutation: [b, c, a]
511 spec_perm = manager.prepare_load([key_b, key_c, key_a], _EMPTY_REQ_CTX)
512 assert isinstance(spec_perm, CPULoadStoreSpec)
513 - assert [int(x) for x in spec_perm.block_ids] == [
514 - key_to_block_id[key_b],
515 - key_to_block_id[key_c],
516 - key_to_block_id[key_a],
513 + assert [int(x) for x in spec_perm.chunk_ids] == [
514 + key_to_chunk_id[key_b],
515 + key_to_chunk_id[key_c],
516 + key_to_chunk_id[key_a],
517 ]
518 manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant
519
@@ -522,10 +522,10 @@ class TestARCPolicy:
522 """Unit tests for CPUOffloadingManager with ARC eviction policy."""
523
524 def _make_manager(
525 - self, num_blocks: int = 4, enable_events: bool = True
525 + self, num_chunks: int = 4, enable_events: bool = True
526 ) -> tuple[CPUOffloadingManager, ARCCachePolicy]:
527 manager = make_cpu_manager(
528 - num_blocks=num_blocks,
528 + num_chunks=num_chunks,
529 cache_policy="arc",
530 enable_events=enable_events,
531 )
@@ -548,7 +548,7 @@ class TestARCPolicy:
548 prepare_store_output,
549 ExpectedPrepareStoreOutput(
550 keys_to_store=[1, 2],
551 - store_block_ids=[0, 1],
551 + store_chunk_ids=[0, 1],
552 evicted_keys=[],
553 ),
554 )
@@ -569,36 +569,36 @@ class TestARCPolicy:
569 assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT
570 assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS
571
572 - # blocks should be in T1 (recent)
572 + # chunks should be in T1 (recent)
573 assert len(arc_policy.t1) == 2
574 assert len(arc_policy.t2) == 0
575
576 def test_t1_to_t2_promotion(self):
577 """
578 - Tests that accessing a block in T1 promotes it to T2 (frequent).
578 + Tests that accessing a chunk in T1 promotes it to T2 (frequent).
579 This is a key feature of ARC's adaptive behavior.
580 """
581 cpu_manager, arc_policy = self._make_manager(enable_events=False)
582
583 - # store and complete block 1
583 + # store and complete chunk 1
584 cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
585 cpu_manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX)
586
587 - # block 1 starts in T1 (recent)
587 + # chunk 1 starts in T1 (recent)
588 assert to_keys([1])[0] in arc_policy.t1
589 assert to_keys([1])[0] not in arc_policy.t2
590
591 - # touch block 1 (simulate second access)
591 + # touch chunk 1 (simulate second access)
592 cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
593
594 - # block 1 should now be in T2 (frequent)
594 + # chunk 1 should now be in T2 (frequent)
595 assert to_keys([1])[0] not in arc_policy.t1
596 assert to_keys([1])[0] in arc_policy.t2
597
598 def test_eviction_with_load(self):
599 """
600 Tests ARC eviction behavior similar to LRU test.
601 - Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted.
601 + Verifies that chunks being loaded (ref_cnt > 0) cannot be evicted.
602 """
603 cpu_manager, _ = self._make_manager()
604
@@ -610,7 +610,7 @@ class TestARCPolicy:
610 prepare_store_output,
611 ExpectedPrepareStoreOutput(
612 keys_to_store=[1, 2, 3, 4],
613 - store_block_ids=[0, 1, 2, 3],
613 + store_chunk_ids=[0, 1, 2, 3],
614 evicted_keys=[],
615 ),
616 )
@@ -628,36 +628,36 @@ class TestARCPolicy:
628 cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
629
630 # now prepare store [5, 6, 7] should succeed
631 - # ARC will evict blocks one at a time from T1 as needed
631 + # ARC will evict chunks one at a time from T1 as needed
632 prepare_store_output = cpu_manager.prepare_store(
633 to_keys([5, 6, 7]), _EMPTY_REQ_CTX
634 )
635 assert prepare_store_output is not None
636 - # Should successfully evict enough blocks to make room (at least 1)
636 + # Should successfully evict enough chunks to make room (at least 1)
637 assert len(prepare_store_output.evicted_keys) >= 1
638
639 def test_adaptive_target(self):
640 """
641 Tests ARC's adaptive target adjustment via ghost lists.
642 - When a block in B1 (ghost list) is accessed, target_t1_size increases.
643 - When a block in B2 is accessed, target_t1_size decreases.
642 + When a chunk in B1 (ghost list) is accessed, target_t1_size increases.
643 + When a chunk in B2 is accessed, target_t1_size decreases.
644 """
645 - cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
645 + cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False)
646
647 - # store blocks 1, 2 (fills cache)
647 + # store chunks 1, 2 (fills cache)
648 cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
649 cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
650
651 initial_target = arc_policy.target_t1_size
652
653 - # store block 3, evicting block 1 (moves to B1 ghost list)
653 + # store chunk 3, evicting chunk 1 (moves to B1 ghost list)
654 cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX)
655 cpu_manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX)
656
657 - # block 1 should be in B1 (ghost list)
657 + # chunk 1 should be in B1 (ghost list)
658 assert to_keys([1])[0] in arc_policy.b1
659
660 - # touch block 1 (cache miss, but in B1)
660 + # touch chunk 1 (cache miss, but in B1)
661 # this should increase target_t1_size (favor recency)
662 cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
663
@@ -671,11 +671,11 @@ class TestARCPolicy:
671 """
672 cpu_manager, arc_policy = self._make_manager(enable_events=False)
673
674 - # store blocks 1, 2, 3, 4
674 + # store chunks 1, 2, 3, 4
675 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
676 cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
677
678 - # promote blocks 3, 4 to T2 by touching them
678 + # promote chunks 3, 4 to T2 by touching them
679 cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX)
680
681 # now: T1 = {1, 2}, T2 = {3, 4}
@@ -686,22 +686,22 @@ class TestARCPolicy:
686 # (when |T1| >= target, evict from T1)
687 arc_policy.target_t1_size = 1
688
689 - # store block 5, should evict from T1 (block 1, LRU in T1)
689 + # store chunk 5, should evict from T1 (chunk 1, LRU in T1)
690 output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
691 assert output is not None
692 assert to_keys([1]) == output.evicted_keys
693
694 cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX)
695
696 - # block 1 should be in B1 (ghost list)
696 + # chunk 1 should be in B1 (ghost list)
697 assert to_keys([1])[0] in arc_policy.b1
698 - # block 5 should be in T1
698 + # chunk 5 should be in T1
699 assert to_keys([5])[0] in arc_policy.t1
700
701 def test_batch_eviction_scans_t1_and_t2_once(self):
702 """ARC batch eviction must preserve order without restarting scans."""
703 cpu_manager, arc_policy = self._make_manager(
704 - num_blocks=256, enable_events=False
704 + num_chunks=256, enable_events=False
705 )
706 keys = to_keys(list(range(256)))
707 cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
@@ -721,13 +721,13 @@ class TestARCPolicy:
721 t2_order = list(arc_policy.t2)
722 expected_t1 = [
723 key
724 - for key, block in arc_policy.t1.items()
725 - if block.ref_cnt == 0 and key not in protected
724 + for key, chunk in arc_policy.t1.items()
725 + if chunk.ref_cnt == 0 and key not in protected
726 ][:num_t1_evictions]
727 expected_t2 = [
728 key
729 - for key, block in arc_policy.t2.items()
730 - if block.ref_cnt == 0 and key not in protected
729 + for key, chunk in arc_policy.t2.items()
730 + if chunk.ref_cnt == 0 and key not in protected
731 ][:num_t2_evictions]
732 expected_t1_scans = t1_order.index(expected_t1[-1]) + 1
733 expected_t2_scans = t2_order.index(expected_t2[-1]) + 1
@@ -746,7 +746,7 @@ class TestARCPolicy:
746
747 def test_batch_eviction_falls_back_after_t1_iterator_exhausted(self):
748 """An exhausted T1 scan must keep falling back to T2."""
749 - cpu_manager, arc_policy = self._make_manager(num_blocks=8, enable_events=False)
749 + cpu_manager, arc_policy = self._make_manager(num_chunks=8, enable_events=False)
750 keys = to_keys(list(range(8)))
751 cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
752 cpu_manager.complete_store(keys, _EMPTY_REQ_CTX)
@@ -764,8 +764,8 @@ class TestARCPolicy:
764 # the target, so each remaining selection must retry T1 then use T2.
765 eligible_t1 = [
766 key
767 - for key, block in arc_policy.t1.items()
768 - if block.ref_cnt == 0 and key not in protected
767 + for key, chunk in arc_policy.t1.items()
768 + if chunk.ref_cnt == 0 and key not in protected
769 ]
770 assert eligible_t1 == t1_order[:1]
771 assert len(t1_order) - 1 >= int(arc_policy.target_t1_size)
@@ -839,7 +839,7 @@ class TestARCPolicy:
839
840 def test_batch_eviction_failure_is_atomic(self):
841 """Finding only some candidates must not partially evict the cache."""
842 - cpu_manager, arc_policy = self._make_manager(num_blocks=4, enable_events=False)
842 + cpu_manager, arc_policy = self._make_manager(num_chunks=4, enable_events=False)
843 keys = to_keys(list(range(4)))
844 cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX)
845 cpu_manager.complete_store(keys, _EMPTY_REQ_CTX)
@@ -858,13 +858,13 @@ class TestARCPolicy:
858 Tests that ghost lists (B1, B2) don't grow unbounded.
859 They should be capped at cache_capacity.
860 """
861 - cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
861 + cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False)
862
863 - # fill cache with blocks 1, 2
863 + # fill cache with chunks 1, 2
864 cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
865 cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
866
867 - # store many blocks to fill ghost lists
867 + # store many chunks to fill ghost lists
868 for i in range(3, 20):
869 cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX)
870 cpu_manager.complete_store(to_keys([i]), _EMPTY_REQ_CTX)
@@ -880,7 +880,7 @@ class TestARCPolicy:
880 """
881 cpu_manager, arc_policy = self._make_manager()
882
883 - # store blocks 1, 2, 3, 4
883 + # store chunks 1, 2, 3, 4
884 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
885 cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
886
@@ -895,13 +895,13 @@ class TestARCPolicy:
895 assert len(arc_policy.t1) == 1
896 assert len(arc_policy.t2) == 3
897
898 - # store block 5, should evict from T1 (block 2, only one in T1)
898 + # store chunk 5, should evict from T1 (chunk 2, only one in T1)
899 prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
900 verify_store_output(
901 prepare_store_output,
902 ExpectedPrepareStoreOutput(
903 keys_to_store=[5],
904 - store_block_ids=[1], # reuses block 2's storage
904 + store_chunk_ids=[1], # reuses chunk 2's storage
905 evicted_keys=[2],
906 ),
907 )
@@ -913,11 +913,11 @@ class TestARCPolicy:
913 """
914 cpu_manager, arc_policy = self._make_manager()
915
916 - # store blocks 1, 2, 3, 4
916 + # store chunks 1, 2, 3, 4
917 cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
918 cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
919
920 - # prepare store block 5 (will evict block 1)
920 + # prepare store chunk 5 (will evict chunk 1)
921 prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
922 assert prepare_store_output is not None
923 assert len(prepare_store_output.evicted_keys) == 1
@@ -925,13 +925,13 @@ class TestARCPolicy:
925 # complete store with failure
926 cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False)
927
928 - # block 5 should not be in cache
928 + # chunk 5 should not be in cache
929 assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.MISS
930 - # block 5 should not be in T1 or T2
930 + # chunk 5 should not be in T1 or T2
931 assert to_keys([5])[0] not in arc_policy.t1
932 assert to_keys([5])[0] not in arc_policy.t2
933
934 - # evicted block should still be gone (in B1 ghost list)
934 + # evicted chunk should still be gone (in B1 ghost list)
935 evicted_hash = prepare_store_output.evicted_keys[0]
936 assert evicted_hash in arc_policy.b1
937
@@ -954,7 +954,7 @@ class TestARCPolicy:
954 assert len(prepare_store_output.evicted_keys) == 1
955 cpu_manager.complete_store(to_keys([3, 4, 5]), _EMPTY_REQ_CTX)
956
957 - # promote some blocks to T2
957 + # promote some chunks to T2
958 cpu_manager.touch(to_keys([2, 3]), _EMPTY_REQ_CTX)
959
960 # T1 has {4, 5}, T2 has {2, 3}
@@ -966,7 +966,7 @@ class TestARCPolicy:
966 assert prepare_store_output is not None
967 cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX)
968
969 - # verify blocks 2, 3 (in T2) are still present
969 + # verify chunks 2, 3 (in T2) are still present
970 assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT
971 assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT
972
@@ -980,7 +980,7 @@ def test_filter_reused_manager():
980 Tests CPUOffloadingManager reuse filtering (store_threshold=2).
981 """
982 manager = make_cpu_manager(
983 - num_blocks=4,
983 + num_chunks=4,
984 cache_policy="lru",
985 enable_events=True,
986 store_threshold=2,
@@ -1023,7 +1023,7 @@ def test_filter_reused_manager():
1023 def test_filter_reused_manager_oversized_offer_makes_progress():
1024 """An offer larger than the tracker capacity must not churn forever."""
1025 manager = make_cpu_manager(
1026 - num_blocks=4,
1026 + num_chunks=4,
1027 cache_policy="lru",
1028 store_threshold=2,
1029 max_tracker_size=3,
@@ -1042,69 +1042,69 @@ def test_filter_reused_manager_oversized_offer_makes_progress():
1042 assert stored_keys == set(keys)
1043
1044
1045 -def test_evictable_cache_block_count():
1045 +def test_evictable_cache_chunk_count():
1046 """
1047 - Verifies _num_evictable_cache_blocks is maintained correctly through the
1047 + Verifies _num_evictable_cache_chunks is maintained correctly through the
1048 full store/load lifecycle, eviction, failed stores, concurrent loads,
1049 reset_cache, and the early-exit fast path in prepare_store.
1050 """
1051 - manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
1051 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
1052
1053 - # Initially no blocks allocated.
1054 - assert manager._num_evictable_cache_blocks == 0
1053 + # Initially no chunks allocated.
1054 + assert manager._num_evictable_cache_chunks == 0
1055
1056 # Initial cache state [x, x, x, x]
1057
1058 - # We get 3 blocks from the cache.
1058 + # We get 3 chunks from the cache.
1059 manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
1060 # cache state [1', 2', 3', x] <- 1', 2', 3' are actively being used.
1061 - assert manager._num_evictable_cache_blocks == 0
1061 + assert manager._num_evictable_cache_chunks == 0
1062
1063 # Completing stores makes them idle.
1064 manager.complete_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
1065 - # cache state [1, 2, 3, x] <- 1, 2, 3 blocks are idle.
1066 - assert manager._num_evictable_cache_blocks == 3
1065 + # cache state [1, 2, 3, x] <- 1, 2, 3 chunks are idle.
1066 + assert manager._num_evictable_cache_chunks == 3
1067
1068 - # prepare_load pins a block: idle count decrements once even if the
1069 - # same block is loaded by two concurrent callers.
1068 + # prepare_load pins a chunk: idle count decrements once even if the
1069 + # same chunk is loaded by two concurrent callers.
1070 manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX)
1071 - # cache state [1', 2, 3, x] <- 2, 3 blocks are idle.
1072 - assert manager._num_evictable_cache_blocks == 2
1071 + # cache state [1', 2, 3, x] <- 2, 3 chunks are idle.
1072 + assert manager._num_evictable_cache_chunks == 2
1073 manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) # 2nd concurrent load
1074 - # cache state [1', 2, 3, x] <- 2, 3 blocks are idle.
1075 - assert manager._num_evictable_cache_blocks == 2 # no double-decrement
1074 + # cache state [1', 2, 3, x] <- 2, 3 chunks are idle.
1075 + assert manager._num_evictable_cache_chunks == 2 # no double-decrement
1076
1077 # First complete_load does not restore idle (ref_cnt still 1).
1078 manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX)
1079 - # cache state [1', 2, 3, x] <- 2, 3 blocks are idle.
1080 - assert manager._num_evictable_cache_blocks == 2
1081 - # Second complete_load drops ref_cnt to 0 -> block becomes idle again.
1079 + # cache state [1', 2, 3, x] <- 2, 3 chunks are idle.
1080 + assert manager._num_evictable_cache_chunks == 2
1081 + # Second complete_load drops ref_cnt to 0 -> chunk becomes idle again.
1082 manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX)
1083 - # cache state [1, 2, 3, x] <- 1, 2, 3 blocks are idle.
1084 - assert manager._num_evictable_cache_blocks == 3
1083 + # cache state [1, 2, 3, x] <- 1, 2, 3 chunks are idle.
1084 + assert manager._num_evictable_cache_chunks == 3
1085
1086 # Eviction decrements idle count.
1087 - # Cache has 3 stored blocks and 1 free slot. Storing 3 new keys needs 2 eviction.
1087 + # Cache has 3 stored chunks and 1 free slot. Storing 3 new keys needs 2 eviction.
1088 manager.prepare_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX)
1089 - # cache state [1, 4', 5', 6'] <- block 1 is idle
1090 - assert manager._num_evictable_cache_blocks == 1
1089 + # cache state [1, 4', 5', 6'] <- chunk 1 is idle
1090 + assert manager._num_evictable_cache_chunks == 1
1091
1092 - # Failed store does not increment idle count (block discarded from cache).
1092 + # Failed store does not increment idle count (chunk discarded from cache).
1093 manager.complete_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX, success=False)
1094 - # cache state [1, x, x, x] <- block 1 is idle. Other returned to cache.
1095 - assert manager._num_evictable_cache_blocks == 1
1094 + # cache state [1, x, x, x] <- chunk 1 is idle. Other returned to cache.
1095 + assert manager._num_evictable_cache_chunks == 1
1096
1097 # reset_cache zeroes the count unconditionally.
1098 manager.reset_cache()
1099 # cache state [x, x, x, x]
1100 - assert manager._num_evictable_cache_blocks == 0
1100 + assert manager._num_evictable_cache_chunks == 0
1101
1102 - # setup 3 blocks with loads so idle count drops to 0.
1102 + # setup 3 chunks with loads so idle count drops to 0.
1103 manager.prepare_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
1104 manager.complete_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
1105 manager.prepare_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
1106 # cache state [10', 11', 12', x]
1107 - assert manager._num_evictable_cache_blocks == 0
1107 + assert manager._num_evictable_cache_chunks == 0
1108
1109 # prepare_store requiring eviction must return None immediately (fast exit).
1110 # Spy on policy.evict to confirm the fast path short-circuits before calling it.
@@ -1120,24 +1120,24 @@ def test_evictable_cache_block_count():
1120 # cache state [10', 11', 12', x] <- cannot evict anything
1121 assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is None
1122 assert not evict_called, (
1123 - "_num_evictable_cache_blocks==0 should short-circuit before evict()"
1123 + "_num_evictable_cache_chunks==0 should short-circuit before evict()"
1124 )
1125
1126 # After releasing the loads, eviction becomes possible again.
1127 manager.complete_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
1128 # cache state [10, 11, 12, x] <- 10, 11, 12 are idle
1129 - assert manager._num_evictable_cache_blocks == 3
1129 + assert manager._num_evictable_cache_chunks == 3
1130 assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is not None
1131 # cache state [10, 11, 14', 15'] <- 10, 11 are idle
1132 - assert manager._num_evictable_cache_blocks == 2
1132 + assert manager._num_evictable_cache_chunks == 2
1133 manager.complete_store(to_keys([14, 15]), _EMPTY_REQ_CTX)
1134 - # cache state [10, 11, 14, 15] <- all blocks idle
1135 - assert manager._num_evictable_cache_blocks == 4
1134 + # cache state [10, 11, 14, 15] <- all chunks idle
1135 + assert manager._num_evictable_cache_chunks == 4
1136
1137
1138 def test_touch_forwards_req_context_to_policy(monkeypatch):
1139 """Regression: CPUOffloadingManager.touch forwards ReqContext to policy."""
1140 - manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
1140 + manager = make_cpu_manager(num_chunks=4, cache_policy="lru")
1141 received = []
1142
1143 def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+55 -55
@@ -38,7 +38,7 @@ def _set_spawn_method(monkeypatch):
38
39 def _make_region(
40 engine_id: str,
41 - num_blocks: int = 4,
41 + num_chunks: int = 4,
42 cpu_page_size: int = PAGE_SIZE,
43 num_workers: int = 1,
44 rank: int = 0,
@@ -47,9 +47,9 @@ def _make_region(
47 assert cpu_page_size % PAGE_SIZE == 0
48 return SharedOffloadRegion(
49 engine_id=engine_id,
50 - num_blocks=num_blocks,
50 + num_chunks=num_chunks,
51 rank=rank,
52 - kv_bytes_per_block=num_workers * cpu_page_size,
52 + kv_bytes_per_chunk=num_workers * cpu_page_size,
53 cpu_page_size=cpu_page_size,
54 barrier=barrier,
55 )
@@ -99,16 +99,16 @@ def _region(engine_id: str, **kwargs):
99 def _multi_region(
100 engine_id: str,
101 num_workers: int,
102 - num_blocks: int = 4,
102 + num_chunks: int = 4,
103 cpu_page_size: int = PAGE_SIZE,
104 ):
105 """Context manager: create one SharedOffloadRegion per rank, clean up on exit."""
106 regions = [
107 SharedOffloadRegion(
108 engine_id=engine_id,
109 - num_blocks=num_blocks,
109 + num_chunks=num_chunks,
110 rank=rank,
111 - kv_bytes_per_block=num_workers * cpu_page_size,
111 + kv_bytes_per_chunk=num_workers * cpu_page_size,
112 cpu_page_size=cpu_page_size,
113 )
114 for rank in range(num_workers)
@@ -124,7 +124,7 @@ def _multi_region(
124 def _race_construct(
125 engine_id: str,
126 num_workers: int,
127 - num_blocks: int = 4,
127 + num_chunks: int = 4,
128 cpu_page_size: int = PAGE_SIZE,
129 ) -> tuple[list[SharedOffloadRegion], list[Exception]]:
130 """Spawn num_workers threads that all race to construct SharedOffloadRegion."""
@@ -137,9 +137,9 @@ def _race_construct(
137 try:
138 regions[rank] = SharedOffloadRegion(
139 engine_id=engine_id,
140 - num_blocks=num_blocks,
140 + num_chunks=num_chunks,
141 rank=rank,
142 - kv_bytes_per_block=num_workers * cpu_page_size,
142 + kv_bytes_per_chunk=num_workers * cpu_page_size,
143 cpu_page_size=cpu_page_size,
144 )
145 except Exception as e:
@@ -156,7 +156,7 @@ def _race_construct(
156
157 def _mp_race_construct_and_write(
158 engine_id: str,
159 - num_blocks: int,
159 + num_chunks: int,
160 rank: int,
161 num_workers: int,
162 cpu_page_size: int,
@@ -170,9 +170,9 @@ def _mp_race_construct_and_write(
170 try:
171 region = SharedOffloadRegion(
172 engine_id=engine_id,
173 - num_blocks=num_blocks,
173 + num_chunks=num_chunks,
174 rank=rank,
175 - kv_bytes_per_block=num_workers * cpu_page_size,
175 + kv_bytes_per_chunk=num_workers * cpu_page_size,
176 cpu_page_size=cpu_page_size,
177 )
178 t = region.create_next_worker_view(cpu_page_size)
@@ -198,9 +198,9 @@ def _mp_barrier_construct_and_hold(
198 try:
199 region = SharedOffloadRegion(
200 engine_id=engine_id,
201 - num_blocks=2,
201 + num_chunks=2,
202 rank=rank,
203 - kv_bytes_per_block=num_workers * PAGE_SIZE,
203 + kv_bytes_per_chunk=num_workers * PAGE_SIZE,
204 cpu_page_size=PAGE_SIZE,
205 barrier=lambda: barrier.wait(30),
206 )
@@ -233,9 +233,9 @@ def iid():
233
234
235 def test_create_next_worker_view_shape_and_stride(iid):
236 - """Returned tensor must have shape (num_blocks, tensor_page_size) and
236 + """Returned tensor must have shape (num_chunks, tensor_page_size) and
237 stride (row_stride, 1) where row_stride = cpu_page_size * num_workers."""
238 - with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r:
238 + with _region(iid, num_chunks=4, cpu_page_size=2 * PAGE_SIZE) as r:
239 t = r.create_next_worker_view(PAGE_SIZE)
240 assert t.shape == (4, PAGE_SIZE)
241 # num_workers=1 → row_stride = cpu_page_size
@@ -253,7 +253,7 @@ def test_create_next_worker_view_storage_offset_rank0(iid):
253
254 def test_create_next_worker_view_storage_offset_rank1(iid):
255 """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap."""
256 - with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1):
256 + with _multi_region(iid, num_workers=2, num_chunks=4) as (r0, r1):
257 t1 = r1.create_next_worker_view(PAGE_SIZE)
258 assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE
259 del t1
@@ -261,7 +261,7 @@ def test_create_next_worker_view_storage_offset_rank1(iid):
261
262 def test_create_next_worker_view_row_stride_with_multiple_workers(iid):
263 """With num_workers=4, row_stride must be 4 * cpu_page_size."""
264 - with _region(iid, num_blocks=2, num_workers=4) as r:
264 + with _region(iid, num_chunks=2, num_workers=4) as r:
265 t = r.create_next_worker_view(PAGE_SIZE)
266 assert t.stride(0) == 4 * PAGE_SIZE
267 del t
@@ -332,12 +332,12 @@ def test_create_next_worker_view_overflow_does_not_mutate_cursor(iid):
332 def test_create_next_worker_view_write_visible_in_raw_mmap(iid):
333 """Writes into a create_next_worker_view view must appear at the correct
334 raw mmap offset"""
335 - with _region(iid, num_blocks=4) as r:
335 + with _region(iid, num_chunks=4) as r:
336 t = r.create_next_worker_view(PAGE_SIZE)
337 - t[2, :] = 42 # write to block row 2
337 + t[2, :] = 42 # write to chunk row 2
338
339 raw = memoryview(r.mmap_obj)
340 - # num_workers=1 → row_stride = PAGE_SIZE; block 2 starts at byte 2*PAGE_SIZE
340 + # num_workers=1 → row_stride = PAGE_SIZE; chunk 2 starts at byte 2*PAGE_SIZE
341 chunk = bytes(raw[2 * PAGE_SIZE : 3 * PAGE_SIZE])
342 assert all(b == 42 for b in chunk)
343 del raw, t
@@ -345,7 +345,7 @@ def test_create_next_worker_view_write_visible_in_raw_mmap(iid):
345
346 def test_create_next_worker_view_multi_tensor_layout(iid):
347 """Two tensors from the same worker land at consecutive byte offsets per row."""
348 - with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r:
348 + with _region(iid, num_chunks=2, cpu_page_size=2 * PAGE_SIZE) as r:
349 ta = r.create_next_worker_view(PAGE_SIZE)
350 tb = r.create_next_worker_view(PAGE_SIZE)
351
@@ -353,8 +353,8 @@ def test_create_next_worker_view_multi_tensor_layout(iid):
353 tb[:, :] = 2
354
355 raw = memoryview(r.mmap_obj)
356 - for blk in range(2):
357 - row_offset = blk * 2 * PAGE_SIZE # num_workers=1
356 + for chunk in range(2):
357 + row_offset = chunk * 2 * PAGE_SIZE # num_workers=1
358 assert all(b == 1 for b in raw[row_offset : row_offset + PAGE_SIZE])
359 assert all(
360 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):
366 """Each worker process calls create_next_worker_view and writes distinct data;
367 the parent verifies each slot lands at the correct interleaved offset."""
368 num_workers = 2
369 - num_blocks = 4
369 + num_chunks = 4
370
371 ctx = get_mp_context()
372 done_queue = ctx.Queue()
@@ -375,9 +375,9 @@ def test_create_next_worker_view_multiprocess_slots(iid):
375 # Parent is rank 0 (creator); child is rank 1 (joiner).
376 region = SharedOffloadRegion(
377 engine_id=iid,
378 - num_blocks=num_blocks,
378 + num_chunks=num_chunks,
379 rank=0,
380 - kv_bytes_per_block=num_workers * PAGE_SIZE,
380 + kv_bytes_per_chunk=num_workers * PAGE_SIZE,
381 cpu_page_size=PAGE_SIZE,
382 )
383 try:
@@ -385,7 +385,7 @@ def test_create_next_worker_view_multiprocess_slots(iid):
385 target=_mp_race_construct_and_write,
386 args=(
387 iid,
388 - num_blocks,
388 + num_chunks,
389 1,
390 num_workers,
391 PAGE_SIZE,
@@ -403,12 +403,12 @@ def test_create_next_worker_view_multiprocess_slots(iid):
403 assert result["error"] is None, result["error"]
404
405 raw = memoryview(region.mmap_obj)
406 - for blk in range(num_blocks):
407 - row_start = blk * num_workers * PAGE_SIZE
406 + for chunk in range(num_chunks):
407 + row_start = chunk * num_workers * PAGE_SIZE
408 w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
409 w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
410 - assert all(b == 11 for b in w0), f"block {blk}: rank0 slot wrong"
411 - assert all(b == 22 for b in w1), f"block {blk}: rank1 slot wrong"
410 + assert all(b == 11 for b in w0), f"chunk {chunk}: rank0 slot wrong"
411 + assert all(b == 22 for b in w1), f"chunk {chunk}: rank1 slot wrong"
412
413 del raw, t0 # release before finally triggers cleanup
414 cleanup_queue.put(True)
@@ -422,8 +422,8 @@ def test_create_next_worker_view_multiprocess_slots(iid):
422 def test_create_next_worker_view_worker_isolation(iid):
423 """Writes by worker 0 must not affect worker 1's slot and vice versa."""
424 num_workers = 2
425 - num_blocks = 4
426 - with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions:
425 + num_chunks = 4
426 + with _multi_region(iid, num_workers=num_workers, num_chunks=num_chunks) as regions:
427 t0 = regions[0].create_next_worker_view(PAGE_SIZE)
428 t1 = regions[1].create_next_worker_view(PAGE_SIZE)
429
@@ -431,12 +431,12 @@ def test_create_next_worker_view_worker_isolation(iid):
431 t1[:, :] = 22
432
433 raw = memoryview(regions[0].mmap_obj)
434 - for blk in range(num_blocks):
435 - row_start = blk * num_workers * PAGE_SIZE
434 + for chunk in range(num_chunks):
435 + row_start = chunk * num_workers * PAGE_SIZE
436 w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
437 w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
438 - assert all(b == 11 for b in w0), f"block {blk}: worker0 slot corrupted"
439 - assert all(b == 22 for b in w1), f"block {blk}: worker1 slot corrupted"
438 + assert all(b == 11 for b in w0), f"chunk {chunk}: worker0 slot corrupted"
439 + assert all(b == 22 for b in w1), f"chunk {chunk}: worker1 slot corrupted"
440 del raw, t0, t1 # release before finally triggers cleanup
441
442
@@ -466,7 +466,7 @@ def test_file_exists_after_construction(iid):
466
467 def test_file_has_correct_size(iid):
468 """The mmap file size on disk must equal total_size_bytes."""
469 - with _region(iid, num_blocks=4) as r:
469 + with _region(iid, num_chunks=4) as r:
470 assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE
471
472
@@ -491,11 +491,11 @@ def test_madvise_success_selects_madvise_population(iid, monkeypatch):
491 monkeypatch.setattr(sor, "_madvise_populate_write", _spy_madvise)
492 monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
493
494 - num_blocks = 3
494 + num_chunks = 3
495 num_workers = 2
496 - with _region(iid, num_blocks=num_blocks, num_workers=num_workers, rank=1):
497 - # 1 probe (PAGESIZE) + N populate calls (one per block per worker column).
498 - expected_populate = num_blocks # ranked path: one call per block
496 + with _region(iid, num_chunks=num_chunks, num_workers=num_workers, rank=1):
497 + # 1 probe (PAGESIZE) + N populate calls (one per chunk per worker column).
498 + expected_populate = num_chunks # ranked path: one call per chunk
499 mmap_id = madvise_calls[0][2]
500 assert madvise_calls == [
501 (0, mmap.PAGESIZE, mmap_id),
@@ -510,7 +510,7 @@ def test_madvise_success_selects_madvise_population(iid, monkeypatch):
510
511
512 def test_madvise_einval_selects_fallback_for_ranked_region(iid, monkeypatch):
513 - """An EINVAL probe must select fallback for every ranked block."""
513 + """An EINVAL probe must select fallback for every ranked chunk."""
514 from vllm.v1.kv_offload.cpu import shared_offload_region as sor
515
516 fallback_calls: list[tuple[int, int]] = []
@@ -526,7 +526,7 @@ def test_madvise_einval_selects_fallback_for_ranked_region(iid, monkeypatch):
526 monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval)
527 monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
528
529 - with _region(iid, num_blocks=3, num_workers=2, rank=1):
529 + with _region(iid, num_chunks=3, num_workers=2, rank=1):
530 assert fallback_calls == [
531 (mmap.PAGESIZE, mmap.PAGESIZE),
532 (3 * mmap.PAGESIZE, mmap.PAGESIZE),
@@ -551,7 +551,7 @@ def test_madvise_einval_selects_fallback_for_unranked_region(iid, monkeypatch):
551 monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval)
552 monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
553
554 - with _region(iid, num_blocks=3, num_workers=2, rank=None):
554 + with _region(iid, num_chunks=3, num_workers=2, rank=None):
555 assert fallback_calls == [(0, 6 * mmap.PAGESIZE)]
556
557
@@ -651,7 +651,7 @@ def test_multiprocess_race_construct_and_write(iid):
651 """N processes race to construct the same SharedOffloadRegion, each writes
652 fill_value = rank+1 into their slot; parent verifies interleaved layout."""
653 num_workers = 4
654 - num_blocks = 3
654 + num_chunks = 3
655
656 ctx = get_mp_context()
657 done_queue = ctx.Queue()
@@ -662,7 +662,7 @@ def test_multiprocess_race_construct_and_write(iid):
662 target=_mp_race_construct_and_write,
663 args=(
664 iid,
665 - num_blocks,
665 + num_chunks,
666 rank,
667 num_workers,
668 PAGE_SIZE,
@@ -689,13 +689,13 @@ def test_multiprocess_race_construct_and_write(iid):
689 with open(mmap_path, "rb") as f:
690 raw = f.read()
691
692 - for blk in range(num_blocks):
692 + for chunk in range(num_chunks):
693 for w in range(num_workers):
694 - slot_start = (blk * num_workers + w) * PAGE_SIZE
694 + slot_start = (chunk * num_workers + w) * PAGE_SIZE
695 slot = raw[slot_start : slot_start + PAGE_SIZE]
696 expected = w + 1 # fill_value = rank + 1
697 assert all(b == expected for b in slot), (
698 - f"block {blk}, worker {w}: expected {expected} but got wrong bytes"
698 + f"chunk {chunk}, worker {w}: expected {expected} but got wrong bytes"
699 )
700
701 # Unblock all workers to clean up.
@@ -839,9 +839,9 @@ def test_insufficient_space_raises_clear_error(monkeypatch):
839 with pytest.raises(RuntimeError, match="Insufficient space"):
840 SharedOffloadRegion(
841 engine_id=engine_id,
842 - num_blocks=4,
842 + num_chunks=4,
843 rank=0,
844 - kv_bytes_per_block=PAGE_SIZE,
844 + kv_bytes_per_chunk=PAGE_SIZE,
845 cpu_page_size=PAGE_SIZE,
846 )
847
@@ -870,9 +870,9 @@ def test_ftruncate_failure_cleans_up_creator(monkeypatch):
870 with pytest.raises(OSError, match="ftruncate failed"):
871 SharedOffloadRegion(
872 engine_id=engine_id,
873 - num_blocks=4,
873 + num_chunks=4,
874 rank=0,
875 - kv_bytes_per_block=PAGE_SIZE,
875 + kv_bytes_per_chunk=PAGE_SIZE,
876 cpu_page_size=PAGE_SIZE,
877 )
878
tests/v1/kv_offload/test_factory.py
+11 -11
@@ -136,7 +136,7 @@ def test_get_spec_cls_defaults_to_cpu():
136 def test_create_cpu_offloading_spec():
137 spec = _create_spec()
138 assert isinstance(spec, CPUOffloadingSpec)
139 - assert spec.num_blocks > 0
139 + assert spec.num_chunks > 0
140
141
142 def test_cpu_spec_sizes_normalized_worker_layout():
@@ -156,7 +156,7 @@ def test_cpu_spec_sizes_normalized_worker_layout():
156 assert isinstance(spec, CPUOffloadingSpec)
157 assert spec.cpu_page_size_per_worker == 32
158 assert spec.kv_bytes_per_chunk == alignment
159 - assert spec.num_blocks == 3
159 + assert spec.num_chunks == 3
160
161
162 def test_cpu_spec_zero_worker_bytes_produces_empty_cache():
@@ -165,7 +165,7 @@ def test_cpu_spec_zero_worker_bytes_produces_empty_cache():
165 assert isinstance(spec, CPUOffloadingSpec)
166 assert spec.cpu_page_size_per_worker == 0
167 assert spec.kv_bytes_per_chunk == 0
168 - assert spec.num_blocks == 0
168 + assert spec.num_chunks == 0
169
170
171 def test_tiering_spec_aligns_row_size():
@@ -183,7 +183,7 @@ def test_tiering_spec_aligns_row_size():
183 assert isinstance(spec, TieringOffloadingSpec)
184 assert spec.cpu_page_size_per_worker == 32
185 assert spec.kv_bytes_per_chunk == alignment
186 - assert spec.num_blocks == 3
186 + assert spec.num_chunks == 3
187
188
189 @pytest.mark.parametrize("world_size", [2, 4, 8])
@@ -201,7 +201,7 @@ def test_tiering_spec_replicated_sizing_removes_world_factor(world_size: int):
201 assert spec.replicated_layout is True
202 assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
203 assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block
204 - assert spec.num_blocks == 8
204 + assert spec.num_chunks == 8
205
206
207 def test_tiering_spec_create_worker_uses_single_slot_for_replicated_layout(monkeypatch):
@@ -239,7 +239,7 @@ def test_tiering_spec_create_worker_uses_single_slot_for_replicated_layout(monke
239 spec.create_worker(kv_caches)
240
241 assert region_calls[0]["rank"] == 0
242 - assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block
242 + assert region_calls[0]["kv_bytes_per_chunk"] == worker_kv_bytes_per_block
243 assert worker_calls[0]["kv_caches"] is kv_caches
244 assert worker_calls[0]["mmap_region"] is region
245
@@ -292,7 +292,7 @@ def test_cpu_spec_replicated_sizing_on_shared_region(monkeypatch, world_size: in
292 assert spec.replicated_layout is True
293 assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
294 assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block
295 - assert spec.num_blocks == 8
295 + assert spec.num_chunks == 8
296
297
298 @pytest.mark.parametrize("world_size", [2, 4, 8])
@@ -320,7 +320,7 @@ def test_cpu_spec_replicated_disabled_without_shared_region(
320 assert spec.replicated_layout is False
321 assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
322 assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block * world_size
323 - assert spec.num_blocks == 2
323 + assert spec.num_chunks == 2
324
325
326 @pytest.mark.parametrize("config_replicated", [True, False])
@@ -381,7 +381,7 @@ def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
381 spec.create_worker(kv_caches)
382
383 assert region_calls[0]["engine_id"] == "test-engine"
384 - assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block * 4
384 + assert region_calls[0]["kv_bytes_per_chunk"] == worker_kv_bytes_per_block * 4
385 assert worker_calls[0]["kv_caches"] is kv_caches
386 assert worker_calls[0]["mmap_region"] is region
387
@@ -419,11 +419,11 @@ def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch):
419 def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
420 import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
421
422 - # worker_kv_bytes_per_block=0 yields num_blocks=0; a zero-byte region cannot
422 + # worker_kv_bytes_per_block=0 yields num_chunks=0; a zero-byte region cannot
423 # be mmap'd, so even on CUDA-alike this must fall back to the tensor path.
424 spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4)
425 assert isinstance(spec, CPUOffloadingSpec)
426 - assert spec.num_blocks == 0
426 + assert spec.num_chunks == 0
427
428 region_calls: list[dict[str, Any]] = []
429 worker_calls: list[dict[str, Any]] = []
tests/v1/kv_offload/tiering/p2p/test_manager.py
+17 -17
@@ -91,17 +91,17 @@ def _req_context(kv_params: dict | None = None) -> ReqContext:
91 def _job_metadata(
92 job_id: int,
93 keys: list[bytes] | None = None,
94 - block_ids: list[int] | None = None,
94 + chunk_ids: list[int] | None = None,
95 kv_params: dict | None = None,
96 ) -> TransferJob:
97 if keys is None:
98 keys = [b"key1"]
99 - if block_ids is None:
100 - block_ids = list(range(len(keys)))
99 + if chunk_ids is None:
100 + chunk_ids = list(range(len(keys)))
101 return TransferJob(
102 job_id=job_id,
103 keys=keys,
104 - block_ids=np.array(block_ids),
104 + chunk_ids=np.array(chunk_ids),
105 is_promotion=False,
106 req_context=_req_context(kv_params),
107 )
@@ -254,9 +254,9 @@ class TestLookup:
254 "kv_params,expected",
255 [
256 (_remote_decoder_kv_params(), OffloadPolicy.REQUEST_LEVEL),
257 - (None, OffloadPolicy.BLOCK_LEVEL),
258 - (_remote_prefiller_kv_params(), OffloadPolicy.BLOCK_LEVEL),
259 - ({"remote_decoder": {}}, OffloadPolicy.BLOCK_LEVEL),
257 + (None, OffloadPolicy.CHUNK_LEVEL),
258 + (_remote_prefiller_kv_params(), OffloadPolicy.CHUNK_LEVEL),
259 + ({"remote_decoder": {}}, OffloadPolicy.CHUNK_LEVEL),
260 ],
261 ids=["producer", "plain", "consumer", "producer_no_id"],
262 )
@@ -366,7 +366,7 @@ class TestSubmitStore:
366 job = _job_metadata(
367 job_id=1,
368 keys=[b"k1", b"k2"],
369 - block_ids=[3, 4],
369 + chunk_ids=[3, 4],
370 kv_params=_remote_decoder_kv_params(kv_request_id="req-1"),
371 )
372 mgr.submit_store(job)
@@ -397,7 +397,7 @@ class TestSubmitStore:
397 job = _job_metadata(
398 job_id=7,
399 keys=[b"k1", b"k2"],
400 - block_ids=[3, 4],
400 + chunk_ids=[3, 4],
401 kv_params=_remote_decoder_kv_params(kv_request_id="req-1"),
402 )
403 mgr.submit_store(job)
@@ -440,7 +440,7 @@ class TestSubmitLoad:
440 """Empty key list succeeds immediately."""
441 mgr = _make_manager()
442 job = _job_metadata(
443 - job_id=1, keys=[], block_ids=[], kv_params=_remote_prefiller_kv_params()
443 + job_id=1, keys=[], chunk_ids=[], kv_params=_remote_prefiller_kv_params()
444 )
445 mgr.submit_load(job)
446 assert mgr._finished_jobs == [JobResult(job_id=1, success=True)]
@@ -463,7 +463,7 @@ class TestSubmitLoad:
463 job = _job_metadata(
464 job_id=42,
465 keys=[b"k1", b"k2"],
466 - block_ids=[5, 6],
466 + chunk_ids=[5, 6],
467 kv_params=_remote_prefiller_kv_params(kv_request_id="req-42"),
468 )
469 mgr.submit_load(job)
@@ -1155,7 +1155,7 @@ class TestBidirectionalManager:
1155 _job_metadata(
1156 job_id=100,
1157 keys=[b"a-block"],
1158 - block_ids=[0],
1158 + chunk_ids=[0],
1159 kv_params=a_prefiller_params,
1160 )
1161 )
@@ -1163,7 +1163,7 @@ class TestBidirectionalManager:
1163 _job_metadata(
1164 job_id=200,
1165 keys=[b"b-block"],
1166 - block_ids=[0],
1166 + chunk_ids=[0],
1167 kv_params=b_prefiller_params,
1168 )
1169 )
@@ -1173,7 +1173,7 @@ class TestBidirectionalManager:
1173 _job_metadata(
1174 job_id=101,
1175 keys=[b"b-block"],
1176 - block_ids=[0],
1176 + chunk_ids=[0],
1177 kv_params=a_decoder_params,
1178 )
1179 )
@@ -1181,7 +1181,7 @@ class TestBidirectionalManager:
1181 _job_metadata(
1182 job_id=201,
1183 keys=[b"a-block"],
1184 - block_ids=[0],
1184 + chunk_ids=[0],
1185 kv_params=b_decoder_params,
1186 )
1187 )
@@ -1573,7 +1573,7 @@ class TestConnectionDeathMidTransfer:
1573 _job_metadata(
1574 job_id=900,
1575 keys=[b"a-block"],
1576 - block_ids=[0],
1576 + chunk_ids=[0],
1577 kv_params=a_prefiller_params,
1578 )
1579 )
@@ -1581,7 +1581,7 @@ class TestConnectionDeathMidTransfer:
1581 _job_metadata(
1582 job_id=901,
1583 keys=[b"b-block"],
1584 - block_ids=[0],
1584 + chunk_ids=[0],
1585 kv_params=a_decoder_params,
1586 )
1587 )
tests/v1/kv_offload/tiering/p2p/test_sessions.py
+2 -2
@@ -259,13 +259,13 @@ class FakeParent:
259 ) -> TransferJob:
260 keys_list = list(keys)
261 self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id))
262 - block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32)
262 + chunk_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32)
263 job_id = self._next_job_id
264 self._next_job_id += 1
265 return TransferJob(
266 job_id=job_id,
267 keys=keys_list,
268 - block_ids=block_ids,
268 + chunk_ids=chunk_ids,
269 is_promotion=False,
270 req_context=ctx,
271 )
tests/v1/kv_offload/tiering/test_fs_tier.py
+13 -13
@@ -105,15 +105,15 @@ def key(n: int) -> OffloadKey:
105 def make_job(
106 job_id: int,
107 keys: list[OffloadKey],
108 - block_ids: list[int] | None = None,
108 + chunk_ids: list[int] | None = None,
109 is_promotion: bool = False,
110 ) -> TransferJob:
111 - if block_ids is None:
112 - block_ids = list(range(len(keys)))
111 + if chunk_ids is None:
112 + chunk_ids = list(range(len(keys)))
113 return TransferJob(
114 job_id=job_id,
115 keys=keys,
116 - block_ids=np.array(block_ids, dtype=np.int64),
116 + chunk_ids=np.array(chunk_ids, dtype=np.int64),
117 is_promotion=is_promotion,
118 req_context=_CTX,
119 )
@@ -382,11 +382,11 @@ def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size):
382 tensor[:] = _page_aligned_rand_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS)
383
384 keys = [key(i) for i in range(batch_size)]
385 - store_block_ids = list(range(batch_size))
386 - load_block_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS))
385 + store_chunk_ids = list(range(batch_size))
386 + load_chunk_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS))
387 expected = tensor[:batch_size].clone()
388
389 - tier.submit_store(make_job(1, keys, store_block_ids))
389 + tier.submit_store(make_job(1, keys, store_chunk_ids))
390 store_results = drain(tier)
391 assert len(store_results) == 1
392 assert store_results[0].success
@@ -396,15 +396,15 @@ def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size):
396 tensor[:] = 0.0
397
398 # Load into a range disjoint by index from the store ids, to also
399 - # exercise loading a block into a different id than it was stored from.
400 - tier.submit_load(make_job(2, keys, load_block_ids, is_promotion=True))
399 + # exercise loading a chunk into a different id than it was stored from.
400 + tier.submit_load(make_job(2, keys, load_chunk_ids, is_promotion=True))
401 load_results = drain(tier)
402 assert len(load_results) == 1
403 assert load_results[0].success
404
405 - for i, bid in enumerate(load_block_ids):
406 - assert torch.allclose(tensor[bid], expected[i]), (
407 - f"Block {bid} data mismatch after store+load"
405 + for i, cid in enumerate(load_chunk_ids):
406 + assert torch.allclose(tensor[cid], expected[i]), (
407 + f"Chunk {cid} data mismatch after store+load"
408 )
409
410
@@ -895,7 +895,7 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path):
895 view = memoryview(tensor.numpy())
896 mock_region = MagicMock()
897 mock_region.create_kv_memoryview.return_value = view
898 - primary = CPUPrimaryTierOffloadingManager(num_blocks=4, mmap_region=mock_region)
898 + primary = CPUPrimaryTierOffloadingManager(num_chunks=4, mmap_region=mock_region)
899 tier = FileSystemTierManager(
900 offloading_spec=_make_offloading_spec(enable_kv_cache_events=True),
901 primary_kv_view=primary.get_kv_memoryview(),
tests/v1/kv_offload/tiering/test_metrics.py
+22 -22
@@ -25,8 +25,8 @@ def to_keys(int_ids: Iterable[int]):
25 def test_tiering_metrics_tracker_records_lookup_metrics():
26 tracker = TieringMetricsTracker(
27 tier_types=["fs", "p2p"],
28 - num_primary_blocks=5,
29 - primary_block_size=16,
28 + num_primary_chunks=5,
29 + primary_chunk_size=16,
30 )
31 tracker.on_new_request(_CTX)
32
@@ -45,18 +45,18 @@ def test_tiering_metrics_tracker_records_lookup_metrics():
45 stats = tracker.take_stats()
46 assert stats is not None
47 values = stats.data["data"]
48 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1
49 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1
50 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("2:p2p",)] == 1
51 - assert values[TieringOffloadingMetrics.BLOCK_HITS][("2:p2p",)] == 1
52 - assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS]
48 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("0:primary",)] == 1
49 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("1:fs",)] == 1
50 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("2:p2p",)] == 1
51 + assert values[TieringOffloadingMetrics.CHUNK_HITS][("2:p2p",)] == 1
52 + assert ("1:fs",) not in values[TieringOffloadingMetrics.CHUNK_HITS]
53
54
55 def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation():
56 tracker = TieringMetricsTracker(
57 tier_types=["fs"],
58 - num_primary_blocks=5,
59 - primary_block_size=16,
58 + num_primary_chunks=5,
59 + primary_chunk_size=16,
60 )
61 tracker.on_new_request(_CTX)
62 key = to_keys([1])[0]
@@ -73,9 +73,9 @@ def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation():
73 stats = tracker.take_stats()
74 assert stats is not None
75 values = stats.data["data"]
76 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1
77 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1
78 - assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:fs",)] == 1
76 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("0:primary",)] == 1
77 + assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("1:fs",)] == 1
78 + assert values[TieringOffloadingMetrics.CHUNK_HITS][("1:fs",)] == 1
79
80 tracker.on_request_allocated(_CTX)
81 tracker.on_lookup(
@@ -90,15 +90,15 @@ def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation():
90 stats = tracker.take_stats()
91 assert stats is not None
92 values = stats.data["data"]
93 - assert TieringOffloadingMetrics.BLOCK_QUERIES not in values
94 - assert TieringOffloadingMetrics.BLOCK_HITS not in values
93 + assert TieringOffloadingMetrics.CHUNK_QUERIES not in values
94 + assert TieringOffloadingMetrics.CHUNK_HITS not in values
95
96
97 def test_tiering_metrics_tracker_records_finished_job_metrics():
98 tracker = TieringMetricsTracker(
99 tier_types=["fs"],
100 - num_primary_blocks=5,
101 - primary_block_size=16,
100 + num_primary_chunks=5,
101 + primary_chunk_size=16,
102 )
103 cascade_key, promotion_key, failed_cascade_key, failed_promotion_key = to_keys(
104 range(4)
@@ -149,8 +149,8 @@ def test_tiering_metrics_tracker_records_finished_job_metrics():
149 def test_tiering_metrics_tracker_records_partial_promotion_success_bytes():
150 tracker = TieringMetricsTracker(
151 tier_types=["fs"],
152 - num_primary_blocks=5,
153 - primary_block_size=16,
152 + num_primary_chunks=5,
153 + primary_chunk_size=16,
154 )
155 keys = to_keys(range(3))
156 job = JobMetadata(
@@ -182,8 +182,8 @@ def test_tiering_metrics_tracker_records_partial_promotion_success_bytes():
182 def test_tiering_metrics_tracker_reports_active_job_and_primary_usage_gauges():
183 tracker = TieringMetricsTracker(
184 tier_types=["fs", "p2p"],
185 - num_primary_blocks=6,
186 - primary_block_size=16,
185 + num_primary_chunks=6,
186 + primary_chunk_size=16,
187 )
188 fs_job = JobMetadata(
189 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():
214 def test_tiering_metrics_tracker_records_promotion_allocation_failures():
215 tracker = TieringMetricsTracker(
216 tier_types=["fs"],
217 - num_primary_blocks=1,
218 - primary_block_size=16,
217 + num_primary_chunks=1,
218 + primary_chunk_size=16,
219 )
220
221 tracker.on_promotion_allocation_failure()
tests/v1/kv_offload/tiering/test_obj_tier.py
+5 -5
@@ -106,14 +106,14 @@ def key(n: int) -> OffloadKey:
106 def make_job(
107 job_id: int,
108 keys: list[OffloadKey],
109 - block_ids: list[int] | None = None,
109 + chunk_ids: list[int] | None = None,
110 ) -> TransferJob:
111 - if block_ids is None:
112 - block_ids = list(range(len(keys)))
111 + if chunk_ids is None:
112 + chunk_ids = list(range(len(keys)))
113 return TransferJob(
114 job_id=job_id,
115 keys=keys,
116 - block_ids=np.array(block_ids, dtype=np.int32),
116 + chunk_ids=np.array(chunk_ids, dtype=np.int32),
117 is_promotion=False,
118 req_context=_CTX,
119 )
@@ -541,7 +541,7 @@ class TestMockObjTierFailures:
541 mmap_region = MagicMock()
542 mmap_region.create_kv_memoryview.return_value = primary_kv_view
543 primary_tier = CPUPrimaryTierOffloadingManager(
544 - num_blocks=num_blocks, mmap_region=mmap_region
544 + num_chunks=num_blocks, mmap_region=mmap_region
545 )
546 obj_tier, agent = _make_tier(
547 num_blocks=num_blocks, primary_kv_view=primary_kv_view
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+204 -203
@@ -5,15 +5,16 @@ Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager.
5
6 These tests verify:
7 1. Basic tiered offloading operations (store, load, lookup)
8 -2. Cascade behavior (blocks stored to all secondary tiers)
9 -3. Promotion behavior (blocks loaded from secondary to primary to GPU)
10 -4. ref_cnt management (blocks protected during async transfers)
8 +2. Cascade behavior (chunks stored to all secondary tiers)
9 +3. Promotion behavior (chunks loaded from secondary to primary to GPU)
10 +4. ref_cnt management (chunks protected during async transfers)
11 5. Eviction coordination between tiers
12 """
13
14 from collections.abc import Iterable
15 from unittest.mock import MagicMock
16
17 +import numpy as np
18 import pytest
19 import torch
20
@@ -57,10 +58,10 @@ _CTX = ReqContext(req_id="test")
58 _MOCK_OFFLOADING_SPEC = MagicMock()
59
60
60 -def _mock_mmap_region(num_blocks: int, row_bytes: int = 16):
61 +def _mock_mmap_region(num_chunks: int, row_bytes: int = 16):
62 """Create a mock SharedOffloadRegion for testing."""
63 mock = MagicMock()
63 - view = memoryview(torch.zeros((num_blocks, row_bytes), dtype=torch.int8).numpy())
64 + view = memoryview(torch.zeros((num_chunks, row_bytes), dtype=torch.int8).numpy())
65 mock.create_kv_memoryview.return_value = view
66 return mock
67
@@ -160,7 +161,7 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch):
161 metadata = metrics[MetricsSecondaryTierManager.MY_TIER_METRIC]
162 assert metadata.documentation == "Number of bytes served by the test tier."
163 assert metadata.labelnames == ("tier",)
163 - assert metrics[TieringOffloadingMetrics.BLOCK_QUERIES].labelnames == ("tier",)
164 + assert metrics[TieringOffloadingMetrics.CHUNK_QUERIES].labelnames == ("tier",)
165 assert metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY].labelnames == ("tier",)
166 assert metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY].labelnames == ("tier",)
167 assert metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES].labelnames == (
@@ -183,7 +184,7 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch):
184 def test_tiering_manager_aggregates_secondary_stats():
185 mock_region = _mock_mmap_region(5)
186 primary_tier = CPUPrimaryTierOffloadingManager(
186 - num_blocks=5, mmap_region=mock_region
187 + num_chunks=5, mmap_region=mock_region
188 )
189 secondary_tier = MetricsSecondaryTierManager(
190 offloading_spec=_MOCK_OFFLOADING_SPEC,
@@ -232,19 +233,19 @@ class TestExampleSecondaryTierManager:
233 )
234
235 # Initially empty
235 - blocks = to_keys(range(3))
236 - assert tier.lookup(blocks[0], _CTX) is LookupResult.MISS
236 + chunks = to_keys(range(3))
237 + assert tier.lookup(chunks[0], _CTX) is LookupResult.MISS
238
238 - # Store blocks (simulate with direct insertion for testing)
239 - tier.blocks[blocks[0]] = True
240 - tier.blocks[blocks[1]] = True
239 + # Store chunks (simulate with direct insertion for testing)
240 + tier.chunks[chunks[0]] = True
241 + tier.chunks[chunks[1]] = True
242
242 - # Lookup should find first two blocks
243 - assert tier.lookup(blocks[0], _CTX) is LookupResult.HIT
244 - assert tier.lookup(blocks[1], _CTX) is LookupResult.HIT
243 + # Lookup should find first two chunks
244 + assert tier.lookup(chunks[0], _CTX) is LookupResult.HIT
245 + assert tier.lookup(chunks[1], _CTX) is LookupResult.HIT
246
246 - # Third block not present
247 - assert tier.lookup(blocks[2], _CTX) is LookupResult.MISS
247 + # Third chunk not present
248 + assert tier.lookup(chunks[2], _CTX) is LookupResult.MISS
249
250
251 # What a request-level cascade does with a key already present in the primary
@@ -260,7 +261,7 @@ _CASCADE_SUPPLY_BEHAVIOR = {
261 # The same test, before it lands.
262 LookupResult.HIT_PENDING: "defer",
263 # test_cascade_drops_deferred_keys_whose_primary_write_failed, where the
263 - # failed write frees the block.
264 + # failed write frees the chunk.
265 LookupResult.MISS: "drop",
266 # test_cascade_rejects_retry_from_primary. The primary tier resolves every
267 # key it holds, so RETRY cannot reach the cascade, and parking on it would
@@ -284,7 +285,7 @@ class TestTieringOffloadingManager:
285 # Create primary tier (CPU-based)
286 mock_region = _mock_mmap_region(5)
287 self.primary_tier = CPUPrimaryTierOffloadingManager(
287 - num_blocks=5, mmap_region=mock_region
288 + num_chunks=5, mmap_region=mock_region
289 )
290
291 mock_view = mock_region.create_kv_memoryview()
@@ -331,7 +332,7 @@ class TestTieringOffloadingManager:
332 TransferJob(
333 job_id=job_id,
334 keys=to_keys([1, 2]),
334 - block_ids=[0, 1],
335 + chunk_ids=np.array([0, 1], dtype=np.int64),
336 is_promotion=True,
337 req_context=_CTX,
338 ),
@@ -358,7 +359,7 @@ class TestTieringOffloadingManager:
359 TransferJob(
360 job_id=job_id,
361 keys=to_keys([1]),
361 - block_ids=[0],
362 + chunk_ids=np.array([0], dtype=np.int64),
363 is_promotion=True,
364 req_context=_CTX,
365 ),
@@ -392,23 +393,23 @@ class TestTieringOffloadingManager:
393
394 def test_basic_store_to_primary(self, manager_setup):
395 """Test basic store operation to primary tier."""
395 - blocks = to_keys(range(3))
396 + chunks = to_keys(range(3))
397
398 # Prepare store
399 self._start_request()
399 - result = self.manager.prepare_store(blocks, _CTX)
400 + result = self.manager.prepare_store(chunks, _CTX)
401 assert result is not None
402 assert len(result.keys_to_store) == 3
403
404 # Complete store
404 - self.manager.complete_store(blocks, _CTX, success=True)
405 + self.manager.complete_store(chunks, _CTX, success=True)
406
406 - # Blocks should be in primary tier
407 - assert count_hits(self.primary_tier, blocks) == 3
407 + # Chunks should be in primary tier
408 + assert count_hits(self.primary_tier, chunks) == 3
409
410 def test_cascade_to_all_secondary_tiers(self, manager_setup):
410 - """Test that blocks are cascaded to ALL secondary tiers."""
411 - blocks = to_keys(range(3))
411 + """Test that chunks are cascaded to ALL secondary tiers."""
412 + chunks = to_keys(range(3))
413
414 self.secondary_tier1.submit_store = MagicMock(
415 wraps=self.secondary_tier1.submit_store
@@ -419,44 +420,44 @@ class TestTieringOffloadingManager:
420
421 # Store to primary
422 self._start_request()
422 - result = self.manager.prepare_store(blocks, _CTX)
423 + result = self.manager.prepare_store(chunks, _CTX)
424 assert result is not None
425
426 # Complete store (triggers cascade via submit_store on each tier)
426 - self.manager.complete_store(blocks, _CTX, success=True)
427 + self.manager.complete_store(chunks, _CTX, success=True)
428
429 # submit_store was called once per secondary tier
430 self.secondary_tier1.submit_store.assert_called_once()
431 self.secondary_tier2.submit_store.assert_called_once()
432
432 - # Blocks should be in both secondary tiers
433 - assert self.secondary_tier1.get_num_blocks() == 3
434 - assert self.secondary_tier2.get_num_blocks() == 3
433 + # Chunks should be in both secondary tiers
434 + assert self.secondary_tier1.get_num_chunks() == 3
435 + assert self.secondary_tier2.get_num_chunks() == 3
436
436 - # Verify blocks are present
437 + # Verify chunks are present
438 assert all(
438 - self.secondary_tier1.lookup(b, _CTX) is LookupResult.HIT for b in blocks
439 + self.secondary_tier1.lookup(c, _CTX) is LookupResult.HIT for c in chunks
440 )
441 assert all(
441 - self.secondary_tier2.lookup(b, _CTX) is LookupResult.HIT for b in blocks
442 + self.secondary_tier2.lookup(c, _CTX) is LookupResult.HIT for c in chunks
443 )
444
445 def test_ref_cnt_protection_during_cascade(self, manager_setup):
445 - """Test that ref_cnt protects blocks during cascade."""
446 - blocks = to_keys(range(3))
446 + """Test that ref_cnt protects chunks during cascade."""
447 + chunks = to_keys(range(3))
448
449 # Store to primary
450 self._start_request()
450 - result = self.manager.prepare_store(blocks, _CTX)
451 + result = self.manager.prepare_store(chunks, _CTX)
452 assert result is not None
452 - self.manager.complete_store(blocks, _CTX, success=True)
453 + self.manager.complete_store(chunks, _CTX, success=True)
454
454 - # After complete_store, blocks should have ref_cnt > 0
455 + # After complete_store, chunks should have ref_cnt > 0
456 # (one for each secondary tier)
456 - for block_hash in blocks:
457 - block = self.primary_tier._policy.get(block_hash)
457 + for key in chunks:
458 + chunk = self.primary_tier._policy.get(key)
459 # ref_cnt should be 2 (one for each secondary tier)
459 - assert block.ref_cnt == 2
460 + assert chunk.ref_cnt == 2
461
462 # End of step 1: _maybe_process_finished_jobs() was already called by
463 # prepare_store() above (setting the per-step flag), so on_schedule_end()
@@ -466,9 +467,9 @@ class TestTieringOffloadingManager:
467
468 # ref_cnt still held: cascade jobs finished (sync tier) but haven't
469 # been polled yet because the per-step guard skipped the second call.
469 - for block_hash in blocks:
470 - block = self.primary_tier._policy.get(block_hash)
471 - assert block.ref_cnt == 2
470 + for key in chunks:
471 + chunk = self.primary_tier._policy.get(key)
472 + assert chunk.ref_cnt == 2
473
474 # Secondary tiers have completed jobs waiting to be drained
475 assert len(self.secondary_tier1.completed_jobs) > 0
@@ -479,37 +480,37 @@ class TestTieringOffloadingManager:
480 self._simulate_on_schedule_end()
481
482 # After cascade completes, ref_cnt should be 0
482 - for block_hash in blocks:
483 - block = self.primary_tier._policy.get(block_hash)
484 - assert block.ref_cnt == 0
483 + for key in chunks:
484 + chunk = self.primary_tier._policy.get(key)
485 + assert chunk.ref_cnt == 0
486
487 # All completed jobs have been drained
488 assert len(self.secondary_tier1.completed_jobs) == 0
489 assert len(self.secondary_tier2.completed_jobs) == 0
490
491 def test_lookup_from_primary(self, manager_setup):
491 - """Test lookup when blocks are in primary tier."""
492 - blocks = to_keys(range(3))
492 + """Test lookup when chunks are in primary tier."""
493 + chunks = to_keys(range(3))
494
494 - # Store blocks
495 + # Store chunks
496 self._start_request()
496 - self.manager.prepare_store(blocks, _CTX)
497 - self.manager.complete_store(blocks, _CTX, success=True)
497 + self.manager.prepare_store(chunks, _CTX)
498 + self.manager.complete_store(chunks, _CTX, success=True)
499
499 - # Lookup should find all blocks in primary
500 - assert count_hits(self.manager, blocks) == 3
500 + # Lookup should find all chunks in primary
501 + assert count_hits(self.manager, chunks) == 3
502
503 def test_promotion_from_secondary(self, manager_setup):
503 - """Test promotion of blocks from secondary to primary tier."""
504 - blocks = to_keys(range(3))
504 + """Test promotion of chunks from secondary to primary tier."""
505 + chunks = to_keys(range(3))
506
506 - # Manually add blocks to secondary tier (simulate previous cascade)
507 - for block in blocks:
508 - self.secondary_tier1.blocks[block] = True
507 + # Manually add chunks to secondary tier (simulate previous cascade)
508 + for chunk in chunks:
509 + self.secondary_tier1.chunks[chunk] = True
510
510 - # Lookup each block to initiate promotion for all of them
511 - for block in blocks:
512 - result = self.manager.lookup(block, _CTX)
511 + # Lookup each chunk to initiate promotion for all of them
512 + for chunk in chunks:
513 + result = self.manager.lookup(chunk, _CTX)
514 assert result is LookupResult.HIT_PENDING # promotion initiated
515
516 # End of step 1: flushes deferred submit_load() calls
@@ -518,11 +519,11 @@ class TestTieringOffloadingManager:
519 # End of step 2: processes the completed promotion jobs
520 self._simulate_on_schedule_end()
521
521 - # Now blocks should be in primary tier
522 - assert count_hits(self.primary_tier, blocks) == 3
522 + # Now chunks should be in primary tier
523 + assert count_hits(self.primary_tier, chunks) == 3
524
525 # Next lookup should succeed
525 - assert count_hits(self.manager, blocks) == 3
526 + assert count_hits(self.manager, chunks) == 3
527
528 @pytest.mark.parametrize(
529 ("successful_indices", "expected_results"),
@@ -538,18 +539,18 @@ class TestTieringOffloadingManager:
539 ],
540 ids=["partial", "legacy-full-failure"],
541 )
541 - def test_failed_promotion_keeps_only_successful_blocks(
542 + def test_failed_promotion_keeps_only_successful_chunks(
543 self, manager_setup, successful_indices, expected_results
544 ):
544 - blocks = to_keys(range(3))
545 - for block in blocks:
546 - self.secondary_tier1.blocks[block] = True
545 + chunks = to_keys(range(3))
546 + for chunk in chunks:
547 + self.secondary_tier1.chunks[chunk] = True
548
549 def submit_partial(job_metadata: TransferJob) -> None:
550 successful_keys = (
551 None
552 if successful_indices is None
552 - else tuple(blocks[i] for i in successful_indices)
553 + else tuple(chunks[i] for i in successful_indices)
554 )
555 self.secondary_tier1.completed_jobs.append(
556 JobResult(
@@ -561,24 +562,24 @@ class TestTieringOffloadingManager:
562
563 self.secondary_tier1.submit_load = submit_partial
564
564 - for block in blocks:
565 - assert self.manager.lookup(block, _CTX) is LookupResult.HIT_PENDING
565 + for chunk in chunks:
566 + assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT_PENDING
567
568 self._simulate_on_schedule_end()
569 self._simulate_on_schedule_end()
570
571 assert [
571 - self.primary_tier.lookup(block, _CTX) for block in blocks
572 + self.primary_tier.lookup(chunk, _CTX) for chunk in chunks
573 ] == expected_results
574
575 def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup):
575 - """Resolved lookups report one sync delay sample per tier and block."""
576 + """Resolved lookups report one sync delay sample per tier and chunk."""
577 self._start_request()
577 - blocks = to_keys(range(2))
578 + chunks = to_keys(range(2))
579
579 - # No tier has these blocks: they resolve immediately as misses.
580 - for block in blocks:
581 - assert self.manager.lookup(block, _CTX) is LookupResult.MISS
580 + # No tier has these chunks: they resolve immediately as misses.
581 + for chunk in chunks:
582 + assert self.manager.lookup(chunk, _CTX) is LookupResult.MISS
583
584 stats = self.manager.get_stats()
585 assert stats is not None
@@ -614,11 +615,11 @@ class TestTieringOffloadingManager:
615 def test_lookup_does_not_report_async_delay_for_promotion(self, manager_setup):
616 """Promotion time is not included in lookup async delay."""
617 self._start_request()
617 - block = to_keys(range(1))[0]
618 - self.secondary_tier1.blocks[block] = True
618 + chunk = to_keys(range(1))[0]
619 + self.secondary_tier1.chunks[chunk] = True
620
620 - # First lookup finds the block in a secondary tier and defers.
621 - assert self.manager.lookup(block, _CTX) is LookupResult.HIT_PENDING
621 + # First lookup finds the chunk in a secondary tier and defers.
622 + assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT_PENDING
623
624 # Promotion is not treated as unresolved secondary lookup time.
625 self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id])
@@ -648,8 +649,8 @@ class TestTieringOffloadingManager:
649 # Promotion completes on the next scheduler step.
650 self._simulate_on_schedule_end()
651
651 - # Next lookup resolves via the now-promoted primary-tier block.
652 - assert self.manager.lookup(block, _CTX) is LookupResult.HIT
652 + # Next lookup resolves via the now-promoted primary-tier chunk.
653 + assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT
654
655 stats = self.manager.get_stats()
656 if stats is not None:
@@ -663,13 +664,13 @@ class TestTieringOffloadingManager:
664 """Async delay is observed when a RETRY lookup later resolves."""
665 ctx = ReqContext(req_id="req_lookup_finish")
666 self._start_request(ctx)
666 - block = to_keys(range(1))[0]
667 + chunk = to_keys(range(1))[0]
668 self.secondary_tier1.lookup = MagicMock(
669 side_effect=[LookupResult.RETRY, LookupResult.HIT]
670 )
671
672 # First lookup is deferred by the secondary tier.
672 - assert self.manager.lookup(block, ctx) is LookupResult.RETRY
673 + assert self.manager.lookup(chunk, ctx) is LookupResult.RETRY
674
675 stats = self.manager.get_stats()
676 if stats is not None:
@@ -677,7 +678,7 @@ class TestTieringOffloadingManager:
678 stats.reduce(), TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY
679 )
680
680 - assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING
681 + assert self.manager.lookup(chunk, ctx) is LookupResult.HIT_PENDING
682
683 stats = self.manager.get_stats()
684 assert stats is not None
@@ -701,68 +702,68 @@ class TestTieringOffloadingManager:
702
703 def test_partial_lookup(self, manager_setup):
704 """Test lookup with partial hits."""
704 - blocks = to_keys(range(5))
705 + chunks = to_keys(range(5))
706
706 - # Store first 3 blocks to primary
707 + # Store first 3 chunks to primary
708 self._start_request()
708 - self.manager.prepare_store(blocks[:3], _CTX)
709 - self.manager.complete_store(blocks[:3], _CTX, success=True)
709 + self.manager.prepare_store(chunks[:3], _CTX)
710 + self.manager.complete_store(chunks[:3], _CTX, success=True)
711
711 - # Lookup all 5 blocks should return 3 (first 3 found)
712 - assert count_hits(self.manager, blocks) == 3
712 + # Lookup all 5 chunks should return 3 (first 3 found)
713 + assert count_hits(self.manager, chunks) == 3
714
715 def test_eviction_in_primary_tier(self, manager_setup):
716 """Test eviction in primary tier when capacity is exceeded."""
716 - # Primary tier has capacity of 5 blocks
717 + # Primary tier has capacity of 5 chunks
718 # First, fill the primary tier
718 - blocks = to_keys(range(5))
719 + chunks = to_keys(range(5))
720 self._start_request()
720 - result = self.manager.prepare_store(blocks, _CTX)
721 + result = self.manager.prepare_store(chunks, _CTX)
722 assert result is not None
723 assert len(result.keys_to_store) == 5
723 - self.manager.complete_store(blocks, _CTX, success=True)
724 + self.manager.complete_store(chunks, _CTX, success=True)
725
726 # End of step: release ref_cnt from cascade
727 self._simulate_on_schedule_end()
728
728 - # Now try to store 2 more blocks (should trigger eviction)
729 - more_blocks = to_keys(range(5, 7))
730 - result = self.manager.prepare_store(more_blocks, _CTX)
729 + # Now try to store 2 more chunks (should trigger eviction)
730 + more_chunks = to_keys(range(5, 7))
731 + result = self.manager.prepare_store(more_chunks, _CTX)
732
732 - # Should evict 2 blocks from primary tier
733 + # Should evict 2 chunks from primary tier
734 assert result is not None
735 assert len(result.evicted_keys) == 2
736 assert len(result.keys_to_store) == 2
737
738 def test_touch_propagates_to_all_tiers(self, manager_setup):
739 """Test that touch() propagates to all tiers."""
739 - blocks = to_keys(range(3))
740 + chunks = to_keys(range(3))
741
741 - # Store blocks
742 + # Store chunks
743 self._start_request()
743 - self.manager.prepare_store(blocks, _CTX)
744 - self.manager.complete_store(blocks, _CTX, success=True)
744 + self.manager.prepare_store(chunks, _CTX)
745 + self.manager.complete_store(chunks, _CTX, success=True)
746 self._simulate_on_schedule_end()
746 - # for secondary tiers to drain jobs, so primary tier's blocks are evictable.
747 + # for secondary tiers to drain jobs, so primary tier's chunks are evictable.
748 self._simulate_on_schedule_end()
749
750 self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch)
751 self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch)
752
752 - # Touch blocks
753 - self.manager.touch(blocks, _CTX)
753 + # Touch chunks
754 + self.manager.touch(chunks, _CTX)
755
756 # Verify touch was called on primary tier (check LRU order)
756 - primary_keys = list(self.primary_tier._policy.evictable_blocks.keys())
757 - assert primary_keys[-3:] == list(reversed(blocks))
757 + primary_keys = list(self.primary_tier._policy.evictable_chunks.keys())
758 + assert primary_keys[-3:] == list(reversed(chunks))
759
760 # Verify touch was propagated to all secondary tiers
760 - self.secondary_tier1.touch.assert_called_once_with(blocks, _CTX)
761 - self.secondary_tier2.touch.assert_called_once_with(blocks, _CTX)
761 + self.secondary_tier1.touch.assert_called_once_with(chunks, _CTX)
762 + self.secondary_tier2.touch.assert_called_once_with(chunks, _CTX)
763
764 def test_failed_store_no_cascade(self, manager_setup):
765 """Test that failed GPU→primary store doesn't cascade."""
765 - blocks = to_keys(range(3))
766 + chunks = to_keys(range(3))
767
768 self.secondary_tier1.submit_store = MagicMock(
769 wraps=self.secondary_tier1.submit_store
@@ -773,11 +774,11 @@ class TestTieringOffloadingManager:
774
775 # Prepare store
776 self._start_request()
776 - result = self.manager.prepare_store(blocks, _CTX)
777 + result = self.manager.prepare_store(chunks, _CTX)
778 assert result is not None
779
780 # Complete store with failure — cascade must not happen
780 - self.manager.complete_store(blocks, _CTX, success=False)
781 + self.manager.complete_store(chunks, _CTX, success=False)
782
783 # submit_store was never called on either secondary tier
784 self.secondary_tier1.submit_store.assert_not_called()
@@ -786,12 +787,12 @@ class TestTieringOffloadingManager:
787 def test_lookup_batches_submit_load_per_request(self, manager_setup):
788 """lookup() defers submit_load until on_schedule_end(), one per request.
789
789 - Blocks from different requests each get their own submit_load call, each
790 + Chunks from different requests each get their own submit_load call, each
791 carrying the correct req_context.
792 """
792 - blocks = to_keys(range(4))
793 - for block in blocks:
794 - self.secondary_tier1.blocks[block] = True
793 + chunks = to_keys(range(4))
794 + for chunk in chunks:
795 + self.secondary_tier1.chunks[chunk] = True
796
797 self.secondary_tier1.submit_load = MagicMock(
798 wraps=self.secondary_tier1.submit_load
@@ -801,10 +802,10 @@ class TestTieringOffloadingManager:
802 ctx_b = ReqContext(req_id="req_b")
803
804 # All lookups return HIT_PENDING: secondary hit triggers promotion
804 - assert self.manager.lookup(blocks[0], ctx_a) is LookupResult.HIT_PENDING
805 - assert self.manager.lookup(blocks[1], ctx_a) is LookupResult.HIT_PENDING
806 - assert self.manager.lookup(blocks[2], ctx_b) is LookupResult.HIT_PENDING
807 - assert self.manager.lookup(blocks[3], ctx_b) is LookupResult.HIT_PENDING
805 + assert self.manager.lookup(chunks[0], ctx_a) is LookupResult.HIT_PENDING
806 + assert self.manager.lookup(chunks[1], ctx_a) is LookupResult.HIT_PENDING
807 + assert self.manager.lookup(chunks[2], ctx_b) is LookupResult.HIT_PENDING
808 + assert self.manager.lookup(chunks[3], ctx_b) is LookupResult.HIT_PENDING
809
810 # submit_load must not fire during lookup - only at end of step
811 self.secondary_tier1.submit_load.assert_not_called()
@@ -816,20 +817,20 @@ class TestTieringOffloadingManager:
817 calls = self.secondary_tier1.submit_load.call_args_list
818 jm_a = calls[0].args[0]
819 jm_b = calls[1].args[0]
819 - assert set(jm_a.keys) == {blocks[0], blocks[1]}
820 + assert set(jm_a.keys) == {chunks[0], chunks[1]}
821 assert jm_a.req_context is ctx_a
821 - assert set(jm_b.keys) == {blocks[2], blocks[3]}
822 + assert set(jm_b.keys) == {chunks[2], chunks[3]}
823 assert jm_b.req_context is ctx_b
824
824 - def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup):
825 - """A block looked up by two requests in the same step is promoted once.
825 + def test_lookup_shared_chunk_no_duplicate_promotion(self, manager_setup):
826 + """A chunk looked up by two requests in the same step is promoted once.
827
828 The first lookup initiates promotion (returns None via secondary hit).
829 The second lookup sees ref_cnt=-1 on the primary slot and returns None
830 via the primary in-flight path — without triggering a second promotion.
831 """
831 - shared_block = to_keys([0])[0]
832 - self.secondary_tier1.blocks[shared_block] = True
832 + shared_chunk = to_keys([0])[0]
833 + self.secondary_tier1.chunks[shared_chunk] = True
834
835 self.secondary_tier1.submit_load = MagicMock(
836 wraps=self.secondary_tier1.submit_load
@@ -838,10 +839,10 @@ class TestTieringOffloadingManager:
839 ctx_a = ReqContext(req_id="req_a")
840 ctx_b = ReqContext(req_id="req_b")
841
841 - result_a = self.manager.lookup(shared_block, ctx_a)
842 - result_b = self.manager.lookup(shared_block, ctx_b)
842 + result_a = self.manager.lookup(shared_chunk, ctx_a)
843 + result_b = self.manager.lookup(shared_chunk, ctx_b)
844
844 - # Both lookups find the shared_block and trigger promotion
845 + # Both lookups find the shared_chunk and trigger promotion
846 # returning HIT_PENDING.
847 assert result_a is LookupResult.HIT_PENDING
848 assert result_b is LookupResult.HIT_PENDING
@@ -851,12 +852,12 @@ class TestTieringOffloadingManager:
852 # Only one submit_load call despite two lookups
853 self.secondary_tier1.submit_load.assert_called_once()
854 job_metadata = self.secondary_tier1.submit_load.call_args.args[0]
854 - assert list(job_metadata.keys) == [shared_block]
855 + assert list(job_metadata.keys) == [shared_chunk]
856 assert job_metadata.req_context is ctx_a
857
858 def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup):
859 """complete_store cascades to secondary tiers with the correct req_context."""
859 - blocks = to_keys(range(2))
860 + chunks = to_keys(range(2))
861
862 self.secondary_tier1.submit_store = MagicMock(
863 wraps=self.secondary_tier1.submit_store
@@ -865,8 +866,8 @@ class TestTieringOffloadingManager:
866 ctx = ReqContext(req_id="req_ctx", kv_transfer_params={"key": "value"})
867
868 self._start_request(ctx)
868 - self.manager.prepare_store(blocks, ctx)
869 - self.manager.complete_store(blocks, ctx, success=True)
869 + self.manager.prepare_store(chunks, ctx)
870 + self.manager.complete_store(chunks, ctx, success=True)
871
872 assert self.secondary_tier1.submit_store.call_count == 1
873 job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
@@ -876,7 +877,7 @@ class TestTieringOffloadingManager:
877 self, manager_setup
878 ):
879 """Manager hook is eager; secondary hooks wait for cascade submission."""
879 - blocks = to_keys(range(2))
880 + chunks = to_keys(range(2))
881 ctx = ReqContext(req_id="req_delayed_secondary")
882 calls: list[tuple[str, str]] = []
883
@@ -911,14 +912,14 @@ class TestTieringOffloadingManager:
912 )
913
914 self._start_request(ctx)
914 - self.manager.prepare_store(blocks, ctx)
915 + self.manager.prepare_store(chunks, ctx)
916 self.manager.on_request_finished(ctx)
917
918 assert calls == [("primary_finish", ctx.req_id)]
919 self.secondary_tier1.on_request_finished.assert_not_called()
920 self.secondary_tier2.on_request_finished.assert_not_called()
921
921 - self.manager.complete_store(blocks, ctx, success=True)
922 + self.manager.complete_store(chunks, ctx, success=True)
923
924 assert calls == [
925 ("primary_finish", ctx.req_id),
@@ -930,7 +931,7 @@ class TestTieringOffloadingManager:
931
932 def test_failed_store_finalizes_finished_request(self, manager_setup):
933 """Failed primary stores still unblock secondary finalization."""
933 - blocks = to_keys(range(2))
934 + chunks = to_keys(range(2))
935 ctx = ReqContext(req_id="req_failed_store_finalize")
936
937 self.secondary_tier1.submit_store = MagicMock(
@@ -947,13 +948,13 @@ class TestTieringOffloadingManager:
948 )
949
950 self._start_request(ctx)
950 - self.manager.prepare_store(blocks, ctx)
951 + self.manager.prepare_store(chunks, ctx)
952 self.manager.on_request_finished(ctx)
953
954 self.secondary_tier1.on_request_finished.assert_not_called()
955 self.secondary_tier2.on_request_finished.assert_not_called()
956
956 - self.manager.complete_store(blocks, ctx, success=False)
957 + self.manager.complete_store(chunks, ctx, success=False)
958
959 self.secondary_tier1.submit_store.assert_not_called()
960 self.secondary_tier2.submit_store.assert_not_called()
@@ -981,7 +982,7 @@ class TestTieringOffloadingManager:
982
983 def test_reset_cache_finalizes_delayed_secondary_request(self, manager_setup):
984 """reset_cache abandons pending primary stores and finalizes secondaries."""
984 - blocks = to_keys(range(2))
985 + chunks = to_keys(range(2))
986 ctx = ReqContext(req_id="req_reset_finalize_secondary")
987
988 self.secondary_tier1.on_request_finished = MagicMock(
@@ -992,7 +993,7 @@ class TestTieringOffloadingManager:
993 )
994
995 self._start_request(ctx)
995 - self.manager.prepare_store(blocks, ctx)
996 + self.manager.prepare_store(chunks, ctx)
997 self.manager.on_request_finished(ctx)
998
999 self.secondary_tier1.on_request_finished.assert_not_called()
@@ -1008,8 +1009,8 @@ class TestTieringOffloadingManager:
1009 self, manager_setup
1010 ):
1011 """reset_cache drops active pending stores so resumed requests finalize."""
1011 - initial_blocks = to_keys(range(2))
1012 - resumed_blocks = to_keys(range(2, 4))
1012 + initial_chunks = to_keys(range(2))
1013 + resumed_chunks = to_keys(range(2, 4))
1014 ctx = ReqContext(req_id="req_reset_resume")
1015
1016 self.secondary_tier1.on_request_finished = MagicMock(
@@ -1020,7 +1021,7 @@ class TestTieringOffloadingManager:
1021 )
1022
1023 self._start_request(ctx)
1023 - self.manager.prepare_store(initial_blocks, ctx)
1024 + self.manager.prepare_store(initial_chunks, ctx)
1025 assert self.manager._req_state[ctx.req_id].pending_primary_stores == 1
1026
1027 self.manager.reset_cache()
@@ -1030,8 +1031,8 @@ class TestTieringOffloadingManager:
1031 self.secondary_tier1.on_request_finished.assert_not_called()
1032 self.secondary_tier2.on_request_finished.assert_not_called()
1033
1033 - self.manager.prepare_store(resumed_blocks, ctx)
1034 - self.manager.complete_store(resumed_blocks, ctx, success=True)
1034 + self.manager.prepare_store(resumed_chunks, ctx)
1035 + self.manager.complete_store(resumed_chunks, ctx, success=True)
1036 self.manager.on_request_finished(ctx)
1037
1038 self.secondary_tier1.on_request_finished.assert_called_once_with(ctx)
@@ -1039,12 +1040,12 @@ class TestTieringOffloadingManager:
1040 assert ctx.req_id not in self.manager._req_state
1041
1042 def test_on_new_request_lifecycle(self, manager_setup):
1042 - """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it,
1043 + """Policy defaults to CHUNK_LEVEL, escalates when a tier requests it,
1044 and is cleaned up on on_request_finished."""
1044 - # Default: all tiers return BLOCK_LEVEL
1045 + # Default: all tiers return CHUNK_LEVEL
1046 ctx = ReqContext(req_id="req_policy_lifecycle")
1047 result = self.manager.on_new_request(ctx)
1047 - assert result.policy == OffloadPolicy.BLOCK_LEVEL
1048 + assert result.policy == OffloadPolicy.CHUNK_LEVEL
1049 assert self.manager._req_state[ctx.req_id].request_level_tiers is None
1050 self.manager.on_request_finished(ctx)
1051 assert ctx.req_id not in self.manager._req_state
@@ -1064,26 +1065,26 @@ class TestTieringOffloadingManager:
1065 assert ctx.req_id not in self.manager._req_state
1066
1067 @pytest.mark.parametrize("new_ids", [(), (3, 4)], ids=["fully_warm", "mixed"])
1067 - def test_prepare_store_cascades_existing_blocks_to_request_level_tiers(
1068 + def test_prepare_store_cascades_existing_chunks_to_request_level_tiers(
1069 self, manager_setup, new_ids
1070 ):
1070 - """prepare_store cascades hit blocks to request-level tiers only.
1071 + """prepare_store cascades hit chunks to request-level tiers only.
1072
1073 The fully-warm case is the producer shape from issue #52808: nothing is
1074 new, so primary returns an empty keys_to_store and no GPU->primary
1074 - transfer runs. A request-level tier must still be handed the blocks,
1075 + transfer runs. A request-level tier must still be handed the chunks,
1076 otherwise a peer waiting on them has nothing to fetch.
1077 """
1077 - # Store some blocks to primary first
1078 - existing_blocks = to_keys(range(3))
1078 + # Store some chunks to primary first
1079 + existing_chunks = to_keys(range(3))
1080 self._start_request()
1080 - result = self.manager.prepare_store(existing_blocks, _CTX)
1081 + result = self.manager.prepare_store(existing_chunks, _CTX)
1082 assert result is not None
1082 - self.manager.complete_store(existing_blocks, _CTX, success=True)
1083 + self.manager.complete_store(existing_chunks, _CTX, success=True)
1084 # Drain cascade completions
1085 self._simulate_on_schedule_end()
1086
1086 - # Make tier1 request-level, tier2 stays block-level
1087 + # Make tier1 request-level, tier2 stays chunk-level
1088 self.secondary_tier1.on_new_request = lambda req_context: (
1089 RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL)
1090 )
@@ -1099,20 +1100,20 @@ class TestTieringOffloadingManager:
1100 wraps=self.secondary_tier2.submit_store
1101 )
1102
1102 - # Call prepare_store with existing + new blocks
1103 - new_blocks = to_keys(new_ids)
1104 - all_blocks = existing_blocks + new_blocks
1105 - result = self.manager.prepare_store(all_blocks, ctx)
1103 + # Call prepare_store with existing + new chunks
1104 + new_chunks = to_keys(new_ids)
1105 + all_chunks = existing_chunks + new_chunks
1106 + result = self.manager.prepare_store(all_chunks, ctx)
1107 assert result is not None
1107 - assert set(result.keys_to_store) == set(new_blocks)
1108 + assert set(result.keys_to_store) == set(new_chunks)
1109
1109 - # Only tier1 (request-level) should get existing blocks cascaded now.
1110 - # New blocks are cascaded to ALL tiers later via complete_store().
1110 + # Only tier1 (request-level) should get existing chunks cascaded now.
1111 + # New chunks are cascaded to ALL tiers later via complete_store().
1112 self.secondary_tier1.submit_store.assert_called_once()
1113 job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
1113 - assert set(job_metadata.keys) == set(existing_blocks)
1114 + assert set(job_metadata.keys) == set(existing_chunks)
1115
1115 - # tier2 (block-level) does not get existing blocks here.
1116 + # tier2 (chunk-level) does not get existing chunks here.
1117 self.secondary_tier2.submit_store.assert_not_called()
1118
1119 def _make_request_level_request(self, req_id: str) -> ReqContext:
@@ -1150,7 +1151,7 @@ class TestTieringOffloadingManager:
1151 self.manager.primary_tier.lookup = lambda key, req_context: (LookupResult.RETRY)
1152
1153 with pytest.raises(AssertionError):
1153 - self.manager._cascade_existing_blocks_to_request_level_tiers(keys, ctx, {0})
1154 + self.manager._cascade_existing_chunks_to_request_level_tiers(keys, ctx, {0})
1155
1156 def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext:
1157 """Leave a GPU->primary write for `keys` open, so they look present to
@@ -1202,7 +1203,7 @@ class TestTieringOffloadingManager:
1203 def test_cascade_drops_deferred_keys_whose_primary_write_failed(
1204 self, manager_setup
1205 ):
1205 - """A failed write frees the block, so the deferred key resolves to MISS
1206 + """A failed write frees the chunk, so the deferred key resolves to MISS
1207 and the request finalizes instead of parking forever."""
1208 keys = to_keys(range(3))
1209 writer_ctx = self._start_in_flight_primary_write(keys)
@@ -1224,21 +1225,21 @@ class TestTieringOffloadingManager:
1225 """reset_cache wipes every kind of orchestrator state and resets
1226 primary tier; pending submissions are dropped without being sent
1227 to the secondary tier. Active request state is retained."""
1227 - # Cascade — populates primary blocks and leaves cascade jobs
1228 + # Cascade — populates primary chunks and leaves cascade jobs
1229 # in _jobs (the synchronous example tier has already
1230 # queued completions); reset_cache's drain loop will pick them up.
1230 - blocks = to_keys(range(3))
1231 + chunks = to_keys(range(3))
1232 self._start_request()
1232 - self.manager.prepare_store(blocks, _CTX)
1233 - self.manager.complete_store(blocks, _CTX, success=True)
1233 + self.manager.prepare_store(chunks, _CTX)
1234 + self.manager.complete_store(chunks, _CTX, success=True)
1235 assert self.manager._jobs
1236
1237 # Pending promotion submission (deferred — no on_schedule_end after
1238 # the lookup that staged it).
1238 - promo_block = to_keys([99])[0]
1239 - self.secondary_tier1.blocks[promo_block] = True
1239 + promo_chunk = to_keys([99])[0]
1240 + self.secondary_tier1.chunks[promo_chunk] = True
1241 assert (
1241 - self.manager.lookup(promo_block, ReqContext(req_id="pending"))
1242 + self.manager.lookup(promo_chunk, ReqContext(req_id="pending"))
1243 is LookupResult.HIT_PENDING
1244 )
1245 assert self.manager._pending_load_submissions
@@ -1268,10 +1269,10 @@ class TestTieringOffloadingManager:
1269 assert self.manager._processed_jobs_this_step is False
1270
1271 # Primary tier reset to a fresh state.
1271 - assert self.primary_tier._num_allocated_blocks == 0
1272 + assert self.primary_tier._num_allocated_chunks == 0
1273 assert self.primary_tier._free_list == []
1273 - for block in blocks:
1274 - assert self.primary_tier.lookup(block, _CTX) is LookupResult.MISS
1274 + for chunk in chunks:
1275 + assert self.primary_tier.lookup(chunk, _CTX) is LookupResult.MISS
1276
1277 # Pending submission was dropped, not submitted.
1278 self.secondary_tier1.submit_load.assert_not_called()
@@ -1291,10 +1292,10 @@ class TestTieringOffloadingManager:
1292 )
1293
1294 # Drive a cascade so a job lands in _jobs.
1294 - blocks = to_keys(range(3))
1295 + chunks = to_keys(range(3))
1296 self._start_request()
1296 - self.manager.prepare_store(blocks, _CTX)
1297 - self.manager.complete_store(blocks, _CTX, success=True)
1297 + self.manager.prepare_store(chunks, _CTX)
1298 + self.manager.complete_store(chunks, _CTX, success=True)
1299 assert self.manager._jobs
1300
1301 self.manager.reset_cache()
@@ -1315,20 +1316,20 @@ class TestTieringOffloadingManager:
1316 self, manager_setup, load_tier_filter
1317 ):
1318 """Filter excluding secondary medium returns MISS from secondaries
1318 - even when they hold the block; primary is unaffected."""
1319 - blocks = to_keys(range(2))
1320 - # Put one block in primary, one only in secondary
1319 + even when they hold the chunk; primary is unaffected."""
1320 + chunks = to_keys(range(2))
1321 + # Put one chunk in primary, one only in secondary
1322 self._start_request()
1322 - self.manager.prepare_store(blocks[:1], _CTX)
1323 - self.manager.complete_store(blocks[:1], _CTX, success=True)
1324 - self.secondary_tier1.blocks[blocks[1]] = True
1323 + self.manager.prepare_store(chunks[:1], _CTX)
1324 + self.manager.complete_store(chunks[:1], _CTX, success=True)
1325 + self.secondary_tier1.chunks[chunks[1]] = True
1326
1327 # Secondaries have medium=CPU, so load_tier_filter skips them.
1328 self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup)
1329
1330 ctx = ReqContext(req_id="r1", load_tier_filter=load_tier_filter)
1330 - assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT
1331 - assert self.manager.lookup(blocks[1], ctx) is LookupResult.MISS
1331 + assert self.manager.lookup(chunks[0], ctx) is LookupResult.HIT
1332 + assert self.manager.lookup(chunks[1], ctx) is LookupResult.MISS
1333 self.secondary_tier1.lookup.assert_not_called()
1334
1335 @pytest.mark.parametrize(
@@ -1344,13 +1345,13 @@ class TestTieringOffloadingManager:
1345 self, manager_setup, load_tier_filter
1346 ):
1347 """Filter that matches the secondary's medium allows lookup."""
1347 - blocks = to_keys(range(1))
1348 - self.secondary_tier1.blocks[blocks[0]] = True
1348 + chunks = to_keys(range(1))
1349 + self.secondary_tier1.chunks[chunks[0]] = True
1350
1351 self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup)
1352
1353 ctx = ReqContext(req_id="r2", load_tier_filter=load_tier_filter)
1353 - assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT_PENDING
1354 + assert self.manager.lookup(chunks[0], ctx) is LookupResult.HIT_PENDING
1355 self.secondary_tier1.lookup.assert_called()
1356
1357
@@ -1360,7 +1361,7 @@ class TestTieringOffloadingWithoutSecondaryTiers:
1361 def test_works_without_secondary_tiers(self):
1362 """Test that manager works with empty secondary_tiers list."""
1363 primary_tier = CPUPrimaryTierOffloadingManager(
1363 - num_blocks=5, mmap_region=_mock_mmap_region(5)
1364 + num_chunks=5, mmap_region=_mock_mmap_region(5)
1365 )
1366
1367 # Create manager with no secondary tiers
@@ -1368,15 +1369,15 @@ class TestTieringOffloadingWithoutSecondaryTiers:
1369 primary_tier=primary_tier, secondary_tiers=[]
1370 )
1371
1371 - blocks = to_keys(range(3))
1372 + chunks = to_keys(range(3))
1373
1374 # Should work like a regular OffloadingManager
1375 manager.on_new_request(_CTX)
1375 - result = manager.prepare_store(blocks, _CTX)
1376 + result = manager.prepare_store(chunks, _CTX)
1377 assert result is not None
1377 - manager.complete_store(blocks, _CTX, success=True)
1378 + manager.complete_store(chunks, _CTX, success=True)
1379
1379 - assert count_hits(manager, blocks) == 3
1380 + assert count_hits(manager, chunks) == 3
1381
1382
1383 @pytest.mark.parametrize(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+1 -1
@@ -1131,7 +1131,7 @@ class OffloadingConnectorScheduler:
1131 # Skip prefix-hit chunks for block-level policy; for
1132 # request-level, next_stored_chunk_idx stays at 0 so all
1133 # chunks (including hits) are offloaded.
1134 - if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL:
1134 + if req_status.offloading_context.policy == OffloadPolicy.CHUNK_LEVEL:
1135 group_state.next_stored_chunk_idx = num_chunks
1136
1137 src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
vllm/v1/kv_offload/base.py
+8 -5
@@ -114,17 +114,17 @@ class LookupResult(Enum):
114
115
116 class OffloadPolicy(Enum):
117 - # Offload only newly-computed blocks as they arrive; prefix-hit
118 - # blocks (already offloaded by a prior request) are skipped.
119 - BLOCK_LEVEL = "block_level"
120 - # Offload all blocks for the request, including prefix hits.
117 + # Offload only newly-computed chunks as they arrive; prefix-hit
118 + # chunks (already offloaded by a prior request) are skipped.
119 + CHUNK_LEVEL = "chunk_level"
120 + # Offload all chunks for the request, including prefix hits.
121 # Used by tiers that need the complete KV context for a request.
122 REQUEST_LEVEL = "request_level"
123
124
125 @dataclass
126 class RequestOffloadingContext:
127 - policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL
127 + policy: OffloadPolicy = OffloadPolicy.CHUNK_LEVEL
128
129
130 class ScheduleEndContext(NamedTuple):
@@ -400,6 +400,9 @@ class OffloadingManager(ABC):
400 class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
401 """
402 Spec for loading/storing KV blocks from given block numbers.
403 +
404 + Subclass semantics differ: GPULoadStoreSpec.block_ids are GPU block
405 + indices; CPULoadStoreSpec.block_ids are CPU cache chunk indices.
406 """
407
408 def __init__(self, block_ids: list[int]):
vllm/v1/kv_offload/cpu/common.py
+11 -2
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import numpy as np
4 +
5 from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec
6
7
@@ -12,6 +14,13 @@ class CPUOffloadingMetrics:
14
15
16 class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
17 + """Spec for loading/storing KV chunks to/from CPU memory.
18 +
19 + The inherited block_ids field holds chunk indices into the
20 + CPU cache (not GPU block IDs). The chunk_ids alias exposes
21 + the same array under the name used by the tiering layer.
22 """
16 - Spec for loading/storing a KV block to CPU memory.
17 - """
23 +
24 + @property
25 + def chunk_ids(self) -> np.ndarray:
26 + return self.block_ids
vllm/v1/kv_offload/cpu/gpu_worker.py
+10 -10
@@ -257,7 +257,7 @@ class SingleDirectionOffloadingHandler:
257 gpu_tensors: list of GPU KV cache tensors.
258 Each of shape (num_gpu_blocks, gpu_page_size_bytes) with dtype int8.
259 cpu_tensors: list of CPU KV cache tensors.
260 - Each of shape (num_cpu_blocks, cpu_page_size_bytes) with dtype int8.
260 + Each of shape (num_cpu_chunks, cpu_page_size_bytes) with dtype int8.
261 Order should match gpu_tensors.
262 layer_refs_per_group: list of CanonicalKVCacheRef per group.
263 gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU.
@@ -527,20 +527,20 @@ class SingleDirectionOffloadingHandler:
527 # 1. GPU -> CPU
528 # 2. CPU -> GPU
529 #
530 - # transfers are also to CPU blocks, EXCEPT MAYBE for the first and last block.
531 - # i.e. the first and last CPU blocks in src_blocks can match against
530 + # transfers are also to CPU chunks, EXCEPT MAYBE for the first and last chunk.
531 + # i.e. the first and last CPU chunks in src_blocks can match against
532 # a smaller (byte-wise) set of GPU blocks in dst_blocks.
533 # In such cases, we may need to skip some gpu-sized sub-blocks,
534 - # and start reading/writing from the middle of the first CPU block.
534 + # and start reading/writing from the middle of the first CPU chunk.
535 # If we have multiple KV cache groups (when using HMA with hybrid models),
536 - # we may have a partial first/last CPU block per each group.
536 + # we may have a partial first/last CPU chunk per each group.
537 # The group_sizes parameter encodes the size of each group of blocks
538 # in the GPU dst_blocks.
539 # If group_sizes is None, we assume all blocks belong to a single group.
540 # The logical_offset parameter maps each group of blocks to its logical
541 # offset inside the request, counting in GPU blocks.
542 # This allows us to find the correct starting position
543 - # in the matching first CPU block.
543 + # in the matching first CPU chunk.
544
545 # extract group_sizes from the GPU spec
546 gpu_spec = src_spec if self.gpu_to_cpu else dst_spec
@@ -751,7 +751,7 @@ class CPUOffloadingWorker(OffloadingWorker):
751 self,
752 kv_caches: CanonicalKVCaches,
753 blocks_per_chunk: int,
754 - num_cpu_blocks: int,
754 + num_cpu_chunks: int,
755 mmap_region: SharedOffloadRegion | None = None,
756 canonical_layout: bool = False,
757 ):
@@ -790,16 +790,16 @@ class CPUOffloadingWorker(OffloadingWorker):
790 else:
791 t0 = time.monotonic()
792 cpu_tensor = torch.zeros(
793 - (num_cpu_blocks, cpu_page_size_bytes),
793 + (num_cpu_chunks, cpu_page_size_bytes),
794 dtype=torch.int8,
795 device="cpu",
796 pin_memory=pin_memory,
797 )
798 logger.debug(
799 "torch.zeros pinned tensor %d×%d (%.2f GB): %.3f s",
800 - num_cpu_blocks,
800 + num_cpu_chunks,
801 cpu_page_size_bytes,
802 - num_cpu_blocks * cpu_page_size_bytes / 1e9,
802 + num_cpu_chunks * cpu_page_size_bytes / 1e9,
803 time.monotonic() - t0,
804 )
805
vllm/v1/kv_offload/cpu/manager.py
+87 -87
@@ -23,7 +23,7 @@ from vllm.v1.kv_offload.cpu.common import (
23 CPULoadStoreSpec,
24 CPUOffloadingMetrics,
25 )
26 -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
26 +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
27 from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory
28
29
@@ -34,14 +34,14 @@ class CPUOffloadingManager(OffloadingManager):
34 register their own or be loaded out-of-tree via cache_policy_module_path).
35
36 The manager owns all shared logic: ref-counting, event emission,
37 - block pool management, and the prepare_store/complete_store skeletons.
38 - Policy-specific block organization and eviction decisions are delegated
37 + chunk pool management, and the prepare_store/complete_store skeletons.
38 + Policy-specific chunk organization and eviction decisions are delegated
39 to the CachePolicy implementation.
40 """
41
42 def __init__(
43 self,
44 - num_blocks: int,
44 + num_chunks: int,
45 cache_policy: str = "lru",
46 cache_policy_module_path: str | None = None,
47 enable_events: bool = False,
@@ -49,59 +49,59 @@ class CPUOffloadingManager(OffloadingManager):
49 max_tracker_size: int = 64_000,
50 ):
51 self.medium: Medium = Medium.CPU
52 - self._num_blocks: int = num_blocks
53 - self._num_allocated_blocks: int = 0
52 + self._num_chunks: int = num_chunks
53 + self._num_allocated_chunks: int = 0
54 self._free_list: list[int] = []
55 self.events: list[OffloadingEvent] | None = [] if enable_events else None
56 policy_cls = CachePolicyFactory.get_cache_policy_cls(
57 cache_policy, cache_policy_module_path
58 )
59 - self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)
60 - # Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0.
61 - self._num_evictable_cache_blocks: int = 0
62 - # Track blocks with an in-flight store (ref_cnt -1, not yet completed).
63 - self._num_write_pending_blocks: int = 0
59 + self._policy: CachePolicy = policy_cls(cache_capacity=num_chunks)
60 + # Track the number of chunks in the cache that are evictable. i.e. ref_cnt 0.
61 + self._num_evictable_cache_chunks: int = 0
62 + # Track chunks with an in-flight store (ref_cnt -1, not yet completed).
63 + self._num_write_pending_chunks: int = 0
64
65 self.store_threshold: int = store_threshold
66 self.max_tracker_size: int = max_tracker_size
67 self.stores_skipped_in_current_batch: int = 0
68 self.allocation_sizes_in_current_batch: list[int] = []
69
70 - # Number of block references. It is ordered so can evict the LRU entry in O(1).
70 + # Number of chunk references. It is ordered so can evict the LRU entry in O(1).
71 self.counts: OrderedDict[OffloadKey, int] | None = (
72 OrderedDict() if store_threshold >= 2 else None
73 )
74
75 - # --- block pool ---
75 + # --- chunk pool ---
76
77 - def _get_num_free_blocks(self) -> int:
78 - return len(self._free_list) + self._num_blocks - self._num_allocated_blocks
77 + def _get_num_free_chunks(self) -> int:
78 + return len(self._free_list) + self._num_chunks - self._num_allocated_chunks
79
80 - def _allocate_blocks(self, keys: list[OffloadKey]) -> list[BlockStatus]:
81 - num_fresh = min(len(keys), self._num_blocks - self._num_allocated_blocks)
80 + def _allocate_chunks(self, keys: list[OffloadKey]) -> list[ChunkStatus]:
81 + num_fresh = min(len(keys), self._num_chunks - self._num_allocated_chunks)
82 num_reused = len(keys) - num_fresh
83 assert len(self._free_list) >= num_reused
84
85 - # allocate fresh blocks
86 - blocks: list[BlockStatus] = []
85 + # allocate fresh chunks
86 + chunks: list[ChunkStatus] = []
87 for _ in range(num_fresh):
88 - blocks.append(BlockStatus(self._num_allocated_blocks))
89 - self._num_allocated_blocks += 1
88 + chunks.append(ChunkStatus(self._num_allocated_chunks))
89 + self._num_allocated_chunks += 1
90
91 - # allocate reused blocks
91 + # allocate reused chunks
92 for _ in range(num_reused):
93 - blocks.append(BlockStatus(self._free_list.pop()))
94 - return blocks
93 + chunks.append(ChunkStatus(self._free_list.pop()))
94 + return chunks
95
96 - def _free_block(self, block: BlockStatus) -> None:
97 - self._free_list.append(block.block_id)
96 + def _free_chunk(self, chunk: ChunkStatus) -> None:
97 + self._free_list.append(chunk.chunk_id)
98
99 def _get_load_store_spec(
100 self,
101 keys: Iterable[OffloadKey],
102 - blocks: Iterable[BlockStatus],
102 + chunks: Iterable[ChunkStatus],
103 ) -> CPULoadStoreSpec:
104 - return CPULoadStoreSpec([block.block_id for block in blocks])
104 + return CPULoadStoreSpec([chunk.chunk_id for chunk in chunks])
105
106 def _record_accesses(self, keys: Collection[OffloadKey]) -> None:
107 """Record an offer without evicting its tracked candidates."""
@@ -132,10 +132,10 @@ class CPUOffloadingManager(OffloadingManager):
132
133 @override
134 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
135 - block = self._policy.get(key)
136 - if block is None:
135 + chunk = self._policy.get(key)
136 + if chunk is None:
137 return LookupResult.MISS
138 - if not block.is_ready:
138 + if not chunk.is_ready:
139 return LookupResult.HIT_PENDING
140 return LookupResult.HIT
141
@@ -145,18 +145,18 @@ class CPUOffloadingManager(OffloadingManager):
145 keys: Collection[OffloadKey],
146 req_context: ReqContext,
147 ) -> LoadStoreSpec:
148 - blocks = []
148 + chunks = []
149 for key in keys:
150 - block = self._policy.get(key)
151 - assert block is not None, f"Block {key!r} not found in cache"
152 - assert block.is_ready, f"Block {key!r} is not ready for reading"
153 - if block.ref_cnt == 0:
150 + chunk = self._policy.get(key)
151 + assert chunk is not None, f"Chunk {key!r} not found in cache"
152 + assert chunk.is_ready, f"Chunk {key!r} is not ready for reading"
153 + if chunk.ref_cnt == 0:
154 self._policy.mark_non_evictable(key)
155 - self._num_evictable_cache_blocks -= 1 # ref_cnt 0 -> 1
156 - assert self._num_evictable_cache_blocks >= 0
157 - block.ref_cnt += 1
158 - blocks.append(block)
159 - return self._get_load_store_spec(keys, blocks)
155 + self._num_evictable_cache_chunks -= 1 # ref_cnt 0 -> 1
156 + assert self._num_evictable_cache_chunks >= 0
157 + chunk.ref_cnt += 1
158 + chunks.append(chunk)
159 + return self._get_load_store_spec(keys, chunks)
160
161 @override
162 def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
@@ -167,12 +167,12 @@ class CPUOffloadingManager(OffloadingManager):
167 self, keys: Collection[OffloadKey], req_context: ReqContext
168 ) -> None:
169 for key in keys:
170 - block = self._policy.get(key)
171 - assert block is not None, f"Block {key!r} not found"
172 - assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0"
173 - block.ref_cnt -= 1
174 - if block.ref_cnt == 0:
175 - self._num_evictable_cache_blocks += 1 # ref_cnt 1 -> 0
170 + chunk = self._policy.get(key)
171 + assert chunk is not None, f"Chunk {key!r} not found"
172 + assert chunk.ref_cnt > 0, f"Chunk {key!r} ref_cnt is already 0"
173 + chunk.ref_cnt -= 1
174 + if chunk.ref_cnt == 0:
175 + self._num_evictable_cache_chunks += 1 # ref_cnt 1 -> 0
176 self._policy.mark_evictable(key)
177
178 @override
@@ -186,7 +186,7 @@ class CPUOffloadingManager(OffloadingManager):
186 self._record_accesses(keys)
187 keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold]
188 self.stores_skipped_in_current_batch += num_keys - len(keys)
189 - # filter out blocks that are already stored
189 + # filter out chunks that are already stored
190 keys_to_store = [k for k in keys if self._policy.get(k) is None]
191
192 if not keys_to_store:
@@ -197,29 +197,29 @@ class CPUOffloadingManager(OffloadingManager):
197 )
198
199 self.allocation_sizes_in_current_batch.append(len(keys_to_store))
200 - num_blocks_to_evict = len(keys_to_store) - self._get_num_free_blocks()
200 + num_chunks_to_evict = len(keys_to_store) - self._get_num_free_chunks()
201
202 to_evict: list[OffloadKey] = []
203 - if num_blocks_to_evict > 0:
204 - if num_blocks_to_evict > self._num_evictable_cache_blocks:
203 + if num_chunks_to_evict > 0:
204 + if num_chunks_to_evict > self._num_evictable_cache_chunks:
205 # Eviction will fail.
206 return None
207 # There is a still a chance for eviction failure as some of the
208 - # idle blocks might be in the protected list.
208 + # idle chunks might be in the protected list.
209
210 - # Blocks from the original input are excluded from eviction candidates:
211 - # a block that was already stored must remain in the cache after this call.
210 + # Chunks from the original input are excluded from eviction candidates:
211 + # a chunk that was already stored must remain in the cache after this call.
212 protected = set(keys)
213 - evicted = self._policy.evict(num_blocks_to_evict, protected)
213 + evicted = self._policy.evict(num_chunks_to_evict, protected)
214 if evicted is None:
215 return None
216
217 - # cache-policy removes only idle blocks.
218 - self._num_evictable_cache_blocks -= len(evicted)
219 - assert self._num_evictable_cache_blocks >= 0
217 + # cache-policy removes only idle chunks.
218 + self._num_evictable_cache_chunks -= len(evicted)
219 + assert self._num_evictable_cache_chunks >= 0
220
221 - for key, block in evicted:
222 - self._free_block(block)
221 + for key, chunk in evicted:
222 + self._free_chunk(chunk)
223 to_evict.append(key)
224
225 if to_evict and self.events is not None:
@@ -231,17 +231,17 @@ class CPUOffloadingManager(OffloadingManager):
231 )
232 )
233
234 - blocks = self._allocate_blocks(keys_to_store)
235 - assert len(blocks) == len(keys_to_store), (
236 - "Block pool did not allocate the expected number of blocks"
234 + chunks = self._allocate_chunks(keys_to_store)
235 + assert len(chunks) == len(keys_to_store), (
236 + "Chunk pool did not allocate the expected number of chunks"
237 )
238
239 - for key, block in zip(keys_to_store, blocks):
240 - self._policy.insert(key, block)
241 - self._num_write_pending_blocks += len(keys_to_store)
239 + for key, chunk in zip(keys_to_store, chunks):
240 + self._policy.insert(key, chunk)
241 + self._num_write_pending_chunks += len(keys_to_store)
242
243 - # build store specs for allocated blocks
244 - store_spec = self._get_load_store_spec(keys_to_store, blocks)
243 + # build store specs for allocated chunks
244 + store_spec = self._get_load_store_spec(keys_to_store, chunks)
245
246 return PrepareStoreOutput(
247 keys_to_store=keys_to_store,
@@ -260,20 +260,20 @@ class CPUOffloadingManager(OffloadingManager):
260
261 if success:
262 for key in keys:
263 - block = self._policy.get(key)
264 - if block is not None and not block.is_ready:
265 - block.ref_cnt = 0
266 - self._num_write_pending_blocks -= 1
267 - self._num_evictable_cache_blocks += 1
263 + chunk = self._policy.get(key)
264 + if chunk is not None and not chunk.is_ready:
265 + chunk.ref_cnt = 0
266 + self._num_write_pending_chunks -= 1
267 + self._num_evictable_cache_chunks += 1
268 self._policy.mark_evictable(key)
269 stored_keys.append(key)
270 else:
271 for key in keys:
272 - block = self._policy.get(key)
273 - if block is not None and not block.is_ready:
274 - self._num_write_pending_blocks -= 1
272 + chunk = self._policy.get(key)
273 + if chunk is not None and not chunk.is_ready:
274 + self._num_write_pending_chunks -= 1
275 self._policy.remove(key)
276 - self._free_block(block)
276 + self._free_chunk(chunk)
277
278 if stored_keys and self.events is not None:
279 self.events.append(
@@ -286,17 +286,17 @@ class CPUOffloadingManager(OffloadingManager):
286
287 @override
288 def reset_cache(self) -> None:
289 - # Clear ALL blocks unconditionally. The scheduler's _stale_job_threshold
289 + # Clear ALL chunks unconditionally. The scheduler's _stale_job_threshold
290 # guarantees that complete_load / complete_store are never called for
291 # pre-reset jobs, so no lazy cleanup is needed. The scheduler also
292 # flushes in-flight load job IDs to the workers before any new stores
293 - # can begin, preventing a cross-direction data race on reused offload block IDs.
293 + # can begin, preventing a cross-direction data race on reused offload chunk IDs.
294 self._policy.clear()
295 - self._num_evictable_cache_blocks = 0
296 - self._num_write_pending_blocks = 0
295 + self._num_evictable_cache_chunks = 0
296 + self._num_write_pending_chunks = 0
297
298 self._free_list.clear()
299 - self._num_allocated_blocks = 0
299 + self._num_allocated_chunks = 0
300
301 @override
302 def take_events(self) -> Iterable[OffloadingEvent]:
@@ -309,11 +309,11 @@ class CPUOffloadingManager(OffloadingManager):
309
310 # Compute cache usage.
311 num_used = (
312 - self._num_allocated_blocks
312 + self._num_allocated_chunks
313 - len(self._free_list)
314 - - self._num_evictable_cache_blocks
314 + - self._num_evictable_cache_chunks
315 )
316 - usage = num_used / self._num_blocks if self._num_blocks > 0 else 0.0
316 + usage = num_used / self._num_chunks if self._num_chunks > 0 else 0.0
317 stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC, usage)
318
319 for allocation_size in self.allocation_sizes_in_current_batch:
@@ -323,8 +323,8 @@ class CPUOffloadingManager(OffloadingManager):
323 self.allocation_sizes_in_current_batch.clear()
324
325 write_usage = (
326 - self._num_write_pending_blocks / self._num_blocks
327 - if self._num_blocks > 0
326 + self._num_write_pending_chunks / self._num_chunks
327 + if self._num_chunks > 0
328 else 0.0
329 )
330 read_usage = max(usage - write_usage, 0.0)
vllm/v1/kv_offload/cpu/policies/arc.py
+33 -33
@@ -6,7 +6,7 @@ from collections.abc import Iterable, Iterator
6 from typing_extensions import override
7
8 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
9 -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
9 +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
10
11
12 class ARCCachePolicy(CachePolicy):
@@ -14,15 +14,15 @@ class ARCCachePolicy(CachePolicy):
14 ARC (Adaptive Replacement Cache) cache policy.
15
16 Data Structures:
17 - T1: Recent cache containing blocks accessed once.
18 - T2: Frequent cache containing blocks accessed multiple times.
19 - B1/B2: Ghost lists tracking recently evicted blocks from T1/T2.
17 + T1: Recent cache containing chunks accessed once.
18 + T2: Frequent cache containing chunks accessed multiple times.
19 + B1/B2: Ghost lists tracking recently evicted chunks from T1/T2.
20 target_t1_size: Adaptive target size for the T1 partition.
21
22 Algorithm Flow:
23 1. Cache lookup (lookup):
24 - Searches T1 and T2 for block hashes and counts consecutive hits
25 - until a miss or non-ready block is encountered.
24 + Searches T1 and T2 for chunk hashes and counts consecutive hits
25 + until a miss or non-ready chunk is encountered.
26
27 2. Cache touch (touch) - Adaptive Learning:
28 For each key (in reverse order):
@@ -31,15 +31,15 @@ class ARCCachePolicy(CachePolicy):
31 - If in B1 ghost list: Increase target_t1_size.
32 - If in B2 ghost list: Decrease target_t1_size.
33
34 - 3. Block eviction (evict) - Adaptive Replacement:
34 + 3. Chunk eviction (evict) - Adaptive Replacement:
35 Determines eviction source based on adaptive target:
36 - If T1 size >= target_t1_size: Evict from T1, add to B1.
37 - Otherwise: Evict from T2, add to B2.
38 Finally, bound each ghost list size.
39
40 - 4. Block insertion (insert):
41 - New blocks are always inserted into T1 and removed from B1/B2 if
42 - present. Blocks may later be promoted to T2 during touch operations.
40 + 4. Chunk insertion (insert):
41 + New chunks are always inserted into T1 and removed from B1/B2 if
42 + present. Chunks may later be promoted to T2 during touch operations.
43
44 Adaptive Behavior:
45 The algorithm self-tunes the recency vs. frequency trade-off:
@@ -50,19 +50,19 @@ class ARCCachePolicy(CachePolicy):
50 def __init__(self, cache_capacity: int):
51 super().__init__(cache_capacity)
52 self.target_t1_size: float = 0.0
53 - self.t1: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
54 - self.t2: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
53 + self.t1: OrderedDict[OffloadKey, ChunkStatus] = OrderedDict()
54 + self.t2: OrderedDict[OffloadKey, ChunkStatus] = OrderedDict()
55 # key -> None (only care about presence)
56 self.b1: OrderedDict[OffloadKey, None] = OrderedDict()
57 self.b2: OrderedDict[OffloadKey, None] = OrderedDict()
58
59 @override
60 - def get(self, key: OffloadKey) -> BlockStatus | None:
60 + def get(self, key: OffloadKey) -> ChunkStatus | None:
61 return self.t1.get(key) or self.t2.get(key)
62
63 @override
64 - def insert(self, key: OffloadKey, block: BlockStatus) -> None:
65 - self.t1[key] = block
64 + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None:
65 + self.t1[key] = chunk
66 self.b1.pop(key, None)
67 self.b2.pop(key, None)
68
@@ -75,13 +75,13 @@ class ARCCachePolicy(CachePolicy):
75 def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
76 for key in reversed(list(keys)):
77 if key in self.t1:
78 - block = self.t1.pop(key)
79 - if not block.is_ready:
80 - # block was just prepared to be stored, not really touched
78 + chunk = self.t1.pop(key)
79 + if not chunk.is_ready:
80 + # chunk was just prepared to be stored, not really touched
81 # twice — keep it in T1 and mark as most recently used
82 - self.t1[key] = block
82 + self.t1[key] = chunk
83 else:
84 - self.t2[key] = block
84 + self.t2[key] = chunk
85
86 elif key in self.t2:
87 self.t2.move_to_end(key)
@@ -111,31 +111,31 @@ class ARCCachePolicy(CachePolicy):
111 @override
112 def evict(
113 self, n: int, protected: set[OffloadKey]
114 - ) -> list[tuple[OffloadKey, BlockStatus]] | None:
114 + ) -> list[tuple[OffloadKey, ChunkStatus]] | None:
115 if n == 0:
116 return []
117
118 # Collect candidates atomically: simulate T1 size changes as we select,
119 # but do not modify actual data structures until all n are found.
120 candidates: list[
121 - tuple[OffloadKey, BlockStatus, bool]
122 - ] = [] # (key, block, from_t1)
121 + tuple[OffloadKey, ChunkStatus, bool]
122 + ] = [] # (key, chunk, from_t1)
123 virtual_t1_size = len(self.t1)
124 # Keep the scans monotonic: restarting from the LRU end after every
125 - # selection makes a batch eviction quadratic in the number of blocks.
125 + # selection makes a batch eviction quadratic in the number of chunks.
126 t1_iter = iter(self.t1.items())
127 t2_iter = iter(self.t2.items())
128
129 def next_candidate(
130 - entries: Iterator[tuple[OffloadKey, BlockStatus]],
131 - ) -> tuple[OffloadKey, BlockStatus] | None:
132 - for key, block in entries:
133 - if block.ref_cnt == 0 and key not in protected:
134 - return key, block
130 + entries: Iterator[tuple[OffloadKey, ChunkStatus]],
131 + ) -> tuple[OffloadKey, ChunkStatus] | None:
132 + for key, chunk in entries:
133 + if chunk.ref_cnt == 0 and key not in protected:
134 + return key, chunk
135 return None
136
137 for _ in range(n):
138 - candidate: tuple[OffloadKey, BlockStatus, bool] | None = None
138 + candidate: tuple[OffloadKey, ChunkStatus, bool] | None = None
139
140 if virtual_t1_size >= int(self.target_t1_size):
141 entry = next_candidate(t1_iter)
@@ -158,15 +158,15 @@ class ARCCachePolicy(CachePolicy):
158 candidates.append(candidate)
159
160 # Apply all evictions now that we know n candidates exist.
161 - result: list[tuple[OffloadKey, BlockStatus]] = []
162 - for key, block, from_t1 in candidates:
161 + result: list[tuple[OffloadKey, ChunkStatus]] = []
162 + for key, chunk, from_t1 in candidates:
163 if from_t1:
164 del self.t1[key]
165 self.b1[key] = None
166 else:
167 del self.t2[key]
168 self.b2[key] = None
169 - result.append((key, block))
169 + result.append((key, chunk))
170
171 # Trim ghost lists to cache_capacity.
172 for ghost in (self.b1, self.b2):
vllm/v1/kv_offload/cpu/policies/base.py
+26 -26
@@ -7,36 +7,36 @@ from collections.abc import Iterable
7 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
8
9
10 -class BlockStatus(ctypes.Structure):
10 +class ChunkStatus(ctypes.Structure):
11 """
12 - Offloading status for a single block of KV data.
12 + Offloading status for a single chunk of KV data.
13 Holds the following information:
14
15 - ref_cnt - the current number of transfers using this block as a source.
16 - A value of -1 indicates the block is not yet ready to be read.
17 - block_id - index of the physical CPU buffer slot.
15 + ref_cnt - the current number of transfers using this chunk as a source.
16 + A value of -1 indicates the chunk is not yet ready to be read.
17 + chunk_id - index of the physical CPU buffer slot.
18 """
19
20 - _fields_ = [("ref_cnt", ctypes.c_int32), ("block_id", ctypes.c_int64)]
20 + _fields_ = [("ref_cnt", ctypes.c_int32), ("chunk_id", ctypes.c_int64)]
21
22 - def __init__(self, block_id: int):
22 + def __init__(self, chunk_id: int):
23 super().__init__()
24 - # initialize block as "not ready" (ref_cnt = -1)
24 + # initialize chunk as "not ready" (ref_cnt = -1)
25 self.ref_cnt = -1
26 - self.block_id = block_id
26 + self.chunk_id = chunk_id
27
28 @property
29 def is_ready(self) -> bool:
30 """
31 - Returns whether the block is ready to be read.
31 + Returns whether the chunk is ready to be read.
32 """
33 return self.ref_cnt >= 0
34
35
36 class CachePolicy(ABC):
37 """
38 - Encapsulates both block organization (data structures) and replacement
39 - decisions (which block to evict). LRU and ARC differ in both dimensions —
38 + Encapsulates both chunk organization (data structures) and replacement
39 + decisions (which chunk to evict). LRU and ARC differ in both dimensions —
40 ARC's ghost lists and target_t1_size live at the intersection of storage
41 and eviction, so they cannot be separated cleanly.
42 """
@@ -45,35 +45,35 @@ class CachePolicy(ABC):
45 self.cache_capacity = cache_capacity
46
47 @abstractmethod
48 - def get(self, key: OffloadKey) -> BlockStatus | None:
49 - """Find block in data structures. Returns None if not present."""
48 + def get(self, key: OffloadKey) -> ChunkStatus | None:
49 + """Find chunk in data structures. Returns None if not present."""
50
51 @abstractmethod
52 - def insert(self, key: OffloadKey, block: BlockStatus) -> None:
53 - """Add a newly allocated block. For ARC: also removes from ghost lists."""
52 + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None:
53 + """Add a newly allocated chunk. For ARC: also removes from ghost lists."""
54
55 @abstractmethod
56 def remove(self, key: OffloadKey) -> None:
57 - """Remove a block (used to clean up after a failed store)."""
57 + """Remove a chunk (used to clean up after a failed store)."""
58
59 @abstractmethod
60 def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
61 """
62 - Mark blocks as recently used.
62 + Mark chunks as recently used.
63
64 Args:
65 - keys: Blocks to mark as recently used.
66 - req_context: Per-request context for the request touching these blocks.
65 + keys: Chunks to mark as recently used.
66 + req_context: Per-request context for the request touching these chunks.
67 """
68
69 @abstractmethod
70 def evict(
71 self, n: int, protected: set[OffloadKey]
72 - ) -> list[tuple[OffloadKey, BlockStatus]] | None:
72 + ) -> list[tuple[OffloadKey, ChunkStatus]] | None:
73 """
74 - Evict exactly n blocks, skipping any in protected.
74 + Evict exactly n chunks, skipping any in protected.
75
76 - Returns a list of (key, block) for the evicted blocks,
76 + Returns a list of (key, chunk) for the evicted chunks,
77 or None if n evictions cannot be satisfied. The operation is atomic:
78 if None is returned, no state changes are made.
79
@@ -84,15 +84,15 @@ class CachePolicy(ABC):
84 @abstractmethod
85 def clear(self) -> None:
86 """
87 - Remove ALL blocks regardless of ref_cnt.
87 + Remove ALL chunks regardless of ref_cnt.
88
89 Ghost lists and adaptive state are also reset.
90 """
91
92 def mark_evictable(self, key: OffloadKey) -> None:
93 - """Called when a block's ref_cnt transitions to 0."""
93 + """Called when a chunk's ref_cnt transitions to 0."""
94 return
95
96 def mark_non_evictable(self, key: OffloadKey) -> None:
97 - """Called when a block's ref_cnt transitions from 0."""
97 + """Called when a chunk's ref_cnt transitions from 0."""
98 return
vllm/v1/kv_offload/cpu/policies/lru.py
+29 -29
@@ -6,7 +6,7 @@ from collections.abc import Iterable
6 from typing_extensions import override
7
8 from vllm.v1.kv_offload.base import OffloadKey, ReqContext
9 -from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
9 +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus
10
11
12 class LRUCachePolicy(CachePolicy):
@@ -20,71 +20,71 @@ class LRUCachePolicy(CachePolicy):
20
21 def __init__(self, cache_capacity: int):
22 super().__init__(cache_capacity)
23 - # Blocks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU
24 - self.evictable_blocks: OrderedDict[OffloadKey, None] = OrderedDict()
25 - self.blocks: dict[OffloadKey, BlockStatus] = {}
23 + # Chunks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU
24 + self.evictable_chunks: OrderedDict[OffloadKey, None] = OrderedDict()
25 + self.chunks: dict[OffloadKey, ChunkStatus] = {}
26
27 @override
28 - def get(self, key: OffloadKey) -> BlockStatus | None:
29 - return self.blocks.get(key)
28 + def get(self, key: OffloadKey) -> ChunkStatus | None:
29 + return self.chunks.get(key)
30
31 @override
32 - def insert(self, key: OffloadKey, block: BlockStatus) -> None:
33 - self.blocks[key] = block
34 - if block.ref_cnt == 0:
35 - self.evictable_blocks[key] = None
32 + def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None:
33 + self.chunks[key] = chunk
34 + if chunk.ref_cnt == 0:
35 + self.evictable_chunks[key] = None
36
37 @override
38 def remove(self, key: OffloadKey) -> None:
39 - del self.blocks[key]
40 - self.evictable_blocks.pop(key, None)
39 + del self.chunks[key]
40 + self.evictable_chunks.pop(key, None)
41
42 @override
43 def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
44 for key in reversed(list(keys)):
45 - if key in self.evictable_blocks:
46 - self.evictable_blocks.move_to_end(key)
47 - # active blocks are untouched as they are non-evictable now. They
48 - # will eventually reach the end of evictable_blocks when they finish.
45 + if key in self.evictable_chunks:
46 + self.evictable_chunks.move_to_end(key)
47 + # active chunks are untouched as they are non-evictable now. They
48 + # will eventually reach the end of evictable_chunks when they finish.
49
50 @override
51 def clear(self) -> None:
52 - self.evictable_blocks.clear()
53 - self.blocks.clear()
52 + self.evictable_chunks.clear()
53 + self.chunks.clear()
54
55 @override
56 def evict(
57 self, n: int, protected: set[OffloadKey]
58 - ) -> list[tuple[OffloadKey, BlockStatus]] | None:
58 + ) -> list[tuple[OffloadKey, ChunkStatus]] | None:
59 if n == 0:
60 return []
61
62 - candidates: list[tuple[OffloadKey, BlockStatus]] = []
63 - for key, _ in self.evictable_blocks.items():
62 + candidates: list[tuple[OffloadKey, ChunkStatus]] = []
63 + for key, _ in self.evictable_chunks.items():
64 if key in protected:
65 continue
66
67 - block = self.blocks[key]
68 - assert block.ref_cnt == 0
69 - candidates.append((key, block))
67 + chunk = self.chunks[key]
68 + assert chunk.ref_cnt == 0
69 + candidates.append((key, chunk))
70 if len(candidates) == n:
71 break
72
73 if len(candidates) < n:
74 return None
75 for key, _ in candidates:
76 - del self.evictable_blocks[key]
77 - del self.blocks[key]
76 + del self.evictable_chunks[key]
77 + del self.chunks[key]
78 return candidates
79
80 @override
81 def mark_evictable(self, key: OffloadKey) -> None:
82 - # blocks can become evictable when,
82 + # chunks can become evictable when,
83 # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list
84 # all loads complete - i.e ref_cnt 1 -> 0 # not in evictable list
85 - self.evictable_blocks[key] = None
85 + self.evictable_chunks[key] = None
86
87 @override
88 def mark_non_evictable(self, key: OffloadKey) -> None:
89 # key must have been in the evictable list.
90 - del self.evictable_blocks[key]
90 + del self.evictable_chunks[key]
vllm/v1/kv_offload/cpu/shared_offload_region.py
+26 -26
@@ -82,24 +82,24 @@ class SharedOffloadRegion:
82 def __init__(
83 self,
84 engine_id: str,
85 - num_blocks: int,
85 + num_chunks: int,
86 rank: int | None,
87 - kv_bytes_per_block: int,
87 + kv_bytes_per_chunk: int,
88 cpu_page_size: int,
89 barrier: Callable[[], None] | None = None,
90 ) -> None:
91 self.page_size = mmap.PAGESIZE
92 - assert kv_bytes_per_block % self.page_size == 0
92 + assert kv_bytes_per_chunk % self.page_size == 0
93
94 - self.num_blocks = num_blocks
95 - self._row_stride = kv_bytes_per_block
96 - self.total_size_bytes = self.num_blocks * self._row_stride
94 + self.num_chunks = num_chunks
95 + self._row_stride = kv_bytes_per_chunk
96 + self.total_size_bytes = self.num_chunks * self._row_stride
97
98 self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"
99 self._creator = False # set True only if this worker creates the file
100 self.rank = rank
101 if rank is not None:
102 - # byte offset to this worker's first slot within each block row
102 + # byte offset to this worker's first slot within each chunk row
103 self._worker_offset = rank * cpu_page_size
104 # exclusive upper bound for this worker's area within each row
105 self._worker_area_end = (rank + 1) * cpu_page_size
@@ -182,19 +182,19 @@ class SharedOffloadRegion:
182 populate_write_fn = _get_populate_write_fn(self.mmap_obj)
183
184 if rank is not None:
185 - # Populate only this worker's pages (one slot per block row).
185 + # Populate only this worker's pages (one slot per chunk row).
186 worker_offset = rank * cpu_page_size
187 _t0 = time.perf_counter()
188 page_size = self.page_size
189 - for block in range(num_blocks):
190 - raw_offset = block * self._row_stride + worker_offset
189 + for chunk in range(num_chunks):
190 + raw_offset = chunk * self._row_stride + worker_offset
191 aligned_offset = (raw_offset // page_size) * page_size
192 end = raw_offset + cpu_page_size
193 aligned_length = end - aligned_offset
194 populate_write_fn(self.mmap_obj, aligned_offset, aligned_length)
195 logger.debug(
196 - "MADV_POPULATE_WRITE loop: %d blocks in %.3f s",
197 - num_blocks,
196 + "MADV_POPULATE_WRITE loop: %d chunks in %.3f s",
197 + num_chunks,
198 time.perf_counter() - _t0,
199 )
200 else:
@@ -215,22 +215,22 @@ class SharedOffloadRegion:
215
216 Must be called once per canonical tensor. The full mmap layout is:
217
218 - worker0_block0 | worker1_block0 | ... | worker{M-1}_block0
219 - worker0_block1 | worker1_block1 | ... | worker{M-1}_block1
218 + worker0_chunk0 | worker1_chunk0 | ... | worker{M-1}_chunk0
219 + worker0_chunk1 | worker1_chunk1 | ... | worker{M-1}_chunk1
220 ...
221
222 - Each worker_block cell is cpu_page_size bytes and holds all canonical
223 - tensors for that worker and block concatenated:
222 + Each worker_chunk cell is cpu_page_size bytes and holds all canonical
223 + tensors for that worker and chunk concatenated:
224 [ tensor0_data | tensor1_data | ... | tensor{L-1}_data ]
225
226 Consecutive rows are separated by row_stride = cpu_page_size * M.
227
228 - Returns an int8 tensor of shape (num_blocks, tensor_page_size) with stride
228 + Returns an int8 tensor of shape (num_chunks, tensor_page_size) with stride
229 (row_stride, 1). Using int8 keeps stride == bytes, so swap_blocks
230 address arithmetic works without any dtype conversion.
231
232 Args:
233 - tensor_page_size: Bytes per block for this tensor.
233 + tensor_page_size: Bytes per chunk for this tensor.
234 """
235 assert self.rank is not None
236 new_offset = self._worker_offset + tensor_page_size
@@ -241,7 +241,7 @@ class SharedOffloadRegion:
241 )
242 worker_layer_view = torch.as_strided(
243 self._base,
244 - size=(self.num_blocks, tensor_page_size),
244 + size=(self.num_chunks, tensor_page_size),
245 stride=(self._row_stride, 1),
246 storage_offset=self._worker_offset,
247 )
@@ -266,8 +266,8 @@ class SharedOffloadRegion:
266 _canonical_offset=0, then advances by each tensor's size
267
268 Each canonical_t{i} cell is that tensor's canonical page for the
269 - block. Canonical areas are carved consecutively from the start of
270 - each block row; consecutive rows are separated by row_stride. Every
269 + chunk. Canonical areas are carved consecutively from the start of
270 + each chunk row; consecutive rows are separated by row_stride. Every
271 worker gets the identical byte ranges and writes only its disjoint
272 bytes within them, as described by its canonical mappings — unlike
273 create_next_worker_view, which gives each worker a private
@@ -279,13 +279,13 @@ class SharedOffloadRegion:
279 where one canonical copy replaces world_size worker copies.
280
281 Args:
282 - tensor_page_size: Canonical bytes per block for this tensor.
282 + tensor_page_size: Canonical bytes per chunk for this tensor.
283 """
284 new_offset = self._canonical_offset + tensor_page_size
285 assert new_offset <= self._row_stride
286 view = torch.as_strided(
287 self._base,
288 - size=(self.num_blocks, tensor_page_size),
288 + size=(self.num_chunks, tensor_page_size),
289 stride=(self._row_stride, 1),
290 storage_offset=self._canonical_offset,
291 )
@@ -296,10 +296,10 @@ class SharedOffloadRegion:
296 def create_kv_memoryview(self) -> memoryview:
297 """Return a zero-copy memoryview over the entire KV buffer.
298
299 - Shape: (num_blocks, row_stride_bytes). Secondary tiers address
300 - block *b* as ``view[b]``.
299 + Shape: (num_chunks, row_stride_bytes). Secondary tiers address
300 + chunk *b* as ``view[b]``.
301 """
302 - kv_tensor = self._base.view(self.num_blocks, self._row_stride)
302 + kv_tensor = self._base.view(self.num_chunks, self._row_stride)
303 np_arr = kv_tensor.numpy()
304 assert np_arr.ctypes.data == self._base.data_ptr(), (
305 "view()/numpy() created a copy instead of sharing the mmap buffer; "
vllm/v1/kv_offload/cpu/spec.py
+14 -14
@@ -73,7 +73,7 @@ class CPUOffloadingSpec(OffloadingSpec):
73 ),
74 CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata(
75 documentation=(
76 - "Histogram of the number of CPU blocks requested by each "
76 + "Histogram of the number of CPU chunks requested by each "
77 "KV offload prepare_store call."
78 ),
79 buckets=(1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144),
@@ -101,7 +101,7 @@ class CPUOffloadingSpec(OffloadingSpec):
101 )
102
103 world_size = config.parallel.world_size
104 - self.num_blocks = 0
104 + self.num_chunks = 0
105 self.kv_bytes_per_chunk = 0
106 self.cpu_page_size_per_worker = 0
107 self.replicated_layout = config.replicated_layout and self._uses_shared_region()
@@ -113,17 +113,17 @@ class CPUOffloadingSpec(OffloadingSpec):
113 # calculate cpu_page_size_per_worker
114 self.cpu_page_size_per_worker = kv_bytes_per_chunk // num_copies
115
116 - # calculate num_blocks
116 + # calculate num_chunks
117 aligned_kv_bytes_per_chunk = round_up(
118 kv_bytes_per_chunk, self.BLOCK_SIZE_ALIGNMENT
119 )
120 - self.num_blocks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk
120 + self.num_chunks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk
121
122 # Expose aligned_kv_bytes_per_chunk as
123 # kv_bytes_per_chunk. Note that this might contain
124 - # some padding. i.e. each offloaded block is of the form,
125 - # |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** |
126 - # or |--- B0 (single copy) ---| *** maybe-pad *** |
124 + # some padding. i.e. each offloaded chunk is of the form,
125 + # |--- W0-C0---|---- W1-C0---| ... |---- Wn-C0---| *** maybe-pad *** |
126 + # or |--- C0 (single copy) ---| *** maybe-pad *** |
127 self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk
128
129 # scheduler-side
@@ -140,7 +140,7 @@ class CPUOffloadingSpec(OffloadingSpec):
140 @override
141 def get_manager(self) -> OffloadingManager:
142 if not self._manager:
143 - # store_threshold: how many times a block must be offered for
143 + # store_threshold: how many times a chunk must be offered for
144 # storage before it is eligible for CPU offloading. Values < 2
145 # disable filtering (a threshold of 1 equals no filter; 0 is the
146 # default).
@@ -150,7 +150,7 @@ class CPUOffloadingSpec(OffloadingSpec):
150 max_tracker_size = int(self.extra_config.get("max_tracker_size", 64_000))
151
152 self._manager = CPUOffloadingManager(
153 - num_blocks=self.num_blocks,
153 + num_chunks=self.num_chunks,
154 cache_policy=self.eviction_policy,
155 cache_policy_module_path=self.cache_policy_module_path,
156 enable_events=self.kv_events_config.enable_kv_cache_events,
@@ -166,9 +166,9 @@ class CPUOffloadingSpec(OffloadingSpec):
166
167 def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
168 mmap_region: SharedOffloadRegion | None = None
169 - # num_blocks == 0 would size the region to zero bytes, which cannot be
169 + # num_chunks == 0 would size the region to zero bytes, which cannot be
170 # mmap'd; fall back to the tensor path (empty tensors) as before.
171 - if self._uses_shared_region() and self.num_blocks > 0:
171 + if self._uses_shared_region() and self.num_chunks > 0:
172 # Replicated layout puts all ranks on slot 0 (single MLA copy);
173 # otherwise each rank takes its own slot by physical device index.
174 if self.replicated_layout:
@@ -178,9 +178,9 @@ class CPUOffloadingSpec(OffloadingSpec):
178 rank = torch.accelerator.current_device_index() % world_size
179 mmap_region = SharedOffloadRegion(
180 engine_id=self.config.engine_id,
181 - num_blocks=self.num_blocks,
181 + num_chunks=self.num_chunks,
182 rank=rank,
183 - kv_bytes_per_block=self.kv_bytes_per_chunk,
183 + kv_bytes_per_chunk=self.kv_bytes_per_chunk,
184 cpu_page_size=self.cpu_page_size_per_worker,
185 barrier=_all_workers_barrier,
186 )
@@ -188,7 +188,7 @@ class CPUOffloadingSpec(OffloadingSpec):
188 return CPUOffloadingWorker(
189 kv_caches=kv_caches,
190 blocks_per_chunk=self.blocks_per_chunk,
191 - num_cpu_blocks=self.num_blocks,
191 + num_cpu_chunks=self.num_chunks,
192 mmap_region=mmap_region,
193 )
194 except Exception:
vllm/v1/kv_offload/tiering/base.py
+23 -23
@@ -45,8 +45,8 @@ class TieringOffloadingMetrics:
45 WRITE_TIME = "vllm:kv_offload_tiering_write_time"
46 PROMOTION_JOB_FAILURES = "vllm:kv_offload_tiering_promotion_job_failures"
47 CASCADE_JOB_FAILURES = "vllm:kv_offload_tiering_cascade_job_failures"
48 - BLOCK_QUERIES = "vllm:kv_offload_tiering_chunk_queries"
49 - BLOCK_HITS = "vllm:kv_offload_tiering_chunk_hits"
48 + CHUNK_QUERIES = "vllm:kv_offload_tiering_chunk_queries"
49 + CHUNK_HITS = "vllm:kv_offload_tiering_chunk_hits"
50 PRIMARY_WRITE_USAGE_PERC = "vllm:kv_offload_tiering_primary_write_usage_perc"
51 PRIMARY_READ_USAGE_PERC = "vllm:kv_offload_tiering_primary_read_usage_perc"
52 PROMOTION_ALLOCATION_FAILURES = (
@@ -62,7 +62,7 @@ class TransferJob:
62
63 job_id: JobId
64 keys: Collection[OffloadKey]
65 - block_ids: np.ndarray
65 + chunk_ids: np.ndarray
66 is_promotion: bool
67 req_context: ReqContext
68
@@ -90,14 +90,14 @@ class ParentManager(ABC):
90
91 Required call sequence for each remote request:
92 1. on_new_request(req_context) — set up per-request state
93 - 2. lookup(key, req_context) — check block availability
94 - (repeat per block)
95 - 3. create_store_job(keys, req_context) — pin blocks and get a
93 + 2. lookup(key, req_context) — check chunk availability
94 + (repeat per chunk)
95 + 3. create_store_job(keys, req_context) — pin chunks and get a
96 job handle
97 4. on_request_finished(req_context) — clean up per-request state
98
99 Steps 2-3 may be interleaved. Step 4 must be called even if no
100 - blocks were found, to avoid leaking async lookup state (e.g. in
100 + chunks were found, to avoid leaking async lookup state (e.g. in
101 the fs tier's AsyncLookupManager).
102 """
103
@@ -155,23 +155,23 @@ class SecondaryTierManager(ABC):
155 @abstractmethod
156 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
157 """
158 - Check whether a block exists in this secondary tier.
158 + Check whether a chunk exists in this secondary tier.
159
160 Args:
161 key: Offload key to look up.
162 req_context: per-request context (e.g. kv_transfer_params).
163
164 Returns:
165 - HIT if the block is present and ready,
165 + HIT if the chunk is present and ready,
166 MISS if not found,
167 - or RETRY if the block is being transferred (retry later).
167 + or RETRY if the chunk is being transferred (retry later).
168 """
169 pass
170
171 @abstractmethod
172 def submit_store(self, job_metadata: TransferJob) -> None:
173 """
174 - Submit an async job to store blocks from the primary tier to this
174 + Submit an async job to store chunks from the primary tier to this
175 secondary tier.
176
177 This method must be lightweight and non-blocking: allocate metadata
@@ -179,19 +179,19 @@ class SecondaryTierManager(ABC):
179 calling thread.
180
181 Preconditions (guaranteed by the framework):
182 - - ``job_metadata.block_ids`` are valid primary-tier slots, pinned
182 + - ``job_metadata.chunk_ids`` are valid primary-tier slots, pinned
183 (ref-counted) for the duration of the transfer.
184
185 The implementation is responsible for:
186 - 1. Filtering out blocks already present in this tier
187 - 2. Evicting blocks if capacity is needed
186 + 1. Filtering out chunks already present in this tier
187 + 2. Evicting chunks if capacity is needed
188 3. Allocating space in this tier
189 - 4. Submitting the async transfer (read from primary via block_ids)
189 + 4. Submitting the async transfer (read from primary via chunk_ids)
190
191 Report completion via ``get_finished_jobs()``.
192
193 Args:
194 - job_metadata: Job metadata including job_id, keys, and block_ids
194 + job_metadata: Job metadata including job_id, keys, and chunk_ids
195 identifying the primary-tier slots to read from.
196 """
197 pass
@@ -199,24 +199,24 @@ class SecondaryTierManager(ABC):
199 @abstractmethod
200 def submit_load(self, job_metadata: TransferJob) -> None:
201 """
202 - Submit an async job to load blocks from this secondary tier to the
202 + Submit an async job to load chunks from this secondary tier to the
203 primary tier.
204
205 - This method must be lightweight and non-blocking: mark blocks as
205 + This method must be lightweight and non-blocking: mark chunks as
206 in-flight and submit the transfer, but do NOT perform the data copy
207 on the calling thread.
208
209 Preconditions (guaranteed by the framework):
210 - - ``job_metadata.block_ids`` are allocated primary-tier slots
210 + - ``job_metadata.chunk_ids`` are allocated primary-tier slots
211 ready to receive data.
212
213 The implementation must copy data from this tier into the
214 - primary-tier slots identified by ``block_ids``.
214 + primary-tier slots identified by ``chunk_ids``.
215
216 Report completion via ``get_finished_jobs()``.
217
218 Args:
219 - job_metadata: Job metadata including job_id, keys, and block_ids
219 + job_metadata: Job metadata including job_id, keys, and chunk_ids
220 identifying the primary-tier slots to write into.
221 """
222 pass
@@ -249,7 +249,7 @@ class SecondaryTierManager(ABC):
249
250 def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
251 """
252 - Mark blocks as recently used for eviction policy.
252 + Mark chunks as recently used for eviction policy.
253
254 Args:
255 keys: Offload keys to mark as recently used.
@@ -263,7 +263,7 @@ class SecondaryTierManager(ABC):
263 Called when a new request is first seen by the scheduler.
264
265 Returns a RequestOffloadingContext expressing this tier's preference
266 - for how blocks should be offloaded for this request.
266 + for how chunks should be offloaded for this request.
267
268 Args:
269 req_context: Per-request context.
vllm/v1/kv_offload/tiering/example/manager.py
+22 -22
@@ -3,7 +3,7 @@
3 """
4 ExampleSecondaryTierManager: A simple in-memory secondary tier.
5
6 -This implementation provides a minimal secondary tier that stores blocks
6 +This implementation provides a minimal secondary tier that stores chunks
7 in memory (using a dictionary) with immediate completion. It serves as a
8 reference for writing new tiers and is useful for testing the
9 TieringOffloadingManager without requiring actual storage or network backends.
@@ -39,7 +39,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
39 A simple in-memory secondary tier.
40
41 This implementation:
42 - - Stores blocks in a dictionary (key -> True)
42 + - Stores chunks in a dictionary (key -> True)
43 - Completes transfers immediately (synchronous)
44 """
45
@@ -69,45 +69,45 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
69 )
70
71 # key -> True (only care about presence)
72 - self.blocks: dict[OffloadKey, bool] = {}
72 + self.chunks: dict[OffloadKey, bool] = {}
73
74 # Completed jobs waiting to be retrieved by get_finished_jobs()
75 self.completed_jobs: list[JobResult] = []
76 assert primary_kv_view.strides is not None
77 - self._block_size = primary_kv_view.strides[0]
77 + self._chunk_size = primary_kv_view.strides[0]
78
79 @override
80 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
81 """
82 - Check whether a block exists in this secondary tier.
82 + Check whether a chunk exists in this secondary tier.
83
84 Args:
85 key: Offload key to look up.
86 req_context: Per-request context.
87
88 Returns:
89 - HIT if the block is present, MISS if not found.
89 + HIT if the chunk is present, MISS if not found.
90 """
91 - return LookupResult.HIT if key in self.blocks else LookupResult.MISS
91 + return LookupResult.HIT if key in self.chunks else LookupResult.MISS
92
93 @override
94 def submit_store(self, job_metadata: TransferJob) -> None:
95 """
96 - Submit a job to store blocks from primary tier to this tier.
96 + Submit a job to store chunks from primary tier to this tier.
97
98 Args:
99 job_metadata: Job metadata including job_id, keys, and
100 - spec for reading blocks from the primary tier.
100 + spec for reading chunks from the primary tier.
101 """
102 keys = job_metadata.keys
103 - block_ids = job_metadata.block_ids
103 + chunk_ids = job_metadata.chunk_ids
104
105 - assert len(keys) == len(block_ids), (
106 - f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
105 + assert len(keys) == len(chunk_ids), (
106 + f"Length mismatch: {len(keys)} keys but {len(chunk_ids)} chunk_ids"
107 )
108
109 for key in keys:
110 - self.blocks[key] = True
110 + self.chunks[key] = True
111 self.completed_jobs.append(
112 JobResult(
113 job_id=job_metadata.job_id,
@@ -119,21 +119,21 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
119 @override
120 def submit_load(self, job_metadata: TransferJob) -> None:
121 """
122 - Submit a job to load blocks from this tier to primary tier.
122 + Submit a job to load chunks from this tier to primary tier.
123
124 Args:
125 job_metadata: Job metadata including job_id, keys, and
126 - spec for writing blocks into the primary tier.
126 + spec for writing chunks into the primary tier.
127 """
128 keys = job_metadata.keys
129 - block_ids = job_metadata.block_ids
129 + chunk_ids = job_metadata.chunk_ids
130
131 - assert len(keys) == len(block_ids), (
132 - f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
131 + assert len(keys) == len(chunk_ids), (
132 + f"Length mismatch: {len(keys)} keys but {len(chunk_ids)} chunk_ids"
133 )
134
135 for key in keys:
136 - if key not in self.blocks:
136 + if key not in self.chunks:
137 self.completed_jobs.append(
138 JobResult(job_id=job_metadata.job_id, success=False)
139 )
@@ -170,6 +170,6 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
170 completes, so there is nothing to wait for."""
171 return
172
173 - def get_num_blocks(self) -> int:
174 - """Get the number of blocks currently stored in this tier."""
175 - return len(self.blocks)
173 + def get_num_chunks(self) -> int:
174 + """Get the number of chunks currently stored in this tier."""
175 + return len(self.chunks)
vllm/v1/kv_offload/tiering/fs/manager.py
+2 -2
@@ -223,7 +223,7 @@ class FileSystemTierManager(SecondaryTierManager):
223 batch_store_block,
224 [self.file_mapper.get_file_name(key) for key in keys],
225 self._primary_kv_view,
226 - [int(bid) * self._block_size for bid in job_metadata.block_ids],
226 + [int(cid) * self._block_size for cid in job_metadata.chunk_ids],
227 self._block_size,
228 self._use_o_direct,
229 )
@@ -237,7 +237,7 @@ class FileSystemTierManager(SecondaryTierManager):
237 keys = list(job_metadata.keys)
238 self._load_job_keys[job_id] = keys
239 paths = [self.file_mapper.get_file_name(key) for key in keys]
240 - offsets = [int(bid) * self._block_size for bid in job_metadata.block_ids]
240 + offsets = [int(cid) * self._block_size for cid in job_metadata.chunk_ids]
241
242 def load_task() -> None:
243 try:
vllm/v1/kv_offload/tiering/manager.py
+64 -64
@@ -8,16 +8,16 @@ and zero or more secondary tiers (Storage, Network, etc.) to provide
8 hierarchical KV cache offloading.
9
10 Key Design Principles:
11 -1. Always offload to all tiers — When a block is stored to the primary tier,
11 +1. Always offload to all tiers — When a chunk is stored to the primary tier,
12 it is cascaded to ALL secondary tiers
13 2. Primary tier is the gateway — Secondary tiers cannot access GPU memory
14 directly; all data flows through the CPU primary tier
15 -3. Staged promotion — Blocks in secondary tiers must be promoted to the
15 +3. Staged promotion — Chunks in secondary tiers must be promoted to the
16 primary tier before GPU can access them
17 4. Transparent retry mechanism — Return None from lookup() to signal
18 "data is being promoted, try later"
19 5. ref_cnt as eviction protection — primary.prepare_read() increments ref_cnt,
20 - protecting blocks from eviction until complete_read() is called
20 + protecting chunks from eviction until complete_read() is called
21 """
22
23 import time
@@ -61,11 +61,11 @@ logger = init_logger(__name__)
61
62 @dataclass
63 class PendingPromotion:
64 - """Accumulator for blocks awaiting submit_load() for one (tier, request)."""
64 + """Accumulator for chunks awaiting submit_load() for one (tier, request)."""
65
66 req_context: ReqContext
67 keys: list[OffloadKey] = field(default_factory=list)
68 - block_ids: list[int] = field(default_factory=list)
68 + chunk_ids: list[int] = field(default_factory=list)
69
70
71 @dataclass(slots=True)
@@ -94,14 +94,14 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
94
95 def __init__(
96 self,
97 - num_blocks: int,
97 + num_chunks: int,
98 mmap_region: SharedOffloadRegion,
99 cache_policy: str = "lru",
100 cache_policy_module_path: str | None = None,
101 enable_events: bool = False,
102 ):
103 super().__init__(
104 - num_blocks=num_blocks,
104 + num_chunks=num_chunks,
105 cache_policy=cache_policy,
106 cache_policy_module_path=cache_policy_module_path,
107 enable_events=enable_events,
@@ -120,9 +120,9 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
120 def get_kv_memoryview(self) -> memoryview:
121 """Return the memoryview over the primary tier's KV cache buffer.
122
123 - The view has shape (num_blocks, row_stride_bytes) and is backed by the
124 - SharedOffloadRegion mmap. Secondary tiers address block *b* as
125 - ``view[b]``.
123 + The view has shape (num_chunks, row_stride_bytes) and is backed by the
124 + SharedOffloadRegion mmap. Secondary tiers address chunk *c* as
125 + ``view[c]``.
126 """
127 return self._kv_memoryview
128
@@ -206,14 +206,14 @@ class TieringOffloadingManager(OffloadingManager):
206 assert primary_view.strides is not None
207 self._metrics = TieringMetricsTracker(
208 tier_types=[tier.tier_type for tier in self.secondary_tiers],
209 - num_primary_blocks=self.primary_tier._num_blocks,
210 - primary_block_size=primary_view.strides[0],
209 + num_primary_chunks=self.primary_tier._num_chunks,
210 + primary_chunk_size=primary_view.strides[0],
211 )
212
213 # Pending promotion requests accumulated during lookup() calls; flushed
214 # as one batched submit_load() per (tier, request) in on_schedule_end().
215 # Outer key: tier index. Inner key: req_context.req_id — the same ReqContext
216 - # object is reused for all block lookups of a given request per engine step.
216 + # object is reused for all chunk lookups of a given request per engine step.
217 self._pending_load_submissions: dict[int, dict[str, PendingPromotion]] = {}
218
219 # Gate for once-per-step execution of _maybe_process_finished_jobs().
@@ -305,7 +305,7 @@ class TieringOffloadingManager(OffloadingManager):
305 2. For completed stores (primary→secondary): calls primary.complete_read()
306 to decrement ref_cnt
307 3. For completed loads (secondary→primary): calls primary.complete_write()
308 - to make blocks available
308 + to make chunks available
309 """
310 for i, tier in enumerate(self.secondary_tiers):
311 for completed_job in tier.get_finished_jobs():
@@ -324,11 +324,11 @@ class TieringOffloadingManager(OffloadingManager):
324
325 if transfer_job.is_promotion:
326 # secondary→primary transfer (promotion) completed.
327 - # Make blocks available in primary tier.
327 + # Make chunks available in primary tier.
328 self._complete_promotion(job_metadata, completed_job)
329 else:
330 # primary→secondary transfer completed.
331 - # Decrement ref_cnt on primary blocks.
331 + # Decrement ref_cnt on primary chunks.
332 self.primary_tier.complete_read(
333 transfer_job.keys, transfer_job.req_context
334 )
@@ -342,7 +342,7 @@ class TieringOffloadingManager(OffloadingManager):
342 exclude_tier_idx: int | None = None,
343 ) -> LookupResult:
344 """
345 - Check whether a single block is offloaded and ready.
345 + Check whether a single chunk is offloaded and ready.
346
347 Algorithm:
348 1. Process any completed async jobs first.
@@ -351,20 +351,20 @@ class TieringOffloadingManager(OffloadingManager):
351 hit and initiate promotion.
352
353 Args:
354 - key: Block hash to look up.
354 + key: Chunk hash to look up.
355 req_context: Per-request context.
356
357 Returns:
358 - HIT — block is ready in the primary tier.
359 - HIT_PENDING — block found but not yet readable (write
358 + HIT — chunk is ready in the primary tier.
359 + HIT_PENDING — chunk found but not yet readable (write
360 in-flight on the primary tier).
361 RETRY — promotion started or a secondary tier is busy.
362 - MISS — block not found in any tier, or primary is full
362 + MISS — chunk not found in any tier, or primary is full
363 and cannot accept a promotion.
364 """
365 # Poll first so a promotion that finished since the last call is
366 # already reflected as HIT (not stale HIT_PENDING/MISS) below, and
367 - # so blocks freed by cascade or promotion completions are evictable
367 + # so chunks freed by cascade or promotion completions are evictable
368 # in time for a promotion this lookup may initiate.
369 self._maybe_process_finished_jobs()
370
@@ -424,30 +424,30 @@ class TieringOffloadingManager(OffloadingManager):
424 req_context: ReqContext,
425 ) -> bool:
426 """
427 - Queue a block for promotion from a secondary tier to the primary tier.
427 + Queue a chunk for promotion from a secondary tier to the primary tier.
428
429 Allocates space in the primary tier immediately (sets ref_cnt=-1 so
430 subsequent lookups within the same step see the slot as in-flight),
431 then defers the actual submit_load() call to _flush_pending_promotions()
432 - so all blocks queued during one engine step are submitted as a single
432 + so all chunks queued during one engine step are submitted as a single
433 batched job.
434
435 Args:
436 tier_idx: The secondary tier index to promote from
437 - key: Block to promote
437 + key: Chunk to promote
438 req_context: Per-request context forwarded to primary.prepare_write().
439
440 Returns:
441 True if promotion was initiated, False if primary tier is full.
442 """
443 - # Allocate space in primary tier for promoted block.
443 + # Allocate space in primary tier for promoted chunk.
444 # Must happen immediately so primary.lookup() returns None (in-flight)
445 # for this key on any subsequent lookup() call within the same step,
446 # preventing duplicate promotion attempts.
447 primary_write_result = self.primary_tier.prepare_write([key], req_context)
448
449 if primary_write_result is None:
450 - # Primary tier is full; caller should treat the block as unavailable
450 + # Primary tier is full; caller should treat the chunk as unavailable
451 # rather than retrying indefinitely.
452 self._metrics.on_promotion_allocation_failure()
453 return False
@@ -455,16 +455,16 @@ class TieringOffloadingManager(OffloadingManager):
455 store_spec = primary_write_result.store_spec
456 assert isinstance(store_spec, CPULoadStoreSpec)
457 # Defer submit_load to on_schedule_end(). Group by (tier, request) so
458 - # each request's blocks are submitted as one batched job per tier.
458 + # each request's chunks are submitted as one batched job per tier.
459 tier_pending = self._pending_load_submissions.setdefault(tier_idx, {})
460 ctx_id = req_context.req_id
461 if ctx_id not in tier_pending:
462 tier_pending[ctx_id] = PendingPromotion(
463 - keys=[], block_ids=[], req_context=req_context
463 + keys=[], chunk_ids=[], req_context=req_context
464 )
465 entry = tier_pending[ctx_id]
466 entry.keys.extend(primary_write_result.keys_to_store)
467 - entry.block_ids.extend(store_spec.block_ids)
467 + entry.chunk_ids.extend(store_spec.chunk_ids)
468 return True
469
470 def _flush_pending_promotions(self) -> None:
@@ -483,7 +483,7 @@ class TieringOffloadingManager(OffloadingManager):
483 job_metadata = TransferJob(
484 job_id=job_id,
485 keys=entry.keys,
486 - block_ids=np.array(entry.block_ids, dtype=np.int32),
486 + chunk_ids=np.array(entry.chunk_ids, dtype=np.int32),
487 is_promotion=True,
488 req_context=entry.req_context,
489 )
@@ -497,16 +497,16 @@ class TieringOffloadingManager(OffloadingManager):
497 self, keys: Collection[OffloadKey], req_context: ReqContext
498 ) -> LoadStoreSpec:
499 """
500 - Prepare blocks to be loaded from primary tier to GPU.
500 + Prepare chunks to be loaded from primary tier to GPU.
501
502 Callers only pass keys already confirmed HIT by lookup() earlier this
503 step.
504
505 - This increments ref_cnt on the blocks in the primary tier, protecting
505 + This increments ref_cnt on the chunks in the primary tier, protecting
506 them from eviction during the transfer.
507
508 Args:
509 - keys: Blocks to prepare for loading.
509 + keys: Chunks to prepare for loading.
510 req_context: Per-request context.
511
512 Returns:
@@ -517,10 +517,10 @@ class TieringOffloadingManager(OffloadingManager):
517 @override
518 def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
519 """
520 - Mark blocks as recently used in all tiers.
520 + Mark chunks as recently used in all tiers.
521
522 Args:
523 - keys: Blocks to mark as recently used.
523 + keys: Chunks to mark as recently used.
524 req_context: Per-request context.
525 """
526 self.primary_tier.touch(keys, req_context)
@@ -530,13 +530,13 @@ class TieringOffloadingManager(OffloadingManager):
530 @override
531 def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
532 """
533 - Mark blocks as done loading from primary tier to GPU.
533 + Mark chunks as done loading from primary tier to GPU.
534
535 - This decrements ref_cnt on the blocks in the primary tier, allowing
535 + This decrements ref_cnt on the chunks in the primary tier, allowing
536 them to be evicted again.
537
538 Args:
539 - keys: Blocks that finished loading.
539 + keys: Chunks that finished loading.
540 req_context: Per-request context.
541 """
542 self.primary_tier.complete_load(keys, req_context)
@@ -546,37 +546,37 @@ class TieringOffloadingManager(OffloadingManager):
546 self, keys: Collection[OffloadKey], req_context: ReqContext
547 ) -> PrepareStoreOutput | None:
548 """
549 - Prepare blocks to be stored from GPU to primary tier.
549 + Prepare chunks to be stored from GPU to primary tier.
550
551 CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure
552 that any completed async transfers have their ref_cnt decremented
553 before the primary tier makes eviction decisions.
554
555 - For request-level tiers, blocks already present in the primary tier
555 + For request-level tiers, chunks already present in the primary tier
556 are immediately cascaded via submit_store().
557
558 Args:
559 - keys: Blocks to prepare for storing.
559 + keys: Chunks to prepare for storing.
560 req_context: Per-request context.
561
562 Returns:
563 - PrepareStoreOutput describing where to store blocks and what was
563 + PrepareStoreOutput describing where to store chunks and what was
564 evicted, or None if store cannot proceed.
565 """
566 # Step 1: Poll for completed async jobs FIRST
567 # _process_finished_jobs() handles two kinds of completions here:
568 # - Cascade completions (store to a secondary tier, either a local
569 # cascade or a store job created for a remote requester via
570 - # create_store_job()): decrements ref_cnt on the primary blocks
570 + # create_store_job()): decrements ref_cnt on the primary chunks
571 # that were read, making them evictable again once ref_cnt hits 0.
572 # - Promotion completions (secondary->primary loads): sets a
573 - # not-yet-ready block's ref_cnt from -1 to 0 via complete_write(),
573 + # not-yet-ready chunk's ref_cnt from -1 to 0 via complete_write(),
574 # making it evictable for the first time.
575 # Both must be accounted for before the eviction decision below.
576 self._maybe_process_finished_jobs()
577
578 - # Step 2: Store to primary tier (new blocks only).
579 - # Cascading of these newly-stored blocks to ALL secondary tiers
578 + # Step 2: Store to primary tier (new chunks only).
579 + # Cascading of these newly-stored chunks to ALL secondary tiers
580 # happens later in complete_store(), after the GPU→Primary transfer
581 # completes.
582 primary_result = self.primary_tier.prepare_store(keys, req_context)
@@ -588,7 +588,7 @@ class TieringOffloadingManager(OffloadingManager):
588 state = self._req_state[req_context.req_id]
589 state.pending_primary_stores += 1
590
591 - # Step 3: For request-level tiers, cascade blocks already in primary
591 + # Step 3: For request-level tiers, cascade chunks already in primary
592 request_level_tiers = self._req_state[req_context.req_id].request_level_tiers
593 if request_level_tiers:
594 keys_to_store_set = set(primary_result.keys_to_store)
@@ -596,13 +596,13 @@ class TieringOffloadingManager(OffloadingManager):
596 k for k in keys if k not in keys_to_store_set
597 )
598 if keys_already_in_primary:
599 - self._cascade_existing_blocks_to_request_level_tiers(
599 + self._cascade_existing_chunks_to_request_level_tiers(
600 keys_already_in_primary, req_context, request_level_tiers
601 )
602
603 return primary_result
604
605 - def _cascade_existing_blocks_to_request_level_tiers(
605 + def _cascade_existing_chunks_to_request_level_tiers(
606 self,
607 keys: Sequence[OffloadKey],
608 req_context: ReqContext,
@@ -610,7 +610,7 @@ class TieringOffloadingManager(OffloadingManager):
610 ) -> None:
611 """
612 For tiers that requested request-level policy, submit_store() for
613 - blocks that are already present in the primary tier.
613 + chunks that are already present in the primary tier.
614
615 A key whose primary write is still in flight (HIT_PENDING) cannot be
616 dropped: prepare_store already excluded it as present, and the
@@ -653,7 +653,7 @@ class TieringOffloadingManager(OffloadingManager):
653 continue
654 assert state.request_level_tiers
655 keys, state.pending_cascade_keys = state.pending_cascade_keys, []
656 - self._cascade_existing_blocks_to_request_level_tiers(
656 + self._cascade_existing_chunks_to_request_level_tiers(
657 keys, state.req_context, state.request_level_tiers
658 )
659 self._maybe_finalize_request(req_id)
@@ -666,30 +666,30 @@ class TieringOffloadingManager(OffloadingManager):
666 success: bool = True,
667 ) -> None:
668 """
669 - Mark blocks as done storing from GPU to primary tier.
669 + Mark chunks as done storing from GPU to primary tier.
670
671 - This is where secondary tier cascading happens — after blocks are
671 + This is where secondary tier cascading happens — after chunks are
672 confirmed to be in the primary tier, they are cascaded to ALL
673 secondary tiers.
674
675 For each secondary tier:
676 1. Call primary.prepare_read() to get LoadStoreSpec AND increment
677 - ref_cnt (protecting blocks during async transfer)
677 + ref_cnt (protecting chunks during async transfer)
678 2. Call tier.submit_store() to start async transfer: primary→secondary
679 3. Track the job in _store_jobs dictionary
680
681 Args:
682 - keys: Blocks that finished storing.
682 + keys: Chunks that finished storing.
683 success: Whether the GPU→primary transfer succeeded.
684 req_context: Per-request context forwarded to primary.prepare_read().
685 """
686 - # Step 1: Complete store in primary tier (makes blocks loadable)
686 + # Step 1: Complete store in primary tier (makes chunks loadable)
687 self.primary_tier.complete_store(keys, req_context, success)
688
689 if success:
690 # Step 2: Cascade to ALL secondary tiers
691 # For each secondary tier, call primary.prepare_read() to get the
692 - # LoadStoreSpec AND to increment ref_cnt (protecting blocks from
692 + # LoadStoreSpec AND to increment ref_cnt (protecting chunks from
693 # eviction during the async transfer). One prepare_read() call per
694 # secondary tier.
695 for tier_idx, tier in enumerate(self.secondary_tiers):
@@ -710,22 +710,22 @@ class TieringOffloadingManager(OffloadingManager):
710 req_context: ReqContext,
711 tier_idx: int = 0,
712 ) -> TransferJob:
713 - """Pin blocks in the primary tier and create a tracked store job.
713 + """Pin chunks in the primary tier and create a tracked store job.
714
715 - Calls prepare_read() to increment ref_cnt (protecting blocks
715 + Calls prepare_read() to increment ref_cnt (protecting chunks
716 from eviction during the async transfer), allocates a job ID,
717 and registers the job in _jobs.
718
719 The caller is responsible for the actual data transfer and
720 reporting completion via get_finished_jobs().
721 """
722 - primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context)
723 - assert isinstance(primary_blocks_spec, CPULoadStoreSpec)
722 + primary_chunks_spec = self.primary_tier.prepare_read(keys, req_context)
723 + assert isinstance(primary_chunks_spec, CPULoadStoreSpec)
724 job_id = self._next_job_id()
725 job_metadata = TransferJob(
726 job_id=job_id,
727 keys=keys,
728 - block_ids=primary_blocks_spec.block_ids,
728 + chunk_ids=primary_chunks_spec.chunk_ids,
729 is_promotion=False,
730 req_context=req_context,
731 )
@@ -760,7 +760,7 @@ class TieringOffloadingManager(OffloadingManager):
760 policy = (
761 OffloadPolicy.REQUEST_LEVEL
762 if state.request_level_tiers
763 - else OffloadPolicy.BLOCK_LEVEL
763 + else OffloadPolicy.CHUNK_LEVEL
764 )
765 return RequestOffloadingContext(policy=policy)
766
vllm/v1/kv_offload/tiering/metrics.py
+23 -23
@@ -39,20 +39,20 @@ class _RequestMetricsState:
39 class _TierState:
40 active_promotion_count: int = 0
41 active_cascade_count: int = 0
42 - primary_write_block_count: int = 0
43 - primary_read_block_count: int = 0
42 + primary_write_chunk_count: int = 0
43 + primary_read_chunk_count: int = 0
44
45
46 class TieringMetricsTracker:
47 def __init__(
48 self,
49 tier_types: list[str],
50 - num_primary_blocks: int,
51 - primary_block_size: int,
50 + num_primary_chunks: int,
51 + primary_chunk_size: int,
52 ) -> None:
53 self._tier_types = tier_types
54 - self._num_primary_blocks = num_primary_blocks
55 - self._primary_block_size = primary_block_size
54 + self._num_primary_chunks = num_primary_chunks
55 + self._primary_chunk_size = primary_chunk_size
56 self._request_states: dict[str, _RequestMetricsState] = {}
57 self._tier_states = [_TierState() for _ in tier_types]
58 self._stats = OffloadingConnectorStats()
@@ -111,13 +111,13 @@ class TieringMetricsTracker:
111 def on_job_registered(self, job_metadata: _JobMetadataLike) -> None:
112 transfer_job = job_metadata.transfer_job
113 state = self._tier_states[job_metadata.tier_idx]
114 - block_count = len(transfer_job.block_ids)
114 + chunk_count = len(transfer_job.chunk_ids)
115 if transfer_job.is_promotion:
116 state.active_promotion_count += 1
117 - state.primary_write_block_count += block_count
117 + state.primary_write_chunk_count += chunk_count
118 else:
119 state.active_cascade_count += 1
120 - state.primary_read_block_count += block_count
120 + state.primary_read_chunk_count += chunk_count
121
122 def on_job_finished(
123 self, job_metadata: _JobMetadataLike, result: JobResult
@@ -149,25 +149,25 @@ class TieringMetricsTracker:
149 assert all(
150 state.active_promotion_count == 0
151 and state.active_cascade_count == 0
152 - and state.primary_write_block_count == 0
153 - and state.primary_read_block_count == 0
152 + and state.primary_write_chunk_count == 0
153 + and state.primary_read_chunk_count == 0
154 for state in self._tier_states
155 )
156
157 def _decrement_tier_state(self, job_metadata: _JobMetadataLike) -> None:
158 transfer_job = job_metadata.transfer_job
159 state = self._tier_states[job_metadata.tier_idx]
160 - block_count = len(transfer_job.block_ids)
160 + chunk_count = len(transfer_job.chunk_ids)
161 if transfer_job.is_promotion:
162 assert state.active_promotion_count > 0
163 state.active_promotion_count -= 1
164 - state.primary_write_block_count -= block_count
165 - assert state.primary_write_block_count >= 0
164 + state.primary_write_chunk_count -= chunk_count
165 + assert state.primary_write_chunk_count >= 0
166 else:
167 assert state.active_cascade_count > 0
168 state.active_cascade_count -= 1
169 - state.primary_read_block_count -= block_count
170 - assert state.primary_read_block_count >= 0
169 + state.primary_read_chunk_count -= chunk_count
170 + assert state.primary_read_chunk_count >= 0
171
172 def _observe_finished_job_stats(
173 self,
@@ -199,7 +199,7 @@ class TieringMetricsTracker:
199 if transfer_job.is_promotion
200 else TieringOffloadingMetrics.WRITE_TIME
201 )
202 - transfer_size = completed_key_count * self._primary_block_size
202 + transfer_size = completed_key_count * self._primary_chunk_size
203 self._stats.increase_counter(bytes_metric, transfer_size, labelvalues)
204 if completed_job.transfer_time is not None:
205 self._stats.increase_counter(
@@ -210,13 +210,13 @@ class TieringMetricsTracker:
210 for tier_idx, state in enumerate(self._tier_states):
211 labelvalues = self.tier_label(tier_idx)
212 write_usage = (
213 - state.primary_write_block_count / self._num_primary_blocks
214 - if self._num_primary_blocks > 0
213 + state.primary_write_chunk_count / self._num_primary_chunks
214 + if self._num_primary_chunks > 0
215 else 0.0
216 )
217 read_usage = (
218 - state.primary_read_block_count / self._num_primary_blocks
219 - if self._num_primary_blocks > 0
218 + state.primary_read_chunk_count / self._num_primary_chunks
219 + if self._num_primary_chunks > 0
220 else 0.0
221 )
222 stats.set_gauge(
@@ -248,12 +248,12 @@ class TieringMetricsTracker:
248 async_start_time: float | None,
249 ) -> None:
250 self._stats.increase_counter(
251 - TieringOffloadingMetrics.BLOCK_QUERIES,
251 + TieringOffloadingMetrics.CHUNK_QUERIES,
252 labelvalues=tier_label,
253 )
254 if result is LookupResult.HIT:
255 self._stats.increase_counter(
256 - TieringOffloadingMetrics.BLOCK_HITS,
256 + TieringOffloadingMetrics.CHUNK_HITS,
257 labelvalues=tier_label,
258 )
259 self._stats.observe_histogram(
vllm/v1/kv_offload/tiering/obj/manager.py
+2 -2
@@ -287,14 +287,14 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
287 self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys)
288 obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
289 self._submit_transfer(
290 - job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE
290 + job_metadata.job_id, job_metadata.chunk_ids, obj_keys, NIXL_WRITE
291 )
292
293 def submit_load(self, job_metadata: TransferJob) -> None:
294 self._load_job_keys[job_metadata.job_id] = list(job_metadata.keys)
295 obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
296 self._submit_transfer(
297 - job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_READ
297 + job_metadata.job_id, job_metadata.chunk_ids, obj_keys, NIXL_READ
298 )
299
300 def on_request_finished(self, req_context: ReqContext) -> None:
vllm/v1/kv_offload/tiering/p2p/manager.py
+2 -2
@@ -417,7 +417,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
417 def submit_store(self, job_metadata: TransferJob) -> None:
418 job_id = job_metadata.job_id
419 keys = list(job_metadata.keys)
420 - block_ids = job_metadata.block_ids.tolist()
420 + block_ids = job_metadata.chunk_ids.tolist()
421
422 assert len(keys) == len(block_ids)
423
@@ -476,7 +476,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
476 def submit_load(self, job_metadata: TransferJob) -> None:
477 job_id = job_metadata.job_id
478 keys = list(job_metadata.keys)
479 - block_ids = job_metadata.block_ids
479 + block_ids = job_metadata.chunk_ids
480
481 source = job_metadata.req_context.get_state(P2PSourceInfo)
482 logger.debug(
vllm/v1/kv_offload/tiering/p2p/session/server.py
+1 -1
@@ -593,7 +593,7 @@ class ServerRole:
593 self.add_stored_blocks(
594 lookup.kv_request_id,
595 list(meta.keys),
596 - meta.block_ids.tolist(),
596 + meta.chunk_ids.tolist(),
597 meta.job_id,
598 round_seq=lookup.round_seq,
599 from_lookup=True,
vllm/v1/kv_offload/tiering/spec.py
+15 -15
@@ -8,7 +8,7 @@ and configurable secondary tiers (e.g., Storage, Network).
8
9 Configuration via kv_connector_extra_config:
10 - cpu_bytes_to_use: (required) Bytes to allocate for CPU primary tier
11 - - block_size: (optional) Block size for offloaded blocks (default: GPU block size)
11 + - block_size: (optional) Tokens per offloaded chunk (default: GPU block size)
12 - eviction_policy: (optional) Primary tier eviction policy: built-in "lru"/
13 "arc", or the name of a policy registered via CachePolicyFactory, or an
14 out-of-tree CachePolicy class name paired with cache_policy_module_path
@@ -105,7 +105,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
105 metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY] = (
106 OffloadingHistogramMetadata(
107 documentation=(
108 - "Histogram of blocking time spent in a per-block tier lookup "
108 + "Histogram of blocking time spent in a per-chunk tier lookup "
109 "that resolved as a hit or miss, labeled by tier, in seconds."
110 ),
111 labelnames=("tier",),
@@ -127,7 +127,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
127 metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY] = (
128 OffloadingHistogramMetadata(
129 documentation=(
130 - "Histogram of wall-clock time from a per-block tier lookup "
130 + "Histogram of wall-clock time from a per-chunk tier lookup "
131 "first returning retry until that same tier lookup resolves "
132 "as a hit or miss, labeled by tier, in seconds."
133 ),
@@ -191,14 +191,14 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
191 labelnames=("tier",),
192 )
193 )
194 - metrics[TieringOffloadingMetrics.BLOCK_QUERIES] = OffloadingCounterMetadata(
194 + metrics[TieringOffloadingMetrics.CHUNK_QUERIES] = OffloadingCounterMetadata(
195 documentation=(
196 - "Number of block lookup queries sent to a tier, labeled by tier."
196 + "Number of chunk lookup queries sent to a tier, labeled by tier."
197 ),
198 labelnames=("tier",),
199 )
200 - metrics[TieringOffloadingMetrics.BLOCK_HITS] = OffloadingCounterMetadata(
201 - documentation="Number of block lookup hits in a tier, labeled by tier.",
200 + metrics[TieringOffloadingMetrics.CHUNK_HITS] = OffloadingCounterMetadata(
201 + documentation="Number of chunk lookup hits in a tier, labeled by tier.",
202 labelnames=("tier",),
203 )
204 metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES] = (
@@ -300,16 +300,16 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
300 # primary tier can eagerly create a memoryview over _base.
301 scheduler_mmap = SharedOffloadRegion(
302 engine_id=self._engine_id,
303 - num_blocks=self.num_blocks,
303 + num_chunks=self.num_chunks,
304 rank=None,
305 - kv_bytes_per_block=self.kv_bytes_per_chunk,
305 + kv_bytes_per_chunk=self.kv_bytes_per_chunk,
306 cpu_page_size=self.cpu_page_size_per_worker,
307 )
308 self._scheduler_mmap = scheduler_mmap
309
310 # Create primary tier (CPU-based)
311 primary_tier = CPUPrimaryTierOffloadingManager(
312 - num_blocks=self.num_blocks,
312 + num_chunks=self.num_chunks,
313 cache_policy=self.eviction_policy,
314 cache_policy_module_path=self.cache_policy_module_path,
315 enable_events=self.kv_events_config.enable_kv_cache_events,
@@ -367,9 +367,9 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
367
368 logger.info(
369 "Created TieringOffloadingManager with primary tier "
370 - "(%s, %s blocks) and %s secondary tier(s)",
370 + "(%s, %s chunks) and %s secondary tier(s)",
371 self.eviction_policy,
372 - self.num_blocks,
372 + self.num_chunks,
373 len(secondary_tiers),
374 )
375
@@ -393,9 +393,9 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
393 rank = torch.accelerator.current_device_index() % world_size
394 worker_mmap = SharedOffloadRegion(
395 engine_id=self._engine_id,
396 - num_blocks=self.num_blocks,
396 + num_chunks=self.num_chunks,
397 rank=rank,
398 - kv_bytes_per_block=self.kv_bytes_per_chunk,
398 + kv_bytes_per_chunk=self.kv_bytes_per_chunk,
399 cpu_page_size=self.cpu_page_size_per_worker,
400 )
401 try:
@@ -404,7 +404,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
404 return CPUOffloadingWorker(
405 kv_caches=kv_caches,
406 blocks_per_chunk=self.blocks_per_chunk,
407 - num_cpu_blocks=self.num_blocks,
407 + num_cpu_chunks=self.num_chunks,
408 mmap_region=worker_mmap,
409 canonical_layout=self.config.canonical_layout,
410 )