[None][feat] Python transceiver support cpp cache manager + offload

[chuangz0]

Summary by CodeRabbit

• New Features

  • Improved KV cache reuse for disaggregated serving with host offload support, enabling efficient memory management and translation across multi-window attention scenarios.

• Tests

  • Added comprehensive test coverage for Python transceiver host offload functionality, including cache reuse verification and memory translation validation.

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.

[chuangz0]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a block-index translation API for disaggregated KV cache management with host offload. The feature adds C++ and Python APIs to translate logical cache block IDs to memory-pool slot indices, integrates this into the cache reuse adapter, and validates correctness through unit and end-to-end tests.

Changes

Block Index Translation and Cache Reuse for Host Offload

Layer / File(s)	Summary
C++ API for block index translation cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h, cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp	BaseKVCacheManager::getMemoryPoolBlockIndicesByBlockIds(...) pure-virtual method translates logical block IDs to current memory-pool indices for a specified window, with optional requirePrimary enforcement. Implementation validates block existence and primary residency before collecting indices.
Python binding and VSWA wrapper cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp, tensorrt_llm/_torch/pyexecutor/resource_manager.py	Nanobind exposes the C++ API as get_memory_pool_block_indices. Python wrapper in KVCacheManager validates and provides window_size for variable-sized attention window configurations, defaulting to the first window only for single-window setups.
Cache reuse adapter migration tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py	CacheReuseAdapter.get_block_ids docstring clarifies that returned values are primary memory-pool slot indices required for pointer arithmetic. _CacheReuseAdapterV1 now translates block IDs to primary indices via the new API with require_primary=True. _CacheReuseAdapterV2 documentation confirms it already provides correct per-cache-level offsets.
Unit tests for translation correctness tests/unittest/disaggregated/test_extractor.py	Three CUDA tests validate translation behavior: VSWA multi-window dispatch and parameter validation; identity mapping when host offload is disabled and pointer arithmetic correctness; offload/eviction/onboard cycles confirming index divergence under pressure and error handling for non-primary blocks.
Integration test configuration and validation tests/integration/defs/disaggregated/test_configs/disagg_config_python_transceiver_host_offload.yaml, tests/integration/defs/disaggregated/test_disaggregated.py, tests/integration/test_lists/test-db/l0_dgx_h100.yml	New test configuration specifies deterministic primary pool sizing, block/partial reuse, and host cache allocation. _verify_python_transceiver_under_host_offload helper validates cache reuse and offload behavior via concurrent replay prompts across multiple passes. Parametrized test integrates verification into disaggregated test suite and adds entry to H100 test matrix.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

• venkywonka
• ruodil
• pcastonguay
• moraxu
• thorjohnsen
• Wanli-Jiang

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.42% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The PR description is only the template and lacks concrete details about the change, tests, and checklist items.	Fill in the Description and Test Coverage sections with the problem, solution, and relevant tests; complete any applicable checklist items.

✅ Passed checks (3 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title matches the main change and uses the required [None][feat] format.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

tests/unittest/disaggregated/test_extractor.py (1)

646-646: 💤 Low value

Add strict=True to zip() calls for defensive validation.

Both lists should be equal length by construction, but adding strict=True would catch unexpected mismatches during test execution rather than silently iterating over partial data.

Suggested fix

Line 646:

-        primary_diverged = [(b, s) for b, s in zip(block_ids_b, slots_b) if b != s]
+        primary_diverged = [(b, s) for b, s in zip(block_ids_b, slots_b, strict=True) if b != s]

Line 723:

-        for ptr, idx in zip(spec_region.memory.ptrs, translated_a2):
+        for ptr, idx in zip(spec_region.memory.ptrs, translated_a2, strict=True):

Also applies to: 723-723

🤖 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/disaggregated/test_extractor.py` at line 646, The zip() calls
used to pair lists like block_ids_b and slots_b (e.g. where primary_diverged is
computed) should be made strict to catch length mismatches; update the zip(...)
invocations in test_extractor.py (including the one that produces
primary_diverged and the other zip usage around the later comparison) to
zip(..., strict=True) so the test will raise on unexpected list-length
differences rather than silently truncating.

Source: Linters/SAST tools

🤖 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 `@cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp`:
- Around line 633-635: The trampoline for BaseKVCacheManager is missing a
pure-virtual override for getMemoryPoolBlockIndicesByBlockIds so Python
subclasses cannot implement it; update PyKvCacheManager to either increase the
NB_TRAMPOLINE(...) count to include the 61st pure virtual or add an explicit
NB_OVERRIDE_PURE forwarding for getMemoryPoolBlockIndicesByBlockIds (matching
the C++ signature used by the binding) and ensure the binding
.def("get_memory_pool_block_indices",
&BaseKVCacheManager::getMemoryPoolBlockIndicesByBlockIds, ...) will dispatch to
the Python override.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 1542-1544: The check that raises "window_size must be provided for
VSWA" should be guarded by the VSWA predicate rather than by the vector length;
update the conditional to use self.is_vswa (e.g., if self.is_vswa and
len(self.max_attention_window_vec) > 1: raise ValueError(...)) so non-VSWA
configs with multiple identical entries don't trigger the error, and keep
assigning window_size = self.max_attention_window_vec[0] afterward.

---

Nitpick comments:
In `@tests/unittest/disaggregated/test_extractor.py`:
- Line 646: The zip() calls used to pair lists like block_ids_b and slots_b
(e.g. where primary_diverged is computed) should be made strict to catch length
mismatches; update the zip(...) invocations in test_extractor.py (including the
one that produces primary_diverged and the other zip usage around the later
comparison) to zip(..., strict=True) so the test will raise on unexpected
list-length differences rather than silently truncating.

🪄 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: f7de916e-43a4-4125-a371-3901319a2c21

📥 Commits

Reviewing files that changed from the base of the PR and between 205920d and 45fa295.

📒 Files selected for processing (9)

• cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
• cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
• tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py
• tests/integration/defs/disaggregated/test_configs/disagg_config_python_transceiver_host_offload.yaml
• tests/integration/defs/disaggregated/test_disaggregated.py
• tests/integration/test_lists/test-db/l0_dgx_h100.yml
• tests/unittest/disaggregated/test_extractor.py

[tensorrt-cicd]

PR_Github #53493 [ run ] triggered by Bot. Commit: 04a8bfd Link to invocation

[tensorrt-cicd]

PR_Github #53493 [ run ] completed with state SUCCESS. Commit: 04a8bfd
/LLM/main/L0_MergeRequest_PR pipeline #42652 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

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53761 [ run ] triggered by Bot. Commit: f61f4ed Link to invocation

[tensorrt-cicd]

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

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54219 [ run ] triggered by Bot. Commit: 531d745 Link to invocation

[tensorrt-cicd]

PR_Github #54219 [ run ] completed with state SUCCESS. Commit: 531d745
/LLM/main/L0_MergeRequest_PR pipeline #43296 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

[fredricz-20070104]

Approved

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54414 [ run ] triggered by Bot. Commit: 0f9c0c0 Link to invocation

[tensorrt-cicd]

PR_Github #54414 [ run ] completed with state FAILURE. Commit: 0f9c0c0
/LLM/main/L0_MergeRequest_PR pipeline #43478 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

[chuangz0]

/bot run --disable-fail-fast

[chuangz0]

/bot skip --comment "all tests have passed"

[tensorrt-cicd]

PR_Github #59414 [ skip ] triggered by Bot. Commit: 5db5c93 Link to invocation

[tensorrt-cicd]

PR_Github #59414 [ skip ] completed with state SUCCESS. Commit: 5db5c93
Skipping testing for commit 5db5c93

Link to invocation