[Hybrid KV Cache] Retain finalized Mamba decode checkpoints for prefix caching

tests/v1/core/test_prefix_caching.py

@@ -3565,6 +3565,142 @@ def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary():
     assert len(computed_blocks.blocks[1]) == 0
 
 
+def _make_decode_checkpoint_manager(monkeypatch):
+    monkeypatch.setenv("VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS", "1")
+    block_size = 4
+    manager = make_kv_cache_manager(
+        _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba_align"]),
+        max_model_len=1024,
+        enable_caching=True,
+        retention_interval=0,
+        hash_block_size=block_size,
+    )
+    return manager, block_size
+
+
+def _materialize_checkpoint_test_request(manager, block_size, num_decode_blocks=3):
+    request = make_request("producer", list(range(2 * block_size)), block_size, sha256)
+
+    def compute(num_tokens):
+        manager.new_step_starts()
+        blocks = manager.allocate_slots(request, num_tokens)
+        assert blocks is not None
+        request.num_computed_tokens += num_tokens
+        manager.update_decode_checkpoint_candidates(request)
+
+    compute(block_size)
+    compute(block_size)
+    for block_idx in range(num_decode_blocks):
+        start = request.num_tokens
+        request.append_output_token_ids(list(range(start, start + block_size)))
+        compute(block_size)
+    return request
+
+
+@pytest.mark.parametrize("keep", [True, False])
+def test_mamba_decode_checkpoints_publish_latest_on_finish(monkeypatch, keep):
+    """The latest materialized decode state becomes reusable after finish."""
+    manager, block_size = _make_decode_checkpoint_manager(monkeypatch)
+    request = _materialize_checkpoint_test_request(manager, block_size)
+
+    assert manager.block_pool.get_cached_block(request.block_hashes[4], [1]) is None
+    manager.finalize_decode_checkpoints(request, keep=keep)
+    manager.free(request)
+
+    # Only a kept checkpoint extends reuse beyond the prompt replay boundary.
+    full_replay = make_request(
+        "full-replay",
+        list(request.all_token_ids) + [100, 101, 102, 103],
+        block_size,
+        sha256,
+    )
+    _, full_hit, _ = manager.get_computed_blocks(full_replay)
+    assert full_hit == (20 if keep else 4)
+
+
+def test_mamba_decode_checkpoint_pin_survives_state_rotation(monkeypatch):
+    """A private pin keeps an old state out of the allocator after rotation."""
+    manager, block_size = _make_decode_checkpoint_manager(monkeypatch)
+    request = _materialize_checkpoint_test_request(
+        manager, block_size, num_decode_blocks=2
+    )
+    mamba_manager = manager.coordinator.single_type_managers[1]
+    candidate = mamba_manager._decode_checkpoint_candidates[request.request_id]
+    assert candidate.num_tokens == 16
+    assert candidate.block.ref_cnt == 2  # request owner + private pin
+
+    # Two subsequent running-state rotations remove boundary 16 from the
+    # request table. Do not update the candidate: this isolates the private
+    # pin that must keep its physical block resident.
+    new_blocks = []
+    for _ in range(2):
+        start = request.num_tokens
+        request.append_output_token_ids(list(range(start, start + block_size)))
+        manager.new_step_starts()
+        allocated = manager.allocate_slots(request, block_size)
+        assert allocated is not None
+        new_blocks.extend(allocated.blocks[1])
+        request.num_computed_tokens += block_size
+
+    assert candidate.block.ref_cnt == 1
+    assert all(block is not candidate.block for block in new_blocks)
+
+    manager.free(request)
+    assert candidate.block.ref_cnt == 0
+
+
+def test_mamba_decode_checkpoints_exclude_unmaterialized_boundary(monkeypatch):
+    """Allocated/in-flight tokens and a sampled EOS cannot form a candidate."""
+    manager, block_size = _make_decode_checkpoint_manager(monkeypatch)
+    request = _materialize_checkpoint_test_request(
+        manager, block_size, num_decode_blocks=2
+    )
+
+    # Materialize only three tokens of the next block.
+    start = request.num_tokens
+    request.append_output_token_ids(list(range(start, start + block_size - 1)))
+    blocks = manager.allocate_slots(request, block_size - 1)
+    assert blocks is not None
+    request.num_computed_tokens += block_size - 1
+
+    # Sampling EOS makes request.num_tokens reach the aligned boundary, but EOS
+    # has not entered a forward. Even if its slot is optimistically in flight,
+    # processed_end remains 19 and boundary 20 is ineligible.
+    request.append_output_token_ids(999)
+    request.num_computed_tokens += 1
+    request.num_in_flight_tokens = 1
+    manager.update_decode_checkpoint_candidates(request)
+
+    mamba_manager = manager.coordinator.single_type_managers[1]
+    candidate = mamba_manager._decode_checkpoint_candidates[request.request_id]
+    assert candidate.num_tokens == 16
+    manager.free(request)
+
+
+@pytest.mark.parametrize(
+    ("retention_interval", "enable_caching", "use_eagle", "expected_match"),
+    [
+        (None, True, False, "prefix_cache_retention_interval=0"),
+        (64, True, False, "prefix_cache_retention_interval=0"),
+        (0, False, False, "prefix caching"),
+        (0, True, True, "hidden-state speculative decoding"),
+    ],
+)
+def test_decode_checkpoints_reject_unsupported_config(
+    monkeypatch, retention_interval, enable_caching, use_eagle, expected_match
+):
+    monkeypatch.setenv("VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS", "1")
+    with pytest.raises(ValueError, match=expected_match):
+        make_kv_cache_manager(
+            _make_hybrid_kv_cache_config(4, 100, ["full", "mamba_align"]),
+            max_model_len=1024,
+            enable_caching=enable_caching,
+            retention_interval=retention_interval,
+            use_eagle=use_eagle,
+            hash_block_size=4,
+        )
+
+
 def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary():
     """Verify MTP/EAGLE SWA retention keeps the extra proof block.
 

vllm/envs.py

@@ -327,6 +327,7 @@
     VLLM_NIC_SELECTION_VARS: str = ""
     VLLM_PREFIX_CACHE_RETENTION_INTERVAL: int | None = None
     VLLM_ENABLE_HPC_OPS: bool = False
+    VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS: bool = False
 
 
 def get_default_cache_root():
@@ -1177,6 +1178,11 @@ def _resolve_rust_cli_path() -> str | None:
         if "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in os.environ
         else None
     ),
+    # With latest-only retention, privately pin the latest materialized
+    # scheduler-aligned Mamba decode state and publish it on a stopped finish.
+    "VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS": lambda: bool(
+        int(os.getenv("VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS", "0"))
+    ),
     # a local directory to look in for unrecognized LoRA adapters.
     # only works if plugins are enabled and
     # VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled.

vllm/v1/core/block_pool.py

@@ -371,6 +371,72 @@ def _build_block_stored_event(
             session_id=request.session_id,
         )
 
+    def cache_decode_checkpoint(
+        self,
+        request: Request,
+        block: KVCacheBlock,
+        block_hash_with_group_id: BlockHashWithGroupId,
+        num_tokens: int,
+        block_size: int,
+        kv_cache_group_id: int,
+    ) -> bool:
+        """Publish one exact, already-materialized recurrent-state checkpoint.
+
+        Unlike ``cache_full_blocks``, this inserts only the supplied hash alias.
+        Any other hashes owned by ``block`` remain intact.
+        """
+        assert self.enable_caching
+        assert not block.is_null and block.ref_cnt > 0
+        assert num_tokens > 0 and num_tokens % block_size == 0
+        assert get_group_id(block_hash_with_group_id) == kv_cache_group_id
+
+        block_hashes = resolve_block_hashes(
+            request.block_hashes, self.hash_block_size, block_size
+        )
+        block_idx = num_tokens // block_size - 1
+        assert 0 <= block_idx < len(block_hashes)
+        block_hash = block_hashes[block_idx]
+        assert block_hash_with_group_id == make_block_hash_with_group_id(
+            block_hash, kv_cache_group_id
+        )
+
+        if self.cached_block_hash_to_block.contain(
+            block_hash_with_group_id, block.block_id
+        ):
+            return False
+
+        self._insert_block_hash(
+            block_hash_with_group_id,
+            block,
+            num_tokens=num_tokens,
+        )
+        if self.enable_kv_cache_events:
+            parent_block_hash = (
+                maybe_convert_block_hash(block_hashes[block_idx - 1])
+                if block_idx > 0
+                else None
+            )
+            block_start = num_tokens - block_size
+            extra_keys, _ = generate_block_hash_extra_keys(
+                request,
+                block_start,
+                num_tokens,
+                0,
+            )
+            self.kv_event_queue.append(
+                self._build_block_stored_event(
+                    request,
+                    block_hashes=[maybe_convert_block_hash(block_hash)],
+                    parent_block_hash=parent_block_hash,
+                    start_token_idx=block_start,
+                    end_token_idx=num_tokens,
+                    block_size=block_size,
+                    kv_cache_group_id=kv_cache_group_id,
+                    extra_keys_list=[extra_keys],
+                )
+            )
+        return True
+
     def emit_cached_block_events(
         self,
         request: Request,

vllm/v1/core/kv_cache_coordinator.py

@@ -4,6 +4,7 @@
 from collections.abc import Sequence
 from typing import NamedTuple
 
+from vllm import envs
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv, round_down
 from vllm.v1.core.block_pool import BlockPool
@@ -33,9 +34,43 @@
 
 def _validate_prefix_cache_retention_interval(
     retention_interval: int | None,
+    retain_decode_checkpoints: bool,
+    enable_caching: bool,
     scheduler_block_size: int,
     kv_cache_config: KVCacheConfig,
 ) -> None:
+    if retain_decode_checkpoints:
+        if not enable_caching:
+            raise ValueError(
+                "VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS requires "
+                "prefix caching to be enabled."
+            )
+        if retention_interval != 0:
+            raise ValueError(
+                "VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS requires "
+                "prefix_cache_retention_interval=0."
+            )
+        mamba_specs = [
+            group.kv_cache_spec
+            for group in kv_cache_config.kv_cache_groups
+            if isinstance(group.kv_cache_spec, MambaSpec)
+        ]
+        if not mamba_specs:
+            raise ValueError(
+                "VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS requires a "
+                "Mamba KV cache group."
+            )
+        if any(spec.mamba_cache_mode != "align" for spec in mamba_specs):
+            raise ValueError(
+                "VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS requires all "
+                "Mamba KV cache groups to use mamba_cache_mode='align'."
+            )
+        if any(spec.block_size != scheduler_block_size for spec in mamba_specs):
+            raise ValueError(
+                "VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS requires every "
+                "Mamba block size to equal scheduler_block_size."
+            )
+
     if retention_interval is None:
         return
 
@@ -160,8 +195,20 @@ def __init__(
         # (``scheduler_block_size``) to land on real cache-hit boundaries.
         # 0 = keep only the latest replay boundary; None = dense;
         self.retention_interval = kv_cache_config.prefix_cache_retention_interval
+        self.retain_decode_checkpoints = (
+            envs.VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS
+        )
+        if self.retain_decode_checkpoints and use_eagle:
+            raise ValueError(
+                "VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS is not compatible "
+                "with hidden-state speculative decoding."
+            )
         _validate_prefix_cache_retention_interval(
-            self.retention_interval, self.scheduler_block_size, kv_cache_config
+            self.retention_interval,
+            self.retain_decode_checkpoints,
+            enable_caching,
+            self.scheduler_block_size,
+            kv_cache_config,
         )
 
     def get_num_blocks_to_allocate(
@@ -352,6 +399,31 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
                 replay_boundaries=boundaries,
             )
 
+    def update_decode_checkpoint_candidates(
+        self, request: Request, materialized_tokens: int
+    ) -> None:
+        """Privately retain newly materialized recurrent decode states."""
+        if not self.retain_decode_checkpoints:
+            return
+        for manager in self.single_type_managers:
+            manager.update_decode_checkpoint_candidate(request, materialized_tokens)
+
+    def finalize_decode_checkpoints(
+        self, request: Request, materialized_tokens: int, keep: bool
+    ) -> None:
+        """Publish stopped-finish candidates, or discard their private pins."""
+        managers = iter(self.single_type_managers)
+        try:
+            for manager in managers:
+                manager.finalize_decode_checkpoints(request, materialized_tokens, keep)
+        finally:
+            # A manager releases its own pins in a finally block. If promotion
+            # raises, discard candidates belonging to groups not yet visited.
+            for manager in managers:
+                manager.finalize_decode_checkpoints(
+                    request, materialized_tokens, keep=False
+                )
+
     def free(self, request_id: str) -> None:
         """
         Free the blocks for the request.

vllm/v1/core/kv_cache_manager.py

@@ -590,6 +590,32 @@ def free(self, request: Request) -> None:
         """
         self.coordinator.free(request.request_id)
 
+    def update_decode_checkpoint_candidates(self, request: Request) -> None:
+        """Retain checkpoints only for tokens whose forward has completed."""
+        if not self.coordinator.retain_decode_checkpoints:
+            return
+        if not self.enable_caching:
+            return
+        processed_end = max(
+            0, request.num_computed_tokens - request.num_in_flight_tokens
+        )
+        materialized_tokens = min(processed_end, request.num_tokens)
+        self.coordinator.update_decode_checkpoint_candidates(
+            request, materialized_tokens
+        )
+
+    def finalize_decode_checkpoints(self, request: Request, keep: bool) -> None:
+        """Publish or discard the request's private decode checkpoints."""
+        if not self.coordinator.retain_decode_checkpoints:
+            return
+        if not self.enable_caching:
+            return
+        processed_end = max(
+            0, request.num_computed_tokens - request.num_in_flight_tokens
+        )
+        materialized_tokens = min(processed_end, request.num_tokens)
+        self.coordinator.finalize_decode_checkpoints(request, materialized_tokens, keep)
+
     def remove_skipped_blocks(
         self,
         request_id: str,

vllm/v1/core/sched/scheduler.py

@@ -2066,6 +2066,9 @@ def update_from_output(
                     request.resumable = False
                     stopped = True
 
+            if not output_is_stale:
+                self.kv_cache_manager.update_decode_checkpoint_candidates(request)
+
             routed_experts = None
             if (
                 self.enable_return_routed_experts
@@ -2568,6 +2571,10 @@ def _free_request(
     ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
         assert request.is_finished()
 
+        self.kv_cache_manager.finalize_decode_checkpoints(
+            request,
+            keep=request.status == RequestStatus.FINISHED_STOPPED,
+        )
         self._inflight_prefills.discard(request)
         connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)