[Bugfix][Core] Restore partial prefix hits on long hybrid prompts (27B and Flash-Next) (#1401)

* [Bugfix][Core] Keep sliding-window runs at retained Mamba checkpoints

With --prefix-cache-retention-interval, aligned Mamba groups keep a state
every interval tokens, so a prompt that shares only an earlier part of a
cached one can resume there. A sliding-window group beside them (the
DFlash2 draft on Qwen3.8-27B) released every block its window had passed
with reuse_first, except the run at the prompt's last alignment boundary.
The prefill itself then reused those blocks, the window run before each
retained checkpoint lost its hash, and the hybrid hit collapsed to zero:
a 32K prompt edited after 90% recomputed from scratch.

Sliding-window managers now also hold the run ending at each retained
checkpoint. The held blocks keep normal eviction priority, as the
last-boundary run already did, so admission reservations are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Release] Keep a 27B GDN state every 8K prompt tokens

The 27B recipe kept only the prompt's last Mamba checkpoint, so an edit
anywhere before the last 2K tokens of a long prompt recomputed it all.
With prefix_cache_retention_interval 8192 (and the sliding-window fix
before this commit), a 32K/64K prompt with its last 10% replaced resumes
at 24576/57344 tokens, as 1.5.1 did: TTFT 7.4 -> 2.0 s and 16.3 -> 2.3 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Release] Keep a Flash-Next GDN state every 8K prompt tokens

The same retention as the 27B recipe for Flash-Next. With the last 10%
of a 32K/64K prompt replaced, prefill resumes at 24576/49152 tokens
(TTFT 1.5/3.1 s against 5.5/11.4 s cold); MTP's lookbehind costs one
alignment block at 64K. Both the bundled launcher and the Docker image.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Misc] Fix line length and formatting in the Flash-Next launcher test

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: mac <mac@MacBook-xinyi.local>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

docker/sm70_release/serve.sh

@@ -34,7 +34,8 @@ case "$MODEL" in
     draft=$MODEL_DIR/Qwen3.8-27B-DFlash2
     args+=(--served-model-name "${SERVED_MODEL_NAME:-qwen3.8-27b-dflash2}"
       --kv-cache-dtype fp8_e4m3 --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION:-0.83}"
-      --mamba-cache-mode align --limit-mm-per-prompt '{"video": 0, "audio": 0}'
+      --mamba-cache-mode align --prefix-cache-retention-interval 8192
+      --limit-mm-per-prompt '{"video": 0, "audio": 0}'
       --default-chat-template-kwargs "${CHAT_TEMPLATE_KWARGS:-$thinking_on}"
       --speculative-config "{\"method\": \"dflash\", \"model\": \"$draft\", \"revision\": \"dedf8df68adfb1afeaf7b7480c0a0243108177b4\", \"num_speculative_tokens\": 7, \"kv_cache_dtype\": \"auto\", \"attention_backend\": \"FLASH_ATTN_V100\", \"draft_sample_method\": \"probabilistic\", \"enforce_eager\": false}")
     ;;
@@ -44,6 +45,7 @@ case "$MODEL" in
     args+=(--served-model-name "${SERVED_MODEL_NAME:-qwen3.8-flash-next-nvfp4}"
       --kv-cache-dtype float16 --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION:-0.89}"
       --language-model-only --mamba-ssm-cache-dtype float32
+      --prefix-cache-retention-interval 8192
       --default-chat-template-kwargs "${CHAT_TEMPLATE_KWARGS:-$thinking_off}"
       --kernel-config '{"sm70_hcx": true, "sm70_qsa_prep": true, "sm70_hcx_local_schedule": true, "qsa_host_kv": true, "qsa_host_kv_dtype": "fp8_e4m3", "qsa_host_kv_hot_dtype": "fp8_e4m3", "qsa_host_kv_hot_tokens": 32768, "qsa_host_kv_prefill_stage_tokens": 262144}'
       --speculative-config "{\"method\": \"mtp\", \"model\": \"$draft\", \"num_speculative_tokens\": 4, \"draft_sample_method\": \"greedy\"}")

scripts/serve_flash_next_nvfp4_v100.sh

@@ -64,6 +64,7 @@ exec "$vllm_cli" serve "$model" \
   --gpu-memory-utilization 0.89 --attention-backend FLASH_ATTN_V100 \
   --enable-prefix-caching --partial-prefix-cache \
   --block-size 2048 --mamba-block-size 2048 --mixed-prefill-step-latency-ms 0 \
+  --prefix-cache-retention-interval 8192 \
   --mamba-ssm-cache-dtype float32 --enable-prompt-tokens-details \
   --kernel-config "$kernel_config" \
   --speculative-config "$speculative" "${serve_options[@]}"

tests/config/test_flash_next_release_launcher.py

@@ -39,6 +39,7 @@ def test_flash_launcher_uses_adjacent_cli_and_preserves_overrides(tmp_path):
     assert args[args.index("--kv-cache-dtype") + 1] == "float16"
     assert args[args.index("--max-model-len") + 1] == "262144"
     assert args[args.index("--max-num-seqs") + 1] == "8"
+    assert args[args.index("--prefix-cache-retention-interval") + 1] == "8192"
     assert json.loads(args[args.index("--kernel-config") + 1])["qsa_host_kv"] is True
     assert json.loads(args[args.index("--speculative-config") + 1]) == {
         "method": "mtp",
@@ -52,7 +53,9 @@ def test_flash_launcher_separate_draft(tmp_path):
     shutil.copy(ROOT / "scripts" / script.name, script)
     script.chmod(0o755)
     cli = tmp_path / "vllm"
-    cli.write_text(f"#!{sys.executable}\nimport sys,json\nprint(json.dumps(sys.argv[1:]))\n")
+    cli.write_text(
+        f"#!{sys.executable}\nimport sys,json\nprint(json.dumps(sys.argv[1:]))\n"
+    )
     cli.chmod(0o755)
     draft = tmp_path / 'mtp "draft" dir'
     result = subprocess.run(

tests/config/test_sm70_release_profile.py

@@ -63,6 +63,8 @@ def test_profile_schema():
     assert profile["args"]["max_num_batched_tokens"] == 4096
     assert profile["args"]["block_size"] == 2048
     assert profile["args"]["mamba_block_size"] == 2048
+    # Edited long prompts resume from a GDN state kept every 8K tokens.
+    assert profile["args"]["prefix_cache_retention_interval"] == 8192
     # Prefill beside resident decoders takes the full token budget.
     assert profile["args"]["mixed_prefill_step_latency_ms"] == 0
     spec = profile["args"]["speculative_config"]

tests/v1/core/test_mamba_sparse_retention.py

@@ -20,6 +20,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     KVCacheGroupSpec,
     MambaSpec,
+    SlidingWindowSpec,
 )
 
 # These tests allocate only block metadata; no device/distributed state exists.
@@ -408,3 +409,63 @@ def test_bulk_prefill_materializes_every_admitted_state(eagle, length, interval_
     if interval == 0:
         assert len(ends) <= length // 8000 + 4
     assert cm.block_pool.get_num_free_blocks() == 999
+
+
+def _draft_window_manager(interval):
+    """Main full attention and Mamba groups beside a draft sliding window.
+
+    Mirrors Qwen3.8-27B + DFlash2: 2048-token main pages, a 1024-token draft
+    page and a 2048-token draft window, scaled down by 64.
+    """
+    main = dict(num_kv_heads=1, head_size=1, dtype=torch.float16)
+    config = KVCacheConfig(
+        1000,
+        [],
+        [
+            KVCacheGroupSpec(["full"], FullAttentionSpec(block_size=32, **main)),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=32,
+                    shapes=((1,),),
+                    dtypes=(torch.float16,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["draft"], SlidingWindowSpec(block_size=16, sliding_window=32, **main)
+            ),
+        ],
+    )
+    return KVCacheManager(
+        config, 100000, 8, use_eagle=False, prefix_cache_retention_interval=interval
+    )
+
+
+@pytest.mark.parametrize("interval", [0, 128])
+def test_draft_window_survives_at_retained_mamba_checkpoints(interval):
+    """A prompt sharing only an earlier part resumes at a retained checkpoint.
+
+    Prefill reuses the blocks a sliding window has moved past first. Without
+    keeping the window run at each retained Mamba checkpoint, the draft group
+    has nothing there and the whole hybrid hit collapses to zero.
+    """
+    cm = _draft_window_manager(interval)
+    request = _request("prompt", 400)
+    for start in range(0, request.num_tokens, 64):
+        end = min(start + 64, request.num_tokens)
+        for manager in cm.coordinator.single_type_managers:
+            manager.new_step_starts()
+            manager.remove_skipped_blocks(request.request_id, start, request.num_tokens)
+            manager.allocate_new_blocks(request.request_id, end, end)
+        cm.coordinator.cache_blocks(request, end)
+    cm.coordinator.free(request.request_id)
+
+    # The same first 300 tokens, then a different tail.
+    edited = _request("edited", 400)
+    edited.block_hashes[: 300 // 8] = request.block_hashes[: 300 // 8]
+    _, hit = cm.get_computed_blocks(edited)
+    # Interval 0 keeps only the prompt's last boundary, which the edit changed.
+    assert hit == (256 if interval else 0)
+    _, repeat = cm.get_computed_blocks(_request("prompt", 400))
+    assert repeat == 384

vllm/sm70_profiles/qwen38_27b_nvfp4_dflash2.json

@@ -32,6 +32,7 @@
     "enable_prefix_caching": true,
     "mamba_cache_mode": "align",
     "partial_prefix_cache": true,
+    "prefix_cache_retention_interval": 8192,
     "enable_prompt_tokens_details": true,
     "served_model_name": "qwen3.8-27b-dflash2",
     "limit_mm_per_prompt": {

vllm/v1/core/kv_cache_coordinator.py

@@ -114,6 +114,9 @@ class KVCacheCoordinator(ABC):
                 f"prefix_cache_retention_interval ({interval}) must be a non-negative "
                 f"multiple of the resolved alignment ({alignment})"
             )
+        # Sliding windows keep their run at every retained Mamba checkpoint.
+        for manager in self.single_type_managers:
+            manager.retention_interval = interval
 
     def get_replay_boundaries(
         self, request: Request, alignment: int

vllm/v1/core/single_type_kv_cache_manager.py

@@ -105,6 +105,10 @@ class SingleTypeKVCacheManager(ABC):
         # Cache-hit alignment, set by the coordinator (its ``lcm_block_size``).
         self.alignment_tokens: int | None = None
         self.mamba_replay_lookbehind = False
+        # The coordinator's ``prefix_cache_retention_interval``: a positive
+        # value keeps a Mamba state every that many tokens, which a sliding
+        # window must also keep its run for.
+        self.retention_interval: int | None = None
 
     def take_pending_boundary_state_offloads(
         self,
@@ -479,7 +483,7 @@ class SingleTypeKVCacheManager(ABC):
         request_id: str,
         first_block: int,
         last_block: int,
-        hold: tuple[int, int] | None = None,
+        hold: Sequence[tuple[int, int]] = (),
         reuse_first: bool = False,
     ) -> None:
         """Free blocks in ``[first_block, last_block)`` and replace with null_block.
@@ -488,8 +492,8 @@ class SingleTypeKVCacheManager(ABC):
         earlier blocks in the range were nulled in a prior call.
 
         Args:
-            hold: Block range ``[start, end)`` whose cached entries keep normal
-                eviction priority: the window a repeat of this prompt needs.
+            hold: Block ranges ``[start, end)`` whose cached entries keep normal
+                eviction priority: the windows a later prompt can resume from.
                 Other expired sliding-window blocks can be reused immediately.
             reuse_first: Hand the freed blocks out again before other free
                 blocks (see ``BlockPool.free_blocks``).
@@ -509,7 +513,7 @@ class SingleTypeKVCacheManager(ABC):
                 # should also have been set to null blocks by the previous calls
                 # to this function.
                 break
-            if hold is not None and hold[0] <= i < hold[1]:
+            if any(start <= i < end for start, end in hold):
                 replay.append(blocks[i])
             else:
                 freed.append(blocks[i])
@@ -521,43 +525,50 @@ class SingleTypeKVCacheManager(ABC):
             # the real-held SWA capacity used by admission and pool sizing.
             self.block_pool.free_blocks(replay)
 
-    def _cached_window_at_last_boundary(
-        self, num_prompt_tokens: int | None
-    ) -> tuple[int, int] | None:
-        """Block range a repeat of this prompt will look up, or ``None``.
+    def _cached_windows(self, num_prompt_tokens: int | None) -> list[tuple[int, int]]:
+        """Block ranges a later prompt can resume from.
 
         A hit ends on an ``alignment_tokens`` boundary and, for a sliding
         window, consults only the blocks ``reachable_block_mask`` caches there:
         the ``need``-wide run ending at that boundary, shifted by one when
-        EAGLE peeks past it. Once outside the live attention window, these
-        blocks retain normal cache priority; all other expired blocks are
-        reused first. This preserves the bounded SWA admission reservation.
+        EAGLE peeks past it. Two kinds of boundary can serve a hit: the last
+        one of the prompt, which a repeat looks up, and each Mamba state a
+        positive ``retention_interval`` keeps, which a prompt sharing only an
+        earlier part resumes from. Once outside the live attention window,
+        these blocks retain normal cache priority; all other expired blocks
+        are reused first. This preserves the bounded SWA admission reservation.
         """
         if (
             not self.enable_caching
             or num_prompt_tokens is None
             or self.reuse_window is None
         ):
-            return None
+            return []
         alignment = self.alignment_tokens or self.block_size
         # Lookup recomputes the final token and EAGLE's last matched page.
         # A sparse Mamba checkpoint can require one whole alignment earlier.
         if self.mamba_replay_lookbehind:
-            boundary = max(0, num_prompt_tokens // alignment - 1) * alignment
+            last = max(0, num_prompt_tokens // alignment - 1) * alignment
         else:
             last_match = (
                 num_prompt_tokens - 1 - (self.block_size if self.use_eagle else 0)
             )
-            boundary = max(0, last_match // alignment) * alignment
-        if boundary <= 0:
-            return None
+            last = max(0, last_match // alignment) * alignment
+        boundaries = {last} if last > 0 else set()
+        if self.retention_interval:
+            boundaries.update(
+                range(self.retention_interval, last, self.retention_interval)
+            )
         need = cdiv(self.reuse_window - 1, self.block_size)
         shift = 0
         if self.use_eagle:
             need += 1
             shift = 1
-        end = boundary // self.block_size + shift
-        return (max(0, end - need), end)
+        windows = []
+        for boundary in sorted(boundaries):
+            end = boundary // self.block_size + shift
+            windows.append((max(0, end - need), end))
+        return windows
 
     def remove_skipped_blocks(
         self,
@@ -592,13 +603,13 @@ class SingleTypeKVCacheManager(ABC):
         if self.reuse_window is None:
             self._remove_blocks_in_range(request_id, 0, num_skipped_blocks)
             return
-        # Sliding window: blocks the window has moved past can never serve a
-        # future hit, so they are reused first -- otherwise a long prefill
-        # walks the whole free queue and evicts other requests' prefixes. The
-        # one exception is the cached window at the last alignment boundary
-        # of the prompt: a repeat of this prompt looks exactly that up, so it
-        # keeps normal cache priority while remaining evictable.
-        hold = self._cached_window_at_last_boundary(num_prompt_tokens)
+        # Sliding window: blocks the window has moved past can serve no future
+        # hit, so they are reused first -- otherwise a long prefill walks the
+        # whole free queue and evicts other requests' prefixes. The exceptions
+        # are the cached windows a later prompt resumes from (the last
+        # alignment boundary and retained Mamba checkpoints): they keep normal
+        # cache priority while remaining evictable.
+        hold = self._cached_windows(num_prompt_tokens)
         # `num_skipped_tokens` may include tokens that haven't been allocated yet
         # (e.g., when the attention window moves into the external computed tokens
         # range); `_remove_blocks_in_range` caps to the number of blocks that