[Bugfix][Kimi-K3] Keep transient checkpoints out of prefix-cache eviction (#56794)

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Kimi Code <noreply@moonshot.ai>

tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py

@@ -299,7 +299,14 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing():
     running = request.num_computed_tokens
     assert not mamba.req_to_blocks["r"][running // block_size].is_null
     request.shared_prefix_boundary = running
-    assert mamba._cache_partial_tail_block(request, running) is None
+    assert (
+        mamba._cache_partial_tail_block(
+            request,
+            running,
+            retention_interval=manager.coordinator.retention_interval,
+        )
+        is None
+    )
 
 
 @pytest.mark.parametrize("eagle_group", [None, 0])

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -615,6 +615,77 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle():
     assert num_computed == 112
 
 
+@pytest.mark.parametrize(
+    "retention_interval,transient_published",
+    [(None, True), (0, False)],
+)
+def test_internal_checkpoint_publication_respects_retention(
+    retention_interval, transient_published
+):
+    """With retention_interval=0, a mid-prompt internal checkpoint is
+    request-local and must not enter the prefix cache; the prompt-end
+    checkpoint keeps the existing publication behavior. Dense retention
+    (None) is unchanged.
+    """
+    hash_block_size = 16
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=32,
+        num_prefill_checkpoint_blocks=1,
+    )
+    manager.coordinator.retention_interval = retention_interval
+
+    request = make_request("producer", list(range(240)), hash_block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
+
+    # Mid-prompt chunk: exports a transient checkpoint at 112.
+    assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
+    transient_hash = request.block_hashes[112 // hash_block_size - 1]
+    transient_hit = manager.block_pool.get_cached_block(transient_hash, [1])
+    assert (transient_hit is not None) == transient_published
+
+    request.num_computed_tokens = 128
+    manager.new_step_starts()
+
+    # Final chunk: the prompt-end checkpoint at 224 stays published.
+    assert manager.allocate_slots(request, 112) is not None
+    end_hash = request.block_hashes[224 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(end_hash, [1]) is not None
+
+
+def test_transient_checkpoint_evicts_retained_boundary_hash():
+    """With retention_interval=0, the checkpoint slot can coincide with a
+    reachable boundary block that this step's full-block pass just hashed.
+    The checkpoint state overwrites the slot, so the stale boundary hash
+    must be evicted rather than left pointing at the wrong state.
+    """
+    hash_block_size = 16
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=32,
+        num_prefill_checkpoint_blocks=1,
+    )
+    manager.coordinator.retention_interval = 0
+
+    request = make_request("producer", list(range(240)), hash_block_size, sha256)
+    # Block 2 ends at the shared-prefix junction 96, so the full-block pass
+    # hashes it as a retained boundary.
+    request.shared_prefix_boundary = 96
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
+    assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
+
+    # The same slot is the transient checkpoint block (state@112): both the
+    # boundary hash and the checkpoint itself must stay unpublished.
+    boundary_hash = request.block_hashes[96 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(boundary_hash, [1]) is None
+    checkpoint_hash = request.block_hashes[112 // hash_block_size - 1]
+    assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None
+
+
 def test_eagle_block_aligned_checkpoint_replaces_newer_hash():
     hash_block_size = mamba_block_size = 32
     manager = make_full_mamba_manager(

vllm/v1/core/single_type_kv_cache_manager.py

@@ -2003,7 +2003,9 @@ class MambaManager(SingleTypeKVCacheManager):
         )
         num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
         if self.mamba_cache_mode == "align":
-            partial_hash = self._cache_partial_tail_block(request, num_tokens)
+            partial_hash = self._cache_partial_tail_block(
+                request, num_tokens, retention_interval=retention_interval
+            )
             if partial_hash is not None:
                 self.cached_blocks_this_step.add(partial_hash)
         if num_cached_blocks_after > num_cached_blocks_before:
@@ -2038,6 +2040,7 @@ class MambaManager(SingleTypeKVCacheManager):
         self,
         request: Request,
         num_tokens: int,
+        retention_interval: int | None,
     ) -> BlockHashWithGroupId | None:
         hash_block_size = self.block_pool.hash_block_size
         # Re-key the reserved block at its exported checkpoint boundary.
@@ -2047,6 +2050,17 @@ class MambaManager(SingleTypeKVCacheManager):
             blocks = self.req_to_blocks[request.request_id]
             assert 0 <= checkpoint_idx < len(blocks)
             checkpoint_block = blocks[checkpoint_idx]
+            if (
+                retention_interval == 0
+                and num_tokens < request.num_prompt_tokens
+                and checkpoint_position != request.shared_prefix_boundary
+            ):
+                # retention_interval == 0 keeps this transient checkpoint
+                # request-local. The slot may carry a hash from this step's
+                # full-block pass; that must go too, since the checkpoint
+                # state is about to overwrite the block.
+                self.block_pool._maybe_evict_cached_block(checkpoint_block)
+                return None
             if checkpoint_block.block_hash_num_tokens == checkpoint_position:
                 return None
             return self.block_pool.cache_partial_block(