diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 8f46d88794cd..8e01d3256bda 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2c3aad405890..133136104321 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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]]): diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 7ba99a6d7679..c12f2ca2f42c 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/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 diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index eb8d5ead7600..4ad99286de58 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/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 diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b3de8a0bc4b0..ff75d9553f10 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/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(