diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index ff36dc8f0b45..451778f7e754 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -226,6 +226,7 @@ class _KVCache: def stop_committing(self) -> None: ... def suspend(self) -> None: ... def resume(self, cuda_stream: CudaStream | None = None) -> bool: ... + def prefetch(self, target: CacheLevel) -> bool: ... def get_scratch_desc(self, layer_group_id: LayerGroupId) -> ScratchDesc | None: ... @property def has_scratch_slots(self) -> bool: ... 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 422f5d2794d6..ee8a7e24da0b 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 @@ -54,6 +54,7 @@ BatchedLockTarget, BlockPage, CommittedPage, + Page, ScratchSlotLock, UncommittedPage, _PageHolder, @@ -932,6 +933,49 @@ def resume(self, cuda_stream: CudaStream | None = None) -> bool: self._status = self.Status.ACTIVE return True + def prefetch(self, target: CacheLevel) -> bool: + """Best-effort prefetch active pages to the target cache level. + + The cache must be suspended. Prefetch is only a performance hint: a False + return value means the requested pages could not be recalled due to cache + pressure, but the cache remains functionally valid. + + Args: + target: Destination cache level for active pages in lower tiers. + + Returns: + True if the prefetch was dispatched, False if storage could not reserve enough pages. + """ + assert self.status == self.Status.SUSPENDED + manager = self.manager + storage = manager._storage + num_tiers = storage.num_cache_levels + assert CacheLevel(0) <= target < num_tiers + + num_pool_groups = storage.num_pool_groups + lc2pg = storage.get_pool_group_index + + all_pages = make_typed( + lambda _: make_typed(lambda _: list[Page](), num_tiers), num_pool_groups + ) + + for ordinal, beam_idx, lc_idx in self._active_pages(): + holder = self._page(ordinal, beam_idx, lc_idx) + if holder is None: + continue + page = expect_type(_PageHolder, holder).page + lvl = page.cache_level + if lvl < target: + continue + pg_idx = lc2pg(lc_idx) + all_pages[pg_idx][lvl].append(page) + + try: + storage.prefetch(target, all_pages) + except OutOfPagesError: + return False + return True + def _active_pages(self) -> Iterator[tuple[BlockOrdinal, BeamIndex, LifeCycleId]]: """Yields (ordinal, beam_idx, lc_idx) for all active pages. @@ -975,12 +1019,25 @@ def tokens_per_block(self) -> int: def _page( self, block_ordinal: BlockOrdinal, beam_index: BeamIndex, life_cycle: LifeCycleId ) -> BlockPage: - return self._blocks[block_ordinal].pages[beam_index][life_cycle] + """Return the page holder for an attention block or the SSM block.""" + is_ssm = block_ordinal == BAD_BLOCK_ORDINAL + assert (life_cycle == self.manager._life_cycles.ssm_life_cycle_id) == is_ssm + return ( + self._ssm_blocks[beam_index][life_cycle] + if is_ssm + else self._blocks[block_ordinal].pages[beam_index][life_cycle] + ) def _block( self, block_ordinal: BlockOrdinal, beam_index: BeamIndex ) -> TypedIndexList[LifeCycleId, BlockPage]: - return self._blocks[block_ordinal].pages[beam_index] + """Return the life-cycle page list for an attention block or the SSM block.""" + is_ssm = block_ordinal == BAD_BLOCK_ORDINAL + return ( + self._ssm_blocks[beam_index] + if is_ssm + else self._blocks[block_ordinal].pages[beam_index] + ) def _snapshot_ssm_to_tree_block( self, tree_block: Block, ssm_lc_id: LifeCycleId, beam_idx: BeamIndex diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 0d90c0e92aa3..90f822ba5539 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -898,3 +898,42 @@ def constrain_ratio( total = sum(num_bytes) assert total > 0 return typed_map(num_bytes, lambda x: x / total) + + def prefetch( + self, + dst_lvl: CacheLevel, + pages: TypedIndexList[PoolGroupIndex, TypedIndexList[CacheLevel, list[Page]]], + ) -> None: + """Dispatch page migration to the destination cache level. + + Args: + dst_lvl: Destination cache level for pages currently in lower tiers. + pages: Pages grouped by pool group and current cache level. + + Raises: + OutOfPagesError: If there are not enough pages available for the prefetch hint. + """ + num_slots = filled_list(0, self.num_pool_groups) + scheduled = list[Page]() + try: + for pg_idx, pg_pages in typed_enumerate(pages): + for lvl, lvl_pages in typed_enumerate(pg_pages): + assert lvl >= dst_lvl or not lvl_pages + for p in lvl_pages: + if p.scheduled_for_eviction: + self.exclude_from_eviction(p) + scheduled.append(p) + elif self.is_evictable(p, dst_lvl): + scheduled.append(p) + assert lvl >= dst_lvl + if lvl == dst_lvl: + continue + num_slots[pg_idx] += 1 + self.prepare_free_slots(dst_lvl, num_slots) + for pg_idx, pg_tasks in typed_enumerate(pages): + for lvl in typed_range(CacheLevel(dst_lvl + 1), self.num_cache_levels): + lvl_tasks = pg_tasks[lvl] + self._batched_migrate(pg_idx, dst_lvl, lvl, lvl_tasks, True) + finally: + for p in scheduled: + self.schedule_for_eviction(p) 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 65ccda597d36..649c7d54536c 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 @@ -1275,6 +1275,22 @@ def test_resize_quota(self) -> None: stream_holder = CachedCudaStream() stream = cast(CudaStream, stream_holder.handle) + def count_active_pages_by_level(kv_cache: _KVCache) -> list[int]: + counts = [0] * self.manager._storage.num_cache_levels + for ordinal, beam_idx, lc_idx in kv_cache._active_pages(): + block_page = kv_cache._page(ordinal, beam_idx, lc_idx) + assert block_page is not None + counts[block_page.page.cache_level] += 1 + return counts + + def assert_prefetched_pages_are_evictable(kv_cache: _KVCache) -> None: + for ordinal, beam_idx, lc_idx in kv_cache._active_pages(): + block_page = kv_cache._page(ordinal, beam_idx, lc_idx) + assert block_page is not None + page = block_page.page + if page.cache_level == HOST_LEVEL and self.manager._storage.is_evictable(page): + self.assertTrue(page.scheduled_for_eviction) + # First commit some blocks to fill all levels of cache. This helps test the case where shrinking # the quota will drop some pages from the last-level cache. for _ in range(11): @@ -1329,6 +1345,19 @@ def test_resize_quota(self) -> None: assert success success = self.manager.resize(HOST_LEVEL, 128 << 20) assert success + prefetch_target = kv_cache_lst[1] + prefetch_counts_before = count_active_pages_by_level(prefetch_target) + self.assertGreater(prefetch_counts_before[DISK_LEVEL], 0) + success = prefetch_target.prefetch(HOST_LEVEL) + self.assertEqual(success, True) + prefetch_counts_after = count_active_pages_by_level(prefetch_target) + self.assertEqual(prefetch_counts_after[GPU_LEVEL], prefetch_counts_before[GPU_LEVEL]) + self.assertEqual(prefetch_counts_after[DISK_LEVEL], 0) + self.assertEqual( + prefetch_counts_after[HOST_LEVEL], + prefetch_counts_before[HOST_LEVEL] + prefetch_counts_before[DISK_LEVEL], + ) + assert_prefetched_pages_are_evictable(prefetch_target) # Now both requests can resume for kv_cache in kv_cache_lst: success = kv_cache.resume(stream)