diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index fd003fdce6..059b4005f0 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/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]) 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 8012c82f78..393ec48465 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 @@ -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( diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index d4e68d8adc..fc5dc41048 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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(