@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b28c3e1568

[BugFix] Retain both replay boundaries so an EAGLE resend of a block-aligned prompt still hits (#54713)

Signed-off-by: tobymao <toby.mao@gmail.com>

Toby Mao committed Sep 10, 2026 at 01:39 UTC b28c3e1568bfae930f61d4b24940e47528c85d4a

4 files changed +130 -31

tests/v1/core/test_prefix_caching.py

+98

@@ -3641,6 +3641,104 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary():

3641 assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 8]

3642

3643

3644 +def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt():

3645 + """An identical resend and a longer sibling resume at DIFFERENT positions.

3646 +

3647 + How far a lookup matches depends on who is asking. A resend of the same

3648 + prompt caps its lookup at ``num_tokens - 1`` (the last token is recomputed

3649 + for logits), while a sibling whose prompt merely starts with this one caps

3650 + above the prompt. The two coincide unless the prompt length is an exact

3651 + multiple of the alignment -- there they differ by one alignment unit, and

3652 + under the EAGLE drop BOTH are reachable.

3653 +

3654 + Retaining only the higher one leaves the resend with every retained state

3655 + above every candidate its lookup can produce, and the reconciled hit

3656 + collapses to 0 -- the same zero-hit failure sparse retention already fixes

3657 + at unaligned prompt lengths.

3658 + """

3659 + block_size = 32

3660 + num_spec = 3

3661 + kv_cache_config = KVCacheConfig(

3662 + num_blocks=100,

3663 + kv_cache_tensors=[],

3664 + kv_cache_groups=[

3665 + KVCacheGroupSpec(

3666 + ["full"],

3667 + FullAttentionSpec(

3668 + block_size=block_size,

3669 + num_kv_heads=1,

3670 + head_size=1,

3671 + dtype=torch.float16,

3672 + ),

3673 + ),

3674 + KVCacheGroupSpec(

3675 + ["mamba_mtp"],

3676 + MambaSpec(

3677 + block_size=block_size,

3678 + shapes=((1, 1),),

3679 + dtypes=(torch.float32,),

3680 + mamba_cache_mode="align",

3681 + num_speculative_blocks=num_spec,

3682 + ),

3683 + ),

3684 + ],

3685 + )

3686 + manager = make_kv_cache_manager(

3687 + kv_cache_config=kv_cache_config,

3688 + max_model_len=8192,

3689 + enable_caching=True,

3690 + hash_block_size=block_size,

3691 + retention_interval=0,

3692 + use_eagle=True,

3693 + )

3694 +

3695 + # 128 tokens, an exact multiple of the 32-token alignment. A longer sibling

3696 + # matches 128 and drops to 96; this prompt's own resend caps at 127, matches

3697 + # 96 and drops to 64. Both states must survive retention.

3698 + token_ids = [i for i in range(4) for _ in range(block_size)]

3699 + req0 = make_request("0", token_ids, block_size, sha256)

3700 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

3701 + assert num_computed_tokens == 0

3702 + # Prefill in block-aligned chunks the way the align-mode scheduler does: a

3703 + # state only materializes as a chunk's running-state block, so a

3704 + # single-shot prefill could not retain the lower one.

3705 + for chunk_end in (32, 64, 96, 128):

3706 + blocks = manager.allocate_slots(

3707 + req0,

3708 + chunk_end - req0.num_computed_tokens,

3709 + num_computed_tokens,

3710 + computed_blocks,

3711 + num_lookahead_tokens=num_spec,

3712 + )

3713 + assert blocks is not None

3714 + req0.num_computed_tokens = chunk_end

3715 +

3716 + # Block ``i`` ends at token ``(i + 1) * 32``, so positions 64 and 96 are

3717 + # mamba blocks 1 and 2.

3718 + pool = manager.block_pool

3719 + expected_mamba_cached = {1, 2}

3720 + for i in range(4):

3721 + cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])

3722 + if i in expected_mamba_cached:

3723 + assert cached is not None, f"mamba hash {i} should be cached"

3724 + else:

3725 + assert cached is None, f"mamba hash {i} should not be cached"

3726 + manager.free(req0)

3727 +

3728 + # The identical resend: full attention matches blocks 0-2 (96 tokens, capped

3729 + # by num_tokens - 1) and the EAGLE drop caps the candidate at 64. Without

3730 + # the lower state retained the reconciled hit would be 0.

3731 + req1 = make_request("1", token_ids, block_size, sha256)

3732 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

3733 + assert num_computed_tokens == 2 * block_size

3734 + assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 2]

3735 +

3736 + # The longer sibling resumes one alignment unit higher, off the same prompt.

3737 + longer = make_request("2", token_ids + [9] * block_size, block_size, sha256)

3738 + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(longer)

3739 + assert num_computed_tokens == 3 * block_size

3740 +

3741 +

3742 def test_block_lookup_cache_single_block_per_key():

3743 cache = BlockHashToBlockMap()

3744 key0 = BlockHashWithGroupId(b"hash0")

tests/v1/core/test_single_type_kv_cache_manager.py

+1 -1

@@ -192,7 +192,7 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime():

192 )

193 assert manager.allocate_new_blocks(request_id, num_tokens, num_tokens) == []

194

195 - manager.cache_blocks(request_id, 1024, replay_boundary=0)

195 + manager.cache_blocks(request_id, 1024, replay_boundaries=(0,))

196 manager.remove_skipped_blocks(request_id, 1024)

197 assert manager.get_num_common_prefix_blocks(request_id) == 0

198 assert manager.get_num_skipped_tokens(1024) == 0

vllm/v1/core/kv_cache_coordinator.py

+19 -20

@@ -307,27 +307,26 @@ class KVCacheCoordinator(ABC):

307 for manager in self.single_type_managers

308 )

309

310 - def get_replay_boundary(self, request: Request) -> int:

311 - """Return the position a later request replaying this prompt resumes at.

310 + def get_replay_boundaries(self, request: Request) -> tuple[int, ...]:

311 + """Positions a later request replaying this prompt can resume at.

312

313 - A cache hit is the shortest hit across all groups, so this is a

314 - model-level position: every group has to retain state here, whether or

315 - not it is the group that drops. Groups differ only in how much they

316 - keep around it -- EAGLE groups also keep the block above, which they

313 + A hit is the shortest across all groups, so every group retains state

314 + at each position; EAGLE groups also keep the block above, which they

315 match and drop back from (see ``reachable_block_mask``).

316

319 - Under EAGLE that block must exist, so the boundary sits one alignment

320 - unit below the prompt's last aligned position; every group's block size

321 - divides the alignment, so the block above always fits in the prompt.

317 + Two positions are reachable: a resend of the identical prompt is capped

318 + at ``num_tokens - 1`` (its last token is recomputed for logits), a

319 + longer sibling matches the final aligned block. They differ only on a

320 + block-aligned prompt, where retaining just the higher one collapses the

321 + resend's hit to 0. The alignment is the scheduler block size, not the

322 + finer hash granularity, which would over-estimate the reach.

323 """

324 if not self.eagle_group_ids:

324 - return request.num_prompt_tokens - 1

325 - aligned = (

326 - request.num_prompt_tokens

327 - // self.scheduler_block_size

328 - * self.scheduler_block_size

329 - )

330 - return max(aligned - self.scheduler_block_size, 0)

325 + return (request.num_prompt_tokens - 1,)

326 + block = self.scheduler_block_size

327 + resend = (request.num_prompt_tokens - 1) // block * block

328 + extension = request.num_prompt_tokens // block * block

329 + return tuple(sorted({max(resend - block, 0), max(extension - block, 0)}))

330

331 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:

332 """

@@ -339,7 +338,7 @@ class KVCacheCoordinator(ABC):

338 that need to be cached

339 (including tokens that are already cached).

340 """

342 - replay_boundary = self.get_replay_boundary(request)

341 + boundaries = self.get_replay_boundaries(request)

342 for manager in self.single_type_managers:

343 # Only cache tokens with finalized KV. The last num_reprefillable_tokens

344 # tokens can be re-prefilled during multi-module MTP.

@@ -350,7 +349,7 @@ class KVCacheCoordinator(ABC):

349 request,

350 num_tokens_to_cache,

351 retention_interval=self.retention_interval,

353 - replay_boundary=replay_boundary,

352 + replay_boundaries=boundaries,

353 )

354

355 def free(self, request_id: str) -> None:

@@ -775,7 +774,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

774

775 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:

776 cached_num_computed_tokens = self._align_cacheable(num_computed_tokens)

778 - replay_boundary = self.get_replay_boundary(request)

777 + boundaries = self.get_replay_boundaries(request)

778 for manager in self.single_type_managers:

779 num_tokens_to_cache = cached_num_computed_tokens

780 # EAGLE groups match one block past each aligned boundary and drop

@@ -802,7 +801,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

801 request,

802 num_tokens_to_cache,

803 retention_interval=self.retention_interval,

805 - replay_boundary=replay_boundary,

804 + replay_boundaries=boundaries,

805 )

806

807 def find_longest_cache_hit(

vllm/v1/core/single_type_kv_cache_manager.py

+12 -10

@@ -452,7 +452,7 @@ class SingleTypeKVCacheManager(ABC):

452 num_tokens: int,

453 retention_interval: int | None = None,

454 *,

455 - replay_boundary: int,

455 + replay_boundaries: Sequence[int],

456 ) -> None:

457 """

458 Cache the blocks for the request.

@@ -465,6 +465,8 @@ class SingleTypeKVCacheManager(ABC):

465 keeps dense checkpointing; ``0`` keeps only the latest replay

466 boundary; a positive multiple of ``scheduler_block_size`` keeps

467 a tail once per that-sized segment. Only SWA acts on it.

468 + replay_boundaries: Positions a later request replaying this prompt

469 + can resume at, from ``get_replay_boundaries``.

470 """

471 num_cached_blocks = self.num_cached_block.get(request.request_id, 0)

472 num_full_blocks = num_tokens // self.block_size

@@ -473,9 +475,9 @@ class SingleTypeKVCacheManager(ABC):

475 return

476

477 # Token boundaries whose reachable tail must be retained under sparse

476 - # retention: the replay boundary (``num_prompt - 1``, capped by

477 - # ``get_computed_blocks``) and any detected shared-prefix junction.

478 - reachable_boundaries = [replay_boundary]

478 + # retention: every position a replaying sibling can resume at (see

479 + # ``get_replay_boundaries``) and any detected shared-prefix junction.

480 + reachable_boundaries = [*replay_boundaries]

481 if request.shared_prefix_boundary:

482 reachable_boundaries.append(request.shared_prefix_boundary)

483

@@ -818,13 +820,13 @@ class FullAttentionManager(SingleTypeKVCacheManager):

820 num_tokens: int,

821 retention_interval: int | None = None,

822 *,

821 - replay_boundary: int,

823 + replay_boundaries: Sequence[int],

824 ) -> None:

825 super().cache_blocks(

826 request,

827 num_tokens,

828 retention_interval=retention_interval,

827 - replay_boundary=replay_boundary,

829 + replay_boundaries=replay_boundaries,

830 )

831 hash_block_size = self.block_pool.hash_block_size

832 if self.block_size == hash_block_size:

@@ -1228,7 +1230,7 @@ class CircularBufferManager(FullAttentionManager):

1230 num_tokens: int,

1231 retention_interval: int | None = None,

1232 *,

1231 - replay_boundary: int,

1233 + replay_boundaries: Sequence[int],

1234 ) -> None:

1235 return

1236

@@ -1940,14 +1942,14 @@ class MambaManager(SingleTypeKVCacheManager):

1942 num_tokens: int,

1943 retention_interval: int | None = None,

1944 *,

1943 - replay_boundary: int,

1945 + replay_boundaries: Sequence[int],

1946 ) -> None:

1947 num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0)

1948 super().cache_blocks(

1949 request,

1950 num_tokens,

1951 retention_interval=retention_interval,

1950 - replay_boundary=replay_boundary,

1952 + replay_boundaries=replay_boundaries,

1953 )

1954 num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)

1955 if self.mamba_cache_mode == "align":

@@ -2090,7 +2092,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):

2092 num_tokens: int,

2093 retention_interval: int | None = None,

2094 *,

2093 - replay_boundary: int,

2095 + replay_boundaries: Sequence[int],

2096 ) -> None:

2097 # We do not cache blocks for cross-attention to be shared between

2098 # requests, so this method is not relevant.