[TRTLLM-14660][feat] Early reused cache transfer in Transceiver V2

[athena-nv]

TRTLLM-14660 Early Reused Cache Transfer

Motivation

Agentic workloads commonly have:

cached ISL >> current-turn ISL >> OSL

The reused prefix therefore accounts for most of the KV cache transferred from
the context worker to the generation worker. Previously, that transfer started
after the current turn's prefill completed, placing both the large cached-prefix
transfer and the current-turn KV transfer on the critical path before decode.

This change starts transferring complete KV blocks from the reused prefix before
the first context forward. The cached-prefix transfer can then overlap with
prefill compute for the current turn. Only the KV produced by the current turn,
plus any partial boundary block, remains to be transferred after prefill.

Transfer timeline

Baseline:

sequenceDiagram
    participant C as Context worker
    participant G as Generation worker
    C->>C: Prefill current-turn tokens
    C->>G: Transfer cached-prefix KV
    C->>G: Transfer current-turn KV
    G->>G: Decode output tokens

Loading

With early reused-cache transfer:

sequenceDiagram
    participant C as Context worker
    participant G as Generation worker
    par Transfer reused prefix
        C->>G: Transfer complete cached-prefix blocks
    and Compute current turn
        C->>C: Prefill current-turn tokens
    end
    C->>G: Transfer current-turn and boundary-block KV
    G->>G: Decode output tokens

Loading

Implementation

Before each context forward, _send_kv_cache_early selects context-only,
non-cancelled requests that are on their first context chunk and have a reused
prefix. It rounds prepopulated_prompt_len down to a KV block boundary so that
only fully populated blocks are sent. Partial boundary blocks remain owned by
the normal prefill path because the current forward may still update them.

The request records py_kv_prefix_sent after dispatching the early slice. When
the normal first prefill chunk is later built, the transceiver checks that flag
and starts the chunk at the reuse boundary instead of block zero, preventing the
early prefix from being sent twice.

The early send is inserted into both executor scheduling paths immediately
before the context forward.

Limitation: We skip early reused cache transfer if blocks have been evicted to host memory.
We can handle this in a future PR by synchronizing the CPU Sender worker with CUDA onboard stream's onboardDone event. We did not add this synchronization because for overlap scheduler, the kv transfer of the previous batch will be blocked by waiting on onboardDone of the current batch. Benchmarking is required to understand the perf characteristics of this added synchronization.

Changes

File and change block	What changed	Why it is necessary
tensorrt_llm/_torch/disaggregation/transceiver.py: first-chunk range selection	Added extends_to_prefix = is_first_chunk and not req.py_kv_prefix_sent; block zero is included only when the prefix was not sent early.	Prevents duplicate transfer of reused-prefix blocks while preserving the existing fallback in which the normal first chunk carries the prefix.
tensorrt_llm/_torch/pyexecutor/llm_request.py: request transfer state	Added py_kv_prefix_sent, initialized to False.	Carries the ownership decision from the pre-forward early send to normal post-forward chunk construction.
tensorrt_llm/_torch/pyexecutor/py_executor.py: non-overlap scheduling path	Calls _send_kv_cache_early before collecting batch statistics and running the context forward.	Starts reused-prefix transfer early for the standard executor path.
tensorrt_llm/_torch/pyexecutor/py_executor.py: overlap scheduling path	Calls _send_kv_cache_early when the scheduled batch can be queued.	Provides the same overlap for the overlap-scheduler path.
tensorrt_llm/_torch/pyexecutor/py_executor.py: _send_kv_cache_early	Added eligibility filtering, block-boundary rounding, request-state update, batching, and dispatch through _send_kv_async.	Sends only valid, complete reused-prefix blocks before compute and reuses the established asynchronous transfer pipeline.
tests/unittest/disaggregated/test_chunked_transfer.py: early-send tests	Added tests for request eligibility, aligned and unaligned reused prefixes, sub-block prefixes, cancellation, and the feature-disabled path.	Verifies that early transfer is narrowly gated and never sends incomplete blocks.
tests/unittest/disaggregated/test_chunked_transfer.py: multi-chunk request fixture	Initializes py_kv_prefix_sent=False.	Keeps the existing end-to-end chunk test representative of a request that did not send a prefix early.
tests/unittest/disaggregated/test_chunked_transfer.py: token-coordinate helper	Added a prefix_sent input and assigns it to the mocked request.	Allows chunk-range tests to exercise both early-prefix and baseline behavior.
tests/unittest/disaggregated/test_chunked_transfer.py: block-coordinate helper	Forwards prefix_sent to the token-coordinate helper.	Exposes early-prefix state to existing block-oriented test cases without duplicating setup.
tests/unittest/disaggregated/test_chunked_transfer.py: first regular chunk after early transfer	Added a test asserting that the regular chunk starts at the reuse boundary when the prefix was already sent.	Guards against retransmitting blocks owned by the early slice.
tests/unittest/disaggregated/test_kv_transfer.py: _send_prefill_chunks fixture	Initializes py_kv_prefix_sent=False.	Makes first-chunk behavior explicit for existing multi-chunk transfer tests.
tests/unittest/disaggregated/test_kv_transfer.py: whole-prompt slicing test	Initializes py_kv_prefix_sent=False.	Preserves the baseline range-selection expectation in the existing test.
tests/unittest/disaggregated/test_kv_transfer.py: partial-SWA test	Initializes py_kv_prefix_sent=False.	Preserves explicit baseline state while testing partial sliding-window chunks.

Added test coverage

Test	Coverage
test_send_kv_cache_early_only_sends_reused_prefixes	Sends only eligible first-chunk reuse hits; skips completed/non-first, zero-reuse, and cancelled requests; floors an unaligned hit to complete blocks; skips prefixes smaller than one block.
test_send_kv_cache_early_requires_pipelined_transfer	Verifies that early transfer is disabled unless pipelined KV transfer is enabled.
test_first_chunk_skips_prefix_already_sent_early	Verifies that the normal first prefill chunk begins at the reuse boundary after an early prefix send.

Validation

• Pre-commit checks passed for all five files changed by c2647ec0f6.
• Focused unit tests could not be collected in the local checkout because
  tensorrt_llm/libs/libth_common.so failed to load; a compatible TensorRT-LLM
  build is required.

Dev Engineer Review

• Added early reused-prefix KV-cache transfer for pipelined Transceiver V2 requests.
• Transfers complete, block-aligned prefix blocks before the first context forward.
• Supports both non-overlap and overlap executor scheduling paths.
• Skips early transfer when KV blocks use host/offload storage.
• Tracks py_kv_prefix_sent to prevent duplicate prefix transfer.
• Keeps post-prefill transfer for current-turn KV and partial boundary blocks.
• Changes are consistent across request state, chunk construction, and executor scheduling.
• No public API or configuration changes were identified.
• Pre-commit checks passed.
• Focused unit tests could not be collected locally because libth_common.so failed to load.

QA Engineer Review

• Added coverage for:
  • Early pipelined prefix transfer.
  • Whole-block alignment.
  • Cancellation.
  • Host/offload-tier exclusion.
  • Disabled-pipeline behavior.
  • Duplicate-transfer prevention.
  • Prefix state in prefill chunks.
• Updated mocked requests and multi-chunk fixtures with py_kv_prefix_sent.
• The modified tests are not listed in the provided context as covered by tests/integration/test_lists/, test-db/, or qa/.
• Verdict: needs follow-up to confirm test-list coverage and to run the focused tests in an environment where libth_common.so loads successfully.

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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: aadbc683-59f6-4965-afbc-b4b70b3f3508

📥 Commits

Reviewing files that changed from the base of the PR and between 3503e3f and 9cdd5b6.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/disaggregation/transceiver.py
• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tests/unittest/disaggregated/test_chunked_transfer.py
• tests/unittest/disaggregated/test_kv_transfer.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The change adds early KV-cache prefix transfer for eligible pipelined context requests. It tracks transfer state on each request and prevents already-transferred prefixes from appearing in the first prefill chunk.

Changes

KV-cache prefix transfer

Layer / File(s)	Summary
Prefix transfer state and chunk construction tensorrt_llm/_torch/pyexecutor/llm_request.py, tensorrt_llm/_torch/disaggregation/transceiver.py	Requests now track py_kv_prefix_sent. The first prefill chunk excludes the reused prefix when that state is set.
Early transfer execution tensorrt_llm/_torch/pyexecutor/py_executor.py	Both executor loops start early asynchronous transfer for eligible, block-aligned reusable prefixes. The logic skips canceled requests, offload tiers, non-initial chunks, empty prefixes, and disabled pipelining.
Transfer and chunking validation tests/unittest/disaggregated/test_chunked_transfer.py, tests/unittest/disaggregated/test_kv_transfer.py	Tests cover transfer eligibility, cancellation, alignment, disabled pipelining, request state, and exclusion of transferred prefixes from regular chunks.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 9cdd5

Eligible pipelined context requests can now transfer complete reused KV-cache prefix blocks before prefill while avoiding duplicate prefix transfer. The implemented gating and coverage indicate no actionable merge-blocking risk remains.

Sequence Diagram(s)

sequenceDiagram
  participant py_executor
  participant LlmRequest
  participant transceiver
  py_executor->>LlmRequest: inspect reusable prefix
  py_executor->>transceiver: start asynchronous prefix transfer
  py_executor->>LlmRequest: set py_kv_prefix_sent
  transceiver-->>py_executor: transfer completes

Loading

Suggested reviewers: juney-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 21 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 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 clearly identifies the feature and matches the main change: early reused KV-cache transfer in Transceiver V2. It uses the required ticket and feature format.
Description check	✅ Passed	The description is complete and relevant. It explains the motivation, implementation, limitations, changed files, test coverage, and validation results, including the local test-collection failure.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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