From da8ec2826898d47bd2fb424bc607052a6dc515c1 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:08:46 -0700 Subject: [PATCH] [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> --- .../offloading_connector/test_scheduler.py | 65 ++++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 31 +++++---- 2 files changed, 81 insertions(+), 15 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 040067dc36..8f4b03280d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -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) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 435823eab6..3cdb0731f7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -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