[Bugfix][Core] Fix mamba prefill checkpoint block reservation and prompt-end eviction in align mode (#59175)

Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -13,6 +13,9 @@ import torch
 
 from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request
 from vllm.distributed.kv_transfer.kv_connector.v1.base import SupportsHMA
+from vllm.model_executor.layers.mamba.checkpoint import (
+    compute_mamba_prefill_checkpoints,
+)
 from vllm.utils.hashing import sha256
 from vllm.v1.core.kv_cache_utils import (
     KVCacheBlockCopy,
@@ -28,6 +31,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheGroupSpec,
     MambaSpec,
     SlidingWindowSpec,
+    get_mamba_prefill_checkpoint_position,
 )
 
 
@@ -686,6 +690,81 @@ def test_transient_checkpoint_evicts_retained_boundary_hash():
     assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None
 
 
+@pytest.mark.parametrize("prompt_len", [120, 104, 96])
+def test_prompt_end_checkpoint_survives_block_aligned_cache_count(prompt_len):
+    """At the default match unit the hybrid coordinator hands ``cache_blocks``
+    a block-aligned token count, which is below the prompt length whenever the
+    prompt is not block-aligned. The prompt-end checkpoint is the replay
+    boundary and must stay published under retention_interval=0; only a
+    mid-prompt checkpoint is transient.
+    """
+    hash_block_size = mamba_block_size = 32
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=mamba_block_size,
+        num_prefill_checkpoint_blocks=1,
+    )
+    manager.coordinator.retention_interval = 0
+    assert not manager.coordinator.enable_partial_hash_hits
+
+    request = make_request("producer", list(range(prompt_len)), hash_block_size, sha256)
+    assert manager.allocate_slots(request, request.num_tokens) is not None
+    checkpoint_position = get_mamba_prefill_checkpoint_position(
+        prompt_len, hash_block_size, drop_eagle_block=False
+    )
+    checkpoint_hash = request.block_hashes[checkpoint_position // hash_block_size - 1]
+    hit = manager.block_pool.get_cached_block(checkpoint_hash, [1])
+    assert hit is not None
+    assert hit[0].block_hash_num_tokens == checkpoint_position
+
+    manager.free(request)
+    replay = make_request("replay", list(range(prompt_len)), hash_block_size, sha256)
+    _, num_computed, _ = manager.get_computed_blocks(replay)
+    assert num_computed == checkpoint_position
+
+
+@pytest.mark.parametrize("retention_interval,reserved", [(None, True), (0, False)])
+def test_intermediate_chunk_checkpoint_reserved_only_under_dense_retention(
+    retention_interval, reserved
+):
+    """A checkpoint inside a chunk that does not finish the prefill is only
+    publishable under dense retention. With retention_interval=0 it would be
+    evicted as transient, so the block is not reserved in the first place;
+    the prefill-end checkpoint is reserved either way.
+    """
+    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
+    mamba_manager = manager.coordinator.single_type_managers[1]
+
+    request = make_request("producer", list(range(240)), hash_block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(request)
+    free_before = manager.block_pool.get_num_free_blocks()
+
+    # Mid-prompt chunk: the chunk-keyed position (112) is inside it. Blocks:
+    # the attention blocks, the running mamba block, plus the checkpoint.
+    assert manager.allocate_slots(request, 128, num_computed, computed_blocks)
+    assert (request.request_id in mamba_manager._checkpoints) == reserved
+    allocated = free_before - manager.block_pool.get_num_free_blocks()
+    assert allocated == 128 // hash_block_size + 1 + int(reserved)
+
+    request.num_computed_tokens = 128
+    manager.new_step_starts()
+    # Final chunk: the prefill-end checkpoint at 224 is reserved either way.
+    assert manager.allocate_slots(request, 112) is not None
+    assert mamba_manager._checkpoints[request.request_id][0] == 224
+    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_eagle_block_aligned_checkpoint_replaces_newer_hash():
     hash_block_size = mamba_block_size = 32
     manager = make_full_mamba_manager(
@@ -2565,3 +2644,71 @@ def test_prompt_end_checkpoint_survives_sparse_retention_with_coarse_hash():
     follower = make_request("follower", list(range(240)) + [7] * 40, block_size, sha256)
     _, num_hit, _ = manager.get_computed_blocks(follower)
     assert num_hit == 224
+
+
+@pytest.mark.parametrize(
+    "hash_block_size,mamba_block_size,use_eagle",
+    [(32, 32, False), (16, 32, False), (32, 32, True)],
+)
+@pytest.mark.parametrize("retention_interval", [0, None])
+@pytest.mark.parametrize(
+    "num_prompt_tokens,num_output_tokens",
+    [(104, 0), (200, 0), (40, 24), (40, 25), (150, 43)],
+)
+@pytest.mark.parametrize("chunk_blocks", [2, 100])
+def test_checkpoint_reservation_matches_worker(
+    hash_block_size,
+    mamba_block_size,
+    use_eagle,
+    retention_interval,
+    num_prompt_tokens,
+    num_output_tokens,
+    chunk_blocks,
+):
+    """The worker exports a checkpoint wherever `compute_mamba_prefill_checkpoints`
+    says, into whatever block sits at its column. Each chunk, the manager must
+    either reserve exactly that position and column, or leave the column null so
+    the export is masked; anything else publishes a hash over the wrong state.
+    Under sparse retention, only the prefill-end chunk may reserve one.
+    """
+    manager = make_full_mamba_manager(
+        dcp_world_size=1,
+        hash_block_size=hash_block_size,
+        full_block_size=hash_block_size,
+        mamba_block_size=mamba_block_size,
+        num_blocks=64,
+        use_eagle=use_eagle,
+        num_speculative_blocks=3 if use_eagle else 0,
+        num_prefill_checkpoint_blocks=1,
+    )
+    manager.coordinator.retention_interval = retention_interval
+    mamba_manager = manager.coordinator.single_type_managers[1]
+
+    request = make_request(
+        "req", list(range(num_prompt_tokens)), hash_block_size, sha256
+    )
+    request.append_output_token_ids(list(range(1000, 1000 + num_output_tokens)))
+    num_tokens = request.num_tokens
+    prefill_end = max(num_prompt_tokens, num_tokens - 1)
+
+    start = 0
+    while start < num_tokens:
+        end = min(start + chunk_blocks * mamba_block_size, num_tokens)
+        assert manager.allocate_slots(request, end - start) is not None
+        offsets, cols = compute_mamba_prefill_checkpoints(
+            [end],
+            [end - start],
+            hash_block_size=hash_block_size,
+            mamba_block_size=mamba_block_size,
+            checkpoint_alignment=16,
+            drop_eagle_block=use_eagle,
+        )
+        reserved = mamba_manager._checkpoints.get(request.request_id)
+        if reserved is not None:
+            assert reserved == (start + offsets[0], cols[0])
+            assert retention_interval is None or end >= prefill_end
+        elif offsets[0]:
+            assert mamba_manager.req_to_blocks[request.request_id][cols[0]].is_null
+        request.num_computed_tokens = end
+        manager.new_step_starts()
+        start = end

vllm/v1/core/kv_cache_coordinator.py

@@ -174,6 +174,7 @@ class KVCacheCoordinator(ABC):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         """Get the number of device blocks needed to be allocated for the request.
 
@@ -195,12 +196,19 @@ class KVCacheCoordinator(ABC):
                 per-request admission cap (SWA / chunked-local). Set only by
                 the full-sequence admission gate; per-step allocation must
                 leave it False so the predictor matches `allocate_new_blocks`.
+            prefill_end: The token index the request's prefill ends at, the
+                same value the scheduler splits chunks against. Under sparse
+                retention Mamba reserves a prefill checkpoint only on the chunk
+                reaching it; under dense retention every chunk publishes a
+                state, so it is ignored (0).
 
         Returns:
             The number of blocks to allocate.
 
         """
         num_blocks_to_allocate = 0
+        if self.retention_interval != 0:
+            prefill_end = 0
         for i, manager in enumerate(self.single_type_managers):
             if isinstance(manager, CrossAttentionManager):
                 # For cross-attention, we issue a single static allocation
@@ -223,6 +231,7 @@ class KVCacheCoordinator(ABC):
                     num_local_computed_tokens,
                     num_tokens_main_model,
                     apply_admission_cap=apply_admission_cap,
+                    prefill_end=prefill_end,
                 )
         return num_blocks_to_allocate
 

vllm/v1/core/kv_cache_manager.py

@@ -512,6 +512,10 @@ class KVCacheManager:
         ):
             watermark_blocks = self.watermark_blocks
 
+        # Matches the scheduler's own prefill boundary: `num_tokens - 1`
+        # extends it to resumed requests replaying their output tokens.
+        prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1)
+
         if full_sequence_must_fit:
             # First check and fail if the full request sequence won't fit.
             full_num_tokens = min(request.num_tokens, self.max_model_len)
@@ -525,6 +529,7 @@ class KVCacheManager:
                 num_local_computed_tokens=num_local_computed_tokens,
                 num_tokens_main_model=full_num_tokens,
                 apply_admission_cap=True,
+                prefill_end=prefill_end,
             )
             required_blocks = num_blocks_to_allocate + watermark_blocks
             if required_blocks > self.block_pool.get_num_free_blocks():
@@ -559,6 +564,7 @@ class KVCacheManager:
             + num_external_computed_tokens,
             num_local_computed_tokens=num_local_computed_tokens,
             num_tokens_main_model=num_tokens_main_model,
+            prefill_end=prefill_end,
         )
 
         # Keep `reserved_blocks` free for other in-flight sequences, and an

vllm/v1/core/sched/scheduler.py

@@ -3003,6 +3003,7 @@ class Scheduler(SchedulerInterface):
             num_local_computed_tokens=request.num_computed_tokens,
             num_tokens_main_model=full_num_tokens,
             apply_admission_cap=True,
+            prefill_end=max(request.num_prompt_tokens, request.num_tokens - 1),
         )
         return num_blocks + self._spec_decode_step_blocks()
 

vllm/v1/core/single_type_kv_cache_manager.py

@@ -184,6 +184,7 @@ class SingleTypeKVCacheManager(ABC):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         """Get the number of blocks needed to be allocated for the request.
 
@@ -203,6 +204,10 @@ class SingleTypeKVCacheManager(ABC):
             apply_admission_cap: If True, clamp by `num_required_blocks` by
                 `_max_admission_blocks_per_request`for recycling-aware specs
                 (SWA, chunked-local).
+            prefill_end: The token index the request's prefill ends at, the
+                same value the scheduler splits chunks against. Mamba reserves
+                a prefill checkpoint only on the chunk reaching it; 0 (dense
+                retention) reserves one on every chunk.
 
         Returns:
             The number of blocks to allocate.
@@ -1215,6 +1220,7 @@ class CircularBufferManager(FullAttentionManager):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         return 0 if self.req_to_blocks.get(request_id) else 1
 
@@ -1678,12 +1684,14 @@ class MambaManager(SingleTypeKVCacheManager):
         query_start: int,
         query_end: int,
         checkpoint_position: int,
+        prefill_end: int,
     ) -> bool:
         assert isinstance(self.kv_cache_spec, MambaSpec)
         checkpoint_idx = cdiv(query_end, self.block_size) - 2
         blocks = self.req_to_blocks[request_id]
         return (
             self.has_prefill_checkpoint_blocks
+            and query_end >= prefill_end
             and is_mamba_prefill_checkpoint_valid(
                 query_start=query_start,
                 query_end=query_end,
@@ -1712,6 +1720,7 @@ class MambaManager(SingleTypeKVCacheManager):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         assert isinstance(self.kv_cache_spec, MambaSpec)
         if (
@@ -1738,6 +1747,7 @@ class MambaManager(SingleTypeKVCacheManager):
                 num_local_computed_tokens,
                 num_tokens_main_model,
                 apply_admission_cap=apply_admission_cap,
+                prefill_end=prefill_end,
             )
         else:
             # We don't allocate blocks for lookahead tokens in align mode, because if
@@ -1765,6 +1775,11 @@ class MambaManager(SingleTypeKVCacheManager):
             )
             if has_partial_hit:
                 num_new_blocks = max(num_new_blocks, 0) + 1
+            # Keyed on the chunk end like the worker's
+            # `compute_mamba_prefill_checkpoints`, which writes the state there.
+            # The helper returns a boundary strictly below its input, so every
+            # chunk has one; under sparse retention only the prefill-end chunk's
+            # is kept, so `prefill_end` skips reserving blocks for the others.
             checkpoint_position = get_mamba_prefill_checkpoint_position(
                 num_tokens,
                 self.block_pool.hash_block_size,
@@ -1775,6 +1790,7 @@ class MambaManager(SingleTypeKVCacheManager):
                 total_computed_tokens,
                 num_tokens,
                 checkpoint_position,
+                prefill_end,
             ):
                 checkpoint_position = 0
             checkpoint_block = int(checkpoint_position > 0)
@@ -2272,6 +2288,7 @@ class HiSparseSourceManager(FullAttentionManager):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         required = super().get_num_blocks_to_allocate(
             request_id,
@@ -2407,6 +2424,7 @@ class HiSparseHotManager(_HiSparseAuxiliaryManager):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         assert not new_computed_blocks
         # A hot region is needed to read host-backed history: an external
@@ -2479,6 +2497,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):
         num_local_computed_tokens: int,
         num_tokens_main_model: int,
         apply_admission_cap: bool = False,
+        prefill_end: int = 0,
     ) -> int:
         del num_tokens_main_model
         assert not new_computed_blocks