[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>
Sy03 committed
Sep 14, 2026 at 23:18 UTC
f547c23ec971fd1c78f0d89d15ac1957ae407471
3 files changed
+94
-2
tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py
+8
-1
@@ -299,7 +299,14 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing():
299
running = request.num_computed_tokens
300
assert not mamba.req_to_blocks["r"][running // block_size].is_null
301
request.shared_prefix_boundary = running
302
- assert mamba._cache_partial_tail_block(request, running) is None
302
+ assert (
303
+ mamba._cache_partial_tail_block(
304
+ request,
305
+ running,
306
+ retention_interval=manager.coordinator.retention_interval,
307
+ )
308
+ is None
309
+ )
310
311
312
@pytest.mark.parametrize("eagle_group", [None, 0])
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+71
@@ -615,6 +615,77 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle():
615
assert num_computed == 112
616
617
618
+@pytest.mark.parametrize(
619
+ "retention_interval,transient_published",
620
+ [(None, True), (0, False)],
621
+)
622
+def test_internal_checkpoint_publication_respects_retention(
623
+ retention_interval, transient_published
624
+):
625
+ """With retention_interval=0, a mid-prompt internal checkpoint is
626
+ request-local and must not enter the prefix cache; the prompt-end
627
+ checkpoint keeps the existing publication behavior. Dense retention
628
+ (None) is unchanged.
629
+ """
630
+ hash_block_size = 16
631
+ manager = make_full_mamba_manager(
632
+ dcp_world_size=1,
633
+ hash_block_size=hash_block_size,
634
+ full_block_size=hash_block_size,
635
+ mamba_block_size=32,
636
+ num_prefill_checkpoint_blocks=1,
637
+ )
638
+ manager.coordinator.retention_interval = retention_interval
639
+
640
+ request = make_request("producer", list(range(240)), hash_block_size, sha256)
641
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
642
+
643
+ # Mid-prompt chunk: exports a transient checkpoint at 112.
644
+ assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
645
+ transient_hash = request.block_hashes[112 // hash_block_size - 1]
646
+ transient_hit = manager.block_pool.get_cached_block(transient_hash, [1])
647
+ assert (transient_hit is not None) == transient_published
648
+
649
+ request.num_computed_tokens = 128
650
+ manager.new_step_starts()
651
+
652
+ # Final chunk: the prompt-end checkpoint at 224 stays published.
653
+ assert manager.allocate_slots(request, 112) is not None
654
+ end_hash = request.block_hashes[224 // hash_block_size - 1]
655
+ assert manager.block_pool.get_cached_block(end_hash, [1]) is not None
656
+
657
+
658
+def test_transient_checkpoint_evicts_retained_boundary_hash():
659
+ """With retention_interval=0, the checkpoint slot can coincide with a
660
+ reachable boundary block that this step's full-block pass just hashed.
661
+ The checkpoint state overwrites the slot, so the stale boundary hash
662
+ must be evicted rather than left pointing at the wrong state.
663
+ """
664
+ hash_block_size = 16
665
+ manager = make_full_mamba_manager(
666
+ dcp_world_size=1,
667
+ hash_block_size=hash_block_size,
668
+ full_block_size=hash_block_size,
669
+ mamba_block_size=32,
670
+ num_prefill_checkpoint_blocks=1,
671
+ )
672
+ manager.coordinator.retention_interval = 0
673
+
674
+ request = make_request("producer", list(range(240)), hash_block_size, sha256)
675
+ # Block 2 ends at the shared-prefix junction 96, so the full-block pass
676
+ # hashes it as a retained boundary.
677
+ request.shared_prefix_boundary = 96
678
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
679
+ assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
680
+
681
+ # The same slot is the transient checkpoint block (state@112): both the
682
+ # boundary hash and the checkpoint itself must stay unpublished.
683
+ boundary_hash = request.block_hashes[96 // hash_block_size - 1]
684
+ assert manager.block_pool.get_cached_block(boundary_hash, [1]) is None
685
+ checkpoint_hash = request.block_hashes[112 // hash_block_size - 1]
686
+ assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None
687
+
688
+
689
def test_eagle_block_aligned_checkpoint_replaces_newer_hash():
690
hash_block_size = mamba_block_size = 32
691
manager = make_full_mamba_manager(
vllm/v1/core/single_type_kv_cache_manager.py
+15
-1
@@ -2003,7 +2003,9 @@ class MambaManager(SingleTypeKVCacheManager):
2003
)
2004
num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
2005
if self.mamba_cache_mode == "align":
2006
- partial_hash = self._cache_partial_tail_block(request, num_tokens)
2006
+ partial_hash = self._cache_partial_tail_block(
2007
+ request, num_tokens, retention_interval=retention_interval
2008
+ )
2009
if partial_hash is not None:
2010
self.cached_blocks_this_step.add(partial_hash)
2011
if num_cached_blocks_after > num_cached_blocks_before:
@@ -2038,6 +2040,7 @@ class MambaManager(SingleTypeKVCacheManager):
2040
self,
2041
request: Request,
2042
num_tokens: int,
2043
+ retention_interval: int | None,
2044
) -> BlockHashWithGroupId | None:
2045
hash_block_size = self.block_pool.hash_block_size
2046
# Re-key the reserved block at its exported checkpoint boundary.
@@ -2047,6 +2050,17 @@ class MambaManager(SingleTypeKVCacheManager):
2050
blocks = self.req_to_blocks[request.request_id]
2051
assert 0 <= checkpoint_idx < len(blocks)
2052
checkpoint_block = blocks[checkpoint_idx]
2053
+ if (
2054
+ retention_interval == 0
2055
+ and num_tokens < request.num_prompt_tokens
2056
+ and checkpoint_position != request.shared_prefix_boundary
2057
+ ):
2058
+ # retention_interval == 0 keeps this transient checkpoint
2059
+ # request-local. The slot may carry a hash from this step's
2060
+ # full-block pass; that must go too, since the checkpoint
2061
+ # state is about to overwrite the block.
2062
+ self.block_pool._maybe_evict_cached_block(checkpoint_block)
2063
+ return None
2064
if checkpoint_block.block_hash_num_tokens == checkpoint_position:
2065
return None
2066
return self.block_pool.cache_partial_block(