iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix][KV Offload] Do not let a recurrent group's unhashed block truncate the load boundary (#52807)

Browse Source 
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
This commit is contained in:
Yifan Jiang
2026-09-03 19:08:46 +03:00
committed by GitHub
parent 4ae6228284
commit da8ec28268
2 changed files with 81 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+64 -1
View File
@@ -36,7 +36,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
is_store_reachable_swa_chunk,
)
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock, KVCacheBlockCopy
from vllm.v1.core.kv_cache_utils import (
BlockHash,
KVCacheBlock,
KVCacheBlockCopy,
make_block_hash_with_group_id,
)
from vllm.v1.core.sched.output import (
KVConnectorBlockState,
SchedulerOutput,
@@ -272,6 +277,64 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys():
assert req_status.partial_tail_boundary is None
def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary():
"""A recurrent group may hold unhashed blocks inside the computed prefix.
A recurrent (Mamba / GDN) group has no per-token KV: it carries a fixed-size
state, so most positions point at the null sentinel and the one real block
holding the state is legitimately not full-and-cached. Sliding-window groups
are similar -- neither is required to retain the whole prefix, so a non-null,
unhashed block can sit *below* the locally-computed mark. Scanning the
block list from index 0 treated such a block as the start of the freshly
allocated region, dragging the load boundary below that mark -- which loads the
wrong keys into the wrong destination blocks, and asserts on the way.
Geometry: tokens_per_block = tokens_per_chunk = 16, blocks_per_chunk = 1.
16 tokens are already computed (block 0) and 16 are loaded (block 1), so the
fresh region must start at index 1. The Mamba group's block 0 is non-null and
unhashed and must be ignored rather than taken as the boundary.
"""
scheduler = _make_partial_tail_scheduler()
request = MagicMock()
request.request_id = "req"
request.kv_transfer_params = None
request.num_prompt_tokens = 64
request.num_tokens = 64
request.block_hashes = [BlockHash(f"b{i}".encode()) for i in range(16)]
request.all_token_ids = list(range(64))
request.lora_request = None
request.is_finished.return_value = False
scheduler.on_new_request(request)
req_status = scheduler._req_status["req"]
req_status.num_locally_computed_tokens = 16
req_status.update_offload_keys()
full_blocks = [KVCacheBlock(31), KVCacheBlock(32)]
mamba_blocks = [KVCacheBlock(41), KVCacheBlock(42)]
full_blocks[0].set_block_hash(
make_block_hash_with_group_id(BlockHash(b"h0"), 0), 16
)
# Both Mamba blocks are non-null and unhashed -- the shape a recurrent group
# produces for a prefix it does not retain.
for b in mamba_blocks:
assert b.block_hash is None and not b.is_null
scheduler.update_state_after_alloc(
request,
KVCacheBlocks((full_blocks, mamba_blocks)),
num_external_tokens=16,
)
[load_job] = scheduler._current_batch_load_jobs.values()
loaded = list(load_job.dst_spec.block_ids)
# The fresh region starts at index 1, so the sparse group's block 0 (id 41) is
# not a load destination. Before the fix the boundary collapsed to 0.
assert 41 not in loaded
assert 42 in loaded
def test_partial_lookup_requires_every_cache_group():
scheduler = _make_partial_tail_scheduler()
_make_partial_tail_request(scheduler)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+17 -14
View File
@@ -1025,18 +1025,23 @@ class OffloadingConnectorScheduler:
num_gpu_blocks = cdiv(num_cached_tokens, tokens_per_block)
assert len(group_blocks) >= num_gpu_blocks
num_locally_computed_gpu_blocks = num_gpu_blocks
# Skip null placeholder blocks (used for sliding window or mamba padding).
for i, block in enumerate(group_blocks[:num_gpu_blocks]):
# ``load_start_gpu_block_idx``: the index in ``group_blocks`` where the
# load region begins -- a slice bound, not a count of computed blocks.
# Scan from the computed boundary, not 0: sparse groups (Mamba,
# SWA) legitimately hold non-null unhashed blocks below it.
# Skip nulls (sentinel / out-of-retention padding).
first_fresh_gpu_block_idx = cdiv(
num_locally_computed_tokens, tokens_per_block
)
load_start_gpu_block_idx = num_gpu_blocks
for i in range(first_fresh_gpu_block_idx, num_gpu_blocks):
block = group_blocks[i]
if not block.is_null and block.block_hash is None:
num_locally_computed_gpu_blocks = i
load_start_gpu_block_idx = i
break
assert (
num_locally_computed_tokens
<= num_locally_computed_gpu_blocks * tokens_per_block
)
num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks
assert num_locally_computed_tokens % tokens_per_block == 0
num_pending_gpu_blocks = num_gpu_blocks - load_start_gpu_block_idx
if group_config.sliding_window_size_in_chunks is not None:
assert (
@@ -1049,7 +1054,7 @@ class OffloadingConnectorScheduler:
num_chunks = cdiv(num_cached_tokens, tokens_per_chunk)
if num_pending_gpu_blocks:
start_chunk_idx = (
num_locally_computed_gpu_blocks // self.config.blocks_per_chunk
load_start_gpu_block_idx // self.config.blocks_per_chunk
)
end_chunk_idx = num_chunks - (partial_tail_boundary is not None)
assert len(offload_keys) >= end_chunk_idx
@@ -1063,12 +1068,10 @@ class OffloadingConnectorScheduler:
dst_block_ids.extend(
block.block_id
for block in group_blocks[
num_locally_computed_gpu_blocks:num_gpu_blocks
]
for block in group_blocks[load_start_gpu_block_idx:num_gpu_blocks]
)
group_sizes.append(num_pending_gpu_blocks)
block_indices.append(num_locally_computed_gpu_blocks)
block_indices.append(load_start_gpu_block_idx)
# Skip prefix-hit chunks for block-level policy; for
# request-level, next_stored_chunk_idx stays at 0 so all