From bb9351d3c090cd0fb55c220df43de08fecd7e22e Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:09:39 -0700 Subject: [PATCH] [Bugfix][KV Offload] Do not let a recurrent group's unhashed block truncate the load boundary update_state_after_alloc scans the allocated block list for the first block that is not null and has no hash, and treats that index as the start of the freshly allocated region. That index drives keys_to_load (via start_chunk_idx), dst_block_ids, group_sizes and block_indices, so getting it wrong loads the wrong keys into the wrong destination blocks. The scan started at index 0, inside the region the scheduler had already reported as locally computed. Blocks there are computed by definition, so scanning them can only do harm -- and for a recurrent group it does. A Mamba / GDN group has no per-token KV at all: it carries a fixed-size state, which is why get_sliding_window_size_in_chunks() returns 1 for MambaSpec ("Mamba depends on a single state"). Most positions in its block table point at the shared null sentinel, and the one real block holding the state is legitimately not full-and-cached. The scan took that as the start of the fresh region, collapsed the boundary below the computed mark, and the following assert fired. Start the scan at the locally computed boundary instead. For a full-attention group this is a no-op: every token has KV, every block below the boundary is full-and-cached, so the old scan skipped them anyway. The existing assert now holds by construction (nlcgb >= cdiv(nlct, tpb) implies nlct <= nlcgb * tpb) and is kept as a guard for other causes. Reproduced on a 4-node GB200 deployment (Qwen3.5-family hybrid GDN+GQA, 4 KV cache groups, 3 recurrent, MTP): under sustained partial-hit load the leader crashed after 47,142 requests, exit=1 with this assert while the other three ranks exited 137 on gang teardown. Partial hits immediately preceding: 12768, 6384, 8512, 6384, 4256 offloaded tokens. Reachable only when num_external_tokens > 0 AND num_locally_computed_tokens > 0, i.e. a partial hit. Tests that flood to full eviction before re-sending produce num_computed_tokens == 0, where the assert reads 0 <= X and cannot fire; the same deployment ran clean and idle for 3h47m immediately before the load. Independent of the drafter-annotation defect in #52735/#52771: the crash above was observed with #52771 applied, no group marked as a drafter, and the volatile-tail pop therefore never running. 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 53a3524c25b8..25bf5d6d33b4 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 @@ 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 402bd55dc549..ba1c2a7960d3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1024,18 +1024,23 @@ def update_state_after_alloc( 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 ( @@ -1048,7 +1053,7 @@ def update_state_after_alloc( 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 @@ -1062,12 +1067,10 @@ def update_state_after_alloc( 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