From 3354dba381ce32bb4850acbe98e9355d48c8ed70 Mon Sep 17 00:00:00 2001 From: drakosha Date: Tue, 7 Jul 2026 13:16:52 +0300 Subject: [PATCH] [Bugfix][KV offload] Store interior chunk-boundary blocks under MTP/Eagle (#46972) Signed-off-by: Mikhail Kostryukov Co-authored-by: Claude Co-authored-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 113 ++++++++++++++---- .../kv_connector/v1/offloading/scheduler.py | 45 +++++-- 2 files changed, 131 insertions(+), 27 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 836f93eb59..6c55b91d8d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -2216,15 +2216,17 @@ class TestEagle: # ------------------------------------------------------------------- @pytest.mark.parametrize("async_scheduling", [True, False]) - def test_full_attn_store_excludes_trailing_block( + def test_full_attn_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): - """Eagle full-attention group stores all blocks except the trailing - one. + """Eagle full-attention group excludes the trailing block only while + decoding. Setup: 2 groups — group 0 is normal full-attention, group 1 is - eagle full-attention. With a 3-block prompt, group 1 should store - only blocks 0 and 1, skipping block 2 (the volatile tail). + eagle full-attention. With a 3-block prompt, group 1 stores all 3 + prompt blocks at the end of prefill (the trailing prompt block is + stable), but skips block 3 once it fills with decoded tokens (its + draft-layer KV is volatile until the next block starts). """ block_size = 4 block_size_factor = 1 @@ -2270,23 +2272,27 @@ class TestEagle: runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) + # 4 decoded tokens fill block 3 entirely with decode tokens (one + # extra token so the block is stored under async scheduling too). runner.run( - decoded_tokens=[EOS_TOKEN_ID], + decoded_tokens=[1, 1, 1, 1, 1, EOS_TOKEN_ID], expected_stored=( (0, 0), (0, 1), (0, 2), + (0, 3), (1, 0), (1, 1), + (1, 2), ), ) @pytest.mark.parametrize("async_scheduling", [True, False]) - def test_sw_store_excludes_trailing_block( + def test_sw_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): - """Eagle sliding-window group stores all blocks except the trailing - one.""" + """Eagle sliding-window group stores all prompt blocks but excludes + the trailing block while decoding.""" block_size = 4 sliding_window = 8 num_gpu_blocks = 100 @@ -2321,15 +2327,18 @@ class TestEagle: runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) + # 4 decoded tokens fill block 3 entirely with decode tokens. runner.run( - decoded_tokens=[EOS_TOKEN_ID], - expected_stored=((0, 0), (0, 1)), + decoded_tokens=[1, 1, 1, 1, EOS_TOKEN_ID], + expected_stored=((0, 0), (0, 1), (0, 2)), ) @pytest.mark.parametrize("async_scheduling", [True, False]) - def test_single_block_nothing_stored(self, request_runner, async_scheduling: bool): - """An eagle group with only one block stores nothing: that block is - the tail.""" + def test_single_block_stored_at_end_of_prefill( + self, request_runner, async_scheduling: bool + ): + """An eagle group with a single-block prompt stores it at the end of + prefill: prompt blocks are stable, so no tail is held back.""" block_size = 4 block_size_factor = 1 offloaded_block_size = block_size * block_size_factor @@ -2360,17 +2369,80 @@ class TestEagle: runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) - runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=()) - runner.manager.prepare_store.assert_not_called() + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=((0, 0),)) + + @pytest.mark.parametrize("async_scheduling", [True, False]) + def test_multichunk_store_no_interior_holes( + self, request_runner, async_scheduling: bool + ): + """Eagle store must not drop interior blocks across prefill chunks. + + Regression: the trailing-block exclusion (num_blocks - 1) was applied + when collecting keys, but next_stored_block_idx advanced by the + non-decremented count, so the trailing block of every chunked-prefill + chunk was skipped and never re-considered. With the harness chunk budget + (1000 tokens) and block_size 4, a prompt longer than one chunk lost the + block at the chunk boundary, leaving a permanent gap that caps prefix + reuse at the first hole. Only the trailing decode block may be held + back; all other blocks must be stored exactly once (no duplicates from + next_stored_block_idx regressing at the prefill->decode transition). + """ + block_size = 4 + block_size_factor = 1 + offloaded_block_size = block_size * block_size_factor + num_gpu_blocks = 1000 + + kv_cache_groups = [ + KVCacheGroupSpec( + ["layer0"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + is_eagle_group=True, + ), + ] + runner = request_runner( + block_size=block_size, + num_gpu_blocks=num_gpu_blocks, + async_scheduling=async_scheduling, + kv_cache_groups=kv_cache_groups, + block_size_factor=block_size_factor, + ) + assert runner.connector_scheduler.config.kv_group_configs[0].is_eagle_group + + # Prompt spans more than one prefill chunk (chunk budget 1000 tokens). + num_blocks = 256 + runner.new_request(token_ids=[0] * offloaded_block_size * num_blocks) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + # Decode a few non-EOS tokens so prefill completes across both chunks + # before the request finishes. + runner._run([1, 1, 1, 1, EOS_TOKEN_ID], complete_transfers=True) + + offsets = sorted( + b.request_block_offset + for t in runner.completed_stores + for b in t.gpu_blocks + ) + # The stored blocks must be contiguous from 0: no interior block is + # dropped at a chunk boundary. (The bug left a gap at offloaded block + # 249, the tail of the first 1000-token chunk.) + assert offsets == list(range(len(offsets))), ( + f"interior hole in stored blocks: {offsets}" + ) @pytest.mark.parametrize("async_scheduling", [True, False]) def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool): """Eagle group constrains load: convergence tightens both groups. - Store 3 offloaded blocks per group (eagle group skips tail → stores - 2). Then a new request loads from CPU. The eagle group's post-pop hit - (2) does not tighten below group 0's hit (3), so both groups load - normally. + Store 3 offloaded blocks per group (all prompt blocks, so the eagle + group stores all 3 as well). Then a new request loads from CPU. The + eagle group pops its trailing hit block on load, tightening the hit + to 2 blocks for both groups. """ block_size = 4 block_size_factor = 1 @@ -2419,6 +2491,7 @@ class TestEagle: (0, 2), (1, 0), (1, 1), + (1, 2), ), ) 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 cf0a506a02..75a9696b72 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -298,12 +298,39 @@ class RequestOffloadState: for group_state, new_blocks in zip(self.group_states, new_block_id_groups): group_state.block_ids.extend(new_blocks) + def storable_blocks( + self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int + ) -> int: + """Number of leading offloaded blocks eligible for store. + + For eagle/MTP groups the volatile trailing block of the offloadable + range is excluded while decoding: the draft-layer KV of the last + accepted position may be rewritten after spec-token rejection. During + prefill the trailing block is stable (the draft input for a chunk's + last position is the next prompt token), so it is stored immediately. + The exclusion must be applied consistently everywhere + ``next_stored_block_idx`` is derived: otherwise the trailing block of + each step is skipped on collection but jumped over by + ``next_stored_block_idx``, so it is never re-considered and a + permanent hole breaks prefix-reuse lookup. + """ + num_blocks = num_offloadable_tokens // group_config.offloaded_block_size + is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens + if group_config.is_eagle_group and is_decoding: + num_blocks = max(0, num_blocks - 1) + return num_blocks + def advance_stored_idx(self, num_offloadable_tokens: int) -> None: + # max(): at the prefill->decode transition of a block-aligned prompt, + # storable_blocks drops by one (the eagle exclusion kicks in), and the + # index must not move backwards past already-stored blocks. for group_config, group_state in zip( self.config.kv_group_configs, self.group_states ): - num_blocks = num_offloadable_tokens // group_config.offloaded_block_size - group_state.next_stored_block_idx = num_blocks + group_state.next_stored_block_idx = max( + group_state.next_stored_block_idx, + self.storable_blocks(group_config, num_offloadable_tokens), + ) def update_num_hit_blocks(self, num_cached_tokens: int) -> None: for group_config, group_state in zip( @@ -903,9 +930,9 @@ class OffloadingConnectorScheduler: for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - num_blocks = num_offloadable_tokens // group_config.offloaded_block_size - if group_config.is_eagle_group: - num_blocks = max(0, num_blocks - 1) + num_blocks = req_status.storable_blocks( + group_config, num_offloadable_tokens + ) start_block_idx = group_state.next_stored_block_idx if num_blocks <= start_block_idx: @@ -977,7 +1004,9 @@ class OffloadingConnectorScheduler: is_sliding_window = ( group_config.sliding_window_size_in_blocks is not None ) - num_blocks = num_offloadable_tokens // group_config.offloaded_block_size + num_blocks = req_status.storable_blocks( + group_config, num_offloadable_tokens + ) start_block_idx = group_state.next_stored_block_idx block_ids = group_state.block_ids num_group_blocks = 0 @@ -1010,7 +1039,9 @@ class OffloadingConnectorScheduler: group_sizes.append(num_group_blocks) block_indices.append(start_gpu_block_idx or 0) - group_state.next_stored_block_idx = num_blocks + group_state.next_stored_block_idx = max( + group_state.next_stored_block_idx, num_blocks + ) src_spec = GPULoadStoreSpec( src_block_ids, group_sizes=group_sizes, block_indices=block_indices