From 8044cc6779e8e107221b154c1439d868c1d22008 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 2 Jun 2026 02:24:28 +0000 Subject: [PATCH 1/2] [None][feat] Enable disk cache config for KV cache v2 Expose disk cache size and path through the LLM API config and bridge them into the KV cache manager v2 disk tier. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/pyexecutor/resource_manager.py | 12 +++++++++- tensorrt_llm/llmapi/llm_args.py | 23 +++++++++++++++++++ tests/unittest/llmapi/test_llm_args.py | 15 ++++++++++++ 3 files changed, 49 insertions(+), 1 deletion(-) 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/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( From d7647ac7e6c81d1d38ea47535f400a826cb0210d Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 5 Jun 2026 05:02:56 +0000 Subject: [PATCH 2/2] [None][fix] Add disk_cache_size/path to KvCacheConfigV2 test mocks The disk cache feature added disk_cache_size and disk_cache_path to KvCacheConfig but did not update the KvCacheConfigV2 mock dataclasses in test files, causing AttributeError in all v2 test variants. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../disaggregated/test_cache_transceiver_single_process.py | 2 ++ tests/unittest/disaggregated/test_kv_transfer.py | 2 ++ 2 files changed, 4 insertions(+) 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