@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7e583e615c

[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>

kliuae committed Oct 1, 2026 at 05:26 UTC 7e583e615c20ee4ff0cd82aac592c7d6310aa7c3

5 files changed +182

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

+147

@@ -13,6 +13,9 @@ import torch

14 from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request

15 from vllm.distributed.kv_transfer.kv_connector.v1.base import SupportsHMA

16 +from vllm.model_executor.layers.mamba.checkpoint import (

17 + compute_mamba_prefill_checkpoints,

18 +)

19 from vllm.utils.hashing import sha256

20 from vllm.v1.core.kv_cache_utils import (

21 KVCacheBlockCopy,

@@ -28,6 +31,7 @@ from vllm.v1.kv_cache_interface import (

31 KVCacheGroupSpec,

32 MambaSpec,

33 SlidingWindowSpec,

34 + get_mamba_prefill_checkpoint_position,

35 )

@@ -686,6 +690,81 @@ def test_transient_checkpoint_evicts_retained_boundary_hash():

690 assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None

691

692

693 +@pytest.mark.parametrize("prompt_len", [120, 104, 96])

694 +def test_prompt_end_checkpoint_survives_block_aligned_cache_count(prompt_len):

695 + """At the default match unit the hybrid coordinator hands ``cache_blocks``

696 + a block-aligned token count, which is below the prompt length whenever the

697 + prompt is not block-aligned. The prompt-end checkpoint is the replay

698 + boundary and must stay published under retention_interval=0; only a

699 + mid-prompt checkpoint is transient.

700 + """

701 + hash_block_size = mamba_block_size = 32

702 + manager = make_full_mamba_manager(

703 + dcp_world_size=1,

704 + hash_block_size=hash_block_size,

705 + full_block_size=hash_block_size,

706 + mamba_block_size=mamba_block_size,

707 + num_prefill_checkpoint_blocks=1,

708 + )

709 + manager.coordinator.retention_interval = 0

710 + assert not manager.coordinator.enable_partial_hash_hits

711 +

712 + request = make_request("producer", list(range(prompt_len)), hash_block_size, sha256)

713 + assert manager.allocate_slots(request, request.num_tokens) is not None

714 + checkpoint_position = get_mamba_prefill_checkpoint_position(

715 + prompt_len, hash_block_size, drop_eagle_block=False

716 + )

717 + checkpoint_hash = request.block_hashes[checkpoint_position // hash_block_size - 1]

718 + hit = manager.block_pool.get_cached_block(checkpoint_hash, [1])

719 + assert hit is not None

720 + assert hit[0].block_hash_num_tokens == checkpoint_position

721 +

722 + manager.free(request)

723 + replay = make_request("replay", list(range(prompt_len)), hash_block_size, sha256)

724 + _, num_computed, _ = manager.get_computed_blocks(replay)

725 + assert num_computed == checkpoint_position

726 +

727 +

728 +@pytest.mark.parametrize("retention_interval,reserved", [(None, True), (0, False)])

729 +def test_intermediate_chunk_checkpoint_reserved_only_under_dense_retention(

730 + retention_interval, reserved

731 +):

732 + """A checkpoint inside a chunk that does not finish the prefill is only

733 + publishable under dense retention. With retention_interval=0 it would be

734 + evicted as transient, so the block is not reserved in the first place;

735 + the prefill-end checkpoint is reserved either way.

736 + """

737 + hash_block_size = 16

738 + manager = make_full_mamba_manager(

739 + dcp_world_size=1,

740 + hash_block_size=hash_block_size,

741 + full_block_size=hash_block_size,

742 + mamba_block_size=32,

743 + num_prefill_checkpoint_blocks=1,

744 + )

745 + manager.coordinator.retention_interval = retention_interval

746 + mamba_manager = manager.coordinator.single_type_managers[1]

747 +

748 + request = make_request("producer", list(range(240)), hash_block_size, sha256)

749 + computed_blocks, num_computed, _ = manager.get_computed_blocks(request)

750 + free_before = manager.block_pool.get_num_free_blocks()

751 +

752 + # Mid-prompt chunk: the chunk-keyed position (112) is inside it. Blocks:

753 + # the attention blocks, the running mamba block, plus the checkpoint.

754 + assert manager.allocate_slots(request, 128, num_computed, computed_blocks)

755 + assert (request.request_id in mamba_manager._checkpoints) == reserved

756 + allocated = free_before - manager.block_pool.get_num_free_blocks()

757 + assert allocated == 128 // hash_block_size + 1 + int(reserved)

758 +

759 + request.num_computed_tokens = 128

760 + manager.new_step_starts()

761 + # Final chunk: the prefill-end checkpoint at 224 is reserved either way.

762 + assert manager.allocate_slots(request, 112) is not None

763 + assert mamba_manager._checkpoints[request.request_id][0] == 224

764 + end_hash = request.block_hashes[224 // hash_block_size - 1]

765 + assert manager.block_pool.get_cached_block(end_hash, [1]) is not None

766 +

767 +

768 def test_eagle_block_aligned_checkpoint_replaces_newer_hash():

769 hash_block_size = mamba_block_size = 32

770 manager = make_full_mamba_manager(

@@ -2565,3 +2644,71 @@ def test_prompt_end_checkpoint_survives_sparse_retention_with_coarse_hash():

2644 follower = make_request("follower", list(range(240)) + [7] * 40, block_size, sha256)

2645 _, num_hit, _ = manager.get_computed_blocks(follower)

2646 assert num_hit == 224

2647 +

2648 +

2649 +@pytest.mark.parametrize(

2650 + "hash_block_size,mamba_block_size,use_eagle",

2651 + [(32, 32, False), (16, 32, False), (32, 32, True)],

2652 +)

2653 +@pytest.mark.parametrize("retention_interval", [0, None])

2654 +@pytest.mark.parametrize(

2655 + "num_prompt_tokens,num_output_tokens",

2656 + [(104, 0), (200, 0), (40, 24), (40, 25), (150, 43)],

2657 +)

2658 +@pytest.mark.parametrize("chunk_blocks", [2, 100])

2659 +def test_checkpoint_reservation_matches_worker(

2660 + hash_block_size,

2661 + mamba_block_size,

2662 + use_eagle,

2663 + retention_interval,

2664 + num_prompt_tokens,

2665 + num_output_tokens,

2666 + chunk_blocks,

2667 +):

2668 + """The worker exports a checkpoint wherever `compute_mamba_prefill_checkpoints`

2669 + says, into whatever block sits at its column. Each chunk, the manager must

2670 + either reserve exactly that position and column, or leave the column null so

2671 + the export is masked; anything else publishes a hash over the wrong state.

2672 + Under sparse retention, only the prefill-end chunk may reserve one.

2673 + """

2674 + manager = make_full_mamba_manager(

2675 + dcp_world_size=1,

2676 + hash_block_size=hash_block_size,

2677 + full_block_size=hash_block_size,

2678 + mamba_block_size=mamba_block_size,

2679 + num_blocks=64,

2680 + use_eagle=use_eagle,

2681 + num_speculative_blocks=3 if use_eagle else 0,

2682 + num_prefill_checkpoint_blocks=1,

2683 + )

2684 + manager.coordinator.retention_interval = retention_interval

2685 + mamba_manager = manager.coordinator.single_type_managers[1]

2686 +

2687 + request = make_request(

2688 + "req", list(range(num_prompt_tokens)), hash_block_size, sha256

2689 + )

2690 + request.append_output_token_ids(list(range(1000, 1000 + num_output_tokens)))

2691 + num_tokens = request.num_tokens

2692 + prefill_end = max(num_prompt_tokens, num_tokens - 1)

2693 +

2694 + start = 0

2695 + while start < num_tokens:

2696 + end = min(start + chunk_blocks * mamba_block_size, num_tokens)

2697 + assert manager.allocate_slots(request, end - start) is not None

2698 + offsets, cols = compute_mamba_prefill_checkpoints(

2699 + [end],

2700 + [end - start],

2701 + hash_block_size=hash_block_size,

2702 + mamba_block_size=mamba_block_size,

2703 + checkpoint_alignment=16,

2704 + drop_eagle_block=use_eagle,

2705 + )

2706 + reserved = mamba_manager._checkpoints.get(request.request_id)

2707 + if reserved is not None:

2708 + assert reserved == (start + offsets[0], cols[0])

2709 + assert retention_interval is None or end >= prefill_end

2710 + elif offsets[0]:

2711 + assert mamba_manager.req_to_blocks[request.request_id][cols[0]].is_null

2712 + request.num_computed_tokens = end

2713 + manager.new_step_starts()

2714 + start = end

vllm/v1/core/kv_cache_coordinator.py

@@ -174,6 +174,7 @@ class KVCacheCoordinator(ABC):

174 num_local_computed_tokens: int,

175 num_tokens_main_model: int,

176 apply_admission_cap: bool = False,

177 + prefill_end: int = 0,

178 ) -> int:

179 """Get the number of device blocks needed to be allocated for the request.

180

@@ -195,12 +196,19 @@ class KVCacheCoordinator(ABC):

196 per-request admission cap (SWA / chunked-local). Set only by

197 the full-sequence admission gate; per-step allocation must

198 leave it False so the predictor matches `allocate_new_blocks`.

199 + prefill_end: The token index the request's prefill ends at, the

200 + same value the scheduler splits chunks against. Under sparse

201 + retention Mamba reserves a prefill checkpoint only on the chunk

202 + reaching it; under dense retention every chunk publishes a

203 + state, so it is ignored (0).

204

205 Returns:

206 The number of blocks to allocate.

207

208 """

209 num_blocks_to_allocate = 0

210 + if self.retention_interval != 0:

211 + prefill_end = 0

212 for i, manager in enumerate(self.single_type_managers):

213 if isinstance(manager, CrossAttentionManager):

214 # For cross-attention, we issue a single static allocation

@@ -223,6 +231,7 @@ class KVCacheCoordinator(ABC):

231 num_local_computed_tokens,

232 num_tokens_main_model,

233 apply_admission_cap=apply_admission_cap,

234 + prefill_end=prefill_end,

235 )

236 return num_blocks_to_allocate

237

vllm/v1/core/kv_cache_manager.py

@@ -512,6 +512,10 @@ class KVCacheManager:

512 ):

513 watermark_blocks = self.watermark_blocks

514

515 + # Matches the scheduler's own prefill boundary: `num_tokens - 1`

516 + # extends it to resumed requests replaying their output tokens.

517 + prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1)

518 +

519 if full_sequence_must_fit:

520 # First check and fail if the full request sequence won't fit.

521 full_num_tokens = min(request.num_tokens, self.max_model_len)

@@ -525,6 +529,7 @@ class KVCacheManager:

529 num_local_computed_tokens=num_local_computed_tokens,

530 num_tokens_main_model=full_num_tokens,

531 apply_admission_cap=True,

532 + prefill_end=prefill_end,

533 )

534 required_blocks = num_blocks_to_allocate + watermark_blocks

535 if required_blocks > self.block_pool.get_num_free_blocks():

@@ -559,6 +564,7 @@ class KVCacheManager:

564 + num_external_computed_tokens,

565 num_local_computed_tokens=num_local_computed_tokens,

566 num_tokens_main_model=num_tokens_main_model,

567 + prefill_end=prefill_end,

568 )

569

570 # Keep `reserved_blocks` free for other in-flight sequences, and an

vllm/v1/core/sched/scheduler.py

@@ -3003,6 +3003,7 @@ class Scheduler(SchedulerInterface):

3003 num_local_computed_tokens=request.num_computed_tokens,

3004 num_tokens_main_model=full_num_tokens,

3005 apply_admission_cap=True,

3006 + prefill_end=max(request.num_prompt_tokens, request.num_tokens - 1),

3007 )

3008 return num_blocks + self._spec_decode_step_blocks()

3009

vllm/v1/core/single_type_kv_cache_manager.py

+19

@@ -184,6 +184,7 @@ class SingleTypeKVCacheManager(ABC):

184 num_local_computed_tokens: int,

185 num_tokens_main_model: int,

186 apply_admission_cap: bool = False,

187 + prefill_end: int = 0,

188 ) -> int:

189 """Get the number of blocks needed to be allocated for the request.

190

@@ -203,6 +204,10 @@ class SingleTypeKVCacheManager(ABC):

204 apply_admission_cap: If True, clamp by `num_required_blocks` by

205 `_max_admission_blocks_per_request`for recycling-aware specs

206 (SWA, chunked-local).

207 + prefill_end: The token index the request's prefill ends at, the

208 + same value the scheduler splits chunks against. Mamba reserves

209 + a prefill checkpoint only on the chunk reaching it; 0 (dense

210 + retention) reserves one on every chunk.

211

212 Returns:

213 The number of blocks to allocate.

@@ -1215,6 +1220,7 @@ class CircularBufferManager(FullAttentionManager):

1220 num_local_computed_tokens: int,

1221 num_tokens_main_model: int,

1222 apply_admission_cap: bool = False,

1223 + prefill_end: int = 0,

1224 ) -> int:

1225 return 0 if self.req_to_blocks.get(request_id) else 1

1226

@@ -1678,12 +1684,14 @@ class MambaManager(SingleTypeKVCacheManager):

1684 query_start: int,

1685 query_end: int,

1686 checkpoint_position: int,

1687 + prefill_end: int,

1688 ) -> bool:

1689 assert isinstance(self.kv_cache_spec, MambaSpec)

1690 checkpoint_idx = cdiv(query_end, self.block_size) - 2

1691 blocks = self.req_to_blocks[request_id]

1692 return (

1693 self.has_prefill_checkpoint_blocks

1694 + and query_end >= prefill_end

1695 and is_mamba_prefill_checkpoint_valid(

1696 query_start=query_start,

1697 query_end=query_end,

@@ -1712,6 +1720,7 @@ class MambaManager(SingleTypeKVCacheManager):

1720 num_local_computed_tokens: int,

1721 num_tokens_main_model: int,

1722 apply_admission_cap: bool = False,

1723 + prefill_end: int = 0,

1724 ) -> int:

1725 assert isinstance(self.kv_cache_spec, MambaSpec)

1726 if (

@@ -1738,6 +1747,7 @@ class MambaManager(SingleTypeKVCacheManager):

1747 num_local_computed_tokens,

1748 num_tokens_main_model,

1749 apply_admission_cap=apply_admission_cap,

1750 + prefill_end=prefill_end,

1751 )

1752 else:

1753 # We don't allocate blocks for lookahead tokens in align mode, because if

@@ -1765,6 +1775,11 @@ class MambaManager(SingleTypeKVCacheManager):

1775 )

1776 if has_partial_hit:

1777 num_new_blocks = max(num_new_blocks, 0) + 1

1778 + # Keyed on the chunk end like the worker's

1779 + # `compute_mamba_prefill_checkpoints`, which writes the state there.

1780 + # The helper returns a boundary strictly below its input, so every

1781 + # chunk has one; under sparse retention only the prefill-end chunk's

1782 + # is kept, so `prefill_end` skips reserving blocks for the others.

1783 checkpoint_position = get_mamba_prefill_checkpoint_position(

1784 num_tokens,

1785 self.block_pool.hash_block_size,

@@ -1775,6 +1790,7 @@ class MambaManager(SingleTypeKVCacheManager):

1790 total_computed_tokens,

1791 num_tokens,

1792 checkpoint_position,

1793 + prefill_end,

1794 ):

1795 checkpoint_position = 0

1796 checkpoint_block = int(checkpoint_position > 0)

@@ -2272,6 +2288,7 @@ class HiSparseSourceManager(FullAttentionManager):

2288 num_local_computed_tokens: int,

2289 num_tokens_main_model: int,

2290 apply_admission_cap: bool = False,

2291 + prefill_end: int = 0,

2292 ) -> int:

2293 required = super().get_num_blocks_to_allocate(

2294 request_id,

@@ -2407,6 +2424,7 @@ class HiSparseHotManager(_HiSparseAuxiliaryManager):

2424 num_local_computed_tokens: int,

2425 num_tokens_main_model: int,

2426 apply_admission_cap: bool = False,

2427 + prefill_end: int = 0,

2428 ) -> int:

2429 assert not new_computed_blocks

2430 # A hot region is needed to read host-backed history: an external

@@ -2479,6 +2497,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):

2497 num_local_computed_tokens: int,

2498 num_tokens_main_model: int,

2499 apply_admission_cap: bool = False,

2500 + prefill_end: int = 0,

2501 ) -> int:

2502 del num_tokens_main_model

2503 assert not new_computed_blocks