[TRTLLM-13333][feat] Add prefetch_reuse_blocks and configurable prefetch count

[reasonsolo]

• Add KVCacheManager.prefetch_reuse_blocks() to probe radix tree and pre-stage disk-tier pages to host before scheduling
• Make prefetch count configurable via kv_cache_config.disk_prefetch_num_reqs
• Add unit test for prefetch_reuse_blocks disk-to-host migration

Performance Summary

Setup: nvidia/Kimi-K2.5-NVFP4, DEP4 (TP4+attention_dp+EP4), 1 CTX + 1 GEN node, host128g + disk256g, ISL=260K/OSL=8K, batch=16, 30min profiling phase, subagents-256k dataset.

Metric	N=0	N=2	N=4	N=8
Request Throughput (req/s)	0.36	0.42	0.49	0.52
Output Token Throughput (tok/s)	372.3	396.8	486.2	446.4
Input Token Throughput (tok/s)	31,484	35,866	40,677	45,159
Mean TTFT (ms)	118,437	114,918	82,804	84,579
Completed Requests (30min)	661	773	921	996
Error Rate (%)	0	0	2.28	2.71

KV Cache Metrics (RANK 0, cumulative across all 30s intervals)

Metric	N=0	N=2	N=4	N=8
Total Offload Blocks	689,564	700,607	806,932	971,136
Total Onboard Blocks	304,929	347,863	445,153	612,533
Total Disk Onboard Blocks	39,603	38,240	72,377	108,324
Recovery Rate (onboard/offload)	44.2%	49.7%	55.2%	63.1%
Dropped Blocks	0	0	0	0

## Summary by CodeRabbit

• New Features

  • Added disk KV cache block prefetching to improve performance when block reuse is enabled.
  • Introduced a new configuration parameter to control how many queued context requests are prefetched for disk-tier KV cache blocks (defaults to 4).

• Tests

  • Added comprehensive tests for KV cache block prefetch functionality across cache tier transitions.

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

The pull request adds disk-to-host KV cache block prefetch for context requests. A new configuration field controls prefetch scope, core storage implements block matching and migration, the V2 manager wraps and orchestrates, and the executor integrates prefetch into batch preparation.

Changes

KV Cache disk-to-host prefetch for context requests

Layer / File(s)	Summary
Configuration and core storage prefetch mechanism tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py	KvCacheConfig.disk_prefetch_num_reqs (default 4, ≥ 0) configures how many context requests are prefetched. KVCacheManager.prefetch_reuse_blocks() matches reuse-tree blocks and stages matching pages from colder cache tiers to a target level, returning False on OutOfPagesError and True on success or no-op.
KV Cache Manager V2 initialization and wrapper tensorrt_llm/_torch/pyexecutor/resource_manager.py	KVCacheManagerV2.__init__ stores disk_prefetch_num_reqs config. New prefetch_for_context_tokens() augments request tokens for block-reuse matching (excluding final token), builds ReuseScope from request LoRA/salt ids, and calls core prefetch_reuse_blocks() to prefetch blocks to host-level cache when block reuse is enabled.
PyExecutor batch preparation hooks tensorrt_llm/_torch/pyexecutor/py_executor.py	_prefetch_for_context_requests() selects active first-context-chunk requests not yet in the KV cache map and invokes manager prefetch. _commit_kv_cache_stats() forwards scheduled stats to KVCacheManagerV2 when scheduler manages KV suspend. Prefetch hook is called during _prepare_and_schedule_batch() preparation.
Prefetch behavior validation tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py	TestPrefetchReuseBlocks.test_prefetch_reuse_blocks() populates a reuse baseline, forces eviction to disk via quota shrink, restores quota, invokes prefetch_reuse_blocks(), and asserts matching pages migrate from disk to host tier.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#14748: Core prefetch infrastructure (StorageManager.prefetch / _KVCache.prefetch) that this PR builds upon for context-request disk staging.
• NVIDIA/TensorRT-LLM#14845: Disk-tier provisioning via disk_cache_size/disk_cache_path in KVCacheManagerV2, complementary to this PR's prefetch activation.

Suggested labels

KV-Cache Management, api-compatible

Suggested reviewers

• lowsfer
• longlee0622
• Superjomn

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	❓ Inconclusive	The PR description includes bullet points explaining the changes and provides performance metrics, but lacks explicit sections for Test Coverage and omits details about PR checklist completion in the template format.	Expand the description to clearly delineate Test Coverage section listing specific test cases (e.g., TestPrefetchReuseBlocks.test_prefetch_reuse_blocks) and ensure all template sections are explicitly addressed.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly and specifically describes the main changes: adding prefetch_reuse_blocks functionality and making prefetch count configurable, which aligns with the changeset.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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]

Actionable comments posted: 4

🧹 Nitpick comments (1)

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)

1387-1388: 💤 Low value

Consider adding a docstring to explain the test objective.

This test validates a new prefetch feature across cache tiers. A brief docstring would clarify the verification goal: that prefetch_reuse_blocks() correctly migrates matching reusable pages from disk tier to host tier when called with tokens that match committed blocks in the radix tree.

📝 Suggested docstring

 class TestPrefetchReuseBlocks(TestNoBatching):
     def test_prefetch_reuse_blocks(self) -> None:
+        """Verify prefetch_reuse_blocks migrates matching pages from disk to host.
+        
+        Commits known tokens to populate reuse tree, forces eviction to disk via
+        quota shrink, then confirms prefetch moves matching pages from disk to host.
+        """
         self.prepare(64 << 20, 128 << 20, 128 << 20, 36, 128, 1, kv_buf_size=32768)

🤖 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/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py` around
lines 1387 - 1388, Add a concise docstring to the TestPrefetchReuseBlocks class
or its test_prefetch_reuse_blocks method that states the test objective: it
verifies that prefetch_reuse_blocks() migrates matching reusable pages from the
disk tier to the host tier when supplied tokens correspond to committed blocks
in the radix tree, and documents expected outcome (which pages are moved and any
assertions made); reference the prefetch_reuse_blocks() call, the radix
tree/committed block matching, and the cache tiers (disk tier, host tier) so
future readers can quickly understand the verification goal.

🤖 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/_torch/pyexecutor/py_executor.py`:
- Around line 2416-2417: The attribute _scheduler_manages_kv_suspend is
referenced in the scheduling loop and in _prefetch_for_context_requests but
never initialized on PyExecutor, which causes AttributeError at runtime; fix by
adding an explicit initialization (e.g., self._scheduler_manages_kv_suspend =
False) in PyExecutor.__init__ so the flag always exists, or alternatively guard
its usage with hasattr(self, "_scheduler_manages_kv_suspend") before
dereferencing it in the scheduling code that checks
isinstance(self.kv_cache_manager, KVCacheManagerV2) and in
_prefetch_for_context_requests; ensure the initializer approach is preferred so
other methods (like _prefetch_for_context_requests) can rely on a consistent
boolean.
- Around line 2414-2419: The _commit_kv_cache_stats hook is never invoked so
KVCacheManagerV2 never receives commit_scheduled_kv_cache_stats calls and stats
drift; locate where batches are scheduled/committed in the scheduling flow
(where ScheduledRequests instances are finalized/returned) and add a call to
_commit_kv_cache_stats(scheduled_batch) after a batch is scheduled/committed
(guarded by self._scheduler_manages_kv_suspend like the method does). Ensure
this invocation uses the same ScheduledRequests object passed through scheduling
so KVCacheManagerV2.commit_scheduled_kv_cache_stats(scheduled_batch) runs; also
add the same call at the other scheduling exit point corresponding to the second
occurrence noted (the analogous spot referenced around where the second dead
hook would be invoked).

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 3077-3089: The function prefetch_for_context_tokens uses
CacheLevel but the symbol isn't imported; add an import for CacheLevel at the
top of the module (where other cache-related types are imported) so the call
HOST_LEVEL = CacheLevel(1) resolves; update the import block to include
CacheLevel (alongside existing imports such as ReuseScope or related cache
types) so prefetch_for_context_tokens and its call to
self.impl.prefetch_reuse_blocks(reuse_scope, tokens, HOST_LEVEL) no longer raise
NameError or fail linting.

In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Line 1402: The call to kv_cache.resume(stream) should assert its return value
for consistency with other resume() checks; capture the result of
kv_cache.resume(stream) into a variable and add an assertion that it indicates
success (e.g., assert result is True or assert result), updating the test
invocation of kv_cache.resume(stream) so the test explicitly verifies the resume
succeeded using the same success expectation used elsewhere in this file.

---

Nitpick comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 1387-1388: Add a concise docstring to the TestPrefetchReuseBlocks
class or its test_prefetch_reuse_blocks method that states the test objective:
it verifies that prefetch_reuse_blocks() migrates matching reusable pages from
the disk tier to the host tier when supplied tokens correspond to committed
blocks in the radix tree, and documents expected outcome (which pages are moved
and any assertions made); reference the prefetch_reuse_blocks() call, the radix
tree/committed block matching, and the cache tiers (disk tier, host tier) so
future readers can quickly understand the verification goal.

🪄 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: 9cd1fe66-aac2-4137-a824-6bf638338ef5

📥 Commits

Reviewing files that changed from the base of the PR and between 34a94ee and b8557ba.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53203 [ run ] triggered by Bot. Commit: 1b696e8 Link to invocation

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53261 [ run ] triggered by Bot. Commit: e3463b8 Link to invocation

[tensorrt-cicd]

PR_Github #53261 [ run ] completed with state SUCCESS. Commit: e3463b8
/LLM/main/L0_MergeRequest_PR pipeline #42453 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53456 [ run ] triggered by Bot. Commit: e3463b8 Link to invocation

[tensorrt-cicd]

PR_Github #53456 [ run ] completed with state SUCCESS. Commit: e3463b8
/LLM/main/L0_MergeRequest_PR pipeline #42621 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53757 [ run ] triggered by Bot. Commit: f11c63a Link to invocation

[tensorrt-cicd]

PR_Github #53757 [ run ] completed with state SUCCESS. Commit: f11c63a
/LLM/main/L0_MergeRequest_PR pipeline #42877 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]

/bot run --disable-fail-fast

[reasonsolo]

/bot run --disable-fail-fast

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54262 [ run ] triggered by Bot. Commit: 4fc32cd Link to invocation

[reasonsolo]

/bot run --disable-fail-fast

[reasonsolo]

/bot run --disable-fail-fast

[reasonsolo]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54301 [ run ] triggered by Bot. Commit: 7a6c1a7 Link to invocation

[tensorrt-cicd]

PR_Github #54262 [ run ] completed with state ABORTED. Commit: 4fc32cd

Link to invocation

[tensorrt-cicd]

PR_Github #54301 [ run ] completed with state SUCCESS. Commit: 7a6c1a7
/LLM/main/L0_MergeRequest_PR pipeline #43372 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54406 [ run ] triggered by Bot. Commit: 59e4208 Link to invocation

[tensorrt-cicd]

PR_Github #54406 [ run ] completed with state SUCCESS. Commit: 59e4208
/LLM/main/L0_MergeRequest_PR pipeline #43470 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54541 [ run ] triggered by Bot. Commit: 2cb014f Link to invocation

[tensorrt-cicd]

PR_Github #54541 [ run ] completed with state SUCCESS. Commit: 2cb014f
/LLM/main/L0_MergeRequest_PR pipeline #43592 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54674 [ run ] triggered by Bot. Commit: 2cb014f Link to invocation

[tensorrt-cicd]

PR_Github #54674 [ run ] completed with state SUCCESS. Commit: 2cb014f
/LLM/main/L0_MergeRequest_PR pipeline #43705 completed with status: 'SUCCESS'

CI Report

Link to invocation