[TRTLLM-13614][feat] Disaggregated KV-cache bounce transfer

[Shixiaowei02]

Description

This pull request introduces a new "KV bounce buffering" subsystem for VRAM disaggregation in the tensorrt_llm codebase. The main goal is to optimize device-to-device (d2d) KV cache transfers by coalescing scattered per-block transfers into contiguous, high-throughput writes using a "bounce buffer" allocated in fabric-accessible memory. The changes are structured to be modular, pluggable, and testable, and include configuration, core logic, buffer management, and interface exposure. Additionally, minor improvements are made to C++ bindings for performance metric recording.

KV Bounce Buffering Subsystem

• Introduced a new modular subsystem for VRAM d2d KV bounce buffering, allowing opt-in use via configuration. This subsystem coalesces scattered per-block KV transfers into a single contiguous fabric-VMM write, improving transfer reliability and throughput. The default per-block path remains unchanged unless explicitly enabled. (tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py)

Subsystem Core and Buffer Management

• Added the core logic for the bounce subsystem, including the transport contract, per-region lifetime state machine, and detailed transfer state tracking. This is implemented in a testable, CUDA/NIXL-free module. (tensorrt_llm/_torch/disaggregation/native/bounce/core.py)
• Implemented buffer and slot management for fabric-VMM bounce buffers, including thread-safe allocation, quarantine logic for in-doubt regions, and observability features. (tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py)

Configuration and Sizing Policy

• Added a flexible configuration and sizing policy system, allowing the bounce buffer to be enabled, sized, and tuned based on available memory and chunk size. Provides both fixed and pluggable sizing strategies and ensures the buffer fits within a specified fraction of free memory. (tensorrt_llm/_torch/disaggregation/native/bounce/config.py)

C++ Bindings Enhancement

• Enhanced the C++ Python bindings to allow the Python-side transceiver to record per-request KV-transfer timing and size metrics, matching the C++ implementation. Overloaded setter methods now allow either explicit or automatic timestamping. (cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp)

Workload / config	Metric	Non-bounce	Bounce	Gain
Data-plane (raw KV WRITE)
gpt-oss real (aa-agentperf), 2P2D TP2+TP2, conc 8	KV-transfer BW	141.6 GB/s	557.4 GB/s	3.94×
gpt-oss real, conc 8	per-WRITE latency	4.04 ms	1.02 ms	−75% (3.98×)
gpt-oss real, conc 64 (n=27k transfers)	KV-transfer BW	129.6 GB/s	497.2 GB/s	3.84×
Kimi-K2.5 MLA real, conc 64 (n=16.5k)	KV-transfer BW / latency	204 GB/s / 4.95 ms	287 GB/s / 3.46 ms	1.40× / −30%
End-to-end serving
Kimi-K2.5 MLA real, conc 8	TTFT p50	6.02 s	5.39 s	−10.6%
Kimi-K2.5 MLA real, conc 8	output throughput	195 tok/s	212 tok/s	+8.6%
Kimi-K2.5 MLA real, conc 8	e2e latency p50	7.81 s	7.70 s	−1.4%
gpt-oss real, conc 64 (6,531 req, 0 err)	e2e latency p50	7.575 s	7.559 s	parity (−0.2%)

Steady-state p50, warmup dropped; same bytes/config, only kv_cache_bounce_size_mb differs. Raw data-plane KV-WRITE metrics (e2e is at parity — bounce gives ~3.9× transfer bandwidth at zero e2e cost).

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.

[Shixiaowei02]

/bot run --add-multi-gpu-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #55730 [ run ] triggered by Bot. Commit: aa358b4 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds bounce-buffer configuration, allocation, copy, and transport plumbing for KV-cache disaggregation, plus receive-side completion, KV-size accounting, and steady-clock transfer timestamps. It also adds a bounce-size config field, a NIXL worker override, and related tests.

Changes

Bounce transfer flow

Layer / File(s)	Summary
Bounce config surface tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py, tensorrt_llm/_torch/disaggregation/native/bounce/config.py, tensorrt_llm/llmapi/llm_args.py, tests/unittest/disaggregated/test_bounce.py, tests/unittest/llmapi/test_llm_args.py	The bounce package exports sizing policies and the bounce-size config field, and tests cover sizing and config validation.
Bounce buffer allocator tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py	A stable virtual-memory buffer and slot allocator reserve chunk-rounded contiguous regions, release them, and expose one registration descriptor.
Contiguous copy primitives tensorrt_llm/_torch/disaggregation/native/bounce/gather.py	Contiguous gather/scatter uses a Triton batched-copy path when aligned and falls back to per-fragment cudaMemcpyAsync.
Bounce transport protocol tensorrt_llm/_torch/disaggregation/native/bounce/transport.py, tensorrt_llm/_torch/disaggregation/native/transfer.py, tests/unittest/disaggregated/test_bounce.py, tests/integration/test_lists/test-db/l0_*.yml	Bounce transport state, request/result encoding, and transport behavior tests are added.
KV timing and size accounting cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp, tensorrt_llm/_torch/disaggregation/native/transfer.py, tensorrt_llm/_torch/disaggregation/transceiver.py	Result decoding, completion deferral, KV-size bookkeeping, and steady-clock timestamp bindings are added.

NIXL worker count override

Layer / File(s)	Summary
Worker count override tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py	The agent reads TRTLLM_NIXL_NUM_WORKERS and writes it into backend_params["num_workers"] when present.

Sequence Diagram(s)

sequenceDiagram
  participant TransferWorker
  participant Transport
  participant Receiver
  participant Sender
  participant RxSession

  TransferWorker->>Transport: create_bounce(cfg)
  Receiver->>Transport: reserve(recv_req)
  Receiver->>Sender: send bounce request bytes
  Sender->>Receiver: KV_AGENT_RESULT + result tail
  Receiver->>RxSession: process_kv_agent_result(dst_ptrs, sizes, src_base)
  RxSession->>Transport: scatter_write_result(...)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15378: Changes the GenericLlmRequest KV-transfer timestamp bindings in cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp.

Suggested labels

api-compatible

Suggested reviewers

• reasonsolo
• byshiue
• syuoni
• xinhe-nv
• pcastonguay
• VALLIS-NERIA

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.76% which is insufficient. The required threshold is 80.00%.	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 summarizes the main change: the new disaggregated KV-cache bounce transfer path.
Description check	✅ Passed	The description covers the feature, motivation, and checklist, but the Test Coverage section is left empty.

✨ 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: 12

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/llmapi/llm_args.py (1)

3384-3396: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Propagate and constrain kv_cache_bounce_size_mb.

The new field is accepted on CacheTransceiverConfig, but _to_pybind() does not pass it into _CacheTransceiverConfig, so kv_cache_bounce_size_mb=384 is silently lost before the native-disagg transceiver sees it. Also reject negative sizes at the Pydantic boundary.

As per coding guidelines, user-facing Pydantic numeric fields should use built-in constraints such as Field(ge=0).

Proposed fix

-    kv_cache_bounce_size_mb: int = Field(
+    kv_cache_bounce_size_mb: int = Field(
         default=0,
+        ge=0,
         description=
         "Per-region size in MiB of the native-disagg KV-cache bounce buffer (one for send, one for recv). Bounce coalesces a request's scattered per-block KV into one contiguous fabric-VMM buffer and issues a single multi-rail NIXL write. The size doubles as the on/off switch: 0 (default) keeps the per-block path, >0 enables bounce at that capacity. Only used by the Python (v2) transceiver."
     )
@@
             max_tokens_in_buffer=self.max_tokens_in_buffer,
             kv_transfer_timeout_ms=self.kv_transfer_timeout_ms,
             kv_transfer_sender_future_timeout_ms=self.
-            kv_transfer_sender_future_timeout_ms)
+            kv_transfer_sender_future_timeout_ms,
+            kv_cache_bounce_size_mb=self.kv_cache_bounce_size_mb)

🤖 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 `@tensorrt_llm/llmapi/llm_args.py` around lines 3384 - 3396, The new
kv_cache_bounce_size_mb setting is not being propagated and is missing
validation. Update CacheTransceiverConfig so _to_pybind() passes
kv_cache_bounce_size_mb into _CacheTransceiverConfig, and add a Pydantic
non-negative constraint on the Field definition (use a built-in constraint such
as ge=0) so negative values are rejected at the boundary.

Source: Coding guidelines

🤖 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/bindings.cpp`:
- Around line 199-202: The KV-transfer binding API now exposes no-argument
setters in `GenLlmReq`, so the Python test call sites still passing `timedelta`
are out of sync. Update the `set_kv_cache_transfer_start` and
`set_kv_cache_transfer_end` usages in the bindings test to call them with no
arguments, and keep the existing offset-based assertions unchanged so the test
still validates the timing behavior through the recorded values.

In `@tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py`:
- Around line 33-49: The __all__ export list in the bounce package is not
sorted, triggering RUF022. Reorder the symbols in the __all__ definition in the
__init__ module so the exported names are alphabetized while keeping the same
set of exports; this should satisfy the lint gate without changing behavior.

In `@tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py`:
- Line 38: The __slots__ tuple in the buffer class is not sorted, which triggers
RUF023; update the __slots__ definition in the native bounce buffer module so
the slot names are in sorted order, and apply the same ordering fix to the other
__slots__ occurrence mentioned in the review to keep the file lint-clean.
- Around line 81-85: The __del__ cleanup in buffer.py is swallowing all cleanup
errors with a broad Exception handler, which hides failed VMM teardown. Update
the Buffer.__del__ path to catch only the specific exception type expected from
close()/destroy(), and add a log message with the failure details instead of
silently passing. Keep the change localized to the destructor and the cleanup
call it wraps so the failure is observable without masking unrelated errors.
- Around line 131-145: The reservation logic in the allocation path of the
buffer manager is checking only the wrapped head position and can miss a free
hole earlier in the ring, causing unnecessary waits/timeouts. Update the
allocation flow in the buffer class’ reservation method to scan existing free
gaps before blocking on the condition variable, so a request can reuse an
available contiguous hole like the freed region before calling _cv.wait().

In `@tensorrt_llm/_torch/disaggregation/native/bounce/gather.py`:
- Around line 181-184: The _copy_frags helper currently uses plain zip(), which
can silently truncate if pairs and sizes get out of sync. Update the
fragment-copy loop in _copy_frags to use strict zip semantics so mismatches fail
fast, and keep the existing cudaMemcpyAsync/ CUASSERT behavior unchanged.

In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 1743-1751: The failure/cancel handling in process_kv_agent_result
and the related bounce-transfer path does not release the reserved recv slot, so
bounced slices can stay pinned in Transport._reserved forever. Add an explicit
abort/release path on the bounce transport, and invoke it from the FAILED and
cancel flows before or alongside task.fail() so reserved entries are always
freed; use the process_kv_agent_result and
Receiver.dispatch_task/scatter_write_result symbols to locate the affected
logic.
- Around line 568-580: The KV result send path is still using the old ASCII
payload on failure, which breaks Receiver._process_kv_agent_result() because it
now always expects the binary _KV_RESULT_PREFIX frame. Refactor
Sender._deliver_kv_to_agent() and _send_failed_result_to_receiver() to route all
KV result sends through one shared helper that always builds the new binary
KV_AGENT_RESULT message for both success and FAILED/CANCELLED cases, and update
tests in test_bounce.py to cover the failure-path frame format.
- Around line 2152-2162: The bounce transport currently has no explicit shutdown
owner, so TransferWorker.shutdown() must be updated to close the bounce created
by create_bounce and avoid leaking its thread and VMM buffers. Make shutdown in
TransferWorker responsible for calling the bounce close path once, and adjust
registered descriptor handling so the bounce memory descriptors are owned in one
place only and are not deregistered twice. Use the existing self._bounce,
self._registered_mem, and Transport.close behavior as the key symbols to locate
and align the cleanup flow.

In `@tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py`:
- Around line 78-80: The TRTLLM_NIXL_NUM_WORKERS env value is being forwarded
directly in _agent_cpp.py without validation, which can let empty, non-integer,
or non-positive values reach the backend. Update the boundary handling in the
agent setup code that reads os.environ.get("TRTLLM_NIXL_NUM_WORKERS") to parse
it as an integer, reject invalid or <=0 values, and only assign
backend_params["num_workers"] when the value is valid. Keep the fix localized
around the existing nixl_num_workers and backend_params logic so the backend
always receives a safe integer.

In `@tests/unittest/disaggregated/test_bounce.py`:
- Around line 31-36: The import guard in test_bounce.py is too broad and can
hide real bugs in tensorrt_llm._torch.disaggregation.native.bounce.transport by
setting _HAVE_TRANSPORT to False on any exception. Narrow the handling in the
top-level transport import block to only the expected missing-CUDA/import
failure, or move the skip logic into the individual bounce tests with
pytest.importorskip, so genuine import regressions still fail CI. Also extend
the bounce test coverage in test_bounce.py to include a FAILED KV_AGENT_RESULT
frame to validate the wire-format/fan-in path.

In `@tests/unittest/llmapi/test_llm_args.py`:
- Around line 1758-1761: The existing cache bounce size test only checks default
and direct Pydantic construction, so extend the coverage around
CacheTransceiverConfig to verify invalid negative kv_cache_bounce_size_mb values
are rejected and that _to_pybind() passes through the kv_cache_bounce_size_mb
value unchanged; use the CacheTransceiverConfig constructor and _to_pybind() in
the tests so the config handoff path is exercised.

---

Outside diff comments:
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 3384-3396: The new kv_cache_bounce_size_mb setting is not being
propagated and is missing validation. Update CacheTransceiverConfig so
_to_pybind() passes kv_cache_bounce_size_mb into _CacheTransceiverConfig, and
add a Pydantic non-negative constraint on the Field definition (use a built-in
constraint such as ge=0) so negative values are rejected at the boundary.

🪄 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: 82f06d94-cf2f-462a-8ef1-bbe3af79c0e6

📥 Commits

Reviewing files that changed from the base of the PR and between 948db5a and aa358b4.

📒 Files selected for processing (14)

• cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
• tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py
• tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py
• tensorrt_llm/_torch/disaggregation/native/bounce/config.py
• tensorrt_llm/_torch/disaggregation/native/bounce/gather.py
• tensorrt_llm/_torch/disaggregation/native/bounce/transport.py
• tensorrt_llm/_torch/disaggregation/native/transfer.py
• tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py
• tensorrt_llm/_torch/disaggregation/transceiver.py
• tensorrt_llm/llmapi/llm_args.py
• tests/integration/test_lists/test-db/l0_a10.yml
• tests/integration/test_lists/test-db/l0_h100.yml
• tests/unittest/disaggregated/test_bounce.py
• tests/unittest/llmapi/test_llm_args.py

[Shixiaowei02]

/bot run --add-multi-gpu-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #55781 [ run ] triggered by Bot. Commit: 19dc328 Link to invocation

[tensorrt-cicd]

PR_Github #55730 [ run ] completed with state ABORTED. Commit: aa358b4

Link to invocation

[tensorrt-cicd]

PR_Github #55781 [ run ] completed with state FAILURE. Commit: 19dc328
/LLM/main/L0_MergeRequest_PR pipeline #44678 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

[Shixiaowei02]

/bot run --stage-list "A30-AutoDeploy-1, DGX_B200-AutoDeploy-1, H100_PCIe-AutoDeploy-1, GB200-4_GPUs-PyTorch-4, GB200-4_GPUs-PyTorch-5, GB200-4_GPUs-PyTorch-PerfSanity-1, GB200-8_GPUs-2_Nodes-PyTorch-1, GB200-8_GPUs-2_Nodes-PyTorch-2, DGX_B200-2_GPUs-PyTorch-1, DGX_H100-4_GPUs-PyTorch-Others-1, DGX_H100-4_GPUs-PyTorch-Ray-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #55955 [ run ] triggered by Bot. Commit: 19dc328 Link to invocation

[tensorrt-cicd]

PR_Github #55955 [ run ] completed with state FAILURE. Commit: 19dc328
/LLM/main/L0_MergeRequest_PR pipeline #44834 (Partly Tested) 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

[Shixiaowei02]

/bot run --add-multi-gpu-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #56334 [ run ] triggered by Bot. Commit: 91d2482 Link to invocation

[Shixiaowei02]

/bot run --add-multi-gpu-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #57771 [ run ] triggered by Bot. Commit: bf50cae Link to invocation

[tensorrt-cicd]

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

[Shixiaowei02]

/bot run --stage-list "B300-PyTorch-2, DGX_B200-PyTorch-3, H100_PCIe-PyTorch-Ray-1, DGX_B200-4_GPUs-PyTorch-3, DGX_B200-8_GPUs-PyTorch-2, DGX_H100-2_GPUs-PyTorch-Ray-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #57881 [ run ] triggered by Bot. Commit: bf50cae Link to invocation

[tensorrt-cicd]

PR_Github #57881 [ run ] completed with state FAILURE. Commit: bf50cae
/LLM/main/L0_MergeRequest_PR pipeline #46577 (Partly Tested) 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

[Shixiaowei02]

/bot skip --comment "Failing is unrelated to the test."

[tensorrt-cicd]

PR_Github #57941 [ skip ] triggered by Bot. Commit: bf50cae Link to invocation

[tensorrt-cicd]

PR_Github #57941 [ skip ] completed with state SUCCESS. Commit: bf50cae
Skipping testing for commit bf50cae

Link to invocation

[Tabrizian]

Would it be possible to reuse max_tokens_in_buffer argument rather than introducing a separate argument?