[Bugfix][KV Offload] Preserve reachable tails for hybrid SWA groups (#48911)
Signed-off-by: Colton Ottley <colton@ottleyengineering.com> Co-authored-by: Colton Ottley <colton@ottleyengineering.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
coltonottley committed
Jul 20, 2026 at 13:12 UTC
fbfe58133d9fa7c09ad94b21073b43c3accb402b
2 files changed
+82
-12
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+50
@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
19
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
20
OffloadingConnectorScheduler,
21
RequestOffloadState,
22
+ is_store_reachable_swa_chunk,
23
)
24
from vllm.v1.core.kv_cache_utils import BlockHash
25
from vllm.v1.kv_cache_interface import (
@@ -126,6 +127,55 @@ def test_scheduler_reports_lookup_sync_delay(request_runner):
127
assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0
128
129
130
+@pytest.mark.parametrize(
131
+ (
132
+ "absolute_chunk_index",
133
+ "storable_chunk_count",
134
+ "alignment_chunk_count",
135
+ "sliding_window_chunks",
136
+ "is_eagle_group",
137
+ "expected",
138
+ ),
139
+ [
140
+ # Full 64-chunk segment: ordinary SWA keeps 62-63; EAGLE also keeps 61.
141
+ (61, 64, 64, 2, False, False),
142
+ (62, 64, 64, 2, False, True),
143
+ (60, 64, 64, 2, True, False),
144
+ (61, 64, 64, 2, True, True),
145
+ # Partial 48-of-64 segment: the reachable tail ends at chunk 47.
146
+ (45, 48, 64, 2, False, False),
147
+ (46, 48, 64, 2, False, True),
148
+ (44, 48, 64, 2, True, False),
149
+ (45, 48, 64, 2, True, True),
150
+ # A later partial segment uses its own actual end (chunks 64-79).
151
+ (76, 80, 64, 3, False, False),
152
+ (77, 80, 64, 3, False, True),
153
+ # No alignment means no store-pruning optimization.
154
+ (0, 1, None, None, False, True),
155
+ # A tail at least as large as the segment keeps every chunk.
156
+ (0, 2, 64, 2, False, True),
157
+ ],
158
+)
159
+def test_is_store_reachable_swa_chunk(
160
+ absolute_chunk_index: int,
161
+ storable_chunk_count: int,
162
+ alignment_chunk_count: int | None,
163
+ sliding_window_chunks: int | None,
164
+ is_eagle_group: bool,
165
+ expected: bool,
166
+):
167
+ assert (
168
+ is_store_reachable_swa_chunk(
169
+ absolute_chunk_index,
170
+ storable_chunk_count,
171
+ alignment_chunk_count,
172
+ sliding_window_chunks,
173
+ is_eagle_group,
174
+ )
175
+ is expected
176
+ )
177
+
178
+
179
def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
180
"""A deferred lookup reports its async delay once it resolves."""
181
runner = request_runner(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+32
-12
@@ -111,6 +111,26 @@ def get_sliding_window_size_in_chunks(
111
return None
112
113
114
+def is_store_reachable_swa_chunk(
115
+ absolute_chunk_index: int,
116
+ storable_chunk_count: int,
117
+ alignment_chunk_count: int | None,
118
+ sliding_window_chunks: int | None,
119
+ is_eagle_group: bool,
120
+) -> bool:
121
+ """Return whether an SWA chunk can participate in an external-cache hit."""
122
+ if alignment_chunk_count is None:
123
+ return True
124
+ assert sliding_window_chunks is not None
125
+ position_in_segment = absolute_chunk_index % alignment_chunk_count
126
+ segment_start = absolute_chunk_index - position_in_segment
127
+ actual_segment_length = min(
128
+ alignment_chunk_count, storable_chunk_count - segment_start
129
+ )
130
+ reachable_tail = sliding_window_chunks + int(is_eagle_group)
131
+ return position_in_segment >= actual_segment_length - reachable_tail
132
+
133
+
134
def resolve_mamba_align_size(
135
spec: "OffloadingSpec", kv_cache_config: KVCacheConfig
136
) -> int | None:
@@ -990,9 +1010,6 @@ class OffloadingConnectorScheduler:
1010
]
1011
assert len(offload_keys) == len(offload_block_ids)
1012
993
- alignment_chunk_count = group_config.alignment_chunk_count
994
- tail = group_config.sliding_window_size_in_chunks
995
-
1013
for key_idx, (offload_key, block_id) in enumerate(
1014
zip(offload_keys, offload_block_ids)
1015
):
@@ -1000,15 +1017,18 @@ class OffloadingConnectorScheduler:
1017
continue
1018
# Skip SWA chunks that can never serve a load hit:
1019
# within each full-attention alignment segment, only the
1003
- # trailing `tail` chunks are reachable by
1004
- # _sliding_window_lookup. For DeepSeek V4 with 100K
1005
- # tokens this reduces SWA stores by ~78%.
1006
- if alignment_chunk_count is not None:
1007
- assert tail is not None
1008
- abs_chunk_idx = start_chunk_idx + key_idx
1009
- pos_in_segment = abs_chunk_idx % alignment_chunk_count
1010
- if pos_in_segment < alignment_chunk_count - tail:
1011
- continue
1020
+ # trailing chunks queried by _sliding_window_lookup are
1021
+ # reachable. EAGLE/MTP requires one additional chunk that
1022
+ # lookup later drops as its volatile draft tail.
1023
+ abs_chunk_idx = start_chunk_idx + key_idx
1024
+ if not is_store_reachable_swa_chunk(
1025
+ abs_chunk_idx,
1026
+ num_chunks,
1027
+ group_config.alignment_chunk_count,
1028
+ group_config.sliding_window_size_in_chunks,
1029
+ group_config.is_eagle_group,
1030
+ ):
1031
+ continue
1032
new_offload_keys.append(offload_key)
1033
1034
if not new_offload_keys: