[None][feat] Enable disk cache config for KV cache v2

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -33,7 +33,7 @@
 # isort: off
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     DEFAULT_BEAM_INDEX, AttentionLayerConfig, BufferConfig, CacheTierConfig,
-    GpuCacheTierConfig, HostCacheTierConfig, ReuseScope)
+    DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, ReuseScope)
 # isort: on
 from tensorrt_llm.runtime.kv_cache_manager_v2 import \
     KVCacheManager as KVCacheManagerPy
@@ -2512,6 +2512,16 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
             logger.info(
                 f"KV cache manager v2 host cache quota set to {host_quota / (1 << 30):.2f}GiB"
             )
+        disk_cache_size = kv_cache_config.disk_cache_size
+        if disk_cache_size is not None and disk_cache_size > 0:
+            disk_cache_path = kv_cache_config.disk_cache_path
+            assert disk_cache_path is not None
+            cache_tiers.append(
+                DiskCacheTierConfig(quota=disk_cache_size,
+                                    path=disk_cache_path))
+            logger.info(
+                f"KV cache manager v2 disk cache quota set to {disk_cache_size / (1 << 30):.2f}GiB at {disk_cache_path}"
+            )
 
         self.vocab_size = vocab_size
 

tensorrt_llm/llmapi/llm_args.py

@@ -2544,6 +2544,16 @@ class KvCacheConfig(StrictBaseModel, PybindMirror):
         description=
         "Size of the host cache in bytes. If both `max_tokens` and `host_cache_size` are specified, memory corresponding to the minimum will be used."
     )
+    disk_cache_size: Optional[NonNegativeInt] = Field(
+        default=None,
+        description=
+        "Size of the disk cache in bytes. Only used by KV cache manager v2 in the PyTorch backend."
+    )
+    disk_cache_path: Optional[str] = Field(
+        default=None,
+        description=
+        "Directory used for disk KV cache files. Must be set when `disk_cache_size` is positive."
+    )
     cross_kv_cache_fraction: Optional[float] = Field(
         default=None,
         description=
@@ -2692,6 +2702,19 @@ def validate_max_gpu_total_bytes(cls, v: int):
                 "kv_cache_config.max_gpu_total_bytes must be non-negative")
         return v
 
+    @model_validator(mode='after')
+    def validate_disk_cache_config(self):
+        if self.disk_cache_size is not None and self.disk_cache_size > 0:
+            if not self.disk_cache_path:
+                raise ValueError(
+                    "kv_cache_config.disk_cache_path must be set when disk_cache_size is positive"
+                )
+            if not os.path.isdir(self.disk_cache_path):
+                raise ValueError(
+                    f"kv_cache_config.disk_cache_path {self.disk_cache_path} does not exist or is not a directory"
+                )
+        return self
+
     @field_validator('max_attention_window')
     @classmethod
     def validate_max_attention_window(cls, v: Optional[List[int]]):

tests/unittest/disaggregated/test_cache_transceiver_single_process.py

@@ -94,6 +94,8 @@ class KvCacheConfigV2:
     sink_token_length: Optional[int] = None
     free_gpu_memory_fraction: Optional[float] = None
     host_cache_size: Optional[int] = None
+    disk_cache_size: Optional[int] = None
+    disk_cache_path: Optional[str] = None
     onboard_blocks: bool = True
     cross_kv_cache_fraction: Optional[float] = None
     secondary_offload_min_priority: Optional[int] = None

tests/unittest/disaggregated/test_kv_transfer.py

@@ -53,6 +53,8 @@ class KvCacheConfigV2:
     sink_token_length: Optional[int] = None
     free_gpu_memory_fraction: Optional[float] = None
     host_cache_size: Optional[int] = None
+    disk_cache_size: Optional[int] = None
+    disk_cache_path: Optional[str] = None
     cross_kv_cache_fraction: Optional[float] = None
     secondary_offload_min_priority: Optional[int] = None
     event_buffer_max_size: int = 0

tests/unittest/llmapi/test_llm_args.py

@@ -317,6 +317,8 @@ def test_KvCacheConfig_declaration():
                            max_attention_window=[1024, 1024, 1024],
                            free_gpu_memory_fraction=0.5,
                            host_cache_size=1024,
+                           disk_cache_size=2048,
+                           disk_cache_path="/tmp",
                            cross_kv_cache_fraction=0.5,
                            secondary_offload_min_priority=1,
                            event_buffer_max_size=0,
@@ -330,6 +332,8 @@ def test_KvCacheConfig_declaration():
     assert pybind_config.max_attention_window == [1024, 1024, 1024]
     assert pybind_config.free_gpu_memory_fraction == 0.5
     assert pybind_config.host_cache_size == 1024
+    assert config.disk_cache_size == 2048
+    assert config.disk_cache_path == "/tmp"
     assert pybind_config.cross_kv_cache_fraction == 0.5
     assert pybind_config.secondary_offload_min_priority == 1
     assert pybind_config.event_buffer_max_size == 0
@@ -338,6 +342,17 @@ def test_KvCacheConfig_declaration():
     assert pybind_config.attention_dp_events_gather_period_ms == 10
 
 
+def test_KvCacheConfig_disk_cache_validation(tmp_path):
+    config = KvCacheConfig(disk_cache_size=2048, disk_cache_path=str(tmp_path))
+
+    assert config.disk_cache_size == 2048
+    assert config.disk_cache_path == str(tmp_path)
+
+    with pytest.raises(ValidationError) as exc_info:
+        KvCacheConfig(disk_cache_size=2048)
+    assert "disk_cache_path" in str(exc_info.value)
+
+
 def test_CapacitySchedulerPolicy():
     val = CapacitySchedulerPolicy.MAX_UTILIZATION
     assert PybindMirror.maybe_to_pybind(