[Bugfix][Spec Decode] DFlash/DSpark no longer drop the last prefix-cache block (mamba align context recompute)

tests/v1/core/test_mamba_align_chunk_split.py

@@ -141,6 +141,26 @@ def test_disabling_eagle_block_drop_keeps_the_trailing_cache_boundary() -> None:
     assert without_drop == 2 * MAMBA_BLOCK_SIZE
 
 
+def test_dflash_does_not_back_off_last_cache_position() -> None:
+    """DFlash/DSpark never write target blocks, so the split must not back
+    off the last prefix-cache position by a mamba block.
+
+    Regression for #53477: the old `use_eagle` back-off made prompts shorter
+    than two mamba blocks skip the final block-aligned chunk, so the mamba
+    recurrent state was never materialized at a block boundary and the next
+    turn's prefix-cache lookup recomputed the whole context.
+    """
+    (request,) = create_requests(1, num_tokens=PROMPT_LEN, block_size=ATTN_BLOCK_SIZE)
+    assert (
+        _split(request, PROMPT_LEN, use_eagle_block_drop=False)
+        == MAMBA_BLOCK_SIZE
+    )
+    assert (
+        _split(request, PROMPT_LEN, use_eagle_block_drop=True)
+        == PROMPT_LEN
+    )
+
+
 def _run_chunked_prefill(
     manager: KVCacheManager, request: Request, budgets: list[int]
 ) -> dict[int, int]:

tests/v1/core/test_scheduler.py

@@ -3947,6 +3947,7 @@ def test_mamba_align_eagle_schedules_encoder_at_boundary():
     )
     scheduler.need_mamba_block_aligned_split = True
     scheduler.use_eagle = True
+    scheduler.use_eagle_block_drop = True
     scheduler.num_prefill_lookahead = 1
     scheduler.max_num_encoder_input_tokens = 2048
     scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=2048)

vllm/config/speculative.py

@@ -1845,8 +1845,23 @@ def use_eagle(self) -> bool:
         return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark")
 
     def use_eagle_block_drop(self) -> bool:
-        """Whether volatile trailing cache blocks should be discarded."""
-        return self.use_eagle() and not self.disable_eagle_block_drop
+        """Whether volatile trailing cache blocks should be discarded.
+
+        Only eagle-family drafters share (and pollute via lookahead KV write)
+        the target's full-attention KV cache groups; DFlash/DSpark draft from
+        their own KV cache and never write target blocks (#53477), so the
+        drop applies to eagle/eagle3/mtp only, unless explicitly disabled.
+        """
+        return (
+            self.use_eagle_preserves_target_kv_cache()
+            and not self.disable_eagle_block_drop
+        )
+
+    def use_eagle_preserves_target_kv_cache(self) -> bool:
+        # Only eagle-family drafters share (and pollute via lookahead KV
+        # write) the target's full-attention KV cache groups; DFlash/DSpark
+        # draft from their own KV cache and never write target blocks.
+        return self.method in ("eagle", "eagle3", "mtp")
 
     def use_dflash(self) -> bool:
         return self.method == "dflash"

vllm/v1/core/sched/scheduler.py

@@ -281,7 +281,10 @@ def __init__(
                     else 1
                 )
             self.use_eagle_block_drop = speculative_config.use_eagle_block_drop()
-            if self.use_eagle and not self.use_eagle_block_drop:
+            if (
+                speculative_config.use_eagle_preserves_target_kv_cache()
+                and not self.use_eagle_block_drop
+            ):
                 logger.warning(
                     "EAGLE trailing prefix-cache block dropping is disabled. "
                     "This is experimental and may affect speculative-token "
@@ -418,9 +421,11 @@ def _mamba_block_aligned_split(
             return num_new_tokens
 
         block_size = self.cache_config.block_size
-        # The last block-aligned position whose state can be cached. With
-        # Eagle, FullAttn prunes the last matching block, so back off one
-        # block to avoid a Mamba cache miss.
+        # The last block-aligned position whose state can be cached.
+        # Eagle-family drafters pollute the target's last matching
+        # full-attention block with their lookahead KV write, so back off one
+        # block to avoid a Mamba cache miss. DFlash/DSpark draft from their own
+        # KV cache and never write target blocks.
         last_cache_position = request.num_tokens - request.num_tokens % block_size
         if self.use_eagle_block_drop:
             last_cache_position = max(last_cache_position - block_size, 0)