[Bugfix][KV Offload] Restore MTP-retained sliding-window history (#56709)
Signed-off-by: jacklin78911-collab <jacklin78911@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Liqian Lin committed
Sep 16, 2026 at 13:14 UTC
961c5b6b7ea98b765a6d436434d18ef945596f32
2 files changed
+70
-52
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+44
-48
@@ -79,20 +79,31 @@ from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
79
from vllm.v1.request import RequestStatus
80
81
82
-@pytest.mark.parametrize("boundary", [3840, 3904])
82
+@pytest.mark.parametrize(
83
+ "window,extra_retained", [(128, 0), (128, 3), (256, 1), (256, 3)]
84
+)
85
+@pytest.mark.parametrize("boundary", [3840, 3904, 3968])
86
@pytest.mark.parametrize("eagle", [False, True])
87
@pytest.mark.parametrize("left_state", ["present", "missing", "pending", "loading"])
85
-def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
88
+def test_swa_offload_window_covers_unaligned_hit(
89
+ boundary, eagle, left_state, window, extra_retained
90
+):
91
"""A smaller SWA group can move the hit inside another group's CPU chunk."""
92
groups = []
88
- for i, (block_size, window) in enumerate([(256, None), (64, 128), (8, 128)]):
93
+ for i, (block_size, group_window) in enumerate(
94
+ [(256, None), (64, window), (8, 128)]
95
+ ):
96
kwargs = dict(
97
block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32
98
)
99
kv_spec = (
100
FullAttentionSpec(**kwargs)
94
- if window is None
95
- else SlidingWindowSpec(**kwargs, sliding_window=window)
101
+ if group_window is None
102
+ else SlidingWindowSpec(
103
+ **kwargs,
104
+ sliding_window=group_window,
105
+ extra_retained_tokens=extra_retained if i == 1 else 0,
106
+ )
107
)
108
groups.append(
109
KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1)
@@ -153,7 +164,9 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
164
sched._chunks_being_loaded.add(left_key)
165
166
external, load_async = sched.get_num_new_matched_tokens(request, 0)
156
- if boundary == 3904 and left_state != "present":
167
+ first_retained_block = (boundary - window + 1 - extra_retained) // 64
168
+ needs_left = first_retained_block // 4 < end - 1
169
+ if needs_left and left_state != "present":
170
assert external == (0 if left_state == "missing" else None)
171
assert not load_async
172
return
@@ -177,8 +190,8 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
190
sched.update_state_after_alloc(request, blocks, external)
191
assert len(sched._jobs) == 1
192
job = next(iter(sched._jobs.values()))
180
- if boundary == 3904:
181
- assert {swa[14], swa[15]} <= job.keys
193
+ if needs_left:
194
+ assert {left_key, swa[end - 1]} <= job.keys
195
manager.complete_load(job.keys, state.req_context)
196
197
@@ -2629,27 +2642,14 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
2642
assert req_status.transfer_jobs == pending_store_jobs
2643
2644
2645
+@pytest.mark.parametrize("extra_retained", [0, 1, 3])
2646
@pytest.mark.parametrize("async_scheduling", [True, False])
2633
-def test_swa_alignment_skip(request_runner, async_scheduling: bool):
2634
- """SWA blocks unreachable by the load path are skipped during store.
2635
-
2636
- Simulates a DeepSeek V4-like hybrid architecture where SWA groups have
2637
- much smaller block sizes than the full-attention (MLA) group, causing
2638
- most SWA blocks to be unreachable by the alignment-based load path.
2639
-
2640
- Setup:
2641
- - Group 0: full attention (MLA-like), block_size=16
2642
- - Group 1: SWA, block_size=4, sliding_window=8
2643
-
2644
- alignment_tokens = 16, so 16 / 4 = 4 SWA blocks per alignment segment.
2645
- sliding_window_size_in_chunks = ceil(8 / 4) = 2.
2646
- Within each segment of 4 SWA blocks, only the trailing 2 are stored.
2647
-
2648
- With 32 tokens (2 full-attn blocks, 8 SWA blocks):
2649
- - Group 0 stores: blocks 0, 1 (all full-attn blocks)
2650
- - Group 1 stores: blocks 2, 3, 6, 7 (skip 0,1,4,5)
2647
+def test_swa_alignment_skip(request_runner, async_scheduling: bool, extra_retained):
2648
+ """Store the window a cold consumer must restore at each FA boundary.
2649
2652
- For real DeepSeek V4 (100K tokens), this reduces SWA stores by ~78%.
2650
+ With FA blocks of 16 tokens and SWA blocks of 4 tokens, a window of 8
2651
+ needs the last two blocks per segment. Retaining three extra tokens for
2652
+ MTP requires one more block; retaining just one token does not.
2653
"""
2654
full_attn_block_size = 16
2655
swa_block_size = 4
@@ -2674,6 +2674,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
2674
head_size=1,
2675
dtype=torch.float32,
2676
sliding_window=sliding_window,
2677
+ extra_retained_tokens=extra_retained,
2678
),
2679
),
2680
]
@@ -2691,16 +2692,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
2692
# Group 0: full attention
2693
assert kv_group_configs[0].sliding_window_size_in_chunks is None
2694
assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size
2694
- # Group 1: SWA with sliding_window_size of 2 chunks
2695
- assert kv_group_configs[1].sliding_window_size_in_chunks == 2
2695
+ # Group 1: SWA, including any additional retained blocks
2696
+ assert kv_group_configs[1].sliding_window_size_in_chunks == (
2697
+ 3 if extra_retained > 1 else 2
2698
+ )
2699
assert kv_group_configs[1].tokens_per_chunk == swa_block_size
2700
# alignment_tokens = full_attn_block_size = 16
2701
assert runner.connector_scheduler.config.alignment_tokens == full_attn_block_size
2702
2700
- # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks
2701
- # (block_size=4). Decode 1 token to kick off processing (stores are
2702
- # deferred to next step).
2703
- num_tokens = 32
2703
+ # Keep a partial last block so decoding does not reuse the producer's
2704
+ # freed blocks before the harness records its deferred stores.
2705
+ num_tokens = 33
2706
runner.new_request(token_ids=[0] * num_tokens)
2707
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2708
generate_store_output(keys)
@@ -2713,12 +2715,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
2715
)
2716
runner.run(
2717
decoded_tokens=[EOS_TOKEN_ID],
2716
- # Group 0 (full attn, block_size=16): 2 offloaded chunks
2717
- # -> GPU blocks (0, 0) and (0, 1)
2718
- # Group 1 (SWA, block_size=4): 8 offloaded chunks, skip first 2
2719
- # per segment of 4:
2720
- # Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3)
2721
- # Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7)
2718
+ # Retain each segment's tail, including MTP history when needed.
2719
expected_stored=(
2720
(0, 0),
2721
(0, 1),
@@ -2726,27 +2723,26 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
2723
(1, 3),
2724
(1, 6),
2725
(1, 7),
2729
- ),
2726
+ )
2727
+ + (((1, 1), (1, 5)) if extra_retained > 1 else ()),
2728
)
2729
2730
# Verify that loads still work correctly for the stored SWA blocks.
2731
runner.scheduler.reset_prefix_cache()
2732
runner.new_request(token_ids=[0] * num_tokens + [1])
2735
- runner.manager.lookup.return_value = LookupResult.HIT
2736
- runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2
2733
+ runner.manager.lookup.side_effect = lambda key, ctx: (
2734
+ LookupResult.HIT if (key, 0) in runner.offloaded else LookupResult.MISS
2735
+ )
2736
runner.run(
2737
decoded_tokens=[EOS_TOKEN_ID],
2739
- # Group 0: full prefix lookup hits 2 offloaded chunks
2740
- # -> loads GPU blocks (0, 0), (0, 1)
2741
- # Group 1: sliding window lookup finds trailing 2 from last segment
2742
- # (blocks 6, 7 which were stored)
2743
- # -> loads GPU blocks (1, 6), (1, 7)
2738
+ # The consumer restores the full FA prefix and the final SWA tail.
2739
expected_loaded=(
2740
(0, 0),
2741
(0, 1),
2742
(1, 6),
2743
(1, 7),
2749
- ),
2744
+ )
2745
+ + (((1, 5),) if extra_retained > 1 else ()),
2746
)
2747
2748
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+26
-4
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
from collections.abc import Iterable, Sequence
5
-from dataclasses import dataclass, field
5
+from dataclasses import dataclass, field, replace
6
from itertools import chain, islice
7
from typing import Any, NamedTuple
8
@@ -114,7 +114,10 @@ class GroupOffloadConfig(NamedTuple):
114
window = max(
115
window,
116
cdiv(
117
- self.kv_cache_spec.sliding_window - 1 + right_padding,
117
+ self.kv_cache_spec.sliding_window
118
+ - 1
119
+ + self.kv_cache_spec.extra_retained_tokens
120
+ + right_padding,
121
self.tokens_per_chunk,
122
),
123
)
@@ -126,7 +129,13 @@ def get_sliding_window_size_in_chunks(
129
) -> int | None:
130
if isinstance(kv_cache_spec, SlidingWindowSpec):
131
assert kv_cache_spec.sliding_window > 0
129
- return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk)
132
+ return max(
133
+ cdiv(kv_cache_spec.sliding_window, tokens_per_chunk),
134
+ cdiv(
135
+ kv_cache_spec.sliding_window - 1 + kv_cache_spec.extra_retained_tokens,
136
+ tokens_per_chunk,
137
+ ),
138
+ )
139
140
if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
141
# Attention never reaches back past one chunk
@@ -1461,11 +1470,24 @@ class OffloadingConnectorScheduler:
1470
) -> list[bool] | None:
1471
"""Build the block mask for a range of candidate offload chunks."""
1472
blocks_per_chunk = self.config.blocks_per_chunk
1473
+ kv_cache_spec = group_config.kv_cache_spec
1474
+ if isinstance(kv_cache_spec, SlidingWindowSpec) and (
1475
+ kv_cache_spec.extra_retained_tokens
1476
+ ):
1477
+ # Offload restores all allocated history, including MTP re-prefill
1478
+ # tokens. Widen only the store mask, not the model's attention window.
1479
+ kv_cache_spec = replace(
1480
+ kv_cache_spec,
1481
+ sliding_window=(
1482
+ kv_cache_spec.sliding_window + kv_cache_spec.extra_retained_tokens
1483
+ ),
1484
+ extra_retained_tokens=0,
1485
+ )
1486
return group_config.manager_cls.reachable_block_mask(
1487
start_block=start_chunk_idx * blocks_per_chunk,
1488
end_block=end_chunk_idx * blocks_per_chunk,
1489
alignment_tokens=self.config.alignment_tokens,
1468
- kv_cache_spec=group_config.kv_cache_spec,
1490
+ kv_cache_spec=kv_cache_spec,
1491
use_eagle=group_config.is_eagle_group,
1492
retention_interval=self.config.retention_interval,
1493
reachable_boundaries=reachable_boundaries,