[Bugfix][Mamba] Keep the prompt-end prefill checkpoint under sparse retention (#59146)
Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Jared Wen committed
Sep 30, 2026 at 05:00 UTC
d882bddbeab6b4a0d5861dfcb171bf61ce2109d6
2 files changed
+46
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+34
@@ -2531,3 +2531,37 @@ def test_boundary_states_offered_past_prompt_for_resumed_prefill():
2531
offered = [b for _, _, b in drain_boundary_state_offloads(manager).get("0", [])]
2532
assert 2 * block_size in offered
2533
assert req0.num_prompt_tokens < 2 * block_size
2534
+
2535
+
2536
+def test_prompt_end_checkpoint_survives_sparse_retention_with_coarse_hash():
2537
+ """With hash_block_size == mamba block size (no fine-grained matching), the
2538
+ hybrid coordinator caches the final chunk only up to its last full block,
2539
+ so the managers see ``num_tokens < num_prompt_tokens`` there too. The
2540
+ prompt-end checkpoint must still be published under retention_interval=0,
2541
+ otherwise the next request that extends this prompt misses the prefix.
2542
+ """
2543
+ block_size = 32
2544
+ manager = make_full_mamba_manager(
2545
+ dcp_world_size=1,
2546
+ hash_block_size=block_size,
2547
+ full_block_size=block_size,
2548
+ mamba_block_size=block_size,
2549
+ num_prefill_checkpoint_blocks=1,
2550
+ )
2551
+ manager.coordinator.retention_interval = 0
2552
+
2553
+ # 240 tokens: not block aligned, prompt-end checkpoint at 224.
2554
+ request = make_request("producer", list(range(240)), block_size, sha256)
2555
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
2556
+ assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
2557
+ request.num_computed_tokens = 128
2558
+ manager.new_step_starts()
2559
+ assert manager.allocate_slots(request, 112) is not None
2560
+
2561
+ end_hash = request.block_hashes[224 // block_size - 1]
2562
+ assert manager.block_pool.get_cached_block(end_hash, [1]) is not None
2563
+
2564
+ # A request extending the prompt resumes from that checkpoint.
2565
+ follower = make_request("follower", list(range(240)) + [7] * 40, block_size, sha256)
2566
+ _, num_hit, _ = manager.get_computed_blocks(follower)
2567
+ assert num_hit == 224
vllm/v1/core/single_type_kv_cache_manager.py
+12
@@ -2042,9 +2042,21 @@ class MambaManager(SingleTypeKVCacheManager):
2042
blocks = self.req_to_blocks[request.request_id]
2043
assert 0 <= checkpoint_idx < len(blocks)
2044
checkpoint_block = blocks[checkpoint_idx]
2045
+ # The prompt's final chunk caches only up to its last full block,
2046
+ # so ``num_tokens < num_prompt_tokens`` holds there too; its
2047
+ # checkpoint is the prompt's reusable boundary, not a transient one.
2048
+ is_prompt_checkpoint = (
2049
+ checkpoint_position
2050
+ == get_mamba_prefill_checkpoint_position(
2051
+ request.num_prompt_tokens,
2052
+ hash_block_size,
2053
+ self.drop_eagle_checkpoint_block,
2054
+ )
2055
+ )
2056
if (
2057
retention_interval == 0
2058
and num_tokens < request.num_prompt_tokens
2059
+ and not is_prompt_checkpoint
2060
and checkpoint_position != request.shared_prefix_boundary
2061
):
2062
# retention_interval == 0 keeps this transient checkpoint