From 7e583e615c20ee4ff0cd82aac592c7d6310aa7c3 Mon Sep 17 00:00:00 2001 From: kliuae <17350011+kliuae@users.noreply.github.com> Date: Thu, 1 Oct 2026 05:26:54 +0800 Subject: [PATCH] [Bugfix][Core] Fix mamba prefill checkpoint block reservation and prompt-end eviction in align mode (#59175) Signed-off-by: kliuae Co-authored-by: Nick Hill --- .../test_partial_prefix_cache_hits.py | 147 ++++++++++++++++++ vllm/v1/core/kv_cache_coordinator.py | 9 ++ vllm/v1/core/kv_cache_manager.py | 6 + vllm/v1/core/sched/scheduler.py | 1 + vllm/v1/core/single_type_kv_cache_manager.py | 19 +++ 5 files changed, 182 insertions(+) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 3b5fd218a6..76beffce45 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/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 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index d69b2f5fb3..daf0cd12c0 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/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 diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 30bf584f8c..2d39a38651 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/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 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index fc924b876c..494dba5e46 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/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() diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 3d784891c5..c8b9c54435 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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