From ac674a1e89e25c1178500c05a87cdf76daf9c491 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 28 Jul 2026 03:37:39 +0000 Subject: [PATCH 1/5] [None][fix] Handle uncertain scratch pages in KV cache sanity check Signed-off-by: Yao Yao --- .../runtime/kv_cache_manager_v2/_core/_kv_cache.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index da0500dc4c06..fb1c95b0de48 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -1896,9 +1896,15 @@ def get_range(lc: LifeCycle): assert holder is None continue start, end = stale_ranges[lc] + lc_obj = self.manager._life_cycles[lc] if start <= ordinal < end: if is_committed or self._commit_state != self.CommitState.ALLOWED: assert holder is None + elif ordinal in self._get_scratch_range(lc_obj, 0): + # It is uncertain whether this block should hold a page: it may + # have been scratch-allocated by an earlier chunk, but the prior + # history length and capacity are not retained. + pass else: # For the decoder-side disagg case, for the first step, we will skip the # out-of-window blocks. @@ -1908,7 +1914,6 @@ def get_range(lc: LifeCycle): ) else: # Scratch blocks have None pages but valid base_page_indices - lc_obj = self.manager._life_cycles[lc] sr = self._get_scratch_range(lc_obj) is_scratch = ordinal in sr if is_scratch: From 2a397259952510a1cc570ac6fb30ec55c38f903a Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Thu, 30 Jul 2026 17:57:47 +0000 Subject: [PATCH 2/5] [None][fix] KVCacheManagerV2: record per-page token coverage A partial trailing tree block (a rewind endpoint, e.g. 16 tokens of a 32-token block) was destroyed when a longer sibling covering the same tokens was created, and refused when it was created second. The covering block does not necessarily hold a page for every life cycle at that token boundary -- typically a SWA life cycle whose page was never allocated because the block was already outside the sliding window when it was committed -- so the reusable endpoint was lost either way. Generalize the mechanism SSM snapshots already used. num_tokens_in_block moves from SsmCommittedPage up to CommittedPage, meaning the number of leading tokens of the owning block that the page's data is valid for. Attention reads it as a prefix-valid length; SSM keeps its exact-endpoint meaning. Each (block, life cycle) slot keeps only the widest page. Block.__init__ now moves a covered sibling's pages into the new block instead of dropping them, and _snapshot_partial_block_to_tree attaches partial pages to a longer sibling. Prefix matching honours the recorded count for attention as well as SSM, and _commit_block's rebase path no longer adopts a page that covers fewer tokens than the block spans, which would have fed uninitialized KV into a live request. Block.storage reads now go through Block.get_page(); the slot never holds a dangling ref because CommittedPage.__del__ unlinks before invalidating. KVCacheEventManager drops _resolve_page_ref and does not announce a life cycle whose page covers less than the whole block, since the event payload carries the block's full token list. SsmCommittedPage is deleted; it no longer adds any field. Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/_block_radix_tree.py | 206 ++++++++------ .../kv_cache_manager_v2/_core/_kv_cache.py | 120 ++++---- .../kv_cache_manager_v2/_event_manager.py | 25 +- .../kv_cache_manager_v2/_introspection.py | 13 +- .../runtime/kv_cache_manager_v2/_page.py | 73 ++--- .../test_kv_cache_event_manager.py | 24 +- .../test_kv_cache_manager_v2.py | 259 ++++++++++++++++++ 7 files changed, 498 insertions(+), 222 deletions(-) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index c47c591f0d06..66e225ee63ba 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING, Iterable, Iterator, NamedTuple, Sequence, TypeVar, cast +from itertools import chain +from typing import TYPE_CHECKING, Iterator, NamedTuple, Sequence, TypeVar, cast from . import rawref from ._cache_key import ( # noqa: F401 @@ -26,15 +27,7 @@ ) from ._common import NDEBUG, BlockOrdinal, PageStatus, TokenIdExt from ._life_cycle_registry import AttnLifeCycle, LifeCycle, LifeCycleId, LifeCycleRegistry -from ._utils import ( - TypedIndexList, - expect_type, - filled_list, - find_index, - map_optional, - typed_enumerate, - unwrap_rawref, -) +from ._utils import TypedIndexList, filled_list, map_optional, typed_range, unwrap_rawref if TYPE_CHECKING: from ._event_manager import KVCacheEventManager @@ -276,12 +269,20 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N for b in prev.next.values(): if b.tokens[: len(tokens)] == tokens: raise UselessBlockError(b) - # If there are sibling blocks fully covered by this block, remove them. + # A later turn may extend a partial endpoint to this longer block, replacing the + # partial sibling. That turn may not have a committable SWA page for this block: + # commit_min_snapshot releases out-of-window pages, while SWA scratch reuse uses + # temporary shared storage that is not preserved. Adopt the partial sibling's + # pages to keep the shorter endpoint reusable, retaining each page's recorded token + # count (see CommittedPage.num_tokens_in_block). to_remove = [] for k, b in prev.next.items(): if len(b.tokens) < len(tokens) and tokens[: len(b.tokens)] == b.tokens: assert NDEBUG or (not b.is_full and b is not self and b.key == k and not b.next) to_remove.append(k) + # Two covered siblings would be prefixes of each other; the insertion logic + # would already have replaced the shorter one. + assert NDEBUG or len(to_remove) <= 1 event_manager = get_tree(prev).event_manager if to_remove else None # Keep RootBlock attached while covered children are replaced. Adding # the replacement first prevents detach_next() from pruning an @@ -290,11 +291,62 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N for k in to_remove: b = detach_next(prev, k) assert isinstance(b, Block) + self._adopt_pages_from(b) if event_manager is not None: event_manager.add_removed_event(b.key) assert b.is_orphan # _KVCache may still hold it. # prev.next keeps a strong ref to this _Block, so no need to remove self from prev.next in __del__(). + def page_coverage(self, lc_idx: LifeCycleId) -> int: + """Return the page's recorded token count, or zero if the slot is empty. + + For attention this is prefix coverage; for SSM it is an exact checkpoint position. + """ + page = self.get_page(lc_idx) + return page.num_tokens_in_block if page is not None else 0 + + def holds_page(self, page: "CommittedPage") -> bool: + return self.get_page(page.life_cycle) is page + + def can_replace_page(self, lc_idx: LifeCycleId, num_tokens_in_block: int) -> bool: + """Whether a page recording `num_tokens_in_block` may take over slot `lc_idx`. + + A slot keeps only the page with the largest recorded token count. For attention, + greater coverage strictly dominates lesser coverage. For SSM, this deliberately + keeps only the latest checkpoint -- two conversation turns rarely end inside the + same block, and if they do, a reuse miss is acceptable. + + Pure; use replace_page() to install. + """ + existing = self.get_page(lc_idx) + return existing is None or existing.num_tokens_in_block < num_tokens_in_block + + def replace_page(self, lc_idx: LifeCycleId, page: "CommittedPage") -> None: + """Install `page` in slot `lc_idx`, detaching whatever it supersedes. + + The superseded page may outlive this call while a request still holds it, so + unlink_page() must clear its back-pointer: _release_pages() walks `storage`, so + nothing would clear it later and it would dangle once this block dies. + """ + assert NDEBUG or self.can_replace_page(lc_idx, page.num_tokens_in_block) + existing = self.unlink_page(lc_idx) + if existing is not None and existing.scheduled_for_eviction: + existing.manager.exclude_from_eviction(existing) + page.block = rawref.ref(self) + self.storage[lc_idx] = rawref.ref(page) + + def _adopt_pages_from(self, other: "Block") -> None: + """Move `other`'s pages into self without changing their recorded token counts.""" + assert other.ordinal == self.ordinal + for lc_idx in typed_range(self.num_life_cycles): + page = other.get_page(lc_idx) + if page is None or not self.can_replace_page(lc_idx, page.num_tokens_in_block): + continue + # Clear the source slot directly rather than via unlink_page(), which would + # null the back-pointer replace_page() is about to overwrite. + other.storage[lc_idx] = None + self.replace_page(lc_idx, page) + def _release_pages(self) -> None: """Reclaim every page held by this block. @@ -309,9 +361,9 @@ def _release_pages(self) -> None: ``Block`` alive past ``StorageManager`` teardown, after which ``page.manager`` would be a dangling reference. """ - for lc_idx, ref in typed_enumerate(self.storage): - if ref is not None and ref() is not None: - page = unwrap_rawref(ref) + for lc_idx in typed_range(self.num_life_cycles): + page = self.get_page(lc_idx) + if page is not None: self.unlink_page(lc_idx) if page.status == PageStatus.DROPPABLE: if page.scheduled_for_eviction: @@ -339,30 +391,39 @@ def prev(self) -> "Block | RootBlock": return unwrap_rawref(self._prev) def get_page(self, lc_idx: LifeCycleId) -> "CommittedPage | None": + """Return the page in slot `lc_idx`, or None when the slot is empty. + + A non-empty slot always resolves: CommittedPage.__del__ unlinks the page from its + block before invalidating its rawref, so `storage` never retains a dangling ref. + """ return map_optional(self.storage[lc_idx], lambda f: f()) def unlink_page( self, lc_idx: LifeCycleId, expected_page: "CommittedPage | None" = None - ) -> bool: - page_ref = self.storage[lc_idx] - if page_ref is None: - return False + ) -> "CommittedPage | None": + """Detach slot `lc_idx`, returning the page that was there, or None. + + The sole place a block-page link is severed. + """ + # Called from CommittedPage.__del__, which invalidates the page's rawref only + # afterwards, so the dying page is still reachable here. + page = self.get_page(lc_idx) + if page is None: + return None # Only unlink when the slot still holds the expected page. During rebase # another block with the same key may have replaced the stored page, and # unlinking then would clobber the newer page's back-pointer. - if expected_page is not None and page_ref() is not expected_page: - return False - page = page_ref() - if page is not None: - page.block = rawref.NULL + if expected_page is not None and page is not expected_page: + return None + page.block = rawref.NULL self.storage[lc_idx] = None - return True + return page @staticmethod def clear_stale_blocks_after_page_unlink( start: "Block", lc_idx: LifeCycleId, lc: LifeCycle ) -> None: - assert start.storage[lc_idx] is None + assert start.get_page(lc_idx) is None ordinal = start.ordinal tree = try_get_tree(start) event_manager = tree.event_manager if tree is not None else None @@ -375,7 +436,7 @@ def clear_stale_blocks_after_page_unlink( # But for simplicity, we leave it for now. curr = start while ( - (isinstance(curr, Block) and curr.storage[lc_idx] is None) + (isinstance(curr, Block) and curr.get_page(lc_idx) is None) and not curr.next and curr._prev() is not None ): @@ -470,14 +531,6 @@ def _num_matched_tokens(self, matched: list[tuple[Block, int]]) -> int: return 0 return self._tokens_per_block * (len(matched) - 1) + matched[-1][1] - @staticmethod - def _has_pages(block: Block, lc_list: Iterable[LifeCycleId]) -> bool: - return all(block.storage[lc] is not None for lc in lc_list) - - @staticmethod - def _has_page(block: Block, lc: LifeCycleId) -> bool: - return block.storage[lc] is not None - # yields tuples of (block, num_matched_tokens). num_matched_tokens should be equal to # tokens_per_block except the last one. def _match_token_path( @@ -512,49 +565,23 @@ def _prune_match(self, matched: list[tuple[Block, int]]) -> list[tuple[Block, in assert all(b[1] == tokens_per_block for b in matched[:-1]) life_cycles = self._life_cycles - - # check for full attention layers attn_life_cycles = list(life_cycles.attention_life_cycles()) - if any(lc.window_size is None for _, lc in attn_life_cycles): - lc_list = [lc_idx for lc_idx, lc in attn_life_cycles if lc.window_size is None] - - def check_no_pages(b: tuple[Block, int]) -> bool: - return not BlockRadixTree._has_pages(b[0], lc_list) - - n = find_index(matched, check_no_pages) - matched = matched[:n] - - swa_life_cycles = tuple( - (lc_idx, lc) for lc_idx, lc in attn_life_cycles if lc.window_size is not None - ) - # check for SWA sink - for lc_idx, lc in swa_life_cycles: - - def check_no_page_lc(b: tuple[Block, int]) -> bool: - return not BlockRadixTree._has_page(b[0], lc_idx) - - n = find_index(matched[: lc.num_sink_blocks], check_no_page_lc) - if n < lc.num_sink_blocks: - matched = matched[:n] - # Check SSM snapshot availability before SWA window constraints. - # Truncating to the last reusable SSM snapshot can change the matched - # length used by the SWA check. ssm_lc_id = life_cycles.ssm_life_cycle_id - if ssm_lc_id is not None: - from ._page import SsmCommittedPage + # Fixed-point loop: SSM may select an earlier exact snapshot, while attention may + # shorten the match to the coverage of a required page. Every retry strictly + # shortens the match, so the loop terminates. while matched: + # Check SSM snapshot availability first: truncating to the last reusable SSM + # snapshot changes the matched length that all the attention checks use. if ssm_lc_id is not None: ssm_trunc = 0 ssm_match_len = 0 for i in reversed(range(len(matched))): - block = matched[i][0] - page = map_optional(block.storage[ssm_lc_id], lambda f: f()) - if page is None: - continue - page = expect_type(SsmCommittedPage, page) - snapshot_len = page.num_tokens_in_block - if matched[i][1] >= snapshot_len: + # An SSM page holds the recurrent state after exactly this many tokens, + # so reuse must stop there instead of anywhere inside the block. + snapshot_len = matched[i][0].page_coverage(ssm_lc_id) + if snapshot_len > 0 and matched[i][1] >= snapshot_len: ssm_trunc = i + 1 ssm_match_len = snapshot_len break @@ -562,29 +589,30 @@ def check_no_page_lc(b: tuple[Block, int]) -> bool: if not matched: break matched[-1] = (matched[-1][0], ssm_match_len) - # SWA window check - num_tokens = self._num_matched_tokens(matched) - for lc_idx, lc in swa_life_cycles: - if lc.window_size is None: - continue - def check_has_page_lc(b: tuple[Block, int]) -> bool: - return BlockRadixTree._has_page(b[0], lc_idx) - - n = find_index(reversed(matched), check_has_page_lc) - if n != 0: - matched = matched[:-n] + # Only pages that are active at this candidate endpoint constrain attention + # reuse. Full attention requires every block. SWA requires sink blocks and the + # trailing window, but not the stale blocks between them. In particular, at an + # exact block boundary with window_size=1, every historical block is stale. + num_tokens = self._num_matched_tokens(matched) + shortened = False + for lc_idx, lc in attn_life_cycles: + stale = lc.get_stale_range(num_tokens, tokens_per_block) + for i in chain(range(stale.beg), range(stale.end, len(matched))): + block, num_matched = matched[i] + coverage = block.page_coverage(lc_idx) + if coverage >= num_matched: + continue + if coverage > 0: + matched = matched[: i + 1] + matched[-1] = (block, coverage) + else: + matched = matched[:i] + shortened = True break - _, stale_end = lc.get_stale_range(num_tokens, tokens_per_block) - - def has_no_page(b: tuple[Block, int]) -> bool: - return not BlockRadixTree._has_page(b[0], lc_idx) - - n = find_index(reversed(matched[stale_end:]), has_no_page) - if len(matched) - n > stale_end: - matched = matched[: len(matched) - n - 1] + if shortened: break - else: + if not shortened: break return matched diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index fb1c95b0de48..d71bf3db6acb 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -57,7 +57,6 @@ CommittedPage, Page, ScratchSlotLock, - SsmCommittedPage, UncommittedPage, _PageHolder, _SharedPageLock, @@ -75,7 +74,6 @@ filled_list, intersect, make_typed, - map_optional, stream_wait_events, to_typed, typed_enumerate, @@ -83,7 +81,6 @@ typed_map, typed_range, unwrap_optional, - unwrap_rawref, value_or, ) from ._moving_average import Average @@ -1441,13 +1438,10 @@ def _copy_page_to_tree_block( tree_block: Block, lc_idx: LifeCycleId, src_page: Page, - ssm_num_tokens_in_block: int | None = None, + num_tokens_in_block: int, ) -> CommittedPage | None: - existing_page = map_optional(tree_block.storage[lc_idx], lambda p: p()) - if existing_page is not None: - return existing_page - is_ssm = lc_idx == self.manager._life_cycles.ssm_life_cycle_id - assert is_ssm == (ssm_num_tokens_in_block is not None) + if not tree_block.can_replace_page(lc_idx, num_tokens_in_block): + return tree_block.get_page(lc_idx) storage = self.manager._storage pg_idx = storage.get_pool_group_index(lc_idx) @@ -1471,19 +1465,12 @@ def _copy_page_to_tree_block( ) new_slot.ready_event = CachedCudaEvent(cuda_stream) priority = self._get_priority(tree_block.ordinal, self.manager._life_cycles[lc_idx]) - if ssm_num_tokens_in_block is None: - committed = CommittedPage(storage, tree_block, lc_idx, lvl, new_slot, priority) - else: - committed = SsmCommittedPage( - storage, - tree_block, - lc_idx, - lvl, - new_slot, - priority, - ssm_num_tokens_in_block, - ) - tree_block.storage[lc_idx] = rawref.ref(committed) + committed = CommittedPage( + storage, tree_block, lc_idx, lvl, new_slot, num_tokens_in_block, priority + ) + # Drops the superseded page, deferred until the copy is issued: an + # OutOfPagesError above must not destroy a usable shorter snapshot. + tree_block.replace_page(lc_idx, committed) storage.schedule_for_eviction(committed) return committed return None @@ -1495,17 +1482,8 @@ def _snapshot_ssm_to_tree_block( tokens_per_block = self.tokens_per_block num_tokens_in_block = num_tokens - tree_block.ordinal * tokens_per_block assert 0 < num_tokens_in_block <= tokens_per_block - existing_page = map_optional(tree_block.storage[ssm_lc_id], lambda p: p()) - existing_num_tokens = 0 - if existing_page is not None: - existing_ssm_page = expect_type(SsmCommittedPage, existing_page) - existing_num_tokens = existing_ssm_page.num_tokens_in_block - if existing_num_tokens >= num_tokens_in_block: + if tree_block.page_coverage(ssm_lc_id) >= num_tokens_in_block: return - if existing_page is not None: - tree_block.storage[ssm_lc_id] = None - if existing_page.scheduled_for_eviction: - existing_page.manager.exclude_from_eviction(existing_page) ssm_block = self._ssm_blocks[DEFAULT_BEAM_INDEX] ssm_lock = expect_type(_SharedPageLock, ssm_block[ssm_lc_id]) @@ -1513,19 +1491,15 @@ def _snapshot_ssm_to_tree_block( if move: src_page = expect_type(UncommittedPage, ssm_lock.unlock()) ssm_block[ssm_lc_id] = None - committed = src_page.convert_to_ssm_committed( + # convert_to_committed() reserves the slot, dropping any shorter snapshot. + committed = src_page.convert_to_committed( tree_block, self.finish_event, num_tokens_in_block ) storage = self.manager._storage storage.schedule_for_eviction(committed) return - self._copy_page_to_tree_block( - tree_block, - ssm_lc_id, - src_page, - ssm_num_tokens_in_block=num_tokens_in_block, - ) + self._copy_page_to_tree_block(tree_block, ssm_lc_id, src_page, num_tokens_in_block) def _snapshot_partial_block_to_tree(self, ordinal: BlockOrdinal, commit_ssm: bool) -> None: tokens_per_block = self.tokens_per_block @@ -1549,23 +1523,33 @@ def _snapshot_partial_block_to_tree(self, ordinal: BlockOrdinal, commit_ssm: boo beam_idx = DEFAULT_BEAM_INDEX beam_block = self._blocks[ordinal].pages[beam_idx] ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - # Only attach partial attention pages to a tree block whose token span is - # exactly the partial snapshot. A longer existing sibling may already - # have full attention pages; if not, a partial attention page would make - # that longer block look more reusable than it is. - if len(tree_block.tokens) == num_tokens: - for lc_idx, _ in self.manager._life_cycles.attention_life_cycles(): - holder = beam_block[lc_idx] - if holder is None or tree_block.storage[lc_idx] is not None: - continue - self._copy_page_to_tree_block(tree_block, lc_idx, holder.page) + # `tree_block` may be a longer existing sibling that covers these tokens. Attaching + # the partial attention pages to it is still correct because each page records the + # token span it covers, and prefix matching honours that span. + attached_lcs = list[LifeCycleId]() + for lc_idx, _ in self.manager._life_cycles.attention_life_cycles(): + holder = beam_block[lc_idx] + if holder is None or tree_block.page_coverage(lc_idx) >= num_tokens: + continue + if ( + self._copy_page_to_tree_block(tree_block, lc_idx, holder.page, num_tokens) + is not None + ): + attached_lcs.append(lc_idx) if commit_ssm: assert ssm_lc_id is not None self._snapshot_ssm_to_tree_block(tree_block, ssm_lc_id, start + num_tokens) - if is_new: - event_manager = self.manager.event_manager - if event_manager is not None: + event_manager = self.manager.event_manager + if event_manager is not None: + if is_new: event_manager.add_stored_block_event_from_block(tree_block) + else: + # The block was already announced, so report just the life cycles this + # snapshot added. The event manager itself drops the ones whose page does + # not span the whole block: the payload carries the block's full token + # list and cannot express a shorter valid prefix. + for lc_idx in attached_lcs: + event_manager.add_stored_life_cycle_event_from_block(tree_block, int(lc_idx)) def _commit_block( self, @@ -1626,8 +1610,7 @@ def _commit_block( for lc, (page, locked) in typed_enumerate(uncommitted_pages): if page is None: continue - p = page.convert_to_committed(tree_block, self.finish_event) - tree_block.storage[lc] = rawref.ref(p) + p = page.convert_to_committed(tree_block, self.finish_event, num_tokens) # The page comes from uncommitted page of self, so safe to skip wait. beam_block[lc] = ( p.lock(self, beam_idx, ordinal, lc, skip_wait=True) if locked else p.hold() @@ -1648,13 +1631,24 @@ def _commit_block( continue # SSM pages are not rebased if beam_block[lc] is None: continue - existing_page = map_optional(tree_block.storage[lc], lambda p: p()) + # A page covering fewer tokens than this block spans (moved in from a + # shorter sibling) is NOT a substitute for our own full page: adopting it + # would feed uninitialized KV for the uncovered tail into a live request. + existing_page = tree_block.get_page(lc) + if existing_page is not None and existing_page.num_tokens_in_block < num_tokens: + existing_page = None locked = isinstance(beam_block[lc], _SharedPageLock) if existing_page is None: # The reusable page is gone. We put our own page into the tree block. - page = cast(UncommittedPage, cast(_SharedPageLock, beam_block[lc]).page) + # Keep this a single expression: a local holding the lock/holder would + # keep it alive past `beam_block[lc] = None`, and convert_to_committed() + # requires the page to be droppable by then. + page = cast( + UncommittedPage, + cast("_SharedPageLock | _PageHolder", beam_block[lc]).page, + ) beam_block[lc] = None - p = page.convert_to_committed(tree_block, self.finish_event) + p = page.convert_to_committed(tree_block, self.finish_event, num_tokens) event_manager = self.manager.event_manager if event_manager is not None: event_manager.add_stored_life_cycle_event_from_block(tree_block, int(lc)) @@ -1840,7 +1834,9 @@ def _get_tree_block(self, ordinal: BlockOrdinal) -> Block: if lc == ssm_lc_id: assert b is None # SSM pages live in _ssm_blocks elif b is not None: - assert isinstance(b.page, CommittedPage) and b.page.block() is ret + # b.page.block() may differ from `ret`: a page can be moved to a longer + # sibling block, or replaced there by one with larger token coverage. + assert isinstance(b.page, CommittedPage) return ret def _take_uncommitted_page( @@ -2049,7 +2045,7 @@ def _setup_for_reuse(self, match: ReuseMatch) -> None: typed_range(stale_start), typed_range(stale_end, BlockOrdinal(len(matched))) ): block = self._block(ordinal, beam_idx) - holder = unwrap_rawref(unwrap_optional(matched[ordinal].storage[lc_idx])).hold() + holder = unwrap_optional(matched[ordinal].get_page(lc_idx)).hold() # For partial blocks (last block, not full), we defer the copy to first resume(). # Just store the holder of the original committed page for now. block[lc_idx] = holder @@ -2073,11 +2069,11 @@ def _setup_for_reuse(self, match: ReuseMatch) -> None: # SSM reuse: hold the snapshot from the last matched block. Copy is deferred to first resume(). if ssm_lc_id is not None and matched: snapshot_block = matched[-1] - snapshot_ref = snapshot_block.storage[ssm_lc_id] - assert snapshot_ref is not None, ( + snapshot_page = snapshot_block.get_page(ssm_lc_id) + assert snapshot_page is not None, ( "Last matched block must have SSM snapshot after truncation" ) - snapshot_holder = unwrap_rawref(snapshot_ref).hold() + snapshot_holder = snapshot_page.hold() self._ssm_blocks[DEFAULT_BEAM_INDEX][ssm_lc_id] = snapshot_holder if should_record_stats and ssm_lc_id is not None: changed = self._pending_stats.record_ssm_snapshot_lookup( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py index ad1cb3c12012..e1565aa869fc 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py @@ -524,12 +524,6 @@ def _drop_hash_cache(self, block_hash: bytes) -> None: self._v1_hash_compatible_keys.discard(block_hash) self._v1_root_attrs_by_block_key.pop(block_hash, None) - @staticmethod - def _resolve_page_ref(page_ref: Any) -> Any: - if page_ref is None: - return None - return page_ref() if callable(page_ref) else page_ref - @staticmethod def _normalize_token(token: TokenIdExt) -> UniqueToken: if isinstance(token, bytes): @@ -542,13 +536,11 @@ def _stored_block_from_radix_block( cache_level: CacheLevel = GPU_LEVEL priority: Priority = PRIORITY_DEFAULT found_page = False - for life_cycle_id, page_ref in enumerate(block.storage): + for life_cycle_id, _ in enumerate(block.storage): if life_cycle_ids is not None and life_cycle_id not in life_cycle_ids: continue - if page_ref is None: - continue - page = self._resolve_page_ref(page_ref) - if page is None: + page = block.get_page(life_cycle_id) + if page is None or page.num_tokens_in_block < len(block.tokens): continue cache_level = page.cache_level priority = page.priority @@ -568,11 +560,12 @@ def _stored_block_from_radix_block( @staticmethod def _life_cycle_ids_from_radix_block(block: Any) -> set[int]: - return { - life_cycle_id - for life_cycle_id, page_ref in enumerate(block.storage) - if page_ref is not None and KVCacheEventManager._resolve_page_ref(page_ref) is not None - } + life_cycle_ids = set[int]() + for life_cycle_id, _ in enumerate(block.storage): + page = block.get_page(life_cycle_id) + if page is not None and page.num_tokens_in_block >= len(block.tokens): + life_cycle_ids.add(life_cycle_id) + return life_cycle_ids def _parent_hash_from_radix_block(self, block: Any) -> EventBlockHash | None: parent = block.prev diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index 4a5bdc44394f..fda004aa1ebb 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -254,16 +254,13 @@ def reuse_match_pages( pages.append((slot_id, None if num_tokens_in_block < 0 else num_tokens_in_block)) return num_tokens, pages - from ._utils import unwrap_rawref - match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) py_pages: list[tuple[int, int | None] | None] = [] for block in match.blocks: - ref = block.storage[lc_id] - if ref is None: + page = block.get_page(lc_id) + if page is None: py_pages.append(None) else: - page = unwrap_rawref(ref) py_pages.append((page.slot_id, getattr(page, "num_tokens_in_block", None))) return match.num_tokens, py_pages @@ -286,13 +283,11 @@ def reuse_match_planned_drop_counts( manager, reuse_scope, list(tokens), lc_id, enable_partial ) - from ._utils import unwrap_rawref - match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) counts: list[int | None] = [] for block in match.blocks: - ref = block.storage[lc_id] - counts.append(None if ref is None else unwrap_rawref(ref).planned_drop_count) + page = block.get_page(lc_id) + counts.append(None if page is None else page.planned_drop_count) return match.num_tokens, counts diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py index 70a82c090c01..c3cee03270f5 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py @@ -164,48 +164,36 @@ def __init__( ) self.set_slot(slot) - def convert_to_committed(self, block: Block, ready_event: CachedCudaEvent) -> "CommittedPage": + def convert_to_committed( + self, block: Block, ready_event: CachedCudaEvent, num_tokens_in_block: int + ) -> "CommittedPage": """ Moves the slot to a new committed page and add the new page to the block. The uncommitted page becomes invalid. + + `num_tokens_in_block` records the page's token count. See + `CommittedPage.num_tokens_in_block` for its attention and SSM interpretations. """ assert not self.scheduled_for_eviction - assert block.storage[self.life_cycle] is None + # Check before building: replace_page() below drops the superseded page, so a + # failure in between must not lose a usable snapshot. + assert block.can_replace_page(self.life_cycle, num_tokens_in_block) # If you hit this assertion failure, it's likely because you are using debugpy, which delayed GC # for _KVCache._take_uncommitted_page(). Disable breakpoints on exceptions to avoid this issue. assert self.status == PageStatus.DROPPABLE, "Release holder/lock first" self.ready_event = ready_event committed_page = CommittedPage( - self.manager, block, self.life_cycle, self.cache_level, self, self.priority - ) - assert not self.has_valid_slot and self.ready_event is CachedCudaEvent.NULL - assert committed_page.has_valid_slot - block.storage[self.life_cycle] = rawref.ref(committed_page) - return committed_page - - def convert_to_ssm_committed( - self, block: Block, ready_event: CachedCudaEvent, num_tokens_in_block: int - ) -> "SsmCommittedPage": - """ - Moves the slot to a new committed SSM page and add the new page to the block. - The uncommitted page becomes invalid. - """ - assert not self.scheduled_for_eviction - assert block.storage[self.life_cycle] is None - assert self.status == PageStatus.DROPPABLE, "Release holder/lock first" - self.ready_event = ready_event - committed_page = SsmCommittedPage( self.manager, block, self.life_cycle, self.cache_level, self, - self.priority, num_tokens_in_block, + self.priority, ) assert not self.has_valid_slot and self.ready_event is CachedCudaEvent.NULL assert committed_page.has_valid_slot - block.storage[self.life_cycle] = rawref.ref(committed_page) + block.replace_page(self.life_cycle, committed_page) return committed_page def __del__(self) -> None: @@ -241,6 +229,16 @@ class CommittedPage(Page): """ block: rawref.ref["Block"] + # Token count recorded for this page. It is usually len(block.tokens), but a snapshot + # taken at an earlier token boundary may live in a block that spans more tokens -- see + # Block.__init__ and _KVCache._snapshot_partial_block_to_tree. + # + # Attention and SSM life cycles interpret it differently: + # * for attention pages, it is the number of leading tokens with valid per-token KV, + # so the page is reusable for any prefix up to that count (compare with `>=`); + # * for an SSM page, it is the exact recurrent-state checkpoint, so reuse must be + # truncated to exactly that boundary. + num_tokens_in_block: int planned_drop_count: int __rawref__: rawref.ref["CommittedPage"] @@ -254,9 +252,12 @@ def __init__( life_cycle: LifeCycleId, cache_level: CacheLevel, slot: Slot, + num_tokens_in_block: int, priority: Priority, ): + assert 0 < num_tokens_in_block <= len(block.tokens) self.block = rawref.ref(block) + self.num_tokens_in_block = num_tokens_in_block self.planned_drop_count = 0 self.__rawref__ = rawref.NULL Page.__init__( @@ -276,7 +277,7 @@ def __del__(self) -> None: block = self.block() # block may be None when rebase happens, i.e. another block with the same key is committed, # replacing it, but the page is still used by a _KVCache. - if block is not None and block.unlink_page(self.life_cycle, self): + if block is not None and block.unlink_page(self.life_cycle, self) is not None: Block.clear_stale_blocks_after_page_unlink( block, self.life_cycle, @@ -286,25 +287,6 @@ def __del__(self) -> None: self.__rawref__.invalidate() -@dataclass(slots=True) -class SsmCommittedPage(CommittedPage): - num_tokens_in_block: int - - def __init__( - self, - storage: "StorageManager", - block: Block, - life_cycle: LifeCycleId, - cache_level: CacheLevel, - slot: Slot, - priority: Priority, - num_tokens_in_block: int, - ): - assert num_tokens_in_block > 0 - self.num_tokens_in_block = num_tokens_in_block - CommittedPage.__init__(self, storage, block, life_cycle, cache_level, slot, priority) - - @dataclass(slots=True) class _PageHolder: "Prevents pages from being dropped." @@ -329,7 +311,10 @@ def __del__(self) -> None: if not page.scheduled_for_eviction: page.manager.schedule_for_eviction(page) block = page.block() - if block is None or block.is_orphan: + # A page that no longer sits in its block's slot (orphaned block, or replaced + # by a page with a larger recorded token count) is unreachable for reuse, so + # keeping it in the eviction LRU would just pin a slot until memory pressure hits. + if block is None or block.is_orphan or not block.holds_page(page): page.manager.exclude_from_eviction(page) elif page.scheduled_for_eviction: page = cast(UncommittedPage, self.page) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index c3d88dc0433b..576d400451b0 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -87,7 +87,8 @@ class _FakePage: - def __init__(self, cache_level=_DEFAULT_CACHE_LEVEL, priority=0): + def __init__(self, num_tokens_in_block, cache_level=_DEFAULT_CACHE_LEVEL, priority=0): + self.num_tokens_in_block = num_tokens_in_block self.cache_level = cache_level self.priority = priority @@ -118,7 +119,11 @@ def __init__(self, key, tokens, num_life_cycles=1, prev=None): self.tokens = tokens self.prev = prev or _FakeRootBlock() self.ordinal = getattr(self.prev, "ordinal", -1) + 1 - self.storage = [_FakePageRef(_FakePage()) for _ in range(num_life_cycles)] + self.storage = [_FakePageRef(_FakePage(len(tokens))) for _ in range(num_life_cycles)] + + def get_page(self, lc_idx): + page_ref = self.storage[lc_idx] + return None if page_ref is None else page_ref() with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): @@ -1133,6 +1138,21 @@ def test_v2_kv_cache_event_manager_readds_life_cycle_emits_stored_event(): assert "layer_groups" not in events[0]["data"] +def test_v2_kv_cache_event_manager_omits_partial_life_cycle_coverage(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block = _FakeBlock(b"\xab\xd3", [1, 2], num_life_cycles=2) + block.storage[1] = _FakePageRef(_FakePage(num_tokens_in_block=1)) + + event_manager.add_stored_block_event_from_block(block) + events = _flush_serialized_events(event_manager) + + assert [event["layer_group_id"] for event in events] == [0] + assert [token["token_id"] for token in events[0]["data"]["blocks"][0]["tokens"]] == [1, 2] + + event_manager.add_stored_life_cycle_event_from_block(block, 1) + assert _flush_serialized_events(event_manager) == [] + + def test_v2_kv_cache_event_manager_reemits_stored_after_all_life_cycles_were_removed(): event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) block = _FakeBlock(b"\xab\xd0", [1, 2]) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index e5b17c6bece3..27065ccdf634 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -745,6 +745,39 @@ def plan_drop(tokens: list[TokenIdExt]) -> PlannedDropHandle: with self.assertRaisesRegex(ValueError, "already been dropped"): long_handle.drop() + def test_planned_drop_handle_rejects_partial_coverage(self) -> None: + window_size = 8 + tokens_per_block = 8 + self.prepare(16 << 20, 0, 0, 2, window_size, 0, tokens_per_block=tokens_per_block) + tokens = [self.next_token() for _ in range(3 * tokens_per_block)] + + with TemporaryCudaStream([]) as stream_holder: + stream = cast(CudaStream, stream_holder.handle) + kv_cache = self.manager.create_kv_cache(None, tokens) + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(len(tokens))) + kv_cache.commit(tokens) + kv_cache.stop_committing() + + swa_lc_id = next( + lc_id + for lc_id, lc in self.manager._life_cycles.attention_life_cycles() + if lc.window_size is not None + ) + tree_block = kv_cache._blocks[2].tree_block + assert tree_block is not None + page = tree_block.get_page(swa_lc_id) + assert page is not None + self.assertEqual(page.num_tokens_in_block, len(tree_block.tokens)) + + page.num_tokens_in_block -= 1 + try: + self.assertIsNone(kv_cache.plan_committed_block_drop()) + finally: + page.num_tokens_in_block += 1 + kv_cache.close() + stream_holder.take_finish_event().synchronize() + def test_reuse_scope_isolates_reuse(self) -> None: self.prepare(16 << 20, 0, 0, 2, None, 0, tokens_per_block=8) tokens = [TokenId(i) for i in range(64)] @@ -2497,6 +2530,45 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: self.assertEqual(ssm_page[0], ssm_slot) self.assertEqual(ssm_page[1], 16) + def test_ssm_snapshot_moves_to_covering_block(self) -> None: + """A snapshot on a partial block survives the full sibling that replaces it.""" + tokens_per_block = 32 + cfg = self._make_ssm_config(tokens_per_block=tokens_per_block, enable_partial_reuse=True) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(96)] + assert self.manager._life_cycles.ssm_life_cycle_id is not None + ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id + + # Turn 1 ends at 48 tokens, i.e. 16 tokens into block 1. + kv1 = self.manager.create_kv_cache() + kv1.resume(stream) + kv1.capacity = 48 + kv1.history_length = 48 + kv1.commit(prompt[:48]) + kv1.close() + + # Turn 2 fills blocks 0..2. Block 1 becomes a full 32-token block that replaces the + # 16-token one, and its own SSM snapshot is taken on block 2, not block 1 -- so + # without moving the pages over, the 48-token endpoint would be lost. + kv2 = self.manager.create_kv_cache(input_tokens=prompt) + kv2.resume(stream) + kv2.capacity = len(prompt) + kv2.history_length = len(prompt) + kv2.commit(prompt[kv2.num_committed_tokens :]) + kv2.close() + + match = self.manager._radix_tree.match(ReuseScope(), prompt[:48], True) + self.assertEqual(match.num_tokens, 48) + block = match.blocks[-1] + self.assertEqual(len(block.tokens), tokens_per_block) + self.assertEqual(block.page_coverage(ssm_lc_id), 16) + + del match, block, kv1, kv2 + gc.collect() + stream_holder.synchronize() + def test_commit_min_snapshot_requires_history_alignment(self) -> None: """commit_min_snapshot requires commit() to start or end at history length.""" cfg = self._make_ssm_config(tokens_per_block=32) @@ -3780,6 +3852,193 @@ def test_reuse_across_prefill_turns_keeps_only_window_minus_one(self) -> None: self.manager.clear_reusable_blocks() +class TestPartialCoverageReuse(TestKVCacheManagerV2): + """Pages that cover fewer tokens than their block spans. + + A rewind endpoint (a partial trailing block, e.g. 16 tokens of a 32-token block) used + to be destroyed when a longer sibling was created, or refused when it was created + second. The covering block may not have a committable page for every life cycle at + that boundary. For SWA, commit_min_snapshot releases out-of-window + pages, while scratch reuse uses temporary shared storage that is not preserved as a + per-block page. Moving the partial sibling's pages into the covering block keeps the + endpoint reusable, with each page tagged by its recorded token count. + """ + + TOKENS_PER_BLOCK = 32 + WINDOW_SIZE = 16 + + def prepare_partial(self, gpu_quota: int = 64 << 20, window_size: int | None = None) -> None: + kv_buf_size = 8192 + window_size = self.WINDOW_SIZE if window_size is None else window_size + self.cfg = KVCacheManagerConfig( + tokens_per_block=self.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig(role=Role.KEY, size=kv_buf_size), + BufferConfig(role=Role.VALUE, size=kv_buf_size), + ], + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[ + BufferConfig(role=Role.KEY, size=kv_buf_size), + BufferConfig(role=Role.VALUE, size=kv_buf_size), + ], + sliding_window_size=window_size, + ), + ], + enable_partial_reuse=True, + commit_min_snapshot=True, + ) + self.engine = FakeEngine(self.cfg) + self.manager = KVCacheManager(self.cfg) + + @property + def _full_attn_lc_id(self) -> LayerGroupId: + return next( + lc_id + for lc_id, lc in self.manager._life_cycles.attention_life_cycles() + if lc.window_size is None + ) + + @property + def _swa_lc_id(self) -> LayerGroupId: + return next( + lc_id + for lc_id, lc in self.manager._life_cycles.attention_life_cycles() + if lc.window_size is not None + ) + + def run_turn(self, prompt: list[TokenIdExt], refcheck: bool = False) -> int: + """Reuse what we can, prefill the rest, commit, close. Returns the reused count. + + `refcheck` runs FakeEngine, which validates the reused history KV against the + expected values for every layer group, so a clean run proves the reused pages hold + real data rather than uninitialized memory. It requires `resize()` to declare the + pre-prefill history length so that every block the engine writes is in the SWA + window; without it the manager is told the final length up front and skips + allocating out-of-window blocks -- which is exactly the situation that leaves a + partial page behind, so the two cannot be combined in the same turn. + """ + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + kv_cache = self.manager.create_kv_cache(input_tokens=prompt) + num_reused = kv_cache.num_committed_tokens + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(len(prompt), num_reused if refcheck else len(prompt))) + history = list(prompt[:num_reused]) + inp = list(prompt[num_reused:]) + if refcheck: + self.engine.execute([Step(kv_cache, inp, history)], stream) + kv_cache.commit(inp) + kv_cache.close() + s.take_finish_event().synchronize() + return num_reused + + def _partial_block(self, prompt: list[TokenIdExt]): + """The tree block holding the tail of `prompt` (block 2 in these tests).""" + match = self.manager._radix_tree.match(ReuseScope(), prompt, True) + return match.blocks[-1] + + def test_rewind_endpoint_survives_longer_sibling_created_after(self) -> None: + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + extended = base + [TokenId(i) for i in range(1000, 1080)] + rewind = base + [TokenId(2000)] + + self.assertEqual(self.run_turn(base), 0) + self.assertEqual(self.run_turn(extended), len(base)) + # The 16-token endpoint block is gone, but its SWA page moved into the full + # 32-token sibling and still covers the first 16 tokens. + block = self._partial_block(rewind) + self.assertEqual(len(block.tokens), self.TOKENS_PER_BLOCK) + self.assertEqual(block.page_coverage(self._full_attn_lc_id), self.TOKENS_PER_BLOCK) + self.assertEqual(block.page_coverage(self._swa_lc_id), len(base) % self.TOKENS_PER_BLOCK) + del block + self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) + # The partial SWA page is stale at the longer endpoint and must not constrain + # the full-attention lifecycle's reusable prefix. + self.assertEqual(self.manager.probe_reuse(input_tokens=extended), len(extended)) + + def test_rewind_endpoint_attaches_to_longer_sibling_created_before(self) -> None: + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + extended = base + [TokenId(i) for i in range(1000, 1080)] + rewind = base + [TokenId(2000)] + + self.assertEqual(self.run_turn(extended), 0) + # Block 2 already spans 32 tokens, so the 80-token prompt cannot create its own + # 16-token endpoint block; its partial pages are attached to the longer sibling. + self.assertEqual(self.run_turn(base), 0) + self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) + + def test_reused_partial_coverage_kv_is_correct(self) -> None: + """The salvaged endpoint must hold real KV, not uninitialized memory.""" + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + extended = base + [TokenId(i) for i in range(1000, 1080)] + rewind = base + [TokenId(3000 + i) for i in range(40)] + + self.run_turn(base, refcheck=True) + self.run_turn(extended) + # FakeEngine checks every reused history token of both layer groups, including the + # 16 tokens of block 2 that only the salvaged partial SWA page covers. + self.assertEqual(self.run_turn(rewind, refcheck=True), len(base)) + + def test_exact_boundary_ignores_stale_last_block_partial_coverage(self) -> None: + self.prepare_partial(window_size=1) + base = [TokenId(i) for i in range(80)] + boundary = base + [TokenId(i) for i in range(1000, 1016)] + + self.assertEqual(self.run_turn(base), 0) + self.assertEqual(self.run_turn(boundary), len(base)) + + block = self._partial_block(boundary) + self.assertEqual(block.page_coverage(self._full_attn_lc_id), self.TOKENS_PER_BLOCK) + self.assertEqual(block.page_coverage(self._swa_lc_id), 16) + del block + # At the 96-token boundary the input token is the entire size-1 SWA window, so no + # historical SWA block is active. The partial SWA page must not constrain the full + # attention lifecycle's reusable prefix. + self.assertEqual(self.manager.probe_reuse(input_tokens=boundary), len(boundary)) + + def test_page_coverage_only_grows(self) -> None: + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + longer_partial = base + [TokenId(i) for i in range(1000, 1008)] # 88 tokens + rewind = base + [TokenId(2000)] + + self.run_turn(base) + block = self._partial_block(rewind) + self.assertEqual(block.page_coverage(self._swa_lc_id), 16) + del block + + self.run_turn(longer_partial) + block = self._partial_block(rewind) + # A slot keeps only the widest page. The 24-token snapshot supersedes the 16-token + # one, and it still covers the shorter rewind endpoint. + self.assertEqual(block.page_coverage(self._swa_lc_id), 24) + del block + self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) + self.assertEqual( + self.manager.probe_reuse(input_tokens=longer_partial + [TokenId(2000)]), + len(longer_partial), + ) + + # A later shorter snapshot cannot replace the wider page. + self.assertEqual(self.run_turn(base), len(base)) + block = self._partial_block(rewind) + self.assertEqual(block.page_coverage(self._swa_lc_id), 24) + del block + self.assertEqual( + self.manager.probe_reuse(input_tokens=longer_partial + [TokenId(2000)]), + len(longer_partial), + ) + + class TestSlotAllocatorShrink(unittest.TestCase): def test_shrink_underused_pool(self) -> None: # Regression for NVBug 6225866: shrinking a pool whose new size is From 8f91a1b0e206f77953f15751861bb85c7ff5a13c Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Fri, 31 Jul 2026 04:23:38 +0000 Subject: [PATCH 3/5] [None][fix] filter partial KV cache event coverage Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/eventManager.cpp | 16 +++++++++++++--- .../kv_cache_manager_v2/_event_manager.py | 8 ++++---- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp index 73fec6f57cd1..f6f4d4f48979 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp @@ -57,6 +57,16 @@ uint64_t hash64Mix(int64_t input, uint64_t seed) return seed ^ (value + static_cast(kHashCombineConst) + (seed << 6U) + (seed >> 2U)); } +bool pageCoversBlock(CommittedPage const* page, Block const& block) +{ + if (page == nullptr) + { + return false; + } + auto const* ssmPage = dynamic_cast(page); + return ssmPage == nullptr || ssmPage->numTokensInBlock >= static_cast(block.tokens.size()); +} + } // namespace EventManager::EventManager(int maxKvEventEntries, int windowSize, std::optional attentionDpRank, @@ -176,7 +186,7 @@ void EventManager::addStoredBlockUnlocked(Block const& block) std::set lifeCycleIds; for (LifeCycleId lifeCycle{0}; lifeCycle < block.storage.size(); ++lifeCycle) { - if (block.storage[lifeCycle] != nullptr) + if (pageCoversBlock(block.getPage(lifeCycle), block)) { lifeCycleIds.insert(lifeCycle.value()); } @@ -538,8 +548,8 @@ std::optional EventManager::storedBlockFromBlock( { continue; } - auto const* page = block.storage[lifeCycle]; - if (page != nullptr) + auto const* page = block.getPage(lifeCycle); + if (pageCoversBlock(page, block)) { cacheLevel = page->cacheLevel; priority = page->priority; diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py index e1565aa869fc..c493918ca367 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py @@ -335,7 +335,7 @@ def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEven def _add_event( self, - data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData, + data: (KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData), layer_group_id: LayerGroupId = None, ) -> None: if self._max_kv_event_entries <= 0: @@ -416,7 +416,7 @@ def _flush_all_removed_events_unlocked(self) -> None: def _add_event_unlocked( self, - data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData, + data: (KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData), layer_group_id: LayerGroupId = None, ) -> KVCacheEvent: if not isinstance(data, KVCacheRemovedData): @@ -536,7 +536,7 @@ def _stored_block_from_radix_block( cache_level: CacheLevel = GPU_LEVEL priority: Priority = PRIORITY_DEFAULT found_page = False - for life_cycle_id, _ in enumerate(block.storage): + for life_cycle_id in range(len(block.storage)): if life_cycle_ids is not None and life_cycle_id not in life_cycle_ids: continue page = block.get_page(life_cycle_id) @@ -561,7 +561,7 @@ def _stored_block_from_radix_block( @staticmethod def _life_cycle_ids_from_radix_block(block: Any) -> set[int]: life_cycle_ids = set[int]() - for life_cycle_id, _ in enumerate(block.storage): + for life_cycle_id in range(len(block.storage)): page = block.get_page(life_cycle_id) if page is not None and page.num_tokens_in_block >= len(block.tokens): life_cycle_ids.add(life_cycle_id) From f97ed7013b899a345217eb9343e8213a2d4b01cb Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Fri, 31 Jul 2026 09:36:47 +0000 Subject: [PATCH 4/5] [None][fix] KVCacheManagerV2: port per-page token coverage to C++ Mirrors the Python changes on this branch so both backends agree. numTokensInBlock moves from SsmCommittedPage up to CommittedPage, and SsmCommittedPage is deleted -- it no longer adds any field. The count means the number of leading tokens of the owning block the page's data is valid for: a prefix length for attention, an exact checkpoint for SSM. Block gains pageCoverage(), holdsPage(), reservePageSlot() and adoptPagesFrom(). addOrGetExistingBlock() now inserts the new block before detaching the covered sibling and moves that sibling's pages over, so a rewind endpoint survives the longer block that replaces it. pruneMatch() becomes the same fixed-point loop as Python: SSM truncates to an exact snapshot, then every attention life cycle clamps the match to the coverage of each page it still needs, skipping the stale range in between. That subsumes the old full-attention and SWA-sink passes. commitBlock()'s rebase path no longer adopts a tree page covering fewer tokens than the block spans, which would have fed uninitialized KV into a live request. snapshotPartialBlockToTree() drops its exact-span guard and reports the life cycles it attached. eventManager's pageCoversBlock() now checks every page instead of only SSM ones -- with numTokensInBlock on the base class the old dynamic_cast would have silently stopped filtering attention pages. _checkSanity() gains the scratch carve-out from the sibling Python commit, and _getTreeBlock() no longer asserts a page still points at that block, since a page may be moved to a longer sibling or replaced by a wider one. reuse_match_pages() reports the recorded count for all pages, not just SSM. The new coverage tests now go through _introspection so they exercise both backends; test_planned_drop_handle_rejects_partial_coverage stays Python-only because it has to write a page field directly. kvCacheManagerV2StatsTest.cpp is updated for the widened CommittedPage constructor. Without it the google-tests target fails to compile, which fails every *-CPP-* CI stage at fixture setup -- that target is built inside those stages, not by the Build-x86_64/Build-SBSA jobs. Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/blockRadixTree.cpp | 200 ++++++++++-------- .../kv_cache_manager_v2/blockRadixTree.h | 25 +++ .../kv_cache_manager_v2/eventManager.cpp | 10 +- .../kv_cache_manager_v2/introspection.cpp | 16 ++ .../kv_cache_manager_v2/introspection.h | 5 + .../kv_cache_manager_v2/kvCache.cpp | 123 ++++++----- .../kv_cache_manager_v2/kvCache.h | 10 +- .../kv_cache_manager_v2/page.cpp | 62 ++---- .../batch_manager/kv_cache_manager_v2/page.h | 39 ++-- .../batch_manager/kvCacheManagerV2.cpp | 13 +- .../kvCacheManagerV2StatsTest.cpp | 6 +- .../kv_cache_manager_v2/_introspection.py | 29 ++- .../test_kv_cache_manager_v2.py | 195 +++++++++++++---- 13 files changed, 456 insertions(+), 277 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index 212161acfda0..5d0d4bc314c0 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -393,6 +393,55 @@ int Block::partialMatchThisNode(TokenIdExt const* otherTokens, size_t otherCount return count; } +int Block::pageCoverage(LifeCycleId lcIdx) const +{ + auto const* page = getPage(lcIdx); + return page != nullptr ? page->numTokensInBlock : 0; +} + +bool Block::holdsPage(CommittedPage const& page) const +{ + return getPage(page.lifeCycle) == &page; +} + +bool Block::canReplacePage(LifeCycleId lcIdx, int numTokensInBlock) const +{ + auto const* existing = getPage(lcIdx); + return existing == nullptr || existing->numTokensInBlock < numTokensInBlock; +} + +void Block::replacePage(LifeCycleId lcIdx, CommittedPage* page) +{ + TLLM_CHECK_DEBUG(canReplacePage(lcIdx, page->numTokensInBlock)); + // Unlink first: excludeFromEviction() below may drop the eviction list's last + // reference and destroy the page, so the slot must already be empty. + auto* existing = unlinkPage(lcIdx); + if (existing != nullptr && existing->scheduledForEviction()) + { + existing->manager->excludeFromEviction(*existing); + existing = nullptr; // May be dangling now, set to nullptr + } + page->block = this; + storage.at(lcIdx) = page; +} + +void Block::adoptPagesFrom(Block& other) +{ + TLLM_CHECK_DEBUG(other.ordinal() == ordinal()); + for (LifeCycleId lcIdx{0}; lcIdx < storage.size(); ++lcIdx) + { + auto* page = other.getPage(lcIdx); + if (page == nullptr || !canReplacePage(lcIdx, page->numTokensInBlock)) + { + continue; + } + // Clear the source slot directly rather than via unlinkPage(), which would + // null the back-pointer replacePage() is about to overwrite. + other.storage.at(lcIdx) = nullptr; + replacePage(lcIdx, page); + } +} + CommittedPage* Block::unlinkPage(LifeCycleId lcIdx, CommittedPage* expectedPage) { auto& slot = storage.at(lcIdx); @@ -494,7 +543,12 @@ SharedPtr addOrGetExistingBlock( } } - // Remove siblings whose tokens are a strict prefix of ours. + // A later turn may extend a partial endpoint to this longer block, replacing the + // partial sibling. That turn may not have a committable SWA page for this block: + // commitMinSnapshot releases out-of-window pages, while SWA scratch reuse uses + // temporary shared storage that is not preserved. Adopt the partial sibling's pages + // to keep the shorter endpoint reusable, retaining each page's recorded token count + // (see CommittedPage::numTokensInBlock). std::vector toRemove; for (auto const& [k, sibling] : prevNext) { @@ -504,18 +558,25 @@ SharedPtr addOrGetExistingBlock( toRemove.push_back(k); } } + // Two covered siblings would be prefixes of each other; the insertion logic + // would already have replaced the shorter one. + TLLM_CHECK_DEBUG(toRemove.size() <= 1); + + // Create the new block. ordinal and tokensPerBlock are derived from prev inside the Block ctor. + auto block = makeShared(newKey, std::move(tokens), prev, numLifeCycles); + + // Keep the parent attached while covered children are replaced. Adding the replacement + // first prevents detachNext() from pruning an emptied RootBlock out of the tree. + prevNext[newKey] = block; + for (auto const& k : toRemove) { auto erasedBlock = prev->detachNext(k); TLLM_CHECK_DEBUG(erasedBlock); + block->adoptPagesFrom(*erasedBlock); TLLM_CHECK_DEBUG_WITH_INFO(erasedBlock->isOrphan(), "erased sibling must be orphan after removal"); - (void) erasedBlock; } - // Create the new block. ordinal and tokensPerBlock are derived from prev inside the Block ctor. - auto block = makeShared(newKey, std::move(tokens), prev, numLifeCycles); - - prevNext[newKey] = block; if (isNew) *isNew = true; return block; @@ -664,11 +725,6 @@ int numMatchedTokens(std::vector const& matched, in return tokensPerBlock * (static_cast(matched.size()) - 1) + matched.back().numMatchedTokens; } -bool hasPage(Block const& block, LifeCycleId lcId) -{ - return block.storage.at(lcId) != nullptr; -} - } // anonymous namespace std::vector BlockRadixTree::matchTokenPath( @@ -732,70 +788,26 @@ std::vector BlockRadixTree::pruneMatch(std::vector< [this](auto const& m) { return m.numMatchedTokens == mTokensPerBlock; })); auto attnLcs = mLifeCycles.attentionLifeCycles(); - - // Full-attention layers require pages on every matched block. - std::vector fullAttnLcList; - for (auto [lcId, attn] : attnLcs) - { - if (!attn->windowSize.has_value()) - { - fullAttnLcList.push_back(lcId); - } - } - if (!fullAttnLcList.empty()) - { - int n = findIndex(matched.begin(), matched.end(), - [&](auto const& match) - { - return std::any_of(fullAttnLcList.begin(), fullAttnLcList.end(), - [&](LifeCycleId lcId) { return !hasPage(*match.block, lcId); }); - }); - matched.resize(static_cast(n)); - } - - std::vector> swaLcs; - for (auto [lcId, attn] : attnLcs) - { - if (attn->windowSize.has_value()) - { - swaLcs.push_back({lcId, attn}); - } - } - - // SWA sink blocks must all be available. - for (auto [lcId, attn] : swaLcs) - { - int const sinkBlocks = attn->numSinkBlocks; - int const limit = std::min(sinkBlocks, static_cast(matched.size())); - int n = findIndex(matched.begin(), matched.begin() + limit, - [&, lcId = lcId](auto const& match) { return !hasPage(*match.block, lcId); }); - if (n < sinkBlocks) - { - matched.resize(static_cast(n)); - } - } - auto ssmLcId = mLifeCycles.ssmLifeCycleId(); + + // Fixed-point loop: SSM may select an earlier exact snapshot, while attention may + // shorten the match to the coverage of a required page. Every retry strictly + // shortens the match, so the loop terminates. while (!matched.empty()) { + // Check SSM snapshot availability first: truncating to the last reusable SSM + // snapshot changes the matched length that all the attention checks use. if (ssmLcId.has_value()) { - // Truncate to the last block whose SSM snapshot is reusable at that - // block's matched-token count, then clamp the tail entry's matched - // token count to the snapshot length (mirrors _block_radix_tree.py). int ssmTrunc = 0; int ssmMatchLen = 0; for (int i = static_cast(matched.size()) - 1; i >= 0; --i) { - CommittedPage* page = matched[static_cast(i)].block->storage.at(*ssmLcId); - if (page == nullptr) - { - continue; - } - auto* ssmPage = dynamic_cast(page); - TLLM_CHECK_DEBUG(ssmPage != nullptr); - int const snapshotLen = ssmPage->numTokensInBlock; - if (matched[static_cast(i)].numMatchedTokens >= snapshotLen) + auto const& entry = matched[static_cast(i)]; + // An SSM page holds the recurrent state after exactly this many tokens, + // so reuse must stop there instead of anywhere inside the block. + int const snapshotLen = entry.block->pageCoverage(*ssmLcId); + if (snapshotLen > 0 && entry.numMatchedTokens >= snapshotLen) { ssmTrunc = i + 1; ssmMatchLen = snapshotLen; @@ -810,35 +822,49 @@ std::vector BlockRadixTree::pruneMatch(std::vector< matched.back().numMatchedTokens = ssmMatchLen; } + // Only pages that are active at this candidate endpoint constrain attention + // reuse. Full attention requires every block. SWA requires sink blocks and the + // trailing window, but not the stale blocks between them. In particular, at an + // exact block boundary with windowSize=1, every historical block is stale. int const numTok = numMatchedTokens(matched, mTokensPerBlock); - bool trimmed = false; - for (auto [lcId, attn] : swaLcs) + bool shortened = false; + for (auto [lcId, attn] : attnLcs) { - int n = findIndex(matched.rbegin(), matched.rend(), - [&, lcId = lcId](auto const& match) { return hasPage(*match.block, lcId); }); - if (n != 0) + auto const staleRange = attn->getStaleRange(numTok, mTokensPerBlock); + int const staleBeg = staleRange.beg.value(); + int const staleEnd = staleRange.end.value(); + int const numMatchedBlocks = static_cast(matched.size()); + for (int i = 0; i < numMatchedBlocks; ++i) { - matched.resize(matched.size() - static_cast(n)); - trimmed = true; + // Mirrors Python's chain(range(stale.beg), range(stale.end, len(matched))). + if (staleBeg <= i && i < staleEnd) + { + continue; + } + int const numMatched = matched[static_cast(i)].numMatchedTokens; + int const coverage = matched[static_cast(i)].block->pageCoverage(lcId); + if (coverage >= numMatched) + { + continue; + } + if (coverage > 0) + { + matched.resize(static_cast(i) + 1); + matched.back().numMatchedTokens = coverage; + } + else + { + matched.resize(static_cast(i)); + } + shortened = true; break; } - - auto staleRange = attn->getStaleRange(numTok, mTokensPerBlock); - BlockOrdinal const staleEnd = staleRange.end; - if (staleEnd < BlockOrdinal{static_cast(matched.size())}) + if (shortened) { - auto tailBegin = matched.begin() + static_cast(toSizeT(staleEnd)); - int nMissing = findIndex(matched.rbegin(), std::make_reverse_iterator(tailBegin), - [&, lcId = lcId](auto const& match) { return !hasPage(*match.block, lcId); }); - if (BlockOrdinal{static_cast(matched.size()) - nMissing} > staleEnd) - { - matched.resize(matched.size() - static_cast(nMissing) - 1); - trimmed = true; - break; - } + break; } } - if (!trimmed) + if (!shortened) { break; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h index c3ae3035415c..c4a1debcb093 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h @@ -234,6 +234,31 @@ struct Block : NodeBase, EnableSharedFromThis return storage[lcIdx]; } + // Return the page's recorded token count, or zero if the slot is empty. For attention + // this is prefix coverage; for SSM it is an exact checkpoint position. + // Mirrors Python's Block.page_coverage(). + int pageCoverage(LifeCycleId lcIdx) const; + + // True when `page` currently occupies its lifecycle's slot in this block. + // Mirrors Python's Block.holds_page(). + bool holdsPage(CommittedPage const& page) const; + + // Whether a page recording `numTokensInBlock` may take over slot `lcIdx`. A slot + // keeps only the page with the largest recorded token count; for SSM that means only + // the latest checkpoint, and a rare second endpoint in one block is a reuse miss. + // Pure; use replacePage() to install. Mirrors Python's Block.can_replace_page(). + bool canReplacePage(LifeCycleId lcIdx, int numTokensInBlock) const; + + // Install `page` in slot `lcIdx`, detaching whatever it supersedes. The superseded + // page may outlive this call while a request still holds it, so unlinkPage() must + // clear its back-pointer — releasePages() walks `storage` and would never see it. + // Mirrors Python's Block.replace_page(). + void replacePage(LifeCycleId lcIdx, CommittedPage* page); + + // Move `other`'s pages into this block without changing their recorded token counts. + // Mirrors Python's Block._adopt_pages_from(). + void adoptPagesFrom(Block& other); + // Clear stale tree nodes after a lifecycle page has been unlinked. // Returns detached blocks that must stay alive until cleanup completes. static std::vector> clearStaleBlocksAfterPageUnlink( diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp index f6f4d4f48979..301d941d1cfc 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp @@ -57,14 +57,12 @@ uint64_t hash64Mix(int64_t input, uint64_t seed) return seed ^ (value + static_cast(kHashCombineConst) + (seed << 6U) + (seed >> 2U)); } +// A page whose recorded token count is short of the block's span cannot be announced: the +// event payload carries the block's full token list and cannot express a shorter valid +// prefix. bool pageCoversBlock(CommittedPage const* page, Block const& block) { - if (page == nullptr) - { - return false; - } - auto const* ssmPage = dynamic_cast(page); - return ssmPage == nullptr || ssmPage->numTokensInBlock >= static_cast(block.tokens.size()); + return page != nullptr && page->numTokensInBlock >= static_cast(block.tokens.size()); } } // namespace diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp index 835c48c9a13b..4a167d3792f3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp @@ -21,6 +21,7 @@ #include "kv_cache_manager_v2/kvCache.h" #include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/page.h" #include "kv_cache_manager_v2/storageManager.h" #include @@ -79,6 +80,21 @@ KvCacheIntrospection::ActivePageStats KvCacheIntrospection::activePageStats(KvCa return {std::move(counts), std::move(unscheduledEvictable)}; } +std::optional KvCacheIntrospection::committedPageIsLinked(KvCache const& kvCache, int ordinal, int lcId) +{ + auto page = kvCache._page(BlockOrdinal{ordinal}, kDefaultBeamIndex, LifeCycleId{lcId}); + if (!page) + { + return std::nullopt; + } + auto committed = dynamicPointerCast(page); + if (!committed) + { + return std::nullopt; + } + return committed->block != nullptr; +} + TypedVec KvCacheIntrospection::storageStatistics( KvCacheManager& manager, CacheLevel level) { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h index 07f0dd9acef0..3d11ab920221 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h @@ -35,6 +35,11 @@ class KvCacheIntrospection using ActivePageStats = std::tuple, TypedVec>; static ActivePageStats activePageStats(KvCache const& kvCache); + + // Whether the sequence's page at (ordinal, lcId) still points at a tree block; + // nullopt when the slot is empty or holds an uncommitted page. Test hook for the + // back-pointer invariant Block::replacePage() maintains. + static std::optional committedPageIsLinked(KvCache const& kvCache, int ordinal, int lcId); static bool allTreePagesDroppable(KvCacheManager& manager); static TypedVec storageStatistics(KvCacheManager& manager, CacheLevel level); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 1715f44c7463..637deb15d3f9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -751,6 +751,8 @@ void KvCache::_subtractPendingAllocationRange(BlockOrdinal blockBegin, BlockOrdi bool KvCache::_hasReuseSource(BlockPage const& page) { + // `block` is set only while the page occupies its block's slot, so this asks whether + // the tree still offers the page. auto const committedPage = dynamicPointerCast(blockPageGetPage(page)); return committedPage && committedPage->block != nullptr; } @@ -779,16 +781,13 @@ void KvCache::_clearBlocks() // attached to a radix tree block. Mirrors Python's _copy_page_to_tree_block. // --------------------------------------------------------------------------- -void KvCache::_copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, - std::optional ssmNumTokensInBlock) +CommittedPage* KvCache::_copyPageToTreeBlock( + SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, int numTokensInBlock) { - if (treeBlock->storage.at(lcIdx) != nullptr) + if (!treeBlock->canReplacePage(lcIdx, numTokensInBlock)) { - return; // block already holds a page for this lifecycle + return treeBlock->getPage(lcIdx); } - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - bool const isSsm = ssmLcId.has_value() && lcIdx == *ssmLcId; - TLLM_CHECK_DEBUG(isSsm == ssmNumTokensInBlock.has_value()); auto& storageMgr = mManager->storage(); PoolGroupIndex pgIdx = storageMgr.getPoolGroupIndex(lcIdx); @@ -813,22 +812,17 @@ void KvCache::_copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleI CachedCudaEvent readyEv(reinterpret_cast(stream)); auto tempPage = makeShared(*this, treeBlock->ordinal(), lcIdx, lvl, kDefaultBeamIndex); tempPage->setSlot(newSlot); - SharedPtr committed; - if (ssmNumTokensInBlock.has_value()) - { - committed = tempPage->convertToSsmCommitted(treeBlock, std::move(readyEv), *ssmNumTokensInBlock); - } - else - { - committed = tempPage->convertToCommitted(treeBlock, std::move(readyEv)); - } + // Drops the superseded page, deferred until the copy is issued: an + // OutOfPagesError above must not destroy a usable shorter snapshot. + auto committed = tempPage->convertToCommitted(treeBlock, std::move(readyEv), numTokensInBlock); // Schedule for eviction so eviction controller keeps a strong reference, // preventing the page from being destroyed. storageMgr.scheduleForEviction(*committed); - return; // success + return committed.get(); // success } // No pages available in any level, silently skip snapshot (matches Python). + return nullptr; } // --------------------------------------------------------------------------- @@ -841,28 +835,10 @@ void KvCache::_snapshotSsmToTreeBlock(SharedPtr const& treeBlock, LifeCyc int const numTokensInBlock = numTokens - treeBlock->ordinal().value() * mTokensPerBlock; TLLM_CHECK_DEBUG(0 < numTokensInBlock && numTokensInBlock <= mTokensPerBlock); - CommittedPage* existingRaw = treeBlock->storage.at(ssmLcId); - int existingNumTokens = 0; - if (existingRaw != nullptr) - { - auto* existingSsm = dynamic_cast(existingRaw); - TLLM_CHECK_DEBUG(existingSsm != nullptr); - existingNumTokens = existingSsm->numTokensInBlock; - } - if (existingNumTokens >= numTokensInBlock) + if (treeBlock->pageCoverage(ssmLcId) >= numTokensInBlock) { return; } - if (existingRaw != nullptr) - { - // Detach the smaller snapshot; the CommittedPage dtor's expected-page guard - // makes the later unlink a no-op (the slot is already null). - treeBlock->storage.at(ssmLcId) = nullptr; - if (existingRaw->scheduledForEviction()) - { - existingRaw->manager->excludeFromEviction(*existingRaw); - } - } auto& ssmBlock = mSsmBlocks[kDefaultBeamIndex]; auto* ssmLock = std::get_if(&ssmBlock[ssmLcId]); @@ -874,7 +850,8 @@ void KvCache::_snapshotSsmToTreeBlock(SharedPtr const& treeBlock, LifeCyc auto up = dynamicPointerCast(unlocked); TLLM_CHECK_DEBUG(up != nullptr); ssmBlock[ssmLcId] = std::monostate{}; - auto committed = up->convertToSsmCommitted(treeBlock, finishEvent(), numTokensInBlock); + // convertToCommitted() installs via replacePage(), dropping any shorter snapshot. + auto committed = up->convertToCommitted(treeBlock, finishEvent(), numTokensInBlock); mManager->storage().scheduleForEviction(*committed); return; } @@ -923,21 +900,21 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) auto& beamBlock = mBlocks.at(ordinal).pages[kDefaultBeamIndex]; auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - // Only attach partial attention pages to a tree block whose token span is - // exactly the partial snapshot. A longer existing sibling may already have - // full attention pages; otherwise a partial page would make it look more - // reusable than it is. - if (static_cast(treeBlock->tokens.size()) == numTokens) + // `treeBlock` may be a longer existing sibling that covers these tokens. Attaching the + // partial attention pages to it is still correct because each page records the token + // span it covers, and prefix matching honours that span. + std::vector attachedLcs; + for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) { - for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) + (void) attn; + auto& bp = beamBlock[lcIdx]; + if (blockPageIsNull(bp) || treeBlock->pageCoverage(lcIdx) >= numTokens) { - (void) attn; - auto& bp = beamBlock[lcIdx]; - if (blockPageIsNull(bp) || treeBlock->storage.at(lcIdx) != nullptr) - { - continue; - } - _copyPageToTreeBlock(treeBlock, lcIdx, blockPageGetPage(bp)); + continue; + } + if (_copyPageToTreeBlock(treeBlock, lcIdx, blockPageGetPage(bp), numTokens) != nullptr) + { + attachedLcs.push_back(lcIdx); } } if (commitSsm) @@ -945,9 +922,23 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) TLLM_CHECK_DEBUG(ssmLcId.has_value()); _snapshotSsmToTreeBlock(treeBlock, *ssmLcId, start + numTokens); } - if (isNew && treeBlock->eventSink) + if (treeBlock->eventSink) { - treeBlock->eventSink->addStoredBlock(*treeBlock); + if (isNew) + { + treeBlock->eventSink->addStoredBlock(*treeBlock); + } + else + { + // The block was already announced, so report just the life cycles this snapshot + // added. The event manager itself drops the ones whose page does not span the + // whole block: the payload carries the block's full token list and cannot + // express a shorter valid prefix. + for (auto lcIdx : attachedLcs) + { + treeBlock->eventSink->addStoredLifeCycle(*treeBlock, lcIdx); + } + } } } @@ -1519,7 +1510,8 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) int start = ord * mTokensPerBlock; int end = std::min(start + mTokensPerBlock, static_cast(mCommittedTokens.size())); std::vector tokenBlock(mCommittedTokens.begin() + start, mCommittedTokens.begin() + end); - bool isFull = static_cast(tokenBlock.size()) == mTokensPerBlock; + int const numTokens = static_cast(tokenBlock.size()); + bool const isFull = (numTokens == mTokensPerBlock); if (!isLast && !isFull) throw LogicError("Cannot commit block that is not full except last block"); @@ -1568,7 +1560,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) auto& [up, locked] = taken[lc]; if (!up) continue; - auto committed = up->convertToCommitted(newBlock, finishEvent()); + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); if (locked) sb.pages[kDefaultBeamIndex][lc] = committed->lock(*this, kDefaultBeamIndex, static_cast(ord), lc); @@ -1595,7 +1587,14 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) auto& bp = sb.pages[kDefaultBeamIndex][lc]; if (blockPageIsNull(bp)) continue; - auto* existingPage = newBlock->storage.at(lc); + // A page covering fewer tokens than this block spans (moved in from a shorter + // sibling) is NOT a substitute for our own full page: adopting it would feed + // uninitialized KV for the uncovered tail into a live request. + auto* existingPage = newBlock->getPage(lc); + if (existingPage != nullptr && existingPage->numTokensInBlock < numTokens) + { + existingPage = nullptr; + } bool isLocked = std::holds_alternative(bp); if (existingPage == nullptr) { @@ -1606,7 +1605,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) if (up) { bp = std::monostate{}; - auto committed = up->convertToCommitted(newBlock, finishEvent()); + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); if (newBlock->eventSink) { newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); @@ -1624,7 +1623,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) if (up) { bp = std::monostate{}; - auto committed = up->convertToCommitted(newBlock, finishEvent()); + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); if (newBlock->eventSink) { newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); @@ -2126,8 +2125,10 @@ SharedPtr const& KvCache::_getTreeBlock(BlockOrdinal ordinal) const else if (!blockPageIsNull(beamBlock[lcId])) { auto page = blockPageGetPage(beamBlock[lcId]); + // committed->block may differ from `ret`: a page can be moved to a longer + // sibling block, or replaced there by one with larger token coverage. auto committed = dynamicPointerCast(page); - TLLM_CHECK(committed && committed->block == ret.get()); + TLLM_CHECK(committed); } } } @@ -2196,6 +2197,12 @@ bool KvCache::_checkSanity() const { TLLM_CHECK_DEBUG(blockPageIsNull(bp)); } + else if (_getScratchRange(lcs.getLifeCycle(lc), 0).contains(ordinal)) + { + // It is uncertain whether this block should hold a page: it may have + // been scratch-allocated by an earlier chunk, but the prior history + // length and capacity are not retained. + } else { // For the decoder-side disagg case, for the first step, we will skip the diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index bd090197433b..2fca81c8d17e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -457,11 +457,11 @@ class KvCache : public std::enable_shared_from_this std::vector _getMatchedTokens(BlockRadixTree::ReuseMatch const& match) const; void _clearBlocks(); // Copy `srcPage` into a new committed page attached to `treeBlock` for lifecycle - // `lcIdx`. When `ssmNumTokensInBlock` is set, the copy is an SsmCommittedPage - // covering that many tokens; otherwise a plain attention CommittedPage. No-op if - // the block already holds a page for this lifecycle, or on OOM in all levels. - void _copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, - std::optional ssmNumTokensInBlock = std::nullopt); + // `lcIdx`, recording `numTokensInBlock` as the page's token count. Returns the page + // now in the slot, or nullptr on OOM in all levels. No-op (returning the existing + // page) if the block already holds a page covering at least that many tokens. + CommittedPage* _copyPageToTreeBlock( + SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, int numTokensInBlock); // Snapshot live SSM state to `treeBlock` reusable at `numTokens` committed tokens. // If `move`, the live SSM page is moved (not copied) into the tree — the caller diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index 91f9921df948..01e34e589a03 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -22,7 +22,6 @@ #include "kv_cache_manager_v2/storageManager.h" // for StorageManager #include "tensorrt_llm/common/assert.h" -#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -96,18 +95,13 @@ SharedPageLock Page::lock(KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal or // CommittedPage // --------------------------------------------------------------------------- -CommittedPage::CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio) +CommittedPage::CommittedPage( + StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, int numTokensInBlock_, Priority prio) : Page(mgr, lc, level, prio) , block(blk.get()) -{ -} - -SsmCommittedPage::SsmCommittedPage( - StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio, int numTokensInBlock_) - : CommittedPage(mgr, std::move(blk), lc, level, prio) , numTokensInBlock(numTokensInBlock_) { - TLLM_CHECK_DEBUG(numTokensInBlock_ > 0); + TLLM_CHECK_DEBUG(0 < numTokensInBlock_ && numTokensInBlock_ <= static_cast(blk->tokens.size())); } CommittedPage::~CommittedPage() @@ -172,17 +166,20 @@ UncommittedPage::~UncommittedPage() // Delegate slot release to Page::~Page(). } -SharedPtr UncommittedPage::convertToCommitted(SharedPtr blk, CachedCudaEvent readyEv) +SharedPtr UncommittedPage::convertToCommitted( + SharedPtr blk, CachedCudaEvent readyEv, int numTokensInBlock) { TLLM_CHECK_DEBUG(!scheduledForEviction()); - TLLM_CHECK_DEBUG_WITH_INFO( - blk->storage.at(lifeCycle) == nullptr, "Block slot for this lifecycle already has a committed page"); + // Check before building: replacePage() below drops the superseded page, so a failure + // in between must not lose a usable snapshot. + TLLM_CHECK_DEBUG_WITH_INFO(blk->canReplacePage(lifeCycle, numTokensInBlock), + "Block slot for this lifecycle already has a page covering more tokens"); TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE, "Release holder/lock before converting"); // Set the ready event before transfer (matches Python: self.ready_event = ready_event). this->readyEvent = std::move(readyEv); - auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, priority); + auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, numTokensInBlock, priority); // Move slot id to the committed page; invalidate our slot. committed->setSlotId(slotId()); // asserts valid committed->readyEvent = std::move(readyEvent); @@ -193,31 +190,7 @@ SharedPtr UncommittedPage::convertToCommitted(SharedPtr bl TLLM_CHECK_DEBUG_WITH_INFO(committed->hasValidSlot(), "committed page must have a valid slot after transfer"); // Register in block storage. - blk->storage.at(lifeCycle) = committed.get(); - - return committed; -} - -SharedPtr UncommittedPage::convertToSsmCommitted( - SharedPtr blk, CachedCudaEvent readyEv, int numTokensInBlock) -{ - TLLM_CHECK_DEBUG(!scheduledForEviction()); - TLLM_CHECK_DEBUG_WITH_INFO( - blk->storage.at(lifeCycle) == nullptr, "Block slot for this lifecycle already has a committed page"); - TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE, "Release holder/lock before converting"); - - this->readyEvent = std::move(readyEv); - - auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, priority, numTokensInBlock); - committed->setSlotId(slotId()); // asserts valid - committed->readyEvent = std::move(readyEvent); - resetSlot(); - readyEvent = CachedCudaEvent::makeNull(); - - TLLM_CHECK_DEBUG(!hasValidSlot() && readyEvent.isClosed()); - TLLM_CHECK_DEBUG_WITH_INFO(committed->hasValidSlot(), "committed page must have a valid slot after transfer"); - - blk->storage.at(lifeCycle) = committed.get(); + blk->replacePage(lifeCycle, committed.get()); return committed; } @@ -244,13 +217,12 @@ PageHolder::~PageHolder() if (!page->scheduledForEviction()) manager->scheduleForEviction(*page); - // If the block is orphan, exclude from eviction immediately. - auto* cp = dynamic_cast(page.get()); - if (cp) - { - if (cp->block == nullptr || cp->block->isOrphan()) - manager->excludeFromEviction(*page); - } + // A page that no longer sits in its block's slot (orphaned block, or replaced by a + // page with a larger recorded token count) is unreachable for reuse, so keeping it + // in the eviction LRU would just pin a slot until memory pressure hits. + auto* cp = static_cast(page.get()); + if (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp)) + manager->excludeFromEviction(*page); } else { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index 9f52dca1f0ee..fc3a02252471 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -95,11 +95,23 @@ class CommittedPage : public Page public: Block* block; + // Token count recorded for this page. It is usually block->tokens.size(), but a + // snapshot taken at an earlier token boundary may live in a block that spans more + // tokens — see addOrGetExistingBlock() and KvCache::_snapshotPartialBlockToTree(). + // + // Attention and SSM life cycles interpret it differently: + // * for attention pages, it is the number of leading tokens with valid per-token KV, + // so the page is reusable for any prefix up to that count (compare with `>=`); + // * for an SSM page, it is the exact recurrent-state checkpoint, so reuse must be + // truncated to exactly that boundary. + int numTokensInBlock; + // Number of outstanding PlannedDropHandles that intend to drop this page. // Mirrors Python's CommittedPage.planned_drop_count. int plannedDropCount{0}; - CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio); + CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, int numTokensInBlock, + Priority prio); ~CommittedPage() override; @@ -109,22 +121,6 @@ class CommittedPage : public Page } }; -// --------------------------------------------------------------------------- -// SsmCommittedPage — a committed SSM snapshot page. -// -// Unlike attention CommittedPages (which always cover a full block), an SSM -// snapshot may cover only a prefix of its block. `numTokensInBlock` records how -// many tokens of the block this snapshot is reusable for. -// --------------------------------------------------------------------------- -class SsmCommittedPage : public CommittedPage -{ -public: - int numTokensInBlock; - - SsmCommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio, - int numTokensInBlock); -}; - // --------------------------------------------------------------------------- // UncommittedPage — page associated with a live KvCache sequence. // --------------------------------------------------------------------------- @@ -147,11 +143,10 @@ class UncommittedPage : public Page // Convert this UncommittedPage into a CommittedPage and attach to `block`. // The UncommittedPage becomes invalid (slot transferred to CommittedPage). - SharedPtr convertToCommitted(SharedPtr block, CachedCudaEvent readyEvent); - - // Convert this UncommittedPage into an SsmCommittedPage covering - // `numTokensInBlock` tokens and attach to `block`. Invalidates this page. - SharedPtr convertToSsmCommitted( + // + // `numTokensInBlock` records the page's token count. See + // CommittedPage::numTokensInBlock for its attention and SSM interpretations. + SharedPtr convertToCommitted( SharedPtr block, CachedCudaEvent readyEvent, int numTokensInBlock); }; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 7cda4942c104..0ddce2a3a8f7 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1521,6 +1521,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) return std::make_tuple(std::move(counts.raw()), std::move(unscheduledEvictable.raw())); }, nb::arg("kv_cache"), nb::call_guard()); + mIntrospection.def("committed_page_is_linked", &kv::KvCacheIntrospection::committedPageIsLinked, + nb::arg("kv_cache"), nb::arg("ordinal"), nb::arg("lc_id"), nb::call_guard()); mIntrospection.def("all_tree_pages_droppable", &kv::KvCacheIntrospection::allTreePagesDroppable, nb::arg("manager"), nb::call_guard()); mIntrospection.def( @@ -1609,8 +1611,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) }, nb::arg("manager"), nb::call_guard()); // Returns (num_tokens, pages) where pages[i] is None for a block with no page in - // this lifecycle, else (slot_id, num_tokens_in_block) with num_tokens_in_block = -1 - // for a non-SSM (attention) page. + // this lifecycle, else (slot_id, num_tokens_in_block). mIntrospection.def( "reuse_match_pages", [](kv::KvCacheManager& manager, nb::object reuseScope, nb::object tokens, int lcId, bool enablePartial) @@ -1627,17 +1628,13 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) pages.reserve(matchResult.blocks.stdSize()); for (auto* block : matchResult.blocks) { - auto* page = block->storage.at(lc); + auto* page = block->getPage(lc); if (page == nullptr) { pages.emplace_back(std::nullopt); continue; } - int const slotId = page->slotId().value(); - int numTokensInBlock = -1; - if (auto* ssm = dynamic_cast(page)) - numTokensInBlock = ssm->numTokensInBlock; - pages.emplace_back(std::make_pair(slotId, numTokensInBlock)); + pages.emplace_back(std::make_pair(page->slotId().value(), page->numTokensInBlock)); } } return std::make_tuple(numTokens, std::move(pages)); diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index cd91444e6028..2f7169638d29 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -220,7 +220,8 @@ TEST(KvCacheManagerV2StatsTest, PeakBlockStatsResetStartsNextIntervalFromCurrent tokens.emplace_back(TokenId{token++}); } auto block = addOrGetExistingBlock(previous, LifeCycleId{1}, std::move(tokens)); - auto page = makeShared(&storage, block, lifeCycle, kGpuLevel, kPriorityDefault); + auto page = makeShared( + &storage, block, lifeCycle, kGpuLevel, static_cast(block->tokens.size()), kPriorityDefault); page->setSlot(slot); block->storage[lifeCycle] = page.get(); storage.scheduleForEviction(*page); @@ -311,7 +312,8 @@ TEST(KvCacheManagerV2StatsTest, MigrationAndLastTierDropRecordersReceiveExactPag tokens.emplace_back(TokenId{tokenBase++}); } auto block = addOrGetExistingBlock(previous, LifeCycleId{1}, std::move(tokens)); - auto page = makeShared(&storage, block, lifeCycle, kGpuLevel, kPriorityDefault); + auto page = makeShared( + &storage, block, lifeCycle, kGpuLevel, static_cast(block->tokens.size()), kPriorityDefault); page->setSlot(slot); block->storage[lifeCycle] = page.get(); storage.scheduleForEviction(*page); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index fda004aa1ebb..f8f5e3214f6e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -49,6 +49,30 @@ def active_page_stats(kv_cache: Any) -> tuple[list[int], list[int]]: return counts, unscheduled_evictable +def committed_page_is_linked(kv_cache: Any, ordinal: int, lc_id: int) -> bool | None: + """Whether the sequence's page at ``(ordinal, lc_id)`` still points at a tree block. + + ``None`` when the slot is empty or holds an uncommitted page. Test hook: in C++ the + back-pointer is raw, so a page left pointing at a block that dies first is read after + free, and freed-but-mapped memory reads back plausibly enough that only a sanitizer + build catches the fault itself. + """ + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return cpp_introspection.committed_page_is_linked(kv_cache, ordinal, lc_id) + + from ._common import DEFAULT_BEAM_INDEX + from ._page import CommittedPage + + block_page = kv_cache._page(ordinal, DEFAULT_BEAM_INDEX, lc_id) + if block_page is None: + return None + page = block_page.page + if not isinstance(page, CommittedPage): + return None + return page.block() is not None + + def all_tree_pages_droppable(manager: Any) -> bool: """Return whether every page reachable from the radix tree is droppable.""" cpp_introspection = _cpp_introspection_module() @@ -237,8 +261,9 @@ def reuse_match_pages( """Match ``tokens`` against the radix tree and report reusable pages per block. Returns ``(num_tokens, pages)`` where ``pages[i]`` is ``None`` when block ``i`` - holds no page for lifecycle ``lc_id``, otherwise ``(slot_id, num_tokens_in_block)`` - with ``num_tokens_in_block`` set only for SSM pages (``None`` for attention pages). + holds no page for lifecycle ``lc_id``, otherwise ``(slot_id, num_tokens_in_block)``. + See ``CommittedPage.num_tokens_in_block`` for how attention and SSM life cycles + interpret the recorded token count. """ cpp_introspection = _cpp_introspection_module() if cpp_introspection is not None: diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 27065ccdf634..4c93a13c58d5 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -26,7 +26,7 @@ from importlib.util import find_spec from random import randbytes from statistics import median -from typing import TYPE_CHECKING, Iterator, NamedTuple, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, cast, get_type_hints if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: from kv_cache_manager_v2 import ( @@ -148,8 +148,19 @@ from kernels import HostGate, enable_kernel_delay +KV_CACHE_MANAGER_V2_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() + +# Gate for white-box tests that reach into the pure-Python implementation's objects +# (e.g. mutating a CommittedPage field). Prefer `_introspection`, which works on both +# backends; use this only when the behaviour under test cannot be reached through it. +requires_python_backend = unittest.skipIf( + KV_CACHE_MANAGER_V2_BACKEND == "cpp", + "white-box test over pure-Python KVCacheManagerV2 internals", +) + + def get_cached_cuda_event_type(): - backend = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() + backend = KV_CACHE_MANAGER_V2_BACKEND if backend == "cpp": try: from bindings.internal.batch_manager.kv_cache_manager_v2 import CachedCudaEvent @@ -745,7 +756,11 @@ def plan_drop(tokens: list[TokenIdExt]) -> PlannedDropHandle: with self.assertRaisesRegex(ValueError, "already been dropped"): long_handle.drop() + @requires_python_backend def test_planned_drop_handle_rejects_partial_coverage(self) -> None: + # plan_committed_block_drop() rejects via _prune_match, which clamps the match to + # the page's recorded token count, so the endpoint no longer matches exactly. + # Forcing that state needs a direct write to the page, hence the backend gate. window_size = 8 tokens_per_block = 8 self.prepare(16 << 20, 0, 0, 2, window_size, 0, tokens_per_block=tokens_per_block) @@ -2538,8 +2553,8 @@ def test_ssm_snapshot_moves_to_covering_block(self) -> None: stream_holder = CachedCudaStream() stream = cast(CudaStream, stream_holder.handle) prompt = [self.next_token() for _ in range(96)] - assert self.manager._life_cycles.ssm_life_cycle_id is not None - ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id + ssm_lc_id = _introspection.ssm_life_cycle_id(self.manager) + assert ssm_lc_id is not None # Turn 1 ends at 48 tokens, i.e. 16 tokens into block 1. kv1 = self.manager.create_kv_cache() @@ -2559,13 +2574,21 @@ def test_ssm_snapshot_moves_to_covering_block(self) -> None: kv2.commit(prompt[kv2.num_committed_tokens :]) kv2.close() - match = self.manager._radix_tree.match(ReuseScope(), prompt[:48], True) - self.assertEqual(match.num_tokens, 48) - block = match.blocks[-1] - self.assertEqual(len(block.tokens), tokens_per_block) - self.assertEqual(block.page_coverage(ssm_lc_id), 16) + num_tokens, pages = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt[:48], ssm_lc_id, True + ) + self.assertEqual(num_tokens, 48) + self.assertEqual(len(pages), 2) + self.assertIsNotNone(pages[-1]) + self.assertEqual(cast(tuple, pages[-1])[1], 16) + # Block 1 is the full 32-token sibling, not the original 16-token block: the + # 96-token prompt still matches end to end through it. + full_match, _ = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt, ssm_lc_id, True + ) + self.assertEqual(full_match, len(prompt)) - del match, block, kv1, kv2 + del kv1, kv2 gc.collect() stream_holder.synchronize() @@ -3897,20 +3920,17 @@ def prepare_partial(self, gpu_quota: int = 64 << 20, window_size: int | None = N self.manager = KVCacheManager(self.cfg) @property - def _full_attn_lc_id(self) -> LayerGroupId: + def _full_attn_lc_id(self) -> int: + swa = set(_introspection.swa_life_cycle_ids(self.manager)) return next( lc_id - for lc_id, lc in self.manager._life_cycles.attention_life_cycles() - if lc.window_size is None + for lc_id in _introspection.attention_life_cycle_ids(self.manager) + if lc_id not in swa ) @property - def _swa_lc_id(self) -> LayerGroupId: - return next( - lc_id - for lc_id, lc in self.manager._life_cycles.attention_life_cycles() - if lc.window_size is not None - ) + def _swa_lc_id(self) -> int: + return _introspection.swa_life_cycle_ids(self.manager)[0] def run_turn(self, prompt: list[TokenIdExt], refcheck: bool = False) -> int: """Reuse what we can, prefill the rest, commit, close. Returns the reused count. @@ -3938,10 +3958,106 @@ def run_turn(self, prompt: list[TokenIdExt], refcheck: bool = False) -> int: s.take_finish_event().synchronize() return num_reused - def _partial_block(self, prompt: list[TokenIdExt]): - """The tree block holding the tail of `prompt` (block 2 in these tests).""" - match = self.manager._radix_tree.match(ReuseScope(), prompt, True) - return match.blocks[-1] + def _tail_coverage(self, prompt: list[TokenIdExt], lc_id: int) -> int: + """Recorded token count of `lc_id`'s page on the last block matching `prompt`. + + Zero when that block holds no page for the life cycle. This is the tree block + holding the tail of `prompt` (block 2 in these tests). + """ + _, pages = _introspection.reuse_match_pages(self.manager, ReuseScope(), prompt, lc_id, True) + self.assertTrue(pages) + page = pages[-1] + if page is None: + return 0 + self.assertIsNotNone(page[1]) + return cast(int, page[1]) + + def _assert_page_unlinked(self, kv_cache: Any, ordinal: int) -> None: + """Every attention page the sequence holds at `ordinal` must be off the tree. + + A page pushed out of its slot keeps no pointer to the block, so nothing + dereferences that block once it dies. + """ + for lc_id in _introspection.attention_life_cycle_ids(self.manager): + self.assertIs( + _introspection.committed_page_is_linked(kv_cache, ordinal, lc_id), + False, + f"page at ordinal {ordinal} lc {lc_id} still points at a block", + ) + + def test_replaced_page_does_not_outlive_its_block(self) -> None: + """A page pushed out of a slot must not keep a pointer to a block that dies first. + + turn1 ends inside block 2 and stays open, holding the 8-token pages it committed. + turn2 grows that block to 16 tokens and commits over them; turn3 replaces the + block with a 32-token one and destroys it. + """ + self.prepare_partial() + tpb = self.TOKENS_PER_BLOCK + prompt = [TokenId(i) for i in range(3 * tpb)] + turn1, turn2 = prompt[: 2 * tpb + 8], prompt[: 2 * tpb + 16] + + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + kv1 = self.manager.create_kv_cache(input_tokens=turn1) + self.assertTrue(kv1.resume(stream)) + self.assertTrue(kv1.resize(len(turn1), len(turn1))) + kv1.commit(turn1[kv1.num_committed_tokens :], is_end=True) + # kv1 stays open, holding block 2's 8-token pages. + + kv2 = self.manager.create_kv_cache(input_tokens=turn2) + self.assertEqual(kv2.num_committed_tokens, len(turn1)) + self.assertTrue(kv2.resume(stream)) + self.assertTrue(kv2.resize(len(turn2), len(turn2))) + # Block 2 becomes a 16-token block that commits over kv1's 8-token pages, + # pushing them out of the slot while kv1 still holds them. + kv2.commit(turn2[kv2.num_committed_tokens :], is_end=True) + kv2.close() + + try: + # kv1's 8-token pages left the slot but kv1 still holds them. + self._assert_page_unlinked(kv1, 2) + + # A 32-token block replaces the 16-token one, which is then destroyed. + self.run_turn(prompt) + self._assert_page_unlinked(kv1, 2) + finally: + # Must run even on failure: leaving a sequence open makes teardown abort + # the process, which would bury the assertion message. + # Dereferences CommittedPage::block for the surviving 8-token pages. + kv1.close() + s.take_finish_event().synchronize() + + def test_replaced_reused_page_does_not_outlive_its_block(self) -> None: + """Same defect as above, reached through reuse rather than through commit. + + The holder matches the 8-token endpoint and holds the tree's pages. It must not + resume: resume()'s deferred copy swaps a reused partial page for a private one and + drops the holder. + """ + self.prepare_partial() + tpb = self.TOKENS_PER_BLOCK + prompt = [TokenId(i) for i in range(3 * tpb)] + short, mid = prompt[: 2 * tpb + 8], prompt[: 2 * tpb + 16] + + self.assertEqual(self.run_turn(short), 0) + + with TemporaryCudaStream([]) as s: + holder = self.manager.create_kv_cache(input_tokens=short) + self.assertEqual(holder.num_committed_tokens, len(short)) + + # Block 2 grows to 16 tokens, adopting the 8-token pages and then committing + # over them while `holder` still holds them. + try: + self.assertEqual(self.run_turn(mid), len(short)) + self._assert_page_unlinked(holder, 2) + # A 32-token block replaces the 16-token one, which is then destroyed. + self.assertEqual(self.run_turn(prompt), len(mid)) + self._assert_page_unlinked(holder, 2) + finally: + # Must run even on failure -- see the note in the sibling test. + holder.close() + s.take_finish_event().synchronize() def test_rewind_endpoint_survives_longer_sibling_created_after(self) -> None: self.prepare_partial() @@ -3952,12 +4068,13 @@ def test_rewind_endpoint_survives_longer_sibling_created_after(self) -> None: self.assertEqual(self.run_turn(base), 0) self.assertEqual(self.run_turn(extended), len(base)) # The 16-token endpoint block is gone, but its SWA page moved into the full - # 32-token sibling and still covers the first 16 tokens. - block = self._partial_block(rewind) - self.assertEqual(len(block.tokens), self.TOKENS_PER_BLOCK) - self.assertEqual(block.page_coverage(self._full_attn_lc_id), self.TOKENS_PER_BLOCK) - self.assertEqual(block.page_coverage(self._swa_lc_id), len(base) % self.TOKENS_PER_BLOCK) - del block + # 32-token sibling and still covers the first 16 tokens. Full coverage of + # TOKENS_PER_BLOCK also proves the block now spans a whole block, since a page + # never records more tokens than its block holds. + self.assertEqual(self._tail_coverage(rewind, self._full_attn_lc_id), self.TOKENS_PER_BLOCK) + self.assertEqual( + self._tail_coverage(rewind, self._swa_lc_id), len(base) % self.TOKENS_PER_BLOCK + ) self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) # The partial SWA page is stale at the longer endpoint and must not constrain # the full-attention lifecycle's reusable prefix. @@ -3996,10 +4113,10 @@ def test_exact_boundary_ignores_stale_last_block_partial_coverage(self) -> None: self.assertEqual(self.run_turn(base), 0) self.assertEqual(self.run_turn(boundary), len(base)) - block = self._partial_block(boundary) - self.assertEqual(block.page_coverage(self._full_attn_lc_id), self.TOKENS_PER_BLOCK) - self.assertEqual(block.page_coverage(self._swa_lc_id), 16) - del block + self.assertEqual( + self._tail_coverage(boundary, self._full_attn_lc_id), self.TOKENS_PER_BLOCK + ) + self.assertEqual(self._tail_coverage(boundary, self._swa_lc_id), 16) # At the 96-token boundary the input token is the entire size-1 SWA window, so no # historical SWA block is active. The partial SWA page must not constrain the full # attention lifecycle's reusable prefix. @@ -4012,16 +4129,12 @@ def test_page_coverage_only_grows(self) -> None: rewind = base + [TokenId(2000)] self.run_turn(base) - block = self._partial_block(rewind) - self.assertEqual(block.page_coverage(self._swa_lc_id), 16) - del block + self.assertEqual(self._tail_coverage(rewind, self._swa_lc_id), 16) self.run_turn(longer_partial) - block = self._partial_block(rewind) # A slot keeps only the widest page. The 24-token snapshot supersedes the 16-token # one, and it still covers the shorter rewind endpoint. - self.assertEqual(block.page_coverage(self._swa_lc_id), 24) - del block + self.assertEqual(self._tail_coverage(rewind, self._swa_lc_id), 24) self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) self.assertEqual( self.manager.probe_reuse(input_tokens=longer_partial + [TokenId(2000)]), @@ -4030,9 +4143,7 @@ def test_page_coverage_only_grows(self) -> None: # A later shorter snapshot cannot replace the wider page. self.assertEqual(self.run_turn(base), len(base)) - block = self._partial_block(rewind) - self.assertEqual(block.page_coverage(self._swa_lc_id), 24) - del block + self.assertEqual(self._tail_coverage(rewind, self._swa_lc_id), 24) self.assertEqual( self.manager.probe_reuse(input_tokens=longer_partial + [TokenId(2000)]), len(longer_partial), From 53419c668954c764efd4950ff8256fff3d831067 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 3 Aug 2026 08:35:30 +0000 Subject: [PATCH 5/5] [None][fix] KVCacheManagerV2: build committed pages without a temporary _copyPageToTreeBlock() built its CommittedPage by routing through a temporary UncommittedPage. That temporary claims the (kvCache, ordinal, beam, lifecycle) identity which the live sequence's own uncommitted page already holds, so ~UncommittedPage's slot-ownership check failed -- a throwing destructor, hence std::terminate under TLLM_DEBUG_MODE=1. Build the CommittedPage straight from the new slot, as _copy_page_to_tree_block() does in Python. The end state is the same page with the same slot, post-copy ready event, priority and manager, installed through the same replacePage(); the temporary is simply not created. The failure it removes is debug-only -- in release the temporary registered nothing and its destructor released nothing, since the slot had already been transferred -- but the fix is on a non-debug path, so it belongs with the per-page-coverage work that made this path hot rather than in the debug-check cleanup. Signed-off-by: Yao Yao --- .../batch_manager/kv_cache_manager_v2/kvCache.cpp | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 637deb15d3f9..6beb2bd4f9eb 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -809,12 +809,13 @@ CommittedPage* KvCache::_copyPageToTreeBlock( newSlot.readyEvent.waitInStream(reinterpret_cast(stream)); copySlotData(storageMgr, lvl, srcPage->cacheLevel, pgIdx, newSlot.slotId(), srcPage->slotId(), stream); - CachedCudaEvent readyEv(reinterpret_cast(stream)); - auto tempPage = makeShared(*this, treeBlock->ordinal(), lcIdx, lvl, kDefaultBeamIndex); - tempPage->setSlot(newSlot); + newSlot.readyEvent = CachedCudaEvent(reinterpret_cast(stream)); + auto committed = makeShared( + &storageMgr, treeBlock, lcIdx, lvl, numTokensInBlock, getPriority(treeBlock->ordinal(), lcIdx)); + committed->setSlot(newSlot); // Drops the superseded page, deferred until the copy is issued: an // OutOfPagesError above must not destroy a usable shorter snapshot. - auto committed = tempPage->convertToCommitted(treeBlock, std::move(readyEv), numTokensInBlock); + treeBlock->replacePage(lcIdx, committed.get()); // Schedule for eviction so eviction controller keeps a strong reference, // preventing the page from being destroyed.