[None][feat] Add KV cache prefetch

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: ...

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

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)

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)