From 7b9587eb9a6514d24b56f89210592c48bf7a444b Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:58:45 -0700 Subject: [PATCH] kv cache: DFlash2 first reuse hits the full cacheable prefix With DFlash2 and sparse prefix-cache retention (interval 0), the first reuse of a prompt hit one 1,664-token state block less than the full cacheable prefix (32K prompt: 29,952 instead of 31,616 tokens); later reuses hit the full prefix. The drafter's cache group is flagged as an EAGLE group, so get_replay_boundaries rewound the published recurrent-state boundary by one block, although the DFlash2 drafter takes no page-granular drop: its own lookup pops a sub-page block and re-aligns inside the page, and the recurrent groups take no drop at all. The cold prefill computed the state at the full boundary but never published it. The rewind is now owed only to a drop that costs a whole checkpoint page: an engine-level EAGLE/MTP rewind, or a cacheable EAGLE-flagged group whose block size equals the hit alignment. Otherwise both the rewound and the full boundary are published. The rewound one is kept, so the scheduler's prefill chunk stops do not move and the cold computation keeps its chunk geometry. MTP4 (engine-level rewind) and every path without sparse retention return exactly the previous boundaries; models without EAGLE groups are untouched. Tests (CPU): the DFlash2 geometry publishes 31,616 / 98,176 / 129,792 for 32,768 / 99,000 / 131,072-token prompts, and a cold prefill followed by a resend now hits the full prefix (fails on the old logic with exactly the measured deficit); MTP4-style layouts keep the old boundaries and first-reuse hit; cold prefill chunk stops are identical with the old and new boundary sets. Not yet validated on GPU. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- tests/v1/core/test_mamba_align_chunk_split.py | 149 +++++++- tests/v1/core/test_mamba_sparse_retention.py | 337 +++++++++++++++++- vllm/v1/core/kv_cache_coordinator.py | 49 ++- 3 files changed, 524 insertions(+), 11 deletions(-) diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py index 40bd2a9406..f6072d81ff 100644 --- a/tests/v1/core/test_mamba_align_chunk_split.py +++ b/tests/v1/core/test_mamba_align_chunk_split.py @@ -6,13 +6,17 @@ import pytest import torch -from vllm.v1.core.kv_cache_coordinator import KVCacheCoordinator +from vllm.v1.core.kv_cache_coordinator import ( + KVCacheCoordinator, + get_kv_cache_coordinator, +) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, MambaSpec, + SlidingWindowSpec, ) from vllm.v1.structured_output import StructuredOutputManager @@ -300,3 +304,146 @@ def test_periodic_retention_stops_at_all_retained_checkpoints(): def test_dense_retention_still_materializes_every_checkpoint(): length = 3 * MAMBA_BLOCK_SIZE + 30 assert _chunk_ends(length, interval=None) == [816, 1632, 2448, length] + + +# --------------------------------------------------------------------------- +# DFlash2 keeps its own sub-page cache, so its EAGLE flag does not rewind the +# recurrent state (r015). Retaining the unrewound replay boundary must not move +# a single prefill stop of a cold DFlash2 prompt. +# --------------------------------------------------------------------------- + +DFLASH_STATE_BLOCK = 1664 +DFLASH_DRAFT_BLOCK = 128 +# prompt length -> deepest cacheable boundary floor((n - 1) / 1664) * 1664 +DFLASH_PROMPTS = {32_768: 31_616, 99_000: 98_176, 131_072: 129_792} + + +def _dflash2_coordinator( + *, drafter_block: int = DFLASH_DRAFT_BLOCK, use_eagle: bool = False +): + """Target KV + align-mode state on one grid, EAGLE-flagged drafter window.""" + groups = [ + KVCacheGroupSpec( + ["target"], + FullAttentionSpec( + block_size=DFLASH_STATE_BLOCK, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["state"], + MambaSpec( + block_size=DFLASH_STATE_BLOCK, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + KVCacheGroupSpec( + ["drafter"], + SlidingWindowSpec( + block_size=drafter_block, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=3 * drafter_block, + ), + is_eagle_group=True, + ), + ] + coordinator = get_kv_cache_coordinator( + KVCacheConfig(8, [], groups), + max_model_len=max(DFLASH_PROMPTS) + 1, + max_in_flight_tokens=8192, + use_eagle=use_eagle, + enable_caching=True, + enable_kv_cache_events=False, + dcp_world_size=1, + pcp_world_size=1, + hash_block_size=DFLASH_DRAFT_BLOCK, + ) + coordinator.configure_prefix_cache_retention(0) + return coordinator + + +def _prefill_stops(boundary_source, length: int, *, use_eagle: bool) -> list[int]: + """Chunk ends of a cold prompt, driven by the split under test.""" + request = SimpleNamespace( + num_computed_tokens=0, + num_prompt_tokens=length, + num_tokens=length, + shared_prefix_boundary=0, + ) + scheduler = SimpleNamespace( + cache_config=SimpleNamespace(block_size=DFLASH_DRAFT_BLOCK), + mamba_state_block_size=DFLASH_STATE_BLOCK, + mamba_state_retention_interval=0, + use_eagle=use_eagle, + kv_cache_manager=SimpleNamespace(coordinator=boundary_source), + ) + ends = [] + while request.num_computed_tokens < length: + remaining = length - request.num_computed_tokens + chunk = Scheduler._mamba_block_aligned_split( + scheduler, request, min(8192, remaining) + ) + assert 0 < chunk <= min(8192, remaining) + request.num_computed_tokens += chunk + ends.append(request.num_computed_tokens) + return ends + + +def _rewind_only_coordinator(): + """The pre-fix rule: rewind because *some* group carries the EAGLE flag.""" + coordinator = SimpleNamespace(eagle_group_ids={1}) + coordinator.get_replay_boundaries = lambda request, block: ( + KVCacheCoordinator.get_replay_boundaries(coordinator, request, block) + ) + return coordinator + + +def _state_lookbehind(coordinator) -> bool: + """The look-behind in force, also for a coordinator without the new flag.""" + return getattr( + coordinator, "mamba_state_lookbehind", bool(coordinator.eagle_group_ids) + ) + + +@pytest.mark.parametrize("length,deepest", sorted(DFLASH_PROMPTS.items())) +def test_dflash2_deeper_boundary_keeps_the_cold_prefill_stops(length, deepest): + coordinator = _dflash2_coordinator() + assert deepest in coordinator.get_replay_boundaries( + SimpleNamespace(num_tokens=length), DFLASH_STATE_BLOCK + ) + assert _state_lookbehind(coordinator) is False + # DFlash2 does not rewind `last_cache_position` (use_eagle_kv_cache() is + # False for method="dflash"), so that boundary was already a chunk stop. + assert _prefill_stops(coordinator, length, use_eagle=False) == _prefill_stops( + _rewind_only_coordinator(), length, use_eagle=False + ) + assert {deepest - DFLASH_STATE_BLOCK, deepest} <= set( + _prefill_stops(coordinator, length, use_eagle=False) + ) + + +def test_dflash2_cold_prefill_stops_are_pinned(): + ends = _prefill_stops(_dflash2_coordinator(), 32_768, use_eagle=False) + assert ends == [6656, 13312, 19968, 26624, 29952, 31616, 32768] + + +@pytest.mark.parametrize("length,deepest", sorted(DFLASH_PROMPTS.items())) +def test_page_granular_eagle_drop_keeps_the_cold_prefill_stops(length, deepest): + # MTP4-like: the flagged group's block is the whole checkpoint page, and + # the engine rewinds `last_cache_position`, so the last chunk crosses the + # deepest boundary. Both the boundary set and the stops stay as they were. + coordinator = _dflash2_coordinator(drafter_block=DFLASH_STATE_BLOCK, use_eagle=True) + assert _state_lookbehind(coordinator) is True + assert _prefill_stops(coordinator, length, use_eagle=True) == _prefill_stops( + _rewind_only_coordinator(), length, use_eagle=True + ) + assert deepest - DFLASH_STATE_BLOCK in _prefill_stops( + coordinator, length, use_eagle=True + ) + assert deepest not in _prefill_stops(coordinator, length, use_eagle=True) diff --git a/tests/v1/core/test_mamba_sparse_retention.py b/tests/v1/core/test_mamba_sparse_retention.py index 02f8a926b2..fa7e8a23f3 100644 --- a/tests/v1/core/test_mamba_sparse_retention.py +++ b/tests/v1/core/test_mamba_sparse_retention.py @@ -14,21 +14,26 @@ from vllm.v1.core.kv_cache_manager import KVCacheManager from vllm.v1.core.kv_cache_utils import BlockHashListWithBlockSize from vllm.v1.core.sched.scheduler import Scheduler -from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager, MambaManager +from vllm.v1.core.single_type_kv_cache_manager import ( + FullAttentionManager, + MambaManager, + SlidingWindowManager, +) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, MambaSpec, + SlidingWindowSpec, ) # These tests allocate only block metadata; no device/distributed state exists. pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] -def _request(name, length, hash_size=8): +def _request(name, length, hash_size=8, request_id=None): return SimpleNamespace( - request_id=name, + request_id=request_id or name, num_tokens=length, shared_prefix_boundary=0, skip_reading_prefix_cache=False, @@ -408,3 +413,329 @@ def test_bulk_prefill_materializes_every_admitted_state(eagle, length, interval_ if interval == 0: assert len(ends) <= length // 8000 + 4 assert cm.block_pool.get_num_free_blocks() == 999 + + +# --------------------------------------------------------------------------- +# DFlash2: a drafter that owns a *sub-page* cache must not rewind the state. +# +# Measured on 0.1.4: the FIRST reuse of a 32,768 / 99,000 / 131,072-token +# prompt hit 29,952 / 96,512 / 128,128 tokens -- exactly one 1,664-token state +# block below the deepest cacheable boundary 31,616 / 98,176 / 129,792 that +# every later reuse reaches. ``get_replay_boundaries`` rewound whenever *any* +# group was EAGLE-flagged, but DFlash2 flags only its own drafter group +# (kv_cache_utils.py:2446-2462) and that drop pops one sub-page block and +# re-aligns inside the page (single_type_kv_cache_manager.py:947-956), so no +# checkpoint page was ever owed. MTP4 rewinds for real (its EAGLE drop is +# page-granular) and must keep the rewound boundary. +# --------------------------------------------------------------------------- + +STATE_BLOCK = 1664 +"""Recurrent-state grid of the DFlash2 profile (docs/dflash2.md:88-90).""" + +DRAFT_BLOCK = 128 +"""The drafter's own block: a divisor of the state grid, hence sub-page.""" + +DRAFTER_GROUP = 2 +"""Index of the only EAGLE-flagged group in the measured DFlash2 layout.""" + +# prompt length -> deepest cacheable boundary floor((n - 1) / 1664) * 1664. +FIRST_REUSE_DEFICIT = {32_768: 31_616, 99_000: 98_176, 131_072: 129_792} + + +def rewound_replay_boundaries(num_tokens: int, alignment: int) -> tuple[int, ...]: + """The rule before the fix: rewind whenever *any* group is EAGLE-flagged.""" + resend = (num_tokens - 1) // alignment * alignment + extension = num_tokens // alignment * alignment + return tuple(sorted({max(resend - alignment, 0), max(extension - alignment, 0)})) + + +def state_lookbehind(coordinator) -> bool: + """The look-behind a coordinator applies, also before this attribute existed.""" + return getattr( + coordinator, "mamba_state_lookbehind", bool(coordinator.eagle_group_ids) + ) + + +def dflash2_groups( + drafter_block: int = DRAFT_BLOCK, + *, + flag_drafter: bool = True, + state_block: int = STATE_BLOCK, +) -> list[KVCacheGroupSpec]: + """Target KV + align-mode state on one grid, drafter window on a sub-grid.""" + return [ + KVCacheGroupSpec( + ["target"], + FullAttentionSpec( + block_size=state_block, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["state"], + MambaSpec( + block_size=state_block, + shapes=((1,),), + dtypes=(torch.float16,), + mamba_cache_mode="align", + ), + ), + KVCacheGroupSpec( + ["drafter"], + SlidingWindowSpec( + block_size=drafter_block, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + sliding_window=3 * drafter_block, + ), + is_eagle_group=flag_drafter, + ), + ] + + +def dflash2_manager( + *, + length: int, + drafter_block: int = DRAFT_BLOCK, + flag_drafter: bool = True, + use_eagle: bool = False, + interval: int = 0, + num_blocks: int | None = None, +) -> KVCacheManager: + """KVCacheManager as the launch profile builds it: sparse retention, grid B.""" + groups = dflash2_groups(drafter_block, flag_drafter=flag_drafter) + if num_blocks is None: + # Roomy: the checkpoint set under test must never be evicted. + num_blocks = 2 * (length // DRAFT_BLOCK) + 4 * (length // STATE_BLOCK) + 256 + return KVCacheManager( + KVCacheConfig(num_blocks, [], groups), + max_model_len=length + 1, + hash_block_size=DRAFT_BLOCK, + max_in_flight_tokens=8192, + use_eagle=use_eagle, + prefix_cache_retention_interval=interval, + ) + + +def _prompt(name: str, length: int) -> SimpleNamespace: + request = _request(name, length, DRAFT_BLOCK) + request.num_prompt_tokens = length + request.num_computed_tokens = 0 + return request + + +def cold_prefill(cm: KVCacheManager, request, *, use_eagle: bool, budget: int = 8192): + """Drive a cold prefill with the scheduler's own stop rule. + + Returns the chunk ends plus, for every recurrent-state slot, the chunk ends + that could have written it: the align kernel stores the state of a chunk in + the slot of its last whole state block only when the chunk ends there + (worker/gpu/mamba_align.py:243-252). + """ + scheduler = SimpleNamespace( + cache_config=SimpleNamespace(block_size=DRAFT_BLOCK), + mamba_state_block_size=STATE_BLOCK, + mamba_state_retention_interval=cm.coordinator.retention_interval, + use_eagle=use_eagle, + kv_cache_manager=cm, + ) + ends: list[int] = [] + materialized: dict[int, set[int]] = {} + while request.num_computed_tokens < request.num_tokens: + start = request.num_computed_tokens + chunk = Scheduler._mamba_block_aligned_split( + scheduler, request, min(budget, request.num_tokens - start) + ) + assert chunk > 0 + end = start + chunk + for manager in cm.coordinator.single_type_managers: + manager.new_step_starts() + manager.remove_skipped_blocks( + request.request_id, start, request.num_prompt_tokens + ) + # The EAGLE-flagged drafter peeks one of its own blocks past the + # boundary so that block is eligible to be cached, exactly as + # HybridKVCacheCoordinator.cache_blocks does for it. + manager.allocate_new_blocks(request.request_id, end + DRAFT_BLOCK, end) + if isinstance(manager, MambaManager): + slot = (end - 1) // STATE_BLOCK + block = manager.req_to_blocks[request.request_id][slot] + assert not block.is_null + materialized.setdefault(block.block_id, set()).add(end) + cm.coordinator.cache_blocks(request, end) + request.num_computed_tokens = end + ends.append(end) + return ends, materialized + + +def published_state_blocks(cm: KVCacheManager, request) -> list[int]: + """Indices of the state-group blocks that reached the prefix-cache map.""" + hashes = BlockHashListWithBlockSize(request.block_hashes, DRAFT_BLOCK, STATE_BLOCK) + return [ + i + for i, block_hash in enumerate(hashes) + if cm.block_pool.get_cached_block(block_hash, [1]) + ] + + +@pytest.mark.parametrize("length,deepest", sorted(FIRST_REUSE_DEFICIT.items())) +def test_sub_page_drafter_publishes_the_deepest_boundary(length, deepest): + cm = dflash2_manager(length=length, num_blocks=8) + coordinator = cm.coordinator + # The measured geometry: one grid, one sub-page drafter group, flagged alone. + assert coordinator.lcm_block_size == STATE_BLOCK + assert coordinator.eagle_group_ids == {DRAFTER_GROUP} + + boundaries = coordinator.get_replay_boundaries( + _request("doc", length, DRAFT_BLOCK), STATE_BLOCK + ) + # The pre-fix rule published only the rewound boundary -- the reported bug. + rewound = rewound_replay_boundaries(length, STATE_BLOCK) + assert rewound == (deepest - STATE_BLOCK,) + assert deepest not in rewound + # Both checkpoints are retained: the unrewound one a resend lands on, plus + # the rewound one so the cold chunk stops do not move. + assert boundaries == tuple(sorted({deepest - STATE_BLOCK, deepest})) + assert state_lookbehind(coordinator) is False + + mask = MambaManager.reachable_block_mask( + 0, + deepest // STATE_BLOCK + 1, + STATE_BLOCK, + coordinator.single_type_managers[1].kv_cache_spec, + 0, + boundaries, + ) + assert [i for i, keep in enumerate(mask) if keep] == [ + deepest // STATE_BLOCK - 2, + deepest // STATE_BLOCK - 1, + ] + + # The drafter keeps its window at the boundary that is actually admitted. + drafter = coordinator.single_type_managers[DRAFTER_GROUP] + assert drafter.mamba_replay_lookbehind is False + need = SlidingWindowManager._contiguous_blocks_for_hit( + drafter.kv_cache_spec.sliding_window, DRAFT_BLOCK, True + ) + end = deepest // DRAFT_BLOCK + 1 + assert drafter._cached_window_at_last_boundary(length) == (end - need, end) + + +@pytest.mark.parametrize("variant", ["engine_fallback", "page_granular_flag"]) +@pytest.mark.parametrize("length,deepest", sorted(FIRST_REUSE_DEFICIT.items())) +def test_page_granular_eagle_drop_still_rewinds(variant, length, deepest): + # MTP4 takes the drop with a page-granular group, either because nothing is + # labelled and the coordinator flags every group, or because the flagged + # group's block *is* the checkpoint page. + cm = dflash2_manager( + length=length, + num_blocks=8, + drafter_block=STATE_BLOCK if variant == "page_granular_flag" else DRAFT_BLOCK, + flag_drafter=variant == "page_granular_flag", + use_eagle=variant == "engine_fallback", + ) + coordinator = cm.coordinator + assert state_lookbehind(coordinator) is True + assert all(m.mamba_replay_lookbehind for m in coordinator.single_type_managers) + assert ( + coordinator.get_replay_boundaries( + _request("doc", length, DRAFT_BLOCK), STATE_BLOCK + ) + == rewound_replay_boundaries(length, STATE_BLOCK) + == (deepest - STATE_BLOCK,) + ) + + +def test_retention_rules_over_a_sweep_of_prompt_lengths(): + sub_page = dflash2_manager(length=12 * STATE_BLOCK, num_blocks=8).coordinator + page = dflash2_manager( + length=12 * STATE_BLOCK, + drafter_block=STATE_BLOCK, + use_eagle=True, + num_blocks=8, + ).coordinator + lengths = [ + 1, + 2, + STATE_BLOCK - 1, + STATE_BLOCK, + STATE_BLOCK + 1, + 2 * STATE_BLOCK, + 5 * STATE_BLOCK + 3, + 11 * STATE_BLOCK + 7, + ] + for length in lengths: + request = _request("doc", length, DRAFT_BLOCK) + old = rewound_replay_boundaries(length, STATE_BLOCK) + # A page-granular drop keeps today's answer to the token. + assert page.get_replay_boundaries(request, STATE_BLOCK) == old + boundaries = sub_page.get_replay_boundaries(request, STATE_BLOCK) + # The fix only ever adds the unrewound boundary, so the scheduler's + # *earliest* stop -- and with it the cold chunk geometry -- cannot move. + assert set(boundaries) >= set(old) + assert max(boundaries) == length // STATE_BLOCK * STATE_BLOCK + + +def test_sub_page_drafter_first_reuse_hits_the_full_prefix(): + length, deepest = 32_768, FIRST_REUSE_DEFICIT[32_768] + cm = dflash2_manager(length=length) + request = _prompt("doc", length) + ends, materialized = cold_prefill(cm, request, use_eagle=False) + # DFlash2 does not rewind the scheduler's `last_cache_position`, so the + # cold run already ends a chunk at the deepest boundary. + assert ends == [6656, 13312, 19968, 26624, 29952, 31616, 32768] + + published = published_state_blocks(cm, request) + assert published == [deepest // STATE_BLOCK - 2, deepest // STATE_BLOCK - 1] + for index in published: + block = cm.block_pool.get_cached_block( + BlockHashListWithBlockSize(request.block_hashes, DRAFT_BLOCK, STATE_BLOCK)[ + index + ], + [1], + )[0] + # Only a chunk that ended exactly on the boundary writes that state. + assert (index + 1) * STATE_BLOCK in materialized[block.block_id] + + cm.coordinator.free(request.request_id) + _, hit = cm.get_computed_blocks(_prompt("doc", length)) + # Before the fix this was `deepest - STATE_BLOCK` and the junction repair + # cost one extra 1,664-token re-prefill on this very request. + assert hit == deepest + # Every cacheable group now reaches the same boundary: no junction at all. + assert cm.coordinator.shared_prefix_boundary == 0 + # The pool is only metadata here; nothing leaked. + assert cm.block_pool.get_num_free_blocks() == cm.block_pool.num_gpu_blocks - 1 + + +def test_engine_level_rewind_first_reuse_stays_rewound(): + # An engine that rewinds `last_cache_position` (MTP/EAGLE semantics) never + # produces the deepest state, even when only a sub-page group is flagged: + # publishing it would add a chunk stop and change the cold prefill shape. + length, deepest = 32_768, FIRST_REUSE_DEFICIT[32_768] + cm = dflash2_manager(length=length, use_eagle=True) + request = _prompt("doc", length) + ends, materialized = cold_prefill(cm, request, use_eagle=True) + # The last chunk crosses the deepest boundary, so no state exists there. + assert ends == [6656, 13312, 19968, 26624, 29952, 32768] + assert deepest not in ends + published = published_state_blocks(cm, request) + assert published == [deepest // STATE_BLOCK - 2] + for index in published: + block = cm.block_pool.get_cached_block( + BlockHashListWithBlockSize(request.block_hashes, DRAFT_BLOCK, STATE_BLOCK)[ + index + ], + [1], + )[0] + assert (index + 1) * STATE_BLOCK in materialized[block.block_id] + cm.coordinator.free(request.request_id) + reuse = _prompt("doc", length) + _, hit = cm.get_computed_blocks(reuse) + assert hit == deepest - STATE_BLOCK + # Denser groups still reach the deepest boundary, so the junction that + # repairs the gap is recorded -- the behaviour MTP4 keeps today. + assert cm.coordinator.shared_prefix_boundary == deepest diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 9325d9e65e..2e0d14bede 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -51,6 +51,11 @@ def __init__( self.max_model_len = max_model_len self.enable_caching = enable_caching self.retention_interval: int | None = None + # Engine-level EAGLE/MTP rewind. ``Scheduler.use_eagle`` carries the same + # flag (``SpeculativeConfig.use_eagle_kv_cache``): when it is on, the + # scheduler ends the last prefill chunk one page below the deepest + # aligned boundary, so that state is never produced at all. + self.engine_use_eagle = use_eagle self.shared_prefix_boundary = 0 self.block_pool = BlockPool( @@ -68,6 +73,13 @@ def __init__( # Conservatively fall back to flag all groups when no group is flagged. if use_eagle and not self.eagle_group_ids: self.eagle_group_ids = set(range(len(kv_cache_config.kv_cache_groups))) + # Whether a replay of a prompt lands one checkpoint page below the + # deepest aligned boundary its lookup can key. Default to the engine + # rule so a coordinator that never configures sparse retention answers + # exactly as before; ``configure_prefix_cache_retention`` replaces it + # with the drop a page-granular cacheable group really takes, which a + # drafter owning a sub-page cache (DFlash2) does not take. + self.mamba_state_lookbehind = bool(self.eagle_group_ids) self.single_type_managers = tuple( get_manager_for_kv_cache_spec( @@ -97,16 +109,32 @@ def configure_prefix_cache_retention(self, interval: int | None) -> None: "in this backport" ) alignment = getattr(self, "lcm_block_size", mamba[0].block_size) - # SWA must preserve the window at the checkpoint Mamba actually admits, - # rather than the prompt's final (unrestorable) alignment boundary. - sparse_mamba = bool(self.eagle_group_ids) and any( + sparse_mamba = any( isinstance(m.kv_cache_spec, MambaSpec) and m.kv_cache_spec.mamba_cache_mode == "align" and m.block_size == alignment for m in mamba ) + # The lookbehind is only owed to a drop that costs a whole checkpoint + # page, i.e. a drop taken by a cacheable group whose block size equals + # the hit alignment. A drafter that owns a sub-page cache (DFlash2 flags + # only its own group, and ``MambaManager.find_longest_cache_hit`` then + # takes no drop at all) pops one of its own blocks and re-aligns inside + # the page, so the deepest boundary a resend reaches stays + # floor((num_tokens - 1) / alignment) * alignment. An engine that + # rewinds ``last_cache_position`` keeps the rewind: it never produces + # that state. + page_granular_drop = self.engine_use_eagle or any( + manager.kv_cache_group_id in self.eagle_group_ids + and manager.block_size == alignment + for manager in self.single_type_managers + if manager.kv_cache_spec.prefix_cacheable + ) + self.mamba_state_lookbehind = sparse_mamba and page_granular_drop + # SWA must preserve the window at the checkpoint Mamba actually admits, + # rather than the prompt's final (unrestorable) alignment boundary. for manager in self.single_type_managers: - manager.mamba_replay_lookbehind = sparse_mamba + manager.mamba_replay_lookbehind = self.mamba_state_lookbehind if interval < 0 or interval % alignment: raise ValueError( f"prefix_cache_retention_interval ({interval}) must be a non-negative " @@ -130,9 +158,16 @@ def get_replay_boundaries( # Once the cache horizon advances, the next decode token cannot # retroactively admit it. The preceding boundary serves a resend. return tuple(sorted({resend, extension})) - return tuple( - sorted({max(resend - alignment, 0), max(extension - alignment, 0)}) - ) + rewound = {max(resend - alignment, 0), max(extension - alignment, 0)} + lookbehind = getattr(self, "mamba_state_lookbehind", bool(self.eagle_group_ids)) + if not lookbehind: + # Nothing rewinds this group's own lookup, so the unrewound + # boundary is what a resend lands on and must be published. Keep + # the rewound entry too: the scheduler already ends a chunk there, + # so dropping it would move the last prefill stop and change the + # recurrent chunk geometry (and its low bits) for no gain. + return tuple(sorted(rewound | {resend, extension})) + return tuple(sorted(rewound)) def _cache_manager_blocks( self,