[None][feat] KV cache manager v2 + python transceiver bug fix

[chuangz0]

Summary by CodeRabbit

Release Notes

• New Features

  • Added VMM-aware memory descriptor splitting with region metadata support for improved memory efficiency in distributed KV cache operations.
  • Introduced serialization and deserialization APIs for agent descriptions, enabling persistence and transfer of agent metadata across processes.
  • Optimized KV cache index slot management to enable early release for context-only requests.

• Tests

  • Added comprehensive unit tests for agent description serialization and VMM descriptor splitting utilities.

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)

• 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-test

[coderabbitai]

📝 Walkthrough

Walkthrough

These changes introduce VMM-aware descriptor splitting utilities and serialization support for AgentDesc. New public types (VramRegionInfo, VramRegionMap, VmmDescSplitter) provide backend-agnostic VRAM descriptor management with region detection and chunk-aligned splitting. AgentDesc extended with serialize/deserialize APIs and VMM region metadata. NixlTransferAgent refactored to delegate descriptor handling to VmmDescSplitter. Scheduler and resource manager enhanced for coordinated KV cache allocation and early index slot release.

Changes

Cohort / File(s)	Summary
VMM Descriptor Utilities (Public API) cpp/include/tensorrt_llm/executor/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp	Added VramRegionInfo, VramRegionMap, VmmDescSplitter with static methods for chunk lookup, descriptor splitting, region detection, and transfer descriptor coordination. AgentDesc extended with VMM region metadata constructor, accessor, and serialize/deserialize APIs.
NixlTransferAgent Refactoring cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp	Removed internal VMM chunk registry logic; replaced with mLocalVramRegionInfo and mRemoteVramRegionInfo maps. Delegated descriptor splitting to VmmDescSplitter. Added remote agent VMM metadata propagation via loadRemoteAgent and metadata packing in getLocalAgentDesc.
Python Bindings cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp, tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py	Added serialize/deserialize instance and static methods to AgentDesc nanobind wrapper, converting between nb::bytes and internal C++ serialization. Python-side load_remote_agent and get_local_agent_desc updated to use explicit serialization APIs.
Build Configuration cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt	Removed CUDA driver library linkage, retaining only CUDA runtime for cudaMemcpy support in NIXL wrapper.
KV Cache Resource Management tensorrt_llm/_torch/pyexecutor/resource_manager.py, tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py	KVCacheManagerV2 added release_index_slot method and early-freed request tracking to conditionally defer sequence slot removal. Scheduler now performs inline KV allocation (prepare_context, resize_context) during DISAGG_GENERATION_INIT state.
KV Cache Completion Flow tensorrt_llm/_torch/pyexecutor/py_executor.py	Context-only request completion now conditionally releases IndexMapper slot via kv_cache_manager.release_index_slot before async transfer initiation.
Virtual Memory Initialization tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py	Updated cuMemAddressReserve call to pass 0 instead of phys_mem_size as second parameter.
Tests cpp/tests/unit_tests/executor/transferAgentTest.cpp, tests/unittest/bindings/test_transfer_agent_bindings.py	Added comprehensive unit tests for AgentDesc serialization roundtrips (empty/multiple regions, null bytes, unaligned pools, truncation). Added VmmDescSplitter tests covering chunk lookup, descriptor splitting across regions, transfer descriptor coordination.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ❌ 3

❌ Failed checks (2 warnings, 1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is entirely unfilled—only the template is present with placeholder sections and no actual content provided in the Description, Test Coverage, or specific checklist item details.	Complete the PR description by filling in the Description section explaining the issue and solution, detailing relevant tests in Test Coverage, and providing specifics for each checklist item.
Docstring Coverage	⚠️ Warning	Docstring coverage is 28.57% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title is overly broad and vague—it uses generic terms like 'KV cache manager v2' and 'python transceiver bug fix' without clearly specifying the primary change.	Refine the title to be more specific about the main change. For example: '[TRTLLM-xxxx][feat] Add VMM-aware descriptor splitting and index mapper slot release' or use CodeRabbit's summary feature.

_{✏️ 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.}

[tensorrt-cicd]

PR_Github #42695 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: --disable-fail-test

Link to invocation

[coderabbitai]

Actionable comments posted: 4

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp`:
- Around line 653-664: The cached remote region map
(mRemoteVramRegionInfo[name]) must be replaced on agent reload instead of only
updating when agentDesc.getVramRegions() is non-empty; clear or erase the
existing remoteMap for the agent (identified by name) before populating it so
stale entries aren't kept when regions are removed or empty, ensuring
VmmDescSplitter::splitTransferDescsWithRegionMaps and submitTransferRequests()
use the current remote region set.

In `@cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp`:
- Around line 80-98: AgentDesc::deserialize currently calls
regions.reserve(numRegions) using an untrusted numRegions from the payload which
can cause huge allocations; validate numRegions before reserving by checking it
against the remaining stream length or a reasonable maximum (e.g., derive
maxRegions from data.size() or a defined constant) and fail early if it's out of
range, then only call regions.reserve(validatedNumRegions); ensure the
validation logic happens before using numRegions and reference the
variables/functions AgentDesc::deserialize, numRegions, regions.reserve, and the
input stream `is`.
- Around line 257-264: The current global detectedChunkSize enforcement in
splitVmmDescs() (the if block comparing detectedChunkSize and chunkSize)
incorrectly rejects valid mixed-pool registrations; remove or relax that check
so callers can register VMM pools with different chunk sizes and let per-region
metadata in VramRegionMap/AgentDesc carry chunkSize instead. Concretely, delete
the global consistency check around detectedChunkSize and stop writing a single
detectedChunkSize out for registerMemory(); instead ensure chunkSize is recorded
on each AgentDesc/VramRegion entry created by splitVmmDescs() so
registerMemory() and downstream logic read the per-region chunk sizes rather
than relying on a global detectedChunkSize.
- Around line 155-206: The function
VmmDescSplitter::splitTransferDescsWithRegionMaps currently returns early when
srcDescs.getType() != MemoryType::kVRAM, which skips VRAM chunk splitting for
DRAM↔VRAM transfers; change the logic to handle cases where either side is VRAM:
iterate over paired srcVec/dstVec regardless of src type, use lookupChunkInfo
against localRegionMap when splitting VRAM on the source and against
remoteRegionMap when splitting VRAM on the destination, and when only one side
needs splitting produce matching piece entries for the other side (reuse the
original srcVec[i] or dstVec[i] addresses/device ids but with pieceSize
fragments) so splitSrc and splitDst remain aligned; update the code paths around
srcChunkSize/dstChunkSize, srcBase/dstBase, and the while loop to compute
pieceSize from whichever chunk sizes are nonzero and push fragmented MemoryDesc
entries into splitSrc/splitDst accordingly.

🪄 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: Pro

Run ID: df50b36a-7e50-4d83-9694-d9153ef96755

📥 Commits

Reviewing files that changed from the base of the PR and between 4811704 and a890eb6.

📒 Files selected for processing (13)

• cpp/include/tensorrt_llm/executor/transferAgent.h
• cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt
• cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp
• cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
• cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
• cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp
• cpp/tests/unit_tests/executor/transferAgentTest.cpp
• tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py
• tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py
• tests/unittest/bindings/test_transfer_agent_bindings.py

[chuangz0]

/bot --help

[github-actions]

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental) --high-priority]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #42920 [ run ] triggered by Bot. Commit: a890eb6 Link to invocation

[tensorrt-cicd]

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

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #43214 [ run ] triggered by Bot. Commit: d9d271a Link to invocation

[tensorrt-cicd]

PR_Github #43214 [ run ] completed with state DISABLED
Freeze main and open the PR merge only after CI is back to healthy https://nvidia.slack.com/archives/C059LSY62BT/p1776141760843319?thread_ts=1775985925.442509&cid=C059LSY62BT

Link to invocation

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #43374 [ run ] triggered by Bot. Commit: c155506 Link to invocation

[tensorrt-cicd]

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

[chuangz0]

/bot run --disable-fail-fast

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49353 [ run ] triggered by Bot. Commit: 9c6014c Link to invocation

[tensorrt-cicd]

PR_Github #49353 [ run ] completed with state SUCCESS. Commit: 9c6014c
/LLM/main/L0_MergeRequest_PR pipeline #39009 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 #49600 [ run ] triggered by Bot. Commit: 270827d Link to invocation

[tensorrt-cicd]

PR_Github #49600 [ run ] completed with state SUCCESS. Commit: 270827d
/LLM/main/L0_MergeRequest_PR pipeline #39222 completed with status: 'SUCCESS'

CI Report

Link to invocation