[Bugfix][KV Offload] Do not let a recurrent group's unhashed block truncate the load boundary (#52807)
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
Yifan Jiang committed
Sep 3, 2026 at 09:08 UTC
da8ec2826898d47bd2fb424bc607052a6dc515c1
2 files changed
+81
-15
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+64
-1
@@ -36,7 +36,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
36
is_store_reachable_swa_chunk,
37
)
38
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
39
-from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock, KVCacheBlockCopy
39
+from vllm.v1.core.kv_cache_utils import (
40
+ BlockHash,
41
+ KVCacheBlock,
42
+ KVCacheBlockCopy,
43
+ make_block_hash_with_group_id,
44
+)
45
from vllm.v1.core.sched.output import (
46
KVConnectorBlockState,
47
SchedulerOutput,
@@ -272,6 +277,64 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys():
277
assert req_status.partial_tail_boundary is None
278
279
280
+def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary():
281
+ """A recurrent group may hold unhashed blocks inside the computed prefix.
282
+
283
+ A recurrent (Mamba / GDN) group has no per-token KV: it carries a fixed-size
284
+ state, so most positions point at the null sentinel and the one real block
285
+ holding the state is legitimately not full-and-cached. Sliding-window groups
286
+ are similar -- neither is required to retain the whole prefix, so a non-null,
287
+ unhashed block can sit *below* the locally-computed mark. Scanning the
288
+ block list from index 0 treated such a block as the start of the freshly
289
+ allocated region, dragging the load boundary below that mark -- which loads the
290
+ wrong keys into the wrong destination blocks, and asserts on the way.
291
+
292
+ Geometry: tokens_per_block = tokens_per_chunk = 16, blocks_per_chunk = 1.
293
+ 16 tokens are already computed (block 0) and 16 are loaded (block 1), so the
294
+ fresh region must start at index 1. The Mamba group's block 0 is non-null and
295
+ unhashed and must be ignored rather than taken as the boundary.
296
+ """
297
+ scheduler = _make_partial_tail_scheduler()
298
+
299
+ request = MagicMock()
300
+ request.request_id = "req"
301
+ request.kv_transfer_params = None
302
+ request.num_prompt_tokens = 64
303
+ request.num_tokens = 64
304
+ request.block_hashes = [BlockHash(f"b{i}".encode()) for i in range(16)]
305
+ request.all_token_ids = list(range(64))
306
+ request.lora_request = None
307
+ request.is_finished.return_value = False
308
+ scheduler.on_new_request(request)
309
+
310
+ req_status = scheduler._req_status["req"]
311
+ req_status.num_locally_computed_tokens = 16
312
+ req_status.update_offload_keys()
313
+
314
+ full_blocks = [KVCacheBlock(31), KVCacheBlock(32)]
315
+ mamba_blocks = [KVCacheBlock(41), KVCacheBlock(42)]
316
+ full_blocks[0].set_block_hash(
317
+ make_block_hash_with_group_id(BlockHash(b"h0"), 0), 16
318
+ )
319
+ # Both Mamba blocks are non-null and unhashed -- the shape a recurrent group
320
+ # produces for a prefix it does not retain.
321
+ for b in mamba_blocks:
322
+ assert b.block_hash is None and not b.is_null
323
+
324
+ scheduler.update_state_after_alloc(
325
+ request,
326
+ KVCacheBlocks((full_blocks, mamba_blocks)),
327
+ num_external_tokens=16,
328
+ )
329
+
330
+ [load_job] = scheduler._current_batch_load_jobs.values()
331
+ loaded = list(load_job.dst_spec.block_ids)
332
+ # The fresh region starts at index 1, so the sparse group's block 0 (id 41) is
333
+ # not a load destination. Before the fix the boundary collapsed to 0.
334
+ assert 41 not in loaded
335
+ assert 42 in loaded
336
+
337
+
338
def test_partial_lookup_requires_every_cache_group():
339
scheduler = _make_partial_tail_scheduler()
340
_make_partial_tail_request(scheduler)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+17
-14
@@ -1025,18 +1025,23 @@ class OffloadingConnectorScheduler:
1025
num_gpu_blocks = cdiv(num_cached_tokens, tokens_per_block)
1026
1027
assert len(group_blocks) >= num_gpu_blocks
1028
- num_locally_computed_gpu_blocks = num_gpu_blocks
1029
- # Skip null placeholder blocks (used for sliding window or mamba padding).
1030
- for i, block in enumerate(group_blocks[:num_gpu_blocks]):
1028
+ # ``load_start_gpu_block_idx``: the index in ``group_blocks`` where the
1029
+ # load region begins -- a slice bound, not a count of computed blocks.
1030
+ # Scan from the computed boundary, not 0: sparse groups (Mamba,
1031
+ # SWA) legitimately hold non-null unhashed blocks below it.
1032
+ # Skip nulls (sentinel / out-of-retention padding).
1033
+ first_fresh_gpu_block_idx = cdiv(
1034
+ num_locally_computed_tokens, tokens_per_block
1035
+ )
1036
+ load_start_gpu_block_idx = num_gpu_blocks
1037
+ for i in range(first_fresh_gpu_block_idx, num_gpu_blocks):
1038
+ block = group_blocks[i]
1039
if not block.is_null and block.block_hash is None:
1032
- num_locally_computed_gpu_blocks = i
1040
+ load_start_gpu_block_idx = i
1041
break
1042
1035
- assert (
1036
- num_locally_computed_tokens
1037
- <= num_locally_computed_gpu_blocks * tokens_per_block
1038
- )
1039
- num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks
1043
+ assert num_locally_computed_tokens % tokens_per_block == 0
1044
+ num_pending_gpu_blocks = num_gpu_blocks - load_start_gpu_block_idx
1045
1046
if group_config.sliding_window_size_in_chunks is not None:
1047
assert (
@@ -1049,7 +1054,7 @@ class OffloadingConnectorScheduler:
1054
num_chunks = cdiv(num_cached_tokens, tokens_per_chunk)
1055
if num_pending_gpu_blocks:
1056
start_chunk_idx = (
1052
- num_locally_computed_gpu_blocks // self.config.blocks_per_chunk
1057
+ load_start_gpu_block_idx // self.config.blocks_per_chunk
1058
)
1059
end_chunk_idx = num_chunks - (partial_tail_boundary is not None)
1060
assert len(offload_keys) >= end_chunk_idx
@@ -1063,12 +1068,10 @@ class OffloadingConnectorScheduler:
1068
1069
dst_block_ids.extend(
1070
block.block_id
1066
- for block in group_blocks[
1067
- num_locally_computed_gpu_blocks:num_gpu_blocks
1068
- ]
1071
+ for block in group_blocks[load_start_gpu_block_idx:num_gpu_blocks]
1072
)
1073
group_sizes.append(num_pending_gpu_blocks)
1071
- block_indices.append(num_locally_computed_gpu_blocks)
1074
+ block_indices.append(load_start_gpu_block_idx)
1075
1076
# Skip prefix-hit chunks for block-level policy; for
1077
# request-level, next_stored_chunk_idx stays at 0 so all