[TRTLLM-11878][feat] Gen-only sync transfer v2 and manager v2

[Shixiaowei02]

Summary by CodeRabbit

• New Features

  • Implemented synchronous KV cache transfer mechanism for disaggregated serving, enabling blocking KV slice reception with optional auxiliary transfers.

• Tests

  • Added integration tests for disaggregated serving with synchronous transfers, cache manager v2, and asymmetric parallel topologies.
  • Extended unit test coverage for synchronous cache transfer workflows.

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.

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #42493 [ run ] triggered by Bot. Commit: d1ef1f0 Link to invocation

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #42497 [ run ] triggered by Bot. Commit: a2b012d Link to invocation

[tensorrt-cicd]

PR_Github #42493 [ run ] completed with state ABORTED. Commit: d1ef1f0

Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes implement synchronous KV cache transfer in KvCacheTransceiverV2.request_and_receive_sync(), previously a stub. Three new integration tests cover sync transfer and KV cache manager v2 configurations. Unit tests are extended with a synchronous workflow variant. Test lists are updated to run the new tests.

Changes

Cohort / File(s)	Summary
Core Implementation tensorrt_llm/_torch/disaggregation/transceiver.py	Implemented request_and_receive_sync() with synchronous KV cache reception workflow: generates request ID, validates no duplicate RX session, manages request state transitions (IN_PROGRESS → COMPLETED or ERROR), creates/stores RX session and request mapping, waits for completion with blocking semantics, applies auxiliary transfers on success, and cleans up session/request entries.
Integration Tests tests/integration/defs/accuracy/test_disaggregated_serving.py	Added three new test methods: test_gen_only_sync (validates sync KV transfer with TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP), test_kv_cache_manager_v2 (tests KV cache manager v2 on both servers), and test_kv_cache_manager_v2_ctx_tp2pp2_gen_tp4 (asymmetric topology with manager v2). All use Python transceivers and GSM8K dataset.
Unit Tests tests/unittest/disaggregated/test_py_cache_transceiver_mp.py	Extended workflow dispatch to handle new "ctx_first_sync" variant. Added _run_ctx_first_sync_transfer() function that executes context-first transfers with GEN side using synchronous request_and_receive_sync() while CTX side submits asynchronously. Updated test parameterization to include sync workflow.
Test Lists tests/integration/test_lists/test-db/l0_dgx_h100.yml, tests/integration/test_lists/test-db/l0_dgx_h200.yml	Added test selections for new disaggregated serving test cases: test_gen_only_sync and test_kv_cache_manager_v2 on H100; test_kv_cache_manager_v2_ctx_tp2pp2_gen_tp4 on H200.

Sequence Diagram(s)

sequenceDiagram
    participant GEN as GEN Side
    participant Transceiver as KvCacheTransceiverV2
    participant RXSession as RX Session
    participant CTX as CTX Side

    GEN->>Transceiver: request_and_receive_sync(req)
    activate Transceiver
    Transceiver->>Transceiver: compute request_id
    Transceiver->>Transceiver: check _recv_sessions<br/>(prevent duplicates)
    Transceiver->>Transceiver: set state to IN_PROGRESS
    Transceiver->>RXSession: create RX session
    activate RXSession
    Transceiver->>Transceiver: store mapping:<br/>request_id → RXSession
    Transceiver->>RXSession: submit KV slice for reception
    RXSession->>CTX: receive KV data
    activate CTX
    CTX-->>RXSession: KV data transfer
    deactivate CTX
    Transceiver->>RXSession: wait_complete(blocking=True)
    RXSession-->>Transceiver: WaitResult.COMPLETED/<br/>WaitResult.ERROR
    deactivate RXSession
    
    alt WaitResult.COMPLETED
        Transceiver->>Transceiver: _apply_aux() if needed
        Transceiver->>Transceiver: set state to COMPLETED
    else WaitResult.ERROR
        Transceiver->>Transceiver: set state to TRANS_ERROR
    end
    
    Transceiver->>Transceiver: close RX session
    Transceiver->>Transceiver: remove from _recv_sessions
    Transceiver->>Transceiver: remove from _recv_reqs
    Transceiver-->>GEN: return
    deactivate Transceiver

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 58.33% 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 contains only template content with no actual description, test coverage, or checklist details filled in by the author.	Add a detailed description explaining the feature implementation, test coverage details, and complete the PR checklist items beyond just the final checkbox.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main feature: implementing synchronous KV transfer for generation-only disaggregated serving with manager v2 support.

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

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

Inline comments:
In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Around line 322-337: The RX flow around create_rx_session / session.receive /
session.wait_complete can throw and currently leaks entries in _recv_sessions
and _recv_reqs and may leave req.state inconsistent; wrap the session lifecycle
in a try/finally (or try/except/finally) so that session.close() is always
called and the rid entries are always deleted from _recv_sessions and
_recv_reqs, and ensure req.state is set to DISAGG_TRANS_ERROR on exceptions;
specifically modify the code around create_rx_session, session.receive,
session.wait_complete, _apply_aux and the existing state assignments to
guarantee cleanup and error-state assignment even when an exception is raised.

🪄 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: 4d6023a3-741a-4f96-a0b0-96ffa3597667

📥 Commits

Reviewing files that changed from the base of the PR and between ce71620 and a2b012d.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/disaggregation/transceiver.py
• tests/integration/defs/accuracy/test_disaggregated_serving.py
• tests/integration/test_lists/test-db/l0_dgx_h100.yml
• tests/integration/test_lists/test-db/l0_dgx_h200.yml
• tests/unittest/disaggregated/test_py_cache_transceiver_mp.py

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #42536 [ run ] triggered by Bot. Commit: 58a431e Link to invocation

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #42629 [ run ] triggered by Bot. Commit: 281a292 Link to invocation

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #43635 [ run ] triggered by Bot. Commit: 789dda3 Link to invocation

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #43736 [ run ] triggered by Bot. Commit: e7c61af Link to invocation

[tensorrt-cicd]

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

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #43925 [ run ] triggered by Bot. Commit: 6401098 Link to invocation

[tensorrt-cicd]

PR_Github #43925 [ run ] completed with state ABORTED. Commit: 6401098

Link to invocation

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #44105 [ run ] triggered by Bot. Commit: 6401098 Link to invocation

[tensorrt-cicd]

PR_Github #44105 [ run ] completed with state SUCCESS. Commit: 6401098
/LLM/main/L0_MergeRequest_PR pipeline #34532 (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

Link to invocation

[Shixiaowei02]

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

[tensorrt-cicd]

PR_Github #44123 [ run ] triggered by Bot. Commit: 6401098 Link to invocation

[tensorrt-cicd]

PR_Github #44123 [ run ] completed with state SUCCESS. Commit: 6401098
/LLM/main/L0_MergeRequest_PR pipeline #34550 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[Shixiaowei02]

/bot skip --comment "Tests have passed separately."

[github-actions]

👎 Promotion blocked, new vulnerability found

Vulnerability report

Component	Vulnerability	Description	Severity
encode/uvicorn	CVE-2020-7694	This affects all versions of package uvicorn. The request logger provided by the package is vulnerable to ASNI escape sequence injection. Whenever any HTTP request is received, the default behaviour of uvicorn is to log its details to either the console or a log file. When attackers request crafted URLs with percent-encoded escape sequences, the logging component will log the URL after it's been processed with urllib.parse.unquote, therefore converting any percent-encoded characters into their single-character equivalent, which can have special meaning in terminal emulators. By requesting URLs with crafted paths, attackers can: * Pollute uvicorn's access logs, therefore jeopardising the integrity of such files. * Use ANSI sequence codes to attempt to interact with the terminal emulator that's displaying the logs (either in real time or from a file).	HIGH
encode/uvicorn	CVE-2020-7695	Uvicorn before 0.11.7 is vulnerable to HTTP response splitting. CRLF sequences are not escaped in the value of HTTP headers. Attackers can exploit this to add arbitrary headers to HTTP responses, or even return an arbitrary response body, whenever crafted input is used to construct HTTP headers.	MEDIUM

[Shixiaowei02]

/bot skip --comment "CI passed in the assembled form."

[github-actions]

👎 Promotion blocked, new vulnerability found

Vulnerability report

Component	Vulnerability	Description	Severity
encode/uvicorn	CVE-2020-7694	This affects all versions of package uvicorn. The request logger provided by the package is vulnerable to ASNI escape sequence injection. Whenever any HTTP request is received, the default behaviour of uvicorn is to log its details to either the console or a log file. When attackers request crafted URLs with percent-encoded escape sequences, the logging component will log the URL after it's been processed with urllib.parse.unquote, therefore converting any percent-encoded characters into their single-character equivalent, which can have special meaning in terminal emulators. By requesting URLs with crafted paths, attackers can: * Pollute uvicorn's access logs, therefore jeopardising the integrity of such files. * Use ANSI sequence codes to attempt to interact with the terminal emulator that's displaying the logs (either in real time or from a file).	HIGH
encode/uvicorn	CVE-2020-7695	Uvicorn before 0.11.7 is vulnerable to HTTP response splitting. CRLF sequences are not escaped in the value of HTTP headers. Attackers can exploit this to add arbitrary headers to HTTP responses, or even return an arbitrary response body, whenever crafted input is used to construct HTTP headers.	MEDIUM

[Shixiaowei02]

/bot skip --comment "CI passed in the assembled form."

[tensorrt-cicd]

PR_Github #44247 [ skip ] triggered by Bot. Commit: 6401098 Link to invocation

[tensorrt-cicd]

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

Link to invocation