Skip to content

[None][feat] Enable disk cache config for KV cache v2 - #14845

Merged
reasonsolo merged 2 commits into
NVIDIA:mainfrom
reasonsolo:feat/enable_disk_cache
Jun 8, 2026
Merged

[None][feat] Enable disk cache config for KV cache v2#14845
reasonsolo merged 2 commits into
NVIDIA:mainfrom
reasonsolo:feat/enable_disk_cache

Conversation

@reasonsolo

@reasonsolo reasonsolo commented Jun 2, 2026

Copy link
Copy Markdown
Collaborator

Expose disk cache size and path through the LLM API config and bridge them into the KV cache manager v2 disk tier.

Summary by CodeRabbit

Release Notes

  • New Features
    • Added disk-backed KV cache configuration options, including disk cache size and storage path settings.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

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>
@reasonsolo
reasonsolo requested review from a team as code owners June 2, 2026 02:42
@reasonsolo
reasonsolo requested a review from Superjomn June 2, 2026 02:42
@reasonsolo
reasonsolo requested a review from lowsfer June 2, 2026 02:43
@reasonsolo

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented Jun 2, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

This PR adds disk-backed KV-cache tier support to the KV cache system by extending KvCacheConfig with configuration fields, implementing conditional disk tier provisioning in KVCacheManagerV2, and providing test coverage for schema validation and integration.

Changes

Disk KV-cache tier support

Layer / File(s) Summary
KvCacheConfig disk cache schema and validation
tensorrt_llm/llmapi/llm_args.py
KvCacheConfig adds disk_cache_size (bytes) and disk_cache_path (directory) fields; a post-validator enforces that disk_cache_path exists and is set whenever disk_cache_size is positive.
KVCacheManagerV2 disk tier provisioning
tensorrt_llm/_torch/pyexecutor/resource_manager.py
When disk_cache_size > 0, KVCacheManagerV2.__init__ creates and appends a DiskCacheTierConfig to cache_tiers, logs the disk cache quota and path, and updates import formatting.
Disk cache configuration test coverage
tests/unittest/llmapi/test_llm_args.py
test_KvCacheConfig_declaration validates the new fields; test_KvCacheConfig_disk_cache_validation verifies that omitting disk_cache_path when disk_cache_size is set raises ValidationError.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Description check ⚠️ Warning The PR description is incomplete, containing only the template with placeholder text and a boilerplate checklist but missing actual content for Description and Test Coverage sections. Fill in the Description section explaining the issue and solution, and the Test Coverage section listing relevant tests that safeguard these changes.
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly describes the main change: enabling disk cache configuration for KV cache v2.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (1)
tests/unittest/llmapi/test_llm_args.py (1)

345-353: ⚡ Quick win

Add coverage for invalid disk-cache paths.

tests/unittest/llmapi/test_llm_args.py is still insufficient for the full validator contract described in this PR: this test only covers the “size set without path” branch, but not the “path does not exist / is not a directory” branch. Please add at least one negative case for a missing path and one for a regular file so the directory validation stays pinned down.

🧪 Proposed test additions
 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)
+
+    missing_dir = tmp_path / "missing"
+    with pytest.raises(ValidationError) as exc_info:
+        KvCacheConfig(disk_cache_size=2048,
+                      disk_cache_path=str(missing_dir))
+    assert "disk_cache_path" in str(exc_info.value)
+
+    not_a_dir = tmp_path / "cache_file"
+    not_a_dir.write_text("x")
+    with pytest.raises(ValidationError) as exc_info:
+        KvCacheConfig(disk_cache_size=2048,
+                      disk_cache_path=str(not_a_dir))
+    assert "disk_cache_path" in str(exc_info.value)

As per coding guidelines, tests/**: Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/llmapi/test_llm_args.py` around lines 345 - 353, The test
test_KvCacheConfig_disk_cache_validation only covers the "size set without path"
case; add two additional negative test cases for KvCacheConfig: one where
disk_cache_path points to a non-existent path (e.g., tmp_path /
"does_not_exist") and one where disk_cache_path points to a regular file (create
a file under tmp_path), and in each use pytest.raises(ValidationError) to assert
the validator rejects the config and that "disk_cache_path" appears in the error
message; update or add assertions within the same test function or as separate
tests (referencing KvCacheConfig and test_KvCacheConfig_disk_cache_validation)
to ensure both "missing path" and "path is not a directory" branches are
covered.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 2547-2556: Add a fast-fail validation in the Pydantic model that
defines disk_cache_size and disk_cache_path (the same model that has
use_kv_cache_manager_v2) so that if disk_cache_size is set to a positive value
or disk_cache_path is non-empty while use_kv_cache_manager_v2 is False, the
model raises a ValueError with a clear message. Implement this as a
`@root_validator` (or a validator that has access to multiple fields) in the same
class that declares disk_cache_size/disk_cache_path, referencing the symbols
disk_cache_size, disk_cache_path and use_kv_cache_manager_v2, and ensure the
same validation is applied for the other occurrence mentioned (the second block
around lines 2705-2716).
- Around line 2705-2716: The model_validator validate_disk_cache_config should
only enforce that disk_cache_path is set when disk_cache_size > 0 and must not
check filesystem state; remove the os.path.isdir(self.disk_cache_path) existence
check from validate_disk_cache_config (keep the ValueError for missing path) and
move the directory-existence and writability checks into the worker-side
disk-tier initialization in KVCacheManagerV2 (e.g., in its disk
provisioning/init method) so node-local filesystem validation happens where the
path is actually used.

---

Nitpick comments:
In `@tests/unittest/llmapi/test_llm_args.py`:
- Around line 345-353: The test test_KvCacheConfig_disk_cache_validation only
covers the "size set without path" case; add two additional negative test cases
for KvCacheConfig: one where disk_cache_path points to a non-existent path
(e.g., tmp_path / "does_not_exist") and one where disk_cache_path points to a
regular file (create a file under tmp_path), and in each use
pytest.raises(ValidationError) to assert the validator rejects the config and
that "disk_cache_path" appears in the error message; update or add assertions
within the same test function or as separate tests (referencing KvCacheConfig
and test_KvCacheConfig_disk_cache_validation) to ensure both "missing path" and
"path is not a directory" branches are covered.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e6f47868-9bb1-4f1a-b821-2f49a448f5b8

📥 Commits

Reviewing files that changed from the base of the PR and between 6222112 and 8044cc6.

📒 Files selected for processing (3)
  • tensorrt_llm/_torch/pyexecutor/resource_manager.py
  • tensorrt_llm/llmapi/llm_args.py
  • tests/unittest/llmapi/test_llm_args.py

Comment thread tensorrt_llm/llmapi/llm_args.py
Comment thread tensorrt_llm/llmapi/llm_args.py
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #51480 [ run ] triggered by Bot. Commit: 8044cc6 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #51480 [ run ] completed with state SUCCESS. Commit: 8044cc6
/LLM/main/L0_MergeRequest_PR pipeline #40886 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@reasonsolo
reasonsolo enabled auto-merge (squash) June 5, 2026 04:30
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>
@reasonsolo

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52253 [ run ] triggered by Bot. Commit: d7647ac Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52253 [ run ] completed with state SUCCESS. Commit: d7647ac
/LLM/main/L0_MergeRequest_PR pipeline #41567 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@reasonsolo

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52335 [ run ] triggered by Bot. Commit: d7647ac Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52335 [ run ] completed with state SUCCESS. Commit: d7647ac
/LLM/main/L0_MergeRequest_PR pipeline #41640 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@reasonsolo

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52679 [ run ] triggered by Bot. Commit: d7647ac Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52679 [ run ] completed with state SUCCESS. Commit: d7647ac
/LLM/main/L0_MergeRequest_PR pipeline #41947 completed with status: 'SUCCESS'

CI Report

Link to invocation

@Superjomn Superjomn left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@reasonsolo
reasonsolo merged commit c93c63d into NVIDIA:main Jun 8, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants