@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4ae6228284

[Bugfix][KV Connector] Populate SimpleCPUOffload BlockStored metadata (#54325)

mevince committed Sep 3, 2026 at 08:58 UTC 4ae6228284c4e9776f62d0867f543e621864176c
4 files changed +1133 -29
tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py new
+92
@@ -0,0 +1,92 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for SimpleCPUOffloadConnector-level behavior."""
4 +
5 +from __future__ import annotations
6 +
7 +import pytest
8 +
9 +from tests.v1.kv_connector.unit.utils import create_vllm_config
10 +from tests.v1.simple_kv_offload.test_scheduler import (
11 + _BYTES_PER_BLOCK,
12 + BLOCK_SIZE,
13 + DTYPE,
14 + HEAD_SIZE,
15 + NUM_KV_HEADS,
16 +)
17 +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
18 +from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (
19 + SimpleCPUOffloadConnector,
20 +)
21 +from vllm.v1.kv_cache_interface import (
22 + FullAttentionSpec,
23 + KVCacheConfig,
24 + KVCacheGroupSpec,
25 + KVCacheTensor,
26 +)
27 +
28 +
29 +def _make_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig:
30 + fa_layers = ["layer0"]
31 + spec = FullAttentionSpec(
32 + block_size=BLOCK_SIZE,
33 + num_kv_heads=NUM_KV_HEADS,
34 + head_size=HEAD_SIZE,
35 + dtype=DTYPE,
36 + )
37 + groups = [KVCacheGroupSpec(fa_layers, spec)]
38 + size = _BYTES_PER_BLOCK * num_blocks
39 + tensors = [
40 + KVCacheTensor(
41 + size=size,
42 + layers=fa_layers,
43 + layer_stride=size,
44 + block_stride=_BYTES_PER_BLOCK,
45 + )
46 + ]
47 + return KVCacheConfig(
48 + num_blocks=num_blocks,
49 + kv_cache_tensors=tensors,
50 + kv_cache_groups=groups,
51 + )
52 +
53 +
54 +def _make_connector(
55 + extra_config: dict | None = None,
56 + num_gpu_blocks: int = 16,
57 + cpu_bytes: int = _BYTES_PER_BLOCK * 8,
58 +) -> SimpleCPUOffloadConnector:
59 + kv_cache_config = _make_kv_cache_config(num_gpu_blocks)
60 + config_dict = {
61 + "cpu_bytes_to_use": cpu_bytes,
62 + }
63 + if extra_config:
64 + config_dict.update(extra_config)
65 + vllm_config = create_vllm_config(
66 + kv_connector="SimpleCPUOffloadConnector",
67 + kv_role="kv_both",
68 + kv_connector_extra_config=config_dict,
69 + )
70 + connector = SimpleCPUOffloadConnector(
71 + vllm_config=vllm_config,
72 + role=KVConnectorRole.SCHEDULER,
73 + kv_cache_config=kv_cache_config,
74 + )
75 + return connector
76 +
77 +
78 +def test_disk_mode_rejects_non_positive_capacity() -> None:
79 + """Disk mode with disk_capacity_bytes <= 0 raises ValueError.
80 +
81 + Without this guard, the worker operates in disk mode but the scheduler
82 + classifies events as MEDIUM_CPU instead of MEDIUM_STORAGE, silently
83 + mislabeling events.
84 + """
85 + with pytest.raises(ValueError, match="disk_capacity_bytes > 0"):
86 + _make_connector(
87 + extra_config={
88 + "kv_offload_backend": "disk",
89 + "disk_path": "/tmp/fake",
90 + "disk_capacity_bytes": 0,
91 + }
92 + )
tests/v1/simple_kv_offload/test_kv_events.py new
+790
@@ -0,0 +1,790 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""KV cache event emission tests for SimpleCPUOffloadScheduler.
4 +
5 +Verifies that BlockStored/BlockRemoved events carry the correct storage
6 +medium (MEDIUM_CPU / MEDIUM_STORAGE) and per-group metadata when
7 +``enable_kv_cache_events`` is set, and that no events are emitted when it
8 +is unset.
9 +"""
10 +
11 +from __future__ import annotations
12 +
13 +from tests.v1.simple_kv_offload.test_scheduler import (
14 + _BYTES_PER_BLOCK,
15 + BLOCK_SIZE,
16 + DTYPE,
17 + HEAD_SIZE,
18 + NUM_KV_HEADS,
19 + SchedulerFixture,
20 + _alloc_and_register,
21 + _allocate_cp_gpu_blocks,
22 + _make_cp_request,
23 + _make_kv_cache_config,
24 + _make_vllm_config,
25 + make_request,
26 + make_scheduler_output,
27 + simulate_store_completion,
28 +)
29 +from vllm import SamplingParams
30 +from vllm.config.kv_events import KVEventsConfig
31 +from vllm.distributed.kv_events import (
32 + MEDIUM_CPU,
33 + MEDIUM_STORAGE,
34 + BlockRemoved,
35 + BlockStored,
36 +)
37 +from vllm.lora.request import LoRARequest
38 +from vllm.utils.hashing import sha256
39 +from vllm.v1.core.block_pool import BlockPool
40 +from vllm.v1.core.kv_cache_manager import KVCacheBlocks
41 +from vllm.v1.core.kv_cache_utils import (
42 + get_block_hash,
43 + get_request_block_hasher,
44 + make_block_hash_with_group_id,
45 + maybe_convert_block_hash,
46 +)
47 +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
48 +from vllm.v1.kv_cache_interface import (
49 + FullAttentionSpec,
50 + KVCacheConfig,
51 + KVCacheGroupSpec,
52 + KVCacheTensor,
53 + MambaSpec,
54 + SlidingWindowSpec,
55 +)
56 +from vllm.v1.request import Request
57 +from vllm.v1.simple_kv_offload.manager import (
58 + SimpleCPUOffloadScheduler,
59 + StoreRequestState,
60 +)
61 +
62 +
63 +def _make_mixed_kv_cache_config(
64 + num_blocks: int,
65 + sliding_window: int | None = None,
66 + with_mamba: bool = False,
67 + mamba_block_size: int = BLOCK_SIZE,
68 +) -> KVCacheConfig:
69 + """Build a KVCacheConfig with a FullAttention group plus optional
70 + SlidingWindow and Mamba groups for per-group metadata tests."""
71 + register_all_kvcache_specs(vllm_config=None)
72 + groups = []
73 + tensors = []
74 +
75 + # Group 0 is always FullAttention (the scheduler requires one FA group).
76 + fa_layers = ["layer_fa"]
77 + fa_spec = FullAttentionSpec(
78 + block_size=BLOCK_SIZE,
79 + num_kv_heads=NUM_KV_HEADS,
80 + head_size=HEAD_SIZE,
81 + dtype=DTYPE,
82 + )
83 + groups.append(KVCacheGroupSpec(fa_layers, fa_spec))
84 + fa_bytes = _BYTES_PER_BLOCK * num_blocks
85 + tensors.append(
86 + KVCacheTensor(
87 + size=fa_bytes,
88 + layers=fa_layers,
89 + layer_stride=fa_bytes,
90 + block_stride=_BYTES_PER_BLOCK,
91 + )
92 + )
93 +
94 + if sliding_window is not None:
95 + sw_layers = ["layer_sw"]
96 + sw_spec = SlidingWindowSpec(
97 + block_size=BLOCK_SIZE,
98 + num_kv_heads=NUM_KV_HEADS,
99 + head_size=HEAD_SIZE,
100 + dtype=DTYPE,
101 + sliding_window=sliding_window,
102 + )
103 + groups.append(KVCacheGroupSpec(sw_layers, sw_spec))
104 + sw_bytes = _BYTES_PER_BLOCK * num_blocks
105 + tensors.append(
106 + KVCacheTensor(
107 + size=sw_bytes,
108 + layers=sw_layers,
109 + layer_stride=sw_bytes,
110 + block_stride=_BYTES_PER_BLOCK,
111 + )
112 + )
113 +
114 + if with_mamba:
115 + m_layers = ["layer_m"]
116 + m_spec = MambaSpec(
117 + block_size=mamba_block_size,
118 + shapes=((NUM_KV_HEADS, HEAD_SIZE, HEAD_SIZE),),
119 + dtypes=(DTYPE,),
120 + mamba_cache_mode="align",
121 + )
122 + groups.append(KVCacheGroupSpec(m_layers, m_spec))
123 + m_bytes = _BYTES_PER_BLOCK * num_blocks
124 + tensors.append(
125 + KVCacheTensor(
126 + size=m_bytes,
127 + layers=m_layers,
128 + layer_stride=m_bytes,
129 + block_stride=_BYTES_PER_BLOCK,
130 + )
131 + )
132 +
133 + return KVCacheConfig(
134 + num_blocks=num_blocks,
135 + kv_cache_tensors=tensors,
136 + kv_cache_groups=groups,
137 + )
138 +
139 +
140 +def make_events_scheduler(
141 + num_cpu_blocks: int = 8,
142 + num_gpu_blocks: int = 16,
143 + num_groups: int = 1,
144 + lazy: bool = False,
145 + disk_capacity_bytes: int = 0,
146 + dcp_world_size: int = 1,
147 + kv_cache_config: KVCacheConfig | None = None,
148 + enable_events: bool = True,
149 +) -> SchedulerFixture:
150 + """Build a SimpleCPUOffloadScheduler with kv cache events enabled."""
151 + if kv_cache_config is None:
152 + kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)
153 + vllm_config = _make_vllm_config()
154 + if enable_events:
155 + vllm_config.kv_events_config = KVEventsConfig(enable_kv_cache_events=True)
156 + vllm_config.parallel_config.decode_context_parallel_size = dcp_world_size
157 +
158 + virtual_block_size = BLOCK_SIZE * dcp_world_size
159 + cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups
160 +
161 + sched = SimpleCPUOffloadScheduler(
162 + vllm_config=vllm_config,
163 + kv_cache_config=kv_cache_config,
164 + cpu_capacity_bytes=cpu_capacity_bytes,
165 + scheduler_block_size=virtual_block_size,
166 + hash_block_size=virtual_block_size,
167 + lazy_offload=lazy,
168 + disk_capacity_bytes=disk_capacity_bytes,
169 + )
170 +
171 + gpu_block_pool = BlockPool(
172 + num_gpu_blocks=num_gpu_blocks,
173 + enable_caching=True,
174 + hash_block_size=virtual_block_size,
175 + )
176 + sched.bind_gpu_block_pool(gpu_block_pool)
177 + return SchedulerFixture(
178 + scheduler=sched,
179 + gpu_block_pool=gpu_block_pool,
180 + vllm_config=vllm_config,
181 + kv_cache_config=kv_cache_config,
182 + num_groups=num_groups,
183 + )
184 +
185 +
186 +def _do_eager_store(
187 + fix: SchedulerFixture,
188 + num_blocks: int = 2,
189 + lora_request: LoRARequest | None = None,
190 +) -> tuple[object, Request]:
191 + """Run an eager store of ``num_blocks`` and complete it.
192 +
193 + Returns (meta, req) for further assertions.
194 + """
195 + sched = fix.scheduler
196 + if lora_request is not None:
197 + # lora_request must be set in the constructor so block_hashes
198 + # (computed in __init__) include the lora extra_keys.
199 + req = Request(
200 + request_id="req-lora",
201 + prompt_token_ids=list(range(num_blocks * BLOCK_SIZE + 1)),
202 + sampling_params=SamplingParams(max_tokens=1),
203 + pooling_params=None,
204 + mm_features=None,
205 + lora_request=lora_request,
206 + block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256),
207 + )
208 + else:
209 + req = make_request(num_blocks=num_blocks)
210 + kv_blocks = _alloc_and_register(fix, req, num_blocks)
211 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
212 + block_ids = kv_blocks.get_block_ids()
213 + sched_out = make_scheduler_output(
214 + {req.request_id: num_blocks * BLOCK_SIZE},
215 + new_reqs={req.request_id: block_ids},
216 + )
217 + meta = sched.build_connector_meta(sched_out)
218 + assert meta.store_event >= 0
219 + simulate_store_completion(sched, meta.store_event)
220 + return meta, req
221 +
222 +
223 +# ---------------------------------------------------------------------------
224 +# events disabled -> no events
225 +# ---------------------------------------------------------------------------
226 +def test_events_disabled_no_events() -> None:
227 + """When enable_kv_cache_events is False, take_events returns nothing."""
228 + fix = make_events_scheduler(enable_events=False)
229 + sched = fix.scheduler
230 + assert sched.enable_kv_cache_events is False
231 + _do_eager_store(fix, num_blocks=2)
232 + events = list(sched.take_events())
233 + assert events == [], f"expected no events, got {events}"
234 +
235 +
236 +# ---------------------------------------------------------------------------
237 +# eager store emits BlockStored with per-group metadata
238 +# ---------------------------------------------------------------------------
239 +def test_block_stored_per_group_metadata_full_attention() -> None:
240 + """BlockStored carries medium, group_idx, kind, sliding_window, block_size,
241 + and the correct block hash."""
242 + fix = make_events_scheduler()
243 + sched = fix.scheduler
244 + assert sched.kv_event_medium == MEDIUM_CPU
245 + _, req = _do_eager_store(fix, num_blocks=2)
246 + events = list(sched.take_events())
247 + stored = [e for e in events if isinstance(e, BlockStored)]
248 + assert len(stored) == 2
249 + ev = stored[0]
250 + assert ev.medium == MEDIUM_CPU
251 + assert ev.group_idx == 0
252 + assert ev.kv_cache_spec_kind == "full_attention"
253 + assert ev.kv_cache_spec_sliding_window is None
254 + assert ev.block_size == BLOCK_SIZE
255 + assert ev.locality == "LOCAL"
256 + expected_hash_0 = maybe_convert_block_hash(
257 + get_block_hash(make_block_hash_with_group_id(req.block_hashes[0], 0))
258 + )
259 + assert ev.block_hashes == [expected_hash_0]
260 + assert ev.parent_block_hash is None
261 + assert ev.token_ids == req.prompt_token_ids[0:BLOCK_SIZE]
262 +
263 + # Second block's parent_block_hash equals the first block's hash.
264 + ev1 = stored[1]
265 + expected_hash_1 = maybe_convert_block_hash(
266 + get_block_hash(make_block_hash_with_group_id(req.block_hashes[1], 0))
267 + )
268 + assert ev1.block_hashes == [expected_hash_1]
269 + assert ev1.parent_block_hash == expected_hash_0
270 + assert ev1.token_ids == req.prompt_token_ids[BLOCK_SIZE : 2 * BLOCK_SIZE]
271 +
272 +
273 +def test_eager_store_lora_metadata() -> None:
274 + """BlockStored carries lora_id and lora_name from req.lora_request."""
275 + fix = make_events_scheduler()
276 + sched = fix.scheduler
277 +
278 + _do_eager_store(
279 + fix,
280 + num_blocks=1,
281 + lora_request=LoRARequest(
282 + lora_name="test-lora",
283 + lora_int_id=42,
284 + lora_path="/fake/lora/path",
285 + ),
286 + )
287 +
288 + events = list(sched.take_events())
289 + stored = [e for e in events if isinstance(e, BlockStored)]
290 + assert len(stored) == 1
291 + ev = stored[0]
292 + assert ev.lora_id == 42, f"expected lora_id=42, got {ev.lora_id}"
293 + assert ev.lora_name == "test-lora", (
294 + f"expected lora_name='test-lora', got {ev.lora_name!r}"
295 + )
296 + assert ev.extra_keys is not None, "expected extra_keys for lora request"
297 + assert any("test-lora" in keys for keys in ev.extra_keys), (
298 + f"expected 'test-lora' in extra_keys, got {ev.extra_keys}"
299 + )
300 + assert len(ev.token_ids) == BLOCK_SIZE
301 +
302 +
303 +# ---------------------------------------------------------------------------
304 +# disk mode -> MEDIUM_STORAGE
305 +# ---------------------------------------------------------------------------
306 +def test_disk_mode_medium_storage_on_block_stored() -> None:
307 + """disk_capacity_bytes>0 sets kv_event_medium=STORAGE on BlockStored."""
308 + fix = make_events_scheduler(
309 + disk_capacity_bytes=_BYTES_PER_BLOCK * 8,
310 + enable_events=True,
311 + )
312 + sched = fix.scheduler
313 + assert sched.kv_event_medium == MEDIUM_STORAGE
314 + _do_eager_store(fix, num_blocks=1)
315 + events = list(sched.take_events())
316 + stored = [e for e in events if isinstance(e, BlockStored)]
317 + assert len(stored) == 1
318 + assert stored[0].medium == MEDIUM_STORAGE
319 + assert stored[0].locality == "LOCAL"
320 +
321 +
322 +def test_finished_eager_store_emits_all_storage_events() -> None:
323 + """Finished stores preserve and announce every registered hash."""
324 + fix = make_events_scheduler(
325 + disk_capacity_bytes=_BYTES_PER_BLOCK * 8,
326 + enable_events=True,
327 + )
328 + sched = fix.scheduler
329 + req = make_request(num_blocks=2)
330 + kv_blocks = _alloc_and_register(fix, req, num_blocks=2)
331 + gpu_blocks = kv_blocks.blocks[0]
332 + primary_hashes = []
333 + for gpu_block in gpu_blocks:
334 + assert gpu_block.block_hash is not None
335 + primary_hashes.append(gpu_block.block_hash)
336 +
337 + fine_grained_req = Request(
338 + request_id="req-finished-fine-grained-hash",
339 + prompt_token_ids=req.prompt_token_ids,
340 + sampling_params=req.sampling_params,
341 + pooling_params=None,
342 + mm_features=None,
343 + block_hasher=get_request_block_hasher(BLOCK_SIZE // 2, sha256),
344 + )
345 + secondary_hash = make_block_hash_with_group_id(fine_grained_req.block_hashes[0], 0)
346 + fix.gpu_block_pool._insert_block_hash(
347 + secondary_hash, gpu_blocks[0], num_tokens=BLOCK_SIZE // 2
348 + )
349 +
350 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
351 + sched.request_finished_all_groups(req, kv_blocks.get_block_ids())
352 + meta = sched.build_connector_meta(make_scheduler_output({}))
353 + assert meta.store_event >= 0
354 + simulate_store_completion(sched, meta.store_event)
355 +
356 + primary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
357 + primary_hashes[0]
358 + )
359 + secondary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
360 + secondary_hash
361 + )
362 + assert primary_cpu_block is not None
363 + assert secondary_cpu_block is primary_cpu_block
364 +
365 + stored = [e for e in sched.take_events() if isinstance(e, BlockStored)]
366 + primary_event_hashes = [
367 + maybe_convert_block_hash(get_block_hash(block_hash))
368 + for block_hash in primary_hashes
369 + ]
370 + secondary_event_hash = maybe_convert_block_hash(get_block_hash(secondary_hash))
371 + by_hash = {event.block_hashes[0]: event for event in stored}
372 + assert set(by_hash) == {*primary_event_hashes, secondary_event_hash}
373 + assert (
374 + by_hash[primary_event_hashes[0]].token_ids == req.prompt_token_ids[:BLOCK_SIZE]
375 + )
376 + assert (
377 + by_hash[primary_event_hashes[1]].token_ids
378 + == req.prompt_token_ids[BLOCK_SIZE : 2 * BLOCK_SIZE]
379 + )
380 + assert by_hash[primary_event_hashes[1]].parent_block_hash == primary_event_hashes[0]
381 + secondary_event = by_hash[secondary_event_hash]
382 + assert secondary_event.medium == MEDIUM_STORAGE
383 + assert secondary_event.block_size == 0
384 + assert secondary_event.token_ids == []
385 + assert secondary_event.parent_block_hash is None
386 +
387 + sched.cpu_block_pool.evict_blocks({primary_cpu_block.block_id})
388 + removed = [
389 + event for event in sched.take_events() if isinstance(event, BlockRemoved)
390 + ]
391 + assert {event.block_hashes[0] for event in removed} == {
392 + primary_event_hashes[0],
393 + secondary_event_hash,
394 + }
395 + assert all(event.medium == MEDIUM_STORAGE for event in removed)
396 +
397 +
398 +def test_secondary_hash_block_stored_metadata() -> None:
399 + """Secondary hashes carry their own token range and block size."""
400 + fix = make_events_scheduler(
401 + disk_capacity_bytes=_BYTES_PER_BLOCK * 8,
402 + enable_events=True,
403 + )
404 + sched = fix.scheduler
405 + sched.hash_block_size = BLOCK_SIZE // 2
406 + req = Request(
407 + request_id="req-secondary-hash",
408 + prompt_token_ids=list(range(BLOCK_SIZE + 1)),
409 + sampling_params=SamplingParams(max_tokens=1),
410 + pooling_params=None,
411 + mm_features=None,
412 + block_hasher=get_request_block_hasher(BLOCK_SIZE // 2, sha256),
413 + )
414 + req.num_computed_tokens = BLOCK_SIZE
415 + gpu_block = fix.gpu_block_pool.get_new_blocks(1)[0]
416 + primary_hash = make_block_hash_with_group_id(req.block_hashes[1], 0)
417 + secondary_hash = make_block_hash_with_group_id(req.block_hashes[0], 0)
418 + fix.gpu_block_pool._insert_block_hash(
419 + primary_hash, gpu_block, num_tokens=BLOCK_SIZE
420 + )
421 + fix.gpu_block_pool._insert_block_hash(
422 + secondary_hash, gpu_block, num_tokens=BLOCK_SIZE // 2
423 + )
424 + state = StoreRequestState(
425 + request=req,
426 + block_ids=([gpu_block.block_id],),
427 + num_stored_blocks=[0],
428 + )
429 +
430 + gpu_block_ids, advanced_per_group, block_meta = sched._select_eager_blocks_to_store(
431 + state,
432 + ([gpu_block.block_id],),
433 + )
434 + assert gpu_block_ids == [gpu_block.block_id]
435 + assert advanced_per_group == [1]
436 + assert block_meta is not None
437 + assert set(block_meta[0]) == {primary_hash, secondary_hash}
438 +
439 + cpu_block = sched.cpu_block_pool.get_new_blocks(1)[0]
440 + sched._process_store_completion(
441 + gpu_block_ids,
442 + [cpu_block.block_id],
443 + block_meta,
444 + )
445 +
446 + primary_event_hash = maybe_convert_block_hash(get_block_hash(primary_hash))
447 + secondary_event_hash = maybe_convert_block_hash(get_block_hash(secondary_hash))
448 + stored = [event for event in sched.take_events() if isinstance(event, BlockStored)]
449 + by_hash = {event.block_hashes[0]: event for event in stored}
450 + assert set(by_hash) == {primary_event_hash, secondary_event_hash}
451 + assert by_hash[primary_event_hash].block_size == BLOCK_SIZE
452 + assert by_hash[primary_event_hash].token_ids == req.prompt_token_ids[:BLOCK_SIZE]
453 + assert by_hash[secondary_event_hash].block_size == BLOCK_SIZE // 2
454 + assert (
455 + by_hash[secondary_event_hash].token_ids
456 + == req.prompt_token_ids[: BLOCK_SIZE // 2]
457 + )
458 +
459 + primary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
460 + primary_hash
461 + )
462 + secondary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
463 + secondary_hash
464 + )
465 + assert primary_cpu_block is not None
466 + assert secondary_cpu_block is primary_cpu_block
467 +
468 +
469 +# ---------------------------------------------------------------------------
470 +# BlockRemoved relabeled to correct medium
471 +# ---------------------------------------------------------------------------
472 +def test_block_removed_relabeled_medium_cpu() -> None:
473 + """Evicted BlockRemoved (emitted as GPU by BlockPool) is relabeled CPU."""
474 + # 5 total = 4 usable (null_block takes 1). Fill 4, then store 2 more to
475 + # evict 2 LRU blocks, producing BlockRemoved events.
476 + fix = make_events_scheduler(num_cpu_blocks=5, num_gpu_blocks=16)
477 + sched = fix.scheduler
478 +
479 + # Fill CPU with 4 blocks (2 + 2).
480 + _do_eager_store(fix, num_blocks=2)
481 + # Drain store events so they don't clutter the eviction check.
482 + list(sched.take_events())
483 + _do_eager_store(fix, num_blocks=2)
484 + list(sched.take_events())
485 +
486 + # Now store 2 more -> evicts 2 LRU blocks -> BlockRemoved emitted by pool
487 + # with MEDIUM_GPU, then relabeled by scheduler.take_events().
488 + _do_eager_store(fix, num_blocks=2)
489 + events = list(sched.take_events())
490 +
491 + removed = [e for e in events if isinstance(e, BlockRemoved)]
492 + assert len(removed) >= 2, (
493 + f"expected >=2 BlockRemoved from eviction, got {len(removed)}"
494 + )
495 + for ev in removed:
496 + assert ev.medium == MEDIUM_CPU, (
497 + f"BlockRemoved should be relabeled CPU, got {ev.medium}"
498 + )
499 + assert ev.locality == "LOCAL", (
500 + f"BlockRemoved should be LOCAL, got {ev.locality}"
501 + )
502 + stored = [e for e in events if isinstance(e, BlockStored)]
503 + assert len(stored) == 2
504 +
505 +
506 +# ---------------------------------------------------------------------------
507 +# lazy store emits BlockStored
508 +# ---------------------------------------------------------------------------
509 +def test_lazy_store_emits_block_stored() -> None:
510 + """Lazy-mode store completion also emits BlockStored with MEDIUM_CPU."""
511 + fix = make_events_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True)
512 + sched = fix.scheduler
513 + gpu_pool = fix.gpu_block_pool
514 +
515 + # Allocate, hash, free -> hashed blocks in free queue.
516 + req = make_request(num_blocks=2)
517 + gpu_blocks = []
518 + for blk in gpu_pool.get_new_blocks(2):
519 + gpu_blocks.append(blk)
520 + gpu_pool.cache_full_blocks(
521 + request=req,
522 + blocks=gpu_blocks,
523 + num_cached_blocks=0,
524 + num_full_blocks=2,
525 + block_size=BLOCK_SIZE,
526 + kv_cache_group_id=0,
527 + )
528 + gpu_pool.free_blocks(gpu_blocks)
529 +
530 + # Push hashed blocks to LRU head (8 total - 1 null = 7 usable; 2 freed
531 + # hashed + 5 other free -> consume 5 fillers).
532 + fillers = gpu_pool.get_new_blocks(5)
533 +
534 + meta = sched.build_connector_meta(make_scheduler_output({}))
535 + assert meta.store_event >= 0
536 + simulate_store_completion(sched, meta.store_event)
537 + gpu_pool.free_blocks(fillers)
538 +
539 + events = list(sched.take_events())
540 + stored = [e for e in events if isinstance(e, BlockStored)]
541 + assert len(stored) == 2, f"expected 2 BlockStored, got {len(stored)}"
542 + for ev in stored:
543 + assert ev.medium == MEDIUM_CPU
544 + assert ev.group_idx == 0
545 + assert ev.kv_cache_spec_kind == "full_attention"
546 + assert ev.locality == "LOCAL"
547 + assert ev.token_ids == []
548 + assert ev.parent_block_hash is None
549 +
550 +
551 +# ---------------------------------------------------------------------------
552 +# cp_world_size scales block_size for non-Mamba
553 +# ---------------------------------------------------------------------------
554 +def test_cp_world_size_scales_block_size() -> None:
555 + """block_size on BlockStored = spec.block_size * cp_world_size (FA)."""
556 + fix = make_events_scheduler(dcp_world_size=2)
557 + sched = fix.scheduler
558 + gpu_pool = fix.gpu_block_pool
559 + vbs = BLOCK_SIZE * 2
560 + assert sched.cp_world_size == 2
561 +
562 + req = _make_cp_request(num_blocks=1, virtual_block_size=vbs)
563 + gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 1, vbs)
564 + kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
565 + req.num_computed_tokens = vbs
566 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
567 + sched_out = make_scheduler_output(
568 + {req.request_id: vbs},
569 + new_reqs={req.request_id: kv_blocks.get_block_ids()},
570 + )
571 + meta = sched.build_connector_meta(sched_out)
572 + assert meta.store_event >= 0
573 + simulate_store_completion(sched, meta.store_event)
574 +
575 + events = list(sched.take_events())
576 + stored = [e for e in events if isinstance(e, BlockStored)]
577 + assert len(stored) == 1
578 + # spec.block_size=BLOCK_SIZE, cp_world_size=2 -> BLOCK_SIZE * 2
579 + assert stored[0].block_size == BLOCK_SIZE * 2
580 + assert stored[0].locality == "LOCAL"
581 +
582 +
583 +# ---------------------------------------------------------------------------
584 +# Mamba block_size NOT scaled by cp_world_size
585 +# ---------------------------------------------------------------------------
586 +def test_mamba_block_size_not_scaled() -> None:
587 + """Mamba group: block_size = spec.block_size (no cp scaling).
588 +
589 + With cp_world_size=2, scheduler_block_size=hash_block_size=32. The
590 + coordinator constraints force Mamba spec.block_size=32 (must divide
591 + scheduler_block_size and be a multiple of hash_block_size). The event
592 + block_size for Mamba must be 32 (spec.block_size * 1), NOT 64
593 + (spec.block_size * cp_world_size) -- that is the bug this guards.
594 + """
595 + kv_cfg = _make_mixed_kv_cache_config(
596 + num_blocks=16, with_mamba=True, mamba_block_size=32
597 + )
598 + fix = make_events_scheduler(
599 + num_cpu_blocks=8,
600 + num_gpu_blocks=16,
601 + kv_cache_config=kv_cfg,
602 + dcp_world_size=2,
603 + num_groups=2,
604 + )
605 + sched = fix.scheduler
606 +
607 + cpu_pool = sched.cpu_block_pool
608 + gpu_pool = fix.gpu_block_pool
609 +
610 + cpu_blk = cpu_pool.get_new_blocks(1)[0]
611 + gpu_blk = gpu_pool.get_new_blocks(1)[0]
612 +
613 + raw_hash = make_request(num_blocks=1).block_hashes[0]
614 + mamba_hash = make_block_hash_with_group_id(raw_hash, group_id=1)
615 + gpu_blk._block_hash = mamba_hash # type: ignore[attr-defined]
616 +
617 + sched._process_store_completion(
618 + gpu_block_ids=[gpu_blk.block_id],
619 + cpu_block_ids=[cpu_blk.block_id],
620 + )
621 +
622 + events = list(sched.take_events())
623 + stored = [e for e in events if isinstance(e, BlockStored)]
624 + assert len(stored) == 1
625 + ev = stored[0]
626 + assert ev.group_idx == 1
627 + assert ev.kv_cache_spec_kind == "mamba"
628 + assert ev.kv_cache_spec_sliding_window is None
629 + assert ev.locality == "LOCAL"
630 + assert ev.block_size == 32, (
631 + f"Mamba block_size should be unscaled (32), got {ev.block_size}"
632 + )
633 +
634 +
635 +# ---------------------------------------------------------------------------
636 +# multi-group FA+SW stores carry correct group_idx and kind
637 +# ---------------------------------------------------------------------------
638 +def test_multi_group_fa_block_stored_group_idx() -> None:
639 + """Two distinct groups (FA + SW): each BlockStored carries its group_idx
640 + and kind. HybridKVCacheCoordinator requires >=2 distinct spec groups, so
641 + we use FA (group 0) + SlidingWindow (group 1)."""
642 + kv_cfg = _make_mixed_kv_cache_config(num_blocks=16, sliding_window=32)
643 + fix = make_events_scheduler(
644 + num_cpu_blocks=8,
645 + num_gpu_blocks=16,
646 + kv_cache_config=kv_cfg,
647 + num_groups=2,
648 + )
649 + sched = fix.scheduler
650 +
651 + cpu_pool = sched.cpu_block_pool
652 + gpu_pool = fix.gpu_block_pool
653 + raw_hash = make_request(num_blocks=1).block_hashes[0]
654 +
655 + # Store one block for each group via direct _process_store_completion.
656 + for gidx in (0, 1):
657 + cpu_blk = cpu_pool.get_new_blocks(1)[0]
658 + gpu_blk = gpu_pool.get_new_blocks(1)[0]
659 + gpu_blk._block_hash = make_block_hash_with_group_id( # type: ignore
660 + raw_hash, group_id=gidx
661 + )
662 + sched._process_store_completion(
663 + gpu_block_ids=[gpu_blk.block_id],
664 + cpu_block_ids=[cpu_blk.block_id],
665 + )
666 +
667 + events = list(sched.take_events())
668 + stored = [e for e in events if isinstance(e, BlockStored)]
669 + assert len(stored) == 2
670 + by_group = {ev.group_idx: ev for ev in stored}
671 + assert set(by_group) == {0, 1}
672 + assert by_group[0].kv_cache_spec_kind == "full_attention"
673 + assert by_group[0].kv_cache_spec_sliding_window is None
674 + assert by_group[1].kv_cache_spec_kind == "sliding_window"
675 + assert by_group[1].kv_cache_spec_sliding_window == 32
676 + for ev in stored:
677 + assert ev.locality == "LOCAL"
678 +
679 +
680 +def test_disk_mode_block_removed_relabeled_storage() -> None:
681 + """Disk mode: evicted BlockRemoved relabeled to MEDIUM_STORAGE."""
682 + fix = make_events_scheduler(
683 + num_cpu_blocks=5,
684 + num_gpu_blocks=16,
685 + disk_capacity_bytes=_BYTES_PER_BLOCK * 5,
686 + )
687 + sched = fix.scheduler
688 + assert sched.kv_event_medium == MEDIUM_STORAGE
689 +
690 + # Fill disk with 4 blocks (2 + 2).
691 + _do_eager_store(fix, num_blocks=2)
692 + list(sched.take_events())
693 + _do_eager_store(fix, num_blocks=2)
694 + list(sched.take_events())
695 +
696 + # Store 2 more -> evicts 2 LRU blocks -> BlockRemoved relabeled STORAGE.
697 + _do_eager_store(fix, num_blocks=2)
698 + events = list(sched.take_events())
699 +
700 + removed = [e for e in events if isinstance(e, BlockRemoved)]
701 + assert len(removed) >= 2, (
702 + f"expected >=2 BlockRemoved from eviction, got {len(removed)}"
703 + )
704 + for ev in removed:
705 + assert ev.medium == MEDIUM_STORAGE, (
706 + f"BlockRemoved should be relabeled STORAGE, got {ev.medium}"
707 + )
708 + assert ev.locality == "LOCAL", (
709 + f"BlockRemoved should be LOCAL, got {ev.locality}"
710 + )
711 + stored = [e for e in events if isinstance(e, BlockStored)]
712 + assert len(stored) == 2
713 + for ev in stored:
714 + assert ev.medium == MEDIUM_STORAGE
715 +
716 +
717 +def test_mamba_dcp_metadata_guard() -> None:
718 + """Mamba+DCP: token_ids guard emits [] when capture and emission sizes differ.
719 +
720 + The metadata capture at _prepare_eager_store_specs uses
721 + g_block_size = spec.block_size * cp_world_size = 32 * 2 = 64 to slice
722 + tokens. The event emission uses block_size = spec.block_size * 1 = 32
723 + for Mamba (SSM state is not CP-sharded). The guard in
724 + _process_store_completion detects len(token_ids) != block_size and emits
725 + token_ids=[] instead of misleading data. When #49962 fixes the capture
726 + path, the guard stops triggering and correct token_ids flow through.
727 + """
728 + vbs = BLOCK_SIZE * 2
729 + kv_cfg = _make_mixed_kv_cache_config(
730 + num_blocks=16, with_mamba=True, mamba_block_size=32
731 + )
732 + fix = make_events_scheduler(
733 + num_cpu_blocks=8,
734 + num_gpu_blocks=16,
735 + kv_cache_config=kv_cfg,
736 + dcp_world_size=2,
737 + num_groups=2,
738 + )
739 + sched = fix.scheduler
740 + gpu_pool = fix.gpu_block_pool
741 +
742 + # Mamba g_block_size = 32 * 2 = 64, so 2 Mamba blocks need 128 computed
743 + # tokens to be "ready"; FA g_block_size = 16 * 2 = 32 needs only 64.
744 + req = _make_cp_request(num_blocks=4, virtual_block_size=vbs)
745 + fa_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=0)
746 + mamba_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=1)
747 + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, mamba_blocks))
748 + req.num_computed_tokens = 4 * vbs
749 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
750 +
751 + block_ids = kv_blocks.get_block_ids()
752 + sched_out = make_scheduler_output(
753 + {req.request_id: 4 * vbs},
754 + new_reqs={req.request_id: block_ids},
755 + )
756 + meta = sched.build_connector_meta(sched_out)
757 + assert meta.store_event >= 0
758 + simulate_store_completion(sched, meta.store_event)
759 +
760 + events = list(sched.take_events())
761 + stored = [e for e in events if isinstance(e, BlockStored)]
762 + assert len(stored) == 4, (
763 + f"expected 4 BlockStored (2 FA + 2 Mamba), got {len(stored)}"
764 + )
765 +
766 + by_group: dict[int, list[BlockStored]] = {}
767 + for ev in stored:
768 + by_group.setdefault(ev.group_idx, []).append(ev)
769 +
770 + # FA group: capture size (64) == emission size (64) → token_ids present.
771 + fa_events = by_group[0]
772 + assert len(fa_events) == 2
773 + for ev in fa_events:
774 + assert ev.block_size == vbs
775 + assert len(ev.token_ids) == vbs, (
776 + f"FA token_ids should be {vbs}, got {len(ev.token_ids)}"
777 + )
778 +
779 + # Mamba group: capture size (64) != emission size (32) → guard emits [].
780 + mamba_events = by_group[1]
781 + assert len(mamba_events) == 2
782 + for ev in mamba_events:
783 + assert ev.block_size == 32
784 + assert ev.token_ids == [], (
785 + f"Mamba token_ids should be [] (guard), got {ev.token_ids}"
786 + )
787 + assert ev.parent_block_hash is None, (
788 + "Mamba parent_block_hash should be None (guard)"
789 + )
790 + assert ev.extra_keys is None, "Mamba extra_keys should be None (guard)"
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+4
@@ -111,6 +111,10 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
111 raise ValueError(
112 'kv_offload_backend="disk" requires disk_path to be set.'
113 )
114 + if disk_capacity_bytes <= 0:
115 + raise ValueError(
116 + 'kv_offload_backend="disk" requires disk_capacity_bytes > 0.'
117 + )
118 else:
119 ignored = [k for k in _DISK_ONLY_KEYS if k in extra_config]
120 if ignored:
vllm/v1/simple_kv_offload/manager.py
+247 -29
@@ -8,7 +8,13 @@ from dataclasses import dataclass, field, replace
8 from typing import TYPE_CHECKING, Any
9
10 from vllm.config import VllmConfig
11 -from vllm.distributed.kv_events import KVCacheEvent
11 +from vllm.distributed.kv_events import (
12 + MEDIUM_CPU,
13 + MEDIUM_STORAGE,
14 + BlockRemoved,
15 + BlockStored,
16 + KVCacheEvent,
17 +)
18 from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
19 from vllm.logger import init_logger
20 from vllm.utils.math_utils import cdiv
@@ -17,11 +23,23 @@ from vllm.v1.core.kv_cache_coordinator import (
23 KVCacheCoordinator,
24 get_kv_cache_coordinator,
25 )
26 +from vllm.v1.core.kv_cache_utils import (
27 + BlockHashWithGroupId,
28 + ExternalBlockHash,
29 + generate_block_hash_extra_keys,
30 + get_block_hash,
31 + get_group_id,
32 + make_block_hash_with_group_id,
33 + maybe_convert_block_hash,
34 + resolve_block_hashes,
35 +)
36 from vllm.v1.core.sched.output import SchedulerOutput
37 from vllm.v1.kv_cache_interface import (
38 FullAttentionSpec,
39 MambaSpec,
40 SlidingWindowSpec,
41 + get_kv_cache_spec_kind,
42 + get_kv_cache_spec_sliding_window,
43 )
44 from vllm.v1.outputs import KVConnectorOutput
45 from vllm.v1.simple_kv_offload.metadata import (
@@ -38,10 +56,27 @@ if TYPE_CHECKING:
56 logger = init_logger(__name__)
57
58
59 +@dataclass
60 +class BlockStoreMeta:
61 + """Per-block metadata snapshot for BlockStored event emission.
62 +
63 + Captured at store-prep time (when the Request is available) and carried
64 + through TransferMeta to async completion, where the Request is gone.
65 + """
66 +
67 + token_ids: list[int]
68 + parent_block_hash: ExternalBlockHash | None
69 + lora_id: int | None
70 + lora_name: str | None
71 + extra_keys: tuple[Any, ...] | None
72 + block_size: int | None = None
73 +
74 +
75 @dataclass
76 class TransferMeta:
77 gpu_block_ids: list[int]
78 cpu_block_ids: list[int]
79 + block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = None
80
81
82 @dataclass
@@ -98,6 +133,7 @@ class SimpleCPUOffloadScheduler:
133 kv_cache_config, offload_capacity
134 )
135 self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
136 + self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU
137 # Find the full attention kv group for prefix cache matching.
138 self.fa_gidx = -1
139 for g_idx, g in enumerate(self.cpu_kv_cache_config.kv_cache_groups):
@@ -414,17 +450,22 @@ class SimpleCPUOffloadScheduler:
450 ) -> SimpleCPUOffloadMetadata:
451 # --- Stores ---
452 store_event = -1
417 - store_gpu, store_cpu, store_req_ids = self.prepare_store_specs(scheduler_output)
453 + store_gpu, store_cpu, store_req_ids, store_meta = self.prepare_store_specs(
454 + scheduler_output
455 + )
456 for transfer in self._pending_finished_stores:
457 store_gpu.extend(transfer.gpu_block_ids)
458 store_cpu.extend(transfer.cpu_block_ids)
459 + if store_meta is not None:
460 + assert transfer.block_meta is not None
461 + store_meta.extend(transfer.block_meta)
462 self._pending_finished_stores.clear()
463
464 if store_gpu:
465 store_event = self._store_event_counter
466 self._store_event_counter += 1
467 self._store_event_to_blocks[store_event] = TransferMeta(
427 - store_gpu, store_cpu
468 + store_gpu, store_cpu, store_meta
469 )
470 if store_req_ids: # For eager mode only, track req->blocks mapping
471 self._store_event_to_reqs[store_event] = store_req_ids
@@ -469,7 +510,12 @@ class SimpleCPUOffloadScheduler:
510
511 def prepare_store_specs(
512 self, scheduler_output: SchedulerOutput
472 - ) -> tuple[list[int], list[int], list[str]]:
513 + ) -> tuple[
514 + list[int],
515 + list[int],
516 + list[str],
517 + list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
518 + ]:
519 """Prepare store specs for the store event."""
520 if self._lazy_mode:
521 return self._prepare_lazy_store_specs()
@@ -478,7 +524,12 @@ class SimpleCPUOffloadScheduler:
524
525 def _prepare_lazy_store_specs(
526 self,
481 - ) -> tuple[list[int], list[int], list[str]]:
527 + ) -> tuple[
528 + list[int],
529 + list[int],
530 + list[str],
531 + list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
532 + ]:
533 """Single-pass cursor walk: offload cached GPU blocks near eviction.
534
535 Walks the GPU free queue from the cursor, counting blocks that are
@@ -487,7 +538,7 @@ class SimpleCPUOffloadScheduler:
538 """
539 gpu_pool = self._gpu_block_pool
540 if gpu_pool is None or self._target_free <= 0:
490 - return [], [], []
541 + return [], [], [], None
542
543 free_queue = gpu_pool.free_block_queue
544 cpu_pool = self.cpu_block_pool
@@ -525,11 +576,16 @@ class SimpleCPUOffloadScheduler:
576 else:
577 cpu_ids = []
578
528 - return gpu_ids, cpu_ids, []
579 + return gpu_ids, cpu_ids, [], None
580
581 def _prepare_eager_store_specs(
582 self, scheduler_output: SchedulerOutput
532 - ) -> tuple[list[int], list[int], list[str]]:
583 + ) -> tuple[
584 + list[int],
585 + list[int],
586 + list[str],
587 + list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
588 + ]:
589 """Identify newly computed blocks to offload from scheduler requests.
590
591 Only considers blocks whose KV data has been **confirmed computed** by
@@ -537,19 +593,22 @@ class SimpleCPUOffloadScheduler:
593 the finish-time flush if the request completes first.
594
595 Returns:
540 - (gpu_block_ids, cpu_block_ids, req_ids) for the store event.
596 + (gpu_block_ids, cpu_block_ids, req_ids, block_meta) for the store
597 + event. ``block_meta`` is None when kv cache events are disabled.
598 """
599
600 merged_gpu_block_ids: list[int] = []
601 merged_cpu_block_ids: list[int] = []
602 req_ids: list[str] = []
603 + merged_block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = (
604 + [] if self.enable_kv_cache_events else None
605 + )
606
607 gpu_block_pool = self._gpu_block_pool
608 if gpu_block_pool is None:
549 - return [], [], []
609 + return [], [], [], merged_block_meta
610 cpu_block_pool = self.cpu_block_pool
551 - kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups
552 - num_groups = len(kv_cache_groups)
611 + num_groups = len(self.cpu_kv_cache_config.kv_cache_groups)
612 # Dedup against blocks already scheduled.
613 in_flight = self._in_flight_store_gpu_blocks
614
@@ -558,13 +617,13 @@ class SimpleCPUOffloadScheduler:
617 if state is None or state.finished:
618 continue
619
561 - # Accumulate new block IDs.
620 if preempted:
621 state.block_ids = tuple([] for _ in range(num_groups))
622 state.num_stored_blocks = [0] * num_groups
623 if new_block_id_groups:
624 for g in range(min(num_groups, len(new_block_id_groups))):
625 if new_block_id_groups[g] is not None:
626 + # Accumulate new block IDs.
627 state.block_ids[g].extend(new_block_id_groups[g])
628
629 num_new_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0)
@@ -575,8 +634,8 @@ class SimpleCPUOffloadScheduler:
634 if not block_ids_by_group:
635 continue
636
578 - gpu_block_ids, advanced_per_group = self._select_eager_blocks_to_store(
579 - state, block_ids_by_group
637 + gpu_block_ids, advanced_per_group, block_meta = (
638 + self._select_eager_blocks_to_store(state, block_ids_by_group)
639 )
640
641 # --- Phase 2: Batch allocate CPU blocks ---
@@ -592,6 +651,9 @@ class SimpleCPUOffloadScheduler:
651 merged_gpu_block_ids.extend(gpu_block_ids)
652 merged_cpu_block_ids.extend(cpu_block_ids)
653 in_flight.update(gpu_block_ids)
654 + if merged_block_meta is not None:
655 + assert block_meta is not None
656 + merged_block_meta.extend(block_meta)
657
658 # Touch GPU blocks to prevent freeing during async copy
659 gpu_block_pool.touch(
@@ -609,20 +671,26 @@ class SimpleCPUOffloadScheduler:
671 for g in range(num_groups):
672 state.num_stored_blocks[g] += advanced_per_group[g]
673
612 - return merged_gpu_block_ids, merged_cpu_block_ids, req_ids
674 + return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta
675
676 def _select_eager_blocks_to_store(
677 self,
678 state: StoreRequestState,
679 block_ids_by_group: tuple[list[int], ...],
618 - ) -> tuple[list[int], list[int]]:
619 - """Return confirmed, uncached eager blocks and per-group cursor advances."""
680 + ) -> tuple[
681 + list[int],
682 + list[int],
683 + list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None,
684 + ]:
685 + """Return confirmed eager blocks, cursor advances, and event metadata."""
686 assert self._gpu_block_pool is not None
687 gpu_block_ids: list[int] = []
688 advanced_per_group = [0] * len(self.cpu_kv_cache_config.kv_cache_groups)
623 - confirmed_tokens = (
624 - state.request.num_computed_tokens - state.request.num_output_placeholders
689 + block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = (
690 + [] if self.enable_kv_cache_events else None
691 )
692 + request = state.request
693 + confirmed_tokens = request.num_computed_tokens - request.num_output_placeholders
694 aligned_tokens = confirmed_tokens // self.block_size * self.block_size
695 num_free = self.cpu_block_pool.get_num_free_blocks()
696
@@ -637,7 +705,17 @@ class SimpleCPUOffloadScheduler:
705 * self.cp_world_size
706 )
707 ready = min(len(group_gpu_ids), aligned_tokens // group_size)
640 - for gpu_block_id in group_gpu_ids[state.num_stored_blocks[g] : ready]:
708 + resolved_hashes = (
709 + resolve_block_hashes(
710 + request.block_hashes, self.hash_block_size, group_size
711 + )
712 + if block_meta is not None
713 + else None
714 + )
715 + curr_mm_idx = 0
716 + secondary_mm_idx = 0
717 + start = state.num_stored_blocks[g]
718 + for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start):
719 gpu_block = self._gpu_block_pool.blocks[gpu_block_id]
720 if gpu_block.is_null or gpu_block.block_hash is None:
721 advanced_per_group[g] += 1
@@ -655,8 +733,65 @@ class SimpleCPUOffloadScheduler:
733 break
734 gpu_block_ids.append(gpu_block_id)
735 advanced_per_group[g] += 1
658 -
659 - return gpu_block_ids, advanced_per_group
736 + if block_meta is not None:
737 + token_start = i * group_size
738 + token_end = token_start + group_size
739 + assert resolved_hashes is not None
740 + parent_hash = (
741 + None
742 + if i == 0
743 + else maybe_convert_block_hash(resolved_hashes[i - 1])
744 + )
745 + lora_req = request.lora_request
746 + extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
747 + request, token_start, token_end, curr_mm_idx
748 + )
749 + meta_by_hash = {
750 + gpu_block.block_hash: BlockStoreMeta(
751 + token_ids=list(
752 + request.all_token_ids[token_start:token_end]
753 + ),
754 + parent_block_hash=parent_hash,
755 + lora_id=lora_req.adapter_id if lora_req else None,
756 + lora_name=lora_req.name if lora_req else None,
757 + extra_keys=extra_keys,
758 + )
759 + }
760 + first_hash_idx = token_start // self.hash_block_size
761 + last_hash_idx = token_end // self.hash_block_size
762 + for hash_idx in range(first_hash_idx, last_hash_idx):
763 + block_hash = make_block_hash_with_group_id(
764 + request.block_hashes[hash_idx], g
765 + )
766 + if block_hash == gpu_block.block_hash:
767 + continue
768 + hash_start = hash_idx * self.hash_block_size
769 + hash_end = hash_start + self.hash_block_size
770 + secondary_extra_keys, secondary_mm_idx = (
771 + generate_block_hash_extra_keys(
772 + request,
773 + hash_start,
774 + hash_end,
775 + secondary_mm_idx,
776 + )
777 + )
778 + meta_by_hash[block_hash] = BlockStoreMeta(
779 + token_ids=list(request.all_token_ids[hash_start:hash_end]),
780 + parent_block_hash=(
781 + None
782 + if hash_idx == 0
783 + else maybe_convert_block_hash(
784 + request.block_hashes[hash_idx - 1]
785 + )
786 + ),
787 + lora_id=lora_req.adapter_id if lora_req else None,
788 + lora_name=lora_req.name if lora_req else None,
789 + extra_keys=secondary_extra_keys,
790 + block_size=self.hash_block_size,
791 + )
792 + block_meta.append(meta_by_hash)
793 +
794 + return gpu_block_ids, advanced_per_group, block_meta
795
796 def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
797 """Handle async transfer completions from worker.
@@ -695,7 +830,11 @@ class SimpleCPUOffloadScheduler:
830 if not self._lazy_mode:
831 self._in_flight_store_gpu_blocks.difference_update(transfer.gpu_block_ids)
832
698 - self._process_store_completion(transfer.gpu_block_ids, transfer.cpu_block_ids)
833 + self._process_store_completion(
834 + transfer.gpu_block_ids,
835 + transfer.cpu_block_ids,
836 + transfer.block_meta,
837 + )
838 logger.debug(
839 "Store event %d completed: cached %d blocks to CPU",
840 event_idx,
@@ -713,15 +852,19 @@ class SimpleCPUOffloadScheduler:
852 self._cleanup_store_request(req_id)
853
854 def _process_store_completion(
716 - self, gpu_block_ids: list[int], cpu_block_ids: list[int]
855 + self,
856 + gpu_block_ids: list[int],
857 + cpu_block_ids: list[int],
858 + block_meta: list[dict[BlockHashWithGroupId, BlockStoreMeta]] | None = None,
859 ) -> None:
860 """Register copied blocks in the CPU prefix cache and release refs."""
861 assert len(cpu_block_ids) == len(gpu_block_ids)
862 + assert block_meta is None or len(block_meta) == len(gpu_block_ids)
863
864 cpu_blocks = [self.cpu_block_pool.blocks[bid] for bid in cpu_block_ids]
865
866 assert self._gpu_block_pool is not None
724 - for gpu_block_id, cpu_block in zip(gpu_block_ids, cpu_blocks):
867 + for i, (gpu_block_id, cpu_block) in enumerate(zip(gpu_block_ids, cpu_blocks)):
868 gpu_block = self._gpu_block_pool.blocks[gpu_block_id]
869 primary_hash = gpu_block.block_hash
870 assert primary_hash is not None
@@ -737,6 +880,76 @@ class SimpleCPUOffloadScheduler:
880 self.cpu_block_pool._insert_block_hash(
881 block_hash, cpu_block, num_tokens=None
882 )
883 + if self.enable_kv_cache_events:
884 + meta_by_hash = block_meta[i] if block_meta is not None else {}
885 + primary_group_idx = get_group_id(primary_hash)
886 + primary_spec = self.cpu_kv_cache_config.kv_cache_groups[
887 + primary_group_idx
888 + ].kv_cache_spec
889 + primary_block_size = primary_spec.block_size * (
890 + self.cp_world_size if not isinstance(primary_spec, MambaSpec) else 1
891 + )
892 + events_to_emit: list[
893 + tuple[BlockHashWithGroupId, BlockStoreMeta | None, int]
894 + ] = [
895 + (
896 + primary_hash,
897 + meta_by_hash.get(primary_hash),
898 + primary_block_size,
899 + )
900 + ]
901 + for secondary_hash in secondary_hashes:
902 + secondary_meta = meta_by_hash.get(secondary_hash)
903 + events_to_emit.append(
904 + (
905 + secondary_hash,
906 + secondary_meta,
907 + secondary_meta.block_size
908 + if secondary_meta is not None
909 + and secondary_meta.block_size is not None
910 + else 0,
911 + )
912 + )
913 +
914 + for block_hash, meta, event_block_size in events_to_emit:
915 + group_idx = get_group_id(block_hash)
916 + spec = self.cpu_kv_cache_config.kv_cache_groups[
917 + group_idx
918 + ].kv_cache_spec
919 + if meta is not None and len(meta.token_ids) != event_block_size:
920 + # token_ids were sliced with a different g_block_size
921 + # (e.g. Mamba+DCP where capture uses spec.block_size *
922 + # cp_world_size but the event emits spec.block_size * 1).
923 + # Emit empty metadata until #49962 fixes the sizing.
924 + token_ids = []
925 + parent_block_hash = None
926 + extra_keys = None
927 + else:
928 + token_ids = meta.token_ids if meta else []
929 + parent_block_hash = meta.parent_block_hash if meta else None
930 + extra_keys = meta.extra_keys if meta else None
931 + self.cpu_block_pool.kv_event_queue.append(
932 + BlockStored(
933 + block_hashes=[
934 + maybe_convert_block_hash(get_block_hash(block_hash))
935 + ],
936 + parent_block_hash=parent_block_hash,
937 + token_ids=token_ids,
938 + block_size=event_block_size,
939 + lora_id=meta.lora_id if meta else None,
940 + medium=self.kv_event_medium,
941 + lora_name=meta.lora_name if meta else None,
942 + extra_keys=(
943 + [extra_keys] if extra_keys is not None else None
944 + ),
945 + group_idx=group_idx,
946 + kv_cache_spec_kind=get_kv_cache_spec_kind(spec).value,
947 + kv_cache_spec_sliding_window=(
948 + get_kv_cache_spec_sliding_window(spec)
949 + ),
950 + locality="LOCAL",
951 + )
952 + )
953
954 # Free CPU and GPU blocks' ref counts to turn them into prefix cache
955 self.cpu_block_pool.free_blocks(cpu_blocks)
@@ -811,12 +1024,12 @@ class SimpleCPUOffloadScheduler:
1024 gpu_pool = self._gpu_block_pool
1025 if state is None or gpu_pool is None:
1026 return
814 - gpu_ids, _ = self._select_eager_blocks_to_store(state, block_ids)
1027 + gpu_ids, _, block_meta = self._select_eager_blocks_to_store(state, block_ids)
1028 if not gpu_ids:
1029 return
1030 cpu_blocks = self.cpu_block_pool.get_new_blocks(len(gpu_ids))
1031 self._pending_finished_stores.append(
819 - TransferMeta(gpu_ids, [b.block_id for b in cpu_blocks])
1032 + TransferMeta(gpu_ids, [b.block_id for b in cpu_blocks], block_meta)
1033 )
1034 self._in_flight_store_gpu_blocks.update(gpu_ids)
1035 gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])
@@ -882,7 +1095,12 @@ class SimpleCPUOffloadScheduler:
1095 state.store_events.clear()
1096
1097 def take_events(self) -> Iterable[KVCacheEvent]:
885 - return self.cpu_block_pool.take_events()
1098 + events = self.cpu_block_pool.take_events()
1099 + for event in events:
1100 + if isinstance(event, BlockRemoved):
1101 + event.medium = self.kv_event_medium
1102 + event.locality = "LOCAL"
1103 + return events
1104
1105 def reset(self) -> bool:
1106 """Abandon pending transfers and reset the CPU cache when safe.