[Bugfix][KV offload] Store interior chunk-boundary blocks under MTP/Eagle (#46972)
Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
drakosha committed
Jul 7, 2026 at 13:16 UTC
3354dba381ce32bb4850acbe98e9355d48c8ed70
2 files changed
+131
-27
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+93
-20
@@ -2216,15 +2216,17 @@ class TestEagle:
2216
# -------------------------------------------------------------------
2217
2218
@pytest.mark.parametrize("async_scheduling", [True, False])
2219
- def test_full_attn_store_excludes_trailing_block(
2219
+ def test_full_attn_store_excludes_trailing_decode_block(
2220
self, request_runner, async_scheduling: bool
2221
):
2222
- """Eagle full-attention group stores all blocks except the trailing
2223
- one.
2222
+ """Eagle full-attention group excludes the trailing block only while
2223
+ decoding.
2224
2225
Setup: 2 groups — group 0 is normal full-attention, group 1 is
2226
- eagle full-attention. With a 3-block prompt, group 1 should store
2227
- only blocks 0 and 1, skipping block 2 (the volatile tail).
2226
+ eagle full-attention. With a 3-block prompt, group 1 stores all 3
2227
+ prompt blocks at the end of prefill (the trailing prompt block is
2228
+ stable), but skips block 3 once it fills with decoded tokens (its
2229
+ draft-layer KV is volatile until the next block starts).
2230
"""
2231
block_size = 4
2232
block_size_factor = 1
@@ -2270,23 +2272,27 @@ class TestEagle:
2272
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2273
generate_store_output(keys)
2274
)
2275
+ # 4 decoded tokens fill block 3 entirely with decode tokens (one
2276
+ # extra token so the block is stored under async scheduling too).
2277
runner.run(
2274
- decoded_tokens=[EOS_TOKEN_ID],
2278
+ decoded_tokens=[1, 1, 1, 1, 1, EOS_TOKEN_ID],
2279
expected_stored=(
2280
(0, 0),
2281
(0, 1),
2282
(0, 2),
2283
+ (0, 3),
2284
(1, 0),
2285
(1, 1),
2286
+ (1, 2),
2287
),
2288
)
2289
2290
@pytest.mark.parametrize("async_scheduling", [True, False])
2285
- def test_sw_store_excludes_trailing_block(
2291
+ def test_sw_store_excludes_trailing_decode_block(
2292
self, request_runner, async_scheduling: bool
2293
):
2288
- """Eagle sliding-window group stores all blocks except the trailing
2289
- one."""
2294
+ """Eagle sliding-window group stores all prompt blocks but excludes
2295
+ the trailing block while decoding."""
2296
block_size = 4
2297
sliding_window = 8
2298
num_gpu_blocks = 100
@@ -2321,15 +2327,18 @@ class TestEagle:
2327
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2328
generate_store_output(keys)
2329
)
2330
+ # 4 decoded tokens fill block 3 entirely with decode tokens.
2331
runner.run(
2325
- decoded_tokens=[EOS_TOKEN_ID],
2326
- expected_stored=((0, 0), (0, 1)),
2332
+ decoded_tokens=[1, 1, 1, 1, EOS_TOKEN_ID],
2333
+ expected_stored=((0, 0), (0, 1), (0, 2)),
2334
)
2335
2336
@pytest.mark.parametrize("async_scheduling", [True, False])
2330
- def test_single_block_nothing_stored(self, request_runner, async_scheduling: bool):
2331
- """An eagle group with only one block stores nothing: that block is
2332
- the tail."""
2337
+ def test_single_block_stored_at_end_of_prefill(
2338
+ self, request_runner, async_scheduling: bool
2339
+ ):
2340
+ """An eagle group with a single-block prompt stores it at the end of
2341
+ prefill: prompt blocks are stable, so no tail is held back."""
2342
block_size = 4
2343
block_size_factor = 1
2344
offloaded_block_size = block_size * block_size_factor
@@ -2360,17 +2369,80 @@ class TestEagle:
2369
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2370
generate_store_output(keys)
2371
)
2363
- runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=())
2364
- runner.manager.prepare_store.assert_not_called()
2372
+ runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=((0, 0),))
2373
+
2374
+ @pytest.mark.parametrize("async_scheduling", [True, False])
2375
+ def test_multichunk_store_no_interior_holes(
2376
+ self, request_runner, async_scheduling: bool
2377
+ ):
2378
+ """Eagle store must not drop interior blocks across prefill chunks.
2379
+
2380
+ Regression: the trailing-block exclusion (num_blocks - 1) was applied
2381
+ when collecting keys, but next_stored_block_idx advanced by the
2382
+ non-decremented count, so the trailing block of every chunked-prefill
2383
+ chunk was skipped and never re-considered. With the harness chunk budget
2384
+ (1000 tokens) and block_size 4, a prompt longer than one chunk lost the
2385
+ block at the chunk boundary, leaving a permanent gap that caps prefix
2386
+ reuse at the first hole. Only the trailing decode block may be held
2387
+ back; all other blocks must be stored exactly once (no duplicates from
2388
+ next_stored_block_idx regressing at the prefill->decode transition).
2389
+ """
2390
+ block_size = 4
2391
+ block_size_factor = 1
2392
+ offloaded_block_size = block_size * block_size_factor
2393
+ num_gpu_blocks = 1000
2394
+
2395
+ kv_cache_groups = [
2396
+ KVCacheGroupSpec(
2397
+ ["layer0"],
2398
+ FullAttentionSpec(
2399
+ block_size=block_size,
2400
+ num_kv_heads=1,
2401
+ head_size=1,
2402
+ dtype=torch.float32,
2403
+ ),
2404
+ is_eagle_group=True,
2405
+ ),
2406
+ ]
2407
+ runner = request_runner(
2408
+ block_size=block_size,
2409
+ num_gpu_blocks=num_gpu_blocks,
2410
+ async_scheduling=async_scheduling,
2411
+ kv_cache_groups=kv_cache_groups,
2412
+ block_size_factor=block_size_factor,
2413
+ )
2414
+ assert runner.connector_scheduler.config.kv_group_configs[0].is_eagle_group
2415
+
2416
+ # Prompt spans more than one prefill chunk (chunk budget 1000 tokens).
2417
+ num_blocks = 256
2418
+ runner.new_request(token_ids=[0] * offloaded_block_size * num_blocks)
2419
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2420
+ generate_store_output(keys)
2421
+ )
2422
+ # Decode a few non-EOS tokens so prefill completes across both chunks
2423
+ # before the request finishes.
2424
+ runner._run([1, 1, 1, 1, EOS_TOKEN_ID], complete_transfers=True)
2425
+
2426
+ offsets = sorted(
2427
+ b.request_block_offset
2428
+ for t in runner.completed_stores
2429
+ for b in t.gpu_blocks
2430
+ )
2431
+ # The stored blocks must be contiguous from 0: no interior block is
2432
+ # dropped at a chunk boundary. (The bug left a gap at offloaded block
2433
+ # 249, the tail of the first 1000-token chunk.)
2434
+ assert offsets == list(range(len(offsets))), (
2435
+ f"interior hole in stored blocks: {offsets}"
2436
+ )
2437
2438
@pytest.mark.parametrize("async_scheduling", [True, False])
2439
def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool):
2440
"""Eagle group constrains load: convergence tightens both groups.
2441
2370
- Store 3 offloaded blocks per group (eagle group skips tail → stores
2371
- 2). Then a new request loads from CPU. The eagle group's post-pop hit
2372
- (2) does not tighten below group 0's hit (3), so both groups load
2373
- normally.
2442
+ Store 3 offloaded blocks per group (all prompt blocks, so the eagle
2443
+ group stores all 3 as well). Then a new request loads from CPU. The
2444
+ eagle group pops its trailing hit block on load, tightening the hit
2445
+ to 2 blocks for both groups.
2446
"""
2447
block_size = 4
2448
block_size_factor = 1
@@ -2419,6 +2491,7 @@ class TestEagle:
2491
(0, 2),
2492
(1, 0),
2493
(1, 1),
2494
+ (1, 2),
2495
),
2496
)
2497
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+38
-7
@@ -298,12 +298,39 @@ class RequestOffloadState:
298
for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
299
group_state.block_ids.extend(new_blocks)
300
301
+ def storable_blocks(
302
+ self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
303
+ ) -> int:
304
+ """Number of leading offloaded blocks eligible for store.
305
+
306
+ For eagle/MTP groups the volatile trailing block of the offloadable
307
+ range is excluded while decoding: the draft-layer KV of the last
308
+ accepted position may be rewritten after spec-token rejection. During
309
+ prefill the trailing block is stable (the draft input for a chunk's
310
+ last position is the next prompt token), so it is stored immediately.
311
+ The exclusion must be applied consistently everywhere
312
+ ``next_stored_block_idx`` is derived: otherwise the trailing block of
313
+ each step is skipped on collection but jumped over by
314
+ ``next_stored_block_idx``, so it is never re-considered and a
315
+ permanent hole breaks prefix-reuse lookup.
316
+ """
317
+ num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
318
+ is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
319
+ if group_config.is_eagle_group and is_decoding:
320
+ num_blocks = max(0, num_blocks - 1)
321
+ return num_blocks
322
+
323
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
324
+ # max(): at the prefill->decode transition of a block-aligned prompt,
325
+ # storable_blocks drops by one (the eagle exclusion kicks in), and the
326
+ # index must not move backwards past already-stored blocks.
327
for group_config, group_state in zip(
328
self.config.kv_group_configs, self.group_states
329
):
305
- num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
306
- group_state.next_stored_block_idx = num_blocks
330
+ group_state.next_stored_block_idx = max(
331
+ group_state.next_stored_block_idx,
332
+ self.storable_blocks(group_config, num_offloadable_tokens),
333
+ )
334
335
def update_num_hit_blocks(self, num_cached_tokens: int) -> None:
336
for group_config, group_state in zip(
@@ -903,9 +930,9 @@ class OffloadingConnectorScheduler:
930
for group_config, group_state in zip(
931
self.config.kv_group_configs, req_status.group_states
932
):
906
- num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
907
- if group_config.is_eagle_group:
908
- num_blocks = max(0, num_blocks - 1)
933
+ num_blocks = req_status.storable_blocks(
934
+ group_config, num_offloadable_tokens
935
+ )
936
937
start_block_idx = group_state.next_stored_block_idx
938
if num_blocks <= start_block_idx:
@@ -977,7 +1004,9 @@ class OffloadingConnectorScheduler:
1004
is_sliding_window = (
1005
group_config.sliding_window_size_in_blocks is not None
1006
)
980
- num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
1007
+ num_blocks = req_status.storable_blocks(
1008
+ group_config, num_offloadable_tokens
1009
+ )
1010
start_block_idx = group_state.next_stored_block_idx
1011
block_ids = group_state.block_ids
1012
num_group_blocks = 0
@@ -1010,7 +1039,9 @@ class OffloadingConnectorScheduler:
1039
1040
group_sizes.append(num_group_blocks)
1041
block_indices.append(start_gpu_block_idx or 0)
1013
- group_state.next_stored_block_idx = num_blocks
1042
+ group_state.next_stored_block_idx = max(
1043
+ group_state.next_stored_block_idx, num_blocks
1044
+ )
1045
1046
src_spec = GPULoadStoreSpec(
1047
src_block_ids, group_sizes=group_sizes, block_indices=block_indices