[TRTLLM-13617][feat] Inclusive host KV-cache tier (shadow reuse) for KVCacheManager V2

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -1048,6 +1048,8 @@ def _build_cache_config(
             max_util_for_resume=kv_cache_config.max_util_for_resume,
             enable_stats=self.enable_stats,
             layers=layer_configs,
+            enable_inclusive_host_cache=kv_cache_config.
+            enable_inclusive_host_cache,
         )
 
     def _extra_buffers_per_layer(

tensorrt_llm/llmapi/llm_args.py

@@ -3370,6 +3370,15 @@ class KvCacheConfig(StrictBaseModel, PybindMirror):
         "Set to 0 to disable prefetch. Only effective with KV cache manager v2 and block reuse enabled."
     )
 
+    enable_inclusive_host_cache: bool = Field(
+        default=False,
+        status="prototype",
+        description=
+        "Keep a block's clean host copy when it is recalled host->GPU so a later eviction can "
+        "reuse it instead of re-writing to host, avoiding redundant device-to-host copies for "
+        "blocks that bounce GPU<->host. Requires a host cache tier (``host_cache_size`` > 0). "
+        "Default False. Only used with KV cache manager v2 (experimental).")
+
     def _to_pybind(self):
         config = _KvCacheConfig(
             enable_block_reuse=self.enable_block_reuse,

tensorrt_llm/runtime/kv_cache_manager_v2/_config.py

@@ -212,6 +212,12 @@ class KVCacheManagerConfig:
     If True, we will try to reuse tokens from partially matched blocks.
     """
 
+    enable_inclusive_host_cache: bool = False
+    """
+    If True, keep a page's clean host copy after recall so a later eviction reuses it instead of
+    re-copying to host. Reclaimed under host memory pressure.
+    """
+
     constraints: list[BatchDesc] = field(default_factory=list)
     """
     A list of step configurations that must always be supported.

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py

@@ -261,6 +261,7 @@ def __init__(
             typical_batch=config.typical_step,
             constraints=config.constraints,
             event_manager=event_manager,
+            enable_inclusive_host_cache=config.enable_inclusive_host_cache,
         )
         self._living_kv_caches = set[rawref.ref[_KVCache]]()
         decay = 0.9999

tensorrt_llm/runtime/kv_cache_manager_v2/_page.py

@@ -17,6 +17,8 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, NamedTuple, cast
 
+from llist import dllistnode
+
 from . import rawref
 from ._block_radix_tree import Block
 from ._common import (
@@ -205,6 +207,13 @@ def check_page(p: "BlockPage") -> bool:
 class CommittedPage(Page):
     block: rawref.ref["Block"]
     __rawref__: rawref.ref["CommittedPage"]
+    # Inclusive host cache: a retained clean copy of this immutable page at a lower (host/disk)
+    # tier. When set, a later eviction of this page to shadow_level can reuse this slot and skip
+    # the device->host copy. None when no shadow is held (the default / exclusive behavior).
+    # shadow_node is this page's node in the StorageManager shadow reclaim FIFO for shadow_level.
+    shadow_slot: "Slot | None"
+    shadow_level: CacheLevel
+    shadow_node: "dllistnode | None"
 
     def is_committed(self) -> bool:
         return True
@@ -220,6 +229,9 @@ def __init__(
     ):
         self.block = rawref.ref(block)
         self.__rawref__ = rawref.NULL
+        self.shadow_slot = None
+        self.shadow_level = cache_level
+        self.shadow_node = None
         Page.__init__(
             self,
             None,
@@ -234,16 +246,23 @@ def __init__(
         self.set_slot(slot)
 
     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:
-            block.unset_page(
-                self.life_cycle,
-                self.manager._life_cycles.get_life_cycle(self.life_cycle),
-            )
-        Page.__del__(self)
-        self.__rawref__.invalidate()
+        try:
+            # Release any retained host shadow first so its slot returns to the pool.
+            if self.shadow_slot is not None:
+                self.manager.drop_shadow(self)
+            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:
+                block.unset_page(
+                    self.life_cycle,
+                    self.manager._life_cycles.get_life_cycle(self.life_cycle),
+                )
+        finally:
+            try:
+                Page.__del__(self)
+            finally:
+                self.__rawref__.invalidate()
 
 
 @dataclass(slots=True)