[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>
2026-07-20 22:12:36 +03:00
parent 9dd62d80ab
commit fbfe58133d
2 changed files with 82 additions and 12 deletions
@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
    OffloadingConnectorScheduler,
    RequestOffloadState,
+    is_store_reachable_swa_chunk,
 )
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
@@ -126,6 +127,55 @@ def test_scheduler_reports_lookup_sync_delay(request_runner):
    assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0


+@pytest.mark.parametrize(
+    (
+        "absolute_chunk_index",
+        "storable_chunk_count",
+        "alignment_chunk_count",
+        "sliding_window_chunks",
+        "is_eagle_group",
+        "expected",
+    ),
+    [
+        # Full 64-chunk segment: ordinary SWA keeps 62-63; EAGLE also keeps 61.
+        (61, 64, 64, 2, False, False),
+        (62, 64, 64, 2, False, True),
+        (60, 64, 64, 2, True, False),
+        (61, 64, 64, 2, True, True),
+        # Partial 48-of-64 segment: the reachable tail ends at chunk 47.
+        (45, 48, 64, 2, False, False),
+        (46, 48, 64, 2, False, True),
+        (44, 48, 64, 2, True, False),
+        (45, 48, 64, 2, True, True),
+        # A later partial segment uses its own actual end (chunks 64-79).
+        (76, 80, 64, 3, False, False),
+        (77, 80, 64, 3, False, True),
+        # No alignment means no store-pruning optimization.
+        (0, 1, None, None, False, True),
+        # A tail at least as large as the segment keeps every chunk.
+        (0, 2, 64, 2, False, True),
+    ],
+)
+def test_is_store_reachable_swa_chunk(
+    absolute_chunk_index: int,
+    storable_chunk_count: int,
+    alignment_chunk_count: int | None,
+    sliding_window_chunks: int | None,
+    is_eagle_group: bool,
+    expected: bool,
+):
+    assert (
+        is_store_reachable_swa_chunk(
+            absolute_chunk_index,
+            storable_chunk_count,
+            alignment_chunk_count,
+            sliding_window_chunks,
+            is_eagle_group,
+        )
+        is expected
+    )
+
+
 def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
    """A deferred lookup reports its async delay once it resolves."""
    runner = request_runner(
@@ -111,6 +111,26 @@ def get_sliding_window_size_in_chunks(
    return None


+def is_store_reachable_swa_chunk(
+    absolute_chunk_index: int,
+    storable_chunk_count: int,
+    alignment_chunk_count: int | None,
+    sliding_window_chunks: int | None,
+    is_eagle_group: bool,
+) -> bool:
+    """Return whether an SWA chunk can participate in an external-cache hit."""
+    if alignment_chunk_count is None:
+        return True
+    assert sliding_window_chunks is not None
+    position_in_segment = absolute_chunk_index % alignment_chunk_count
+    segment_start = absolute_chunk_index - position_in_segment
+    actual_segment_length = min(
+        alignment_chunk_count, storable_chunk_count - segment_start
+    )
+    reachable_tail = sliding_window_chunks + int(is_eagle_group)
+    return position_in_segment >= actual_segment_length - reachable_tail
+
+
 def resolve_mamba_align_size(
    spec: "OffloadingSpec", kv_cache_config: KVCacheConfig
 ) -> int | None:
@@ -990,9 +1010,6 @@ class OffloadingConnectorScheduler:
                ]
                assert len(offload_keys) == len(offload_block_ids)

-                alignment_chunk_count = group_config.alignment_chunk_count
-                tail = group_config.sliding_window_size_in_chunks
-
                for key_idx, (offload_key, block_id) in enumerate(
                    zip(offload_keys, offload_block_ids)
                ):
@@ -1000,15 +1017,18 @@ class OffloadingConnectorScheduler:
                        continue
                    # Skip SWA chunks that can never serve a load hit:
                    # within each full-attention alignment segment, only the
-                    # trailing `tail` chunks are reachable by
-                    # _sliding_window_lookup. For DeepSeek V4 with 100K
-                    # tokens this reduces SWA stores by ~78%.
-                    if alignment_chunk_count is not None:
-                        assert tail is not None
-                        abs_chunk_idx = start_chunk_idx + key_idx
-                        pos_in_segment = abs_chunk_idx % alignment_chunk_count
-                        if pos_in_segment < alignment_chunk_count - tail:
-                            continue
+                    # trailing chunks queried by _sliding_window_lookup are
+                    # reachable. EAGLE/MTP requires one additional chunk that
+                    # lookup later drops as its volatile draft tail.
+                    abs_chunk_idx = start_chunk_idx + key_idx
+                    if not is_store_reachable_swa_chunk(
+                        abs_chunk_idx,
+                        num_chunks,
+                        group_config.alignment_chunk_count,
+                        group_config.sliding_window_size_in_chunks,
+                        group_config.is_eagle_group,
+                    ):
+                        continue
                    new_offload_keys.append(offload_key)

            if not new_offload_keys: