[TRTLLM-9920][feat] Add support for arbitrary KVCache transfer

[Tabrizian]

Dev Engineer Review

• Adds request-free arbitrary KV-cache transfer using serialized transceiver state, reuse-tree lookup, block pinning, and cleanup on success/failure.
• Extends Python, RPC, worker, LLM, and HTTP APIs with get_data_transceiver_state().
• Adds serialization support and explicit arbitrary-transfer metadata.
• Reuse-tree locking and rollback improve consistency, though the broader locking changes may affect contention/performance and need validation.
• Partial-prefix transfer and avoiding blocks already cached by the receiver remain unsupported.
• Mamba cache support is tracked separately in TRTLLM-13970.
• CI merge-request pipelines consistently reported failure; investigation and a successful rerun are required before approval.

QA Engineer Review

Added:

• test_arbitrary_kv_cache_transfer
• test_arbitrary_kv_cache_transfer_missing_blocks

Both are included in tests/integration/test_lists/test-db/l0_a10.yml for CI coverage. The tests cover successful transfer, cache reuse, missing sender blocks, error handling, and continued worker operation.

Verdict: sufficient

Description

Adds support for arbitrary KV cache transfer between TRT-LLM workers: a generation worker can request KV cache blocks directly from another worker's reuse tree, without requiring a preceding context_only request on the sender.

Use cases:

• Pull cached prompts across workers to amortize prefill cost
• Re-route generation to a different worker after routing decisions change
• Move KV cache during load balancing or failover

def arbitrary_cache_transfer():
     """Test KV cache transfer from the reuse tree.

      Flow:
      1. Worker 0 runs a normal generate to fill the KV cache reuse tree.
      2. Retrieve the data_transceiver_state from worker 0.
      3. Worker 1 sends a generation_only request using that state.
         Sender (worker 0) serves blocks directly from its reuse tree.
      """
      kv_cache_config = KvCacheConfig(max_tokens=2048 * 8, enable_block_reuse=True)
      cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT")

      worker0 = LLM(model=model_path(model),
                    tensor_parallel_size=1,
                    disable_overlap_scheduler=True,
                    cuda_graph_config=CudaGraphConfig(),
                    kv_cache_config=kv_cache_config,
                    cache_transceiver_config=cache_transceiver_config)

      worker1 = LLM(model=model_path(model),
                    tensor_parallel_size=1,
                    disable_overlap_scheduler=not generation_overlap,
                    cuda_graph_config=CudaGraphConfig(),
                    kv_cache_config=kv_cache_config,
                    cache_transceiver_config=cache_transceiver_config)

      prompt = "What is the capital of Germany?"

      # 1. Fill worker 0's reuse tree with a normal generate
      worker0.generate(prompt, SamplingParams(max_tokens=10, ignore_eos=True))

      # 2. Grab the transceiver state (opaque bytes) from worker 0
      state = worker0.get_data_transceiver_state()
      assert isinstance(state, bytes) and len(state) > 0

      # 3. generation_only on worker 1 using state from worker 0.
      #    max_tokens=1 — only one decode step, so KV transfer must
      #    deliver all prompt tokens.
      disagg_params = DisaggregatedParams(
          request_type="generation_only",
          opaque_state=state,
          first_gen_tokens=[0],
          disagg_request_id=42,
      )
      output = worker1.generate(prompt,
                                SamplingParams(max_tokens=1, ignore_eos=True),
                                disaggregated_params=disagg_params)
      assert len(output.token_ids) > 0

      # 4. Normal request on worker 1 still works after the transfer
      output2 = worker1.generate(prompt,
                                 SamplingParams(max_tokens=10, ignore_eos=True))
      assert len(output2.token_ids) > 0
      cached1 = output2.cached_tokens

      # 5. Unrelated prompt on worker 1 → no reuse match
      output3 = worker1.generate("List three primes greater than one hundred",
                                 SamplingParams(max_tokens=5, ignore_eos=True))
      cached2 = output3.cached_tokens
      assert cached2 <= 1            # at most BOS overlap
      assert cached2 < cached1       # transferred prompt cached more

Process overview

sequenceDiagram
    participant Client
    participant W0 as Worker 0 (Sender)
    participant W1 as Worker 1 (Receiver)

    Note over W0: Normal request fills<br/>KV cache reuse tree
    Client->>W0: generate("What is...?")
    W0-->>Client: output tokens + KV stored in reuse tree

    Note over Client: Client queries<br/>DataTransceiverState from W0
    Client->>W0: llm.get_data_transceiver_state()
    W0-->>Client: opaque_state (bytes)

    Note over Client,W1: Arbitrary KV transfer via generation_only
    Client->>W1: generate(prompt, generation_only,<br/>opaque_state=state)
    W1->>W0: requestAndReceive (RequestInfo)
    W0->>W0: BlockRange::fromReuseTree(<br/>pinBlocks=true)
    W0->>W0: pinBlocksById (incRef + claimBlock)
    W0->>W1: send KV cache blocks
    W0->>W0: unpinBlocksById (decRef)
    W1-->>Client: generated tokens

Loading

Key implementation points

• LLM.get_data_transceiver_state() — new API that returns the serialized CacheState + CommState (connection + layout info) needed by the receiver.
• CacheFormatter::getBlockRangeForSending() — when no LlmRequest is provided, the blocks are looked up via BlockRange::fromReuseTree() using the token hashes the receiver has declared.
• Block pinning during transfer — fromReuseTree(pinBlocks=true) increments the ref count (and claims the block from the eviction policy if it had no refs) to prevent the sender's reuse tree from evicting the blocks mid-transfer. sendSyncFromReuseTree() unpins after the send completes.

Test Coverage

Integration test in tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer (registered in l0_a10.yml):

1. Worker 0: normal generate → fills reuse tree
2. Query data_transceiver_state from worker 0
3. Worker 1: generation_only with opaque_state + max_tokens=1 → asserts KV cache was delivered (no prefill recompute path)
4. Post-transfer: normal request on worker 1 works (cached_tokens > 0)
5. Different prompt on worker 1 → cached_tokens ≤ 1 (only BOS matches)

PR Checklist

• PR description clearly explains what and why.

• PR follows TRT-LLM coding guidelines.

• Test cases provided for new code paths.

• No new dependencies.

• Reviewers assigned.

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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #43303 [ run ] triggered by Bot. Commit: bd56626 Link to invocation

[tensorrt-cicd]

PR_Github #43303 [ run ] completed with state FAILURE. Commit: bd56626
/LLM/main/L0_MergeRequest_PR pipeline #33845 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #44376 [ run ] triggered by Bot. Commit: 5ed9f52 Link to invocation

[tensorrt-cicd]

PR_Github #44376 [ run ] completed with state FAILURE. Commit: 5ed9f52
/LLM/main/L0_MergeRequest_PR pipeline #34791 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #44473 [ run ] triggered by Bot. Commit: 5ed9f52 Link to invocation

[tensorrt-cicd]

PR_Github #44473 [ run ] completed with state SUCCESS. Commit: 5ed9f52
/LLM/main/L0_MergeRequest_PR pipeline #34877 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60622 [ run ] completed with state FAILURE. Commit: cda7caa
/LLM/main/L0_MergeRequest_PR pipeline #48931 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60786 [ run ] triggered by Bot. Commit: 045e072 Link to invocation

[tensorrt-cicd]

PR_Github #60786 [ run ] completed with state FAILURE. Commit: 045e072
/LLM/main/L0_MergeRequest_PR pipeline #49066 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60835 [ run ] triggered by Bot. Commit: 4c2c4cb Link to invocation

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (2)

cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp (1)

614-648: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Magic literal for the polling interval.

std::chrono::milliseconds(100) is a bare magic literal. As per coding guidelines, "Avoid magic literals except 0, nullptr, true, and false; initialize named constants instead, using k-prefixed camelCase names."

+    auto constexpr kRecvPollIntervalMs = std::chrono::milliseconds(100);
+
     if (!req->isCompleted())
     {
         auto const& terminate = ctx.getTransferTerminate();
-        while (future.wait_for(std::chrono::milliseconds(100)) != std::future_status::ready)
+        while (future.wait_for(kRecvPollIntervalMs) != std::future_status::ready)

🤖 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
`@cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp`
around lines 614 - 648, Replace the bare 100-millisecond value in the receive
polling loop around future.wait_for with a named k-prefixed camelCase duration
constant, then use that constant for the polling interval while preserving the
existing timeout behavior.

Source: Coding guidelines

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1)

2247-2252: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

refreshBlocks() holds the reuse-tree mutex across a blocking transfer sync.

mTransferManager->syncTransfers() is now executed while holding mLookupTree->getMutex(). Since this PR makes the reuse tree concurrently accessed by transfer threads (pinning/lookup) in addition to the main executor thread, serializing a per-iteration stream-sync call behind this shared lock adds a new contention point on a hot path. Consider releasing the lock before the sync call if mEvictionPolicy->refresh() doesn't need to be atomic with it.

 void WindowBlockManager::refreshBlocks()
 {
-    std::lock_guard<std::recursive_mutex> lock(mLookupTree->getMutex());
-    mEvictionPolicy->refresh();
-    mTransferManager->syncTransfers();
+    {
+        std::lock_guard<std::recursive_mutex> lock(mLookupTree->getMutex());
+        mEvictionPolicy->refresh();
+    }
+    mTransferManager->syncTransfers();
 }

🤖 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 `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp` around lines 2247 - 2252,
Update WindowBlockManager::refreshBlocks() so the lookup-tree mutex is held only
while calling mEvictionPolicy->refresh(), then release it before calling
mTransferManager->syncTransfers(). Preserve the existing refresh and
transfer-sync ordering without holding the tree lock across the potentially
blocking synchronization.

🤖 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/batch_manager/dataTransceiver.cpp`:
- Around line 880-894: Guard the mRequestToSession lookup in pinReuseTreeBlocks
by validating that it is not equal to mRequestToSession.end() before accessing
it->second, matching the existing check in sendSyncFromReuseTree; preserve the
current block lookup and return behavior.
- Around line 885-893: Update pinReuseTreeBlocks() around the cache manager
window metadata lookup to proceed only when getWindowSizesMetadata() contains
exactly one window, returning without pinning for multi-window managers.
Preserve the existing single-window lookup and pinnedIds behavior, or use a
window-qualified handle if the surrounding API supports it.

In `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp`:
- Around line 3031-3061: Update WindowBlockManager::pinBlocks to skip blocks
where block->isPlaceholder() is true before calling pinBlock. Keep pinBlock
behavior unchanged and continue pinning all non-placeholder entries in
allocatedBlocks.

---

Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp`:
- Around line 2247-2252: Update WindowBlockManager::refreshBlocks() so the
lookup-tree mutex is held only while calling mEvictionPolicy->refresh(), then
release it before calling mTransferManager->syncTransfers(). Preserve the
existing refresh and transfer-sync ordering without holding the tree lock across
the potentially blocking synchronization.

In
`@cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp`:
- Around line 614-648: Replace the bare 100-millisecond value in the receive
polling loop around future.wait_for with a named k-prefixed camelCase duration
constant, then use that constant for the polling interval while preserving the
existing timeout behavior.

🪄 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: d0ef3292-f06f-4b64-bc7d-65d2f64ebd65

📥 Commits

Reviewing files that changed from the base of the PR and between 5ed9f52 and 4c2c4cb.

📒 Files selected for processing (13)

• cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
• cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
• cpp/include/tensorrt_llm/executor/dataTransceiverState.h
• cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp
• cpp/tensorrt_llm/batch_manager/cacheFormatter.h
• cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
• cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp
• cpp/tensorrt_llm/batch_manager/dataTransceiver.h
• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
• cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp
• cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp
• cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp
• cpp/tensorrt_llm/executor/serialization.cpp

🚧 Files skipped from review as they are similar to previous changes (4)

• cpp/tensorrt_llm/batch_manager/cacheFormatter.h
• cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
• cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp
• cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

[tensorrt-cicd]

PR_Github #60835 [ run ] completed with state FAILURE. Commit: 4c2c4cb
/LLM/main/L0_MergeRequest_PR pipeline #49112 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61100 [ run ] triggered by Bot. Commit: 4c2c4cb Link to invocation

[tensorrt-cicd]

PR_Github #61100 [ run ] completed with state SUCCESS. Commit: 4c2c4cb
/LLM/main/L0_MergeRequest_PR pipeline #49354 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61328 [ run ] triggered by Bot. Commit: 4c2c4cb Link to invocation

[tensorrt-cicd]

PR_Github #61328 [ run ] completed with state FAILURE. Commit: 4c2c4cb
/LLM/main/L0_MergeRequest_PR pipeline #49558 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61429 [ run ] triggered by Bot. Commit: 4c2c4cb Link to invocation

[tensorrt-cicd]

PR_Github #61429 [ run ] completed with state SUCCESS. Commit: 4c2c4cb
/LLM/main/L0_MergeRequest_PR pipeline #49655 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61462 [ run ] triggered by Bot. Commit: 4c2c4cb Link to invocation

[Shixiaowei02]

Today, we effectively have five tools available:

1. Session affinity
2. Conditional disaggregation
3. P2P direct KV transfer (this PR)
4. KV offload
5. Recomputation (re-prefill)

Roughly speaking:

• Short follow-up turns: (1)
• Tool-oriented agentic workflows (short generation + long tool output): (3)
• Shared system prefixes: (1) + (4)
• Heavy agentic workloads (long generation + large inputs): (2)
• Multi-user shared generation: (4)
• KV not yet committed: (4) or (5)

As the framework evolves, we now have a broader set of tools to address different requirements. For the agentic workflow scenario above, P2P KV transfer seems like a viable approach. While extending the current transceiver model could enable additional use cases, it may also be worth balancing that flexibility with keeping the underlying design simple and maintainable. We can certainly explore this direction in v1. For v2, my current view is that it may fit better as a dedicated abstraction or module rather than an extension of the existing transfer model.

Thanks for your efforts!

[tensorrt-cicd]

PR_Github #61462 [ run ] completed with state SUCCESS. Commit: 4c2c4cb
/LLM/main/L0_MergeRequest_PR pipeline #49683 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation