[Bugfix][KV Offload] Emit self-describing CPU events at KV-group block granularity (#51614)
Signed-off-by: Ziqi Fan <ziqif@nvidia.com>
Ziqi Fan committed
Aug 12, 2026 at 20:28 UTC
15227b934ad807180cc4d3c5dbe316ee424b26fb
2 files changed
+59
-16
tests/v1/kv_connector/unit/offloading_connector/test_events.py
+37
-2
@@ -81,13 +81,18 @@ def _group_config(
81
group_idx: int = 0,
82
block_size: int = 4,
83
blocks_per_chunk: int = 1,
84
+ tokens_per_hash: int | None = None,
85
sliding_window_size_in_chunks: int | None = None,
86
) -> GroupOffloadConfig:
87
+ if tokens_per_hash is None:
88
+ tokens_per_hash = block_size
89
+ tokens_per_chunk = block_size * blocks_per_chunk
90
+ assert tokens_per_chunk % tokens_per_hash == 0
91
return GroupOffloadConfig(
92
group_idx=group_idx,
93
tokens_per_block=block_size,
89
- tokens_per_chunk=block_size * blocks_per_chunk,
90
- hashes_per_chunk=blocks_per_chunk,
94
+ tokens_per_chunk=tokens_per_chunk,
95
+ hashes_per_chunk=tokens_per_chunk // tokens_per_hash,
96
sliding_window_size_in_chunks=sliding_window_size_in_chunks,
97
kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,
98
)
@@ -334,6 +339,36 @@ def test_promotion_emits_full_cpu_stored_event():
339
assert event.kv_cache_spec_sliding_window is None
340
341
342
+@pytest.mark.parametrize(
343
+ ("blocks_per_chunk", "expected_hash_indices"),
344
+ [(1, [63]), (2, [63, 127])],
345
+)
346
+def test_event_hashes_use_group_block_size(
347
+ blocks_per_chunk: int, expected_hash_indices: list[int]
348
+):
349
+ tokens_per_hash = 4
350
+ block_size = 256
351
+ hashes_per_block = block_size // tokens_per_hash
352
+ tracker = _tracker()
353
+ group_config = _group_config(
354
+ block_size=block_size,
355
+ blocks_per_chunk=blocks_per_chunk,
356
+ tokens_per_hash=tokens_per_hash,
357
+ )
358
+ req = _request(
359
+ block_hashes=[_hash(i) for i in range(hashes_per_block * blocks_per_chunk)],
360
+ token_count=block_size * blocks_per_chunk,
361
+ )
362
+ [key] = _record_chunks(tracker, req, group_config, num_chunks=1)
363
+
364
+ [event] = tracker.take_events([_stored_event([key])])
365
+
366
+ assert isinstance(event, BlockStored)
367
+ assert event.block_hashes == [_wire_hash(_hash(i)) for i in expected_hash_indices]
368
+ assert event.block_size == block_size
369
+ assert len(event.token_ids) == block_size * blocks_per_chunk
370
+
371
+
372
def test_lookup_promotion_factor_gt_1_store_and_remove():
373
block_size = 4
374
blocks_per_chunk = 2
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py
+22
-14
@@ -27,7 +27,11 @@ from vllm.distributed.kv_events import (
27
KVCacheEvent,
28
)
29
from vllm.logger import init_logger
30
-from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
30
+from vllm.v1.core.kv_cache_utils import (
31
+ BlockHash,
32
+ maybe_convert_block_hash,
33
+ resolve_block_hashes,
34
+)
35
from vllm.v1.kv_cache_interface import (
36
KVCacheGroupSpec,
37
get_kv_cache_spec_kind,
@@ -255,22 +259,26 @@ class OffloadingEventsTracker:
259
"""Build the payload snapshot for one offloaded chunk: its
260
constituent per-block hashes, the whole chunk's tokens, and the
261
per-block ``block_size``."""
258
- hbf = group_config.hashes_per_chunk
259
- assert hbf > 0
262
+ hashes_per_chunk = group_config.hashes_per_chunk
263
+ assert hashes_per_chunk > 0
264
assert chunk_idx >= 0
261
- # per-block token count (= the GPU/hash block size)
262
- tokens_per_hash = group_config.tokens_per_chunk // hbf
263
- # chunk c covers hash-blocks [c*hbf, (c+1)*hbf); its tail block's hash
264
- # is the chunk's OffloadKey.
265
- first_hash_idx = chunk_idx * hbf
266
- last_hash_idx = first_hash_idx + hbf
265
+ tokens_per_hash = group_config.tokens_per_chunk // hashes_per_chunk
266
+ # Each chunk's final raw hash is its OffloadKey.
267
+ first_hash_idx = chunk_idx * hashes_per_chunk
268
+ last_hash_idx = first_hash_idx + hashes_per_chunk
269
assert first_hash_idx >= 0
270
assert last_hash_idx <= len(req.block_hashes)
269
- chunk_hashes: list[BlockHash] = []
270
- for block_hash in req.block_hashes[first_hash_idx:last_hash_idx]:
271
+ raw_chunk_hashes = req.block_hashes[first_hash_idx:last_hash_idx]
272
+ chunk_hashes = resolve_block_hashes(
273
+ raw_chunk_hashes,
274
+ tokens_per_hash,
275
+ group_config.tokens_per_block,
276
+ )
277
+ for block_hash in chunk_hashes:
278
assert block_hash is not None
272
- chunk_hashes.append(block_hash)
273
- assert len(chunk_hashes) == hbf
279
+ assert len(chunk_hashes) == (
280
+ group_config.tokens_per_chunk // group_config.tokens_per_block
281
+ )
282
283
if group_config.sliding_window_size_in_chunks is not None:
284
# The recording methods filter these out before calling this helper.
@@ -298,7 +306,7 @@ class OffloadingEventsTracker:
306
block_hashes=tuple(chunk_hashes),
307
parent_block_hash=parent_block_hash,
308
token_ids=token_ids,
301
- block_size=tokens_per_hash,
309
+ block_size=group_config.tokens_per_block,
310
lora_id=lora_id,
311
lora_name=lora_name,
312
extra_keys=None,