[TRTLLM-13614][fix] Enhance the dis-agg bounce buffer workflow - #16116
Conversation
📝 WalkthroughWalkthroughThis PR adds GIL release guards to additional NIXL nanobind bindings, introduces a batched consensus allgather and an env-gated idle fast-path in the disaggregation transceiver, and updates bounce-test configuration, a new long prompt fixture, and test logic to exercise the KV-cache bounce path. ChangesNIXL Binding GIL Release
Consensus Fast-Path and Bounce Testing
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Rank as TransceiverRank
participant AllReduce as ReduceOp allreduce
participant AllGather as consensus allgather
Rank->>Rank: check_context_transfer_status
Rank->>AllReduce: sum local terminal transfer counts
AllReduce-->>Rank: total terminal count
alt fastpath enabled and no terminal transfers
Rank-->>Rank: skip consensus allgathers, return
else consensus needed
Rank->>AllGather: batched cancelled/failed/completed ids
AllGather-->>Rank: outcome lists per rank
Rank->>Rank: unpack and apply outcome
end
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/integration/defs/disaggregated/test_disaggregated.py (1)
1197-1223: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftTest wires the right prompt/config but never asserts the bounce path actually ran.
The comment block correctly explains the sizing math, and
prompt_file="long_bounce_prompts.json"makes bounce coalescing likely, but the test still only verifies request success (output verification is even skipped forlong_prompts). If the 96-block gate math is ever off, or the bounce arena silently falls back to the per-block path, this test would still pass — defeating the PR's stated goal of "ensuring the dis-agg bounce buffer is tested."Suggest adding a concrete, positive assertion that the bounce path was taken — e.g., parsing server logs for a bounce/coalesced-transfer marker, or exposing a counter/metric from
TransferWorker/bouncemodule that the test can check. As per path instructions, please confirm whetherrun_disaggregated_testalready performs any such log-based verification internally; if not, this coverage gap should be closed either in this PR or tracked as an explicit follow-up.🤖 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 `@tests/integration/defs/disaggregated/test_disaggregated.py` around lines 1197 - 1223, The test in test_disaggregated_overlap_transceiver_runtime_python_bounce only exercises a long prompt and verifies success, but it never proves the bounce/coalesced transfer path ran. Add a positive assertion in this test (or in run_disaggregated_test if it already centralizes verification) that checks for a bounce-specific signal, such as a server log marker or a metric/counter from TransferWorker or the bounce module. Keep the existing prompt/config setup, but ensure the test explicitly fails if the transfer falls back to the per-block path.Source: Path instructions
🤖 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.
Nitpick comments:
In `@tests/integration/defs/disaggregated/test_disaggregated.py`:
- Around line 1197-1223: The test in
test_disaggregated_overlap_transceiver_runtime_python_bounce only exercises a
long prompt and verifies success, but it never proves the bounce/coalesced
transfer path ran. Add a positive assertion in this test (or in
run_disaggregated_test if it already centralizes verification) that checks for a
bounce-specific signal, such as a server log marker or a metric/counter from
TransferWorker or the bounce module. Keep the existing prompt/config setup, but
ensure the test explicitly fails if the transfer falls back to the per-block
path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8c38d363-59ee-4abf-a399-176d00680ae4
📒 Files selected for processing (5)
cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cppexamples/disaggregated/clients/long_bounce_prompts.jsontensorrt_llm/_torch/disaggregation/transceiver.pytests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yamltests/integration/defs/disaggregated/test_disaggregated.py
cb84f78 to
8600836
Compare
8600836 to
f066a1b
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #58238 [ run ] triggered by Bot. Commit: |
|
PR_Github #58238 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_B200-PyTorch-2, DGX_B200-PyTorch-4, DGX_B200-PyTorch-5, GB200-4_GPUs-PyTorch-1, GB200-4_GPUs-PyTorch-3, DGX_B200-8_GPUs-PyTorch-4, DGX_H100-4_GPUs-PyTorch-DeepSeek-1, DGX_H100-4_GPUs-PyTorch-Others-2" |
Shixiaowei02
changed the title
|
PR_Github #58422 [ run ] triggered by Bot. Commit: |
|
PR_Github #58422 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_B200-PyTorch-2, DGX_B200-PyTorch-4, DGX_B200-PyTorch-5, DGX_B200-8_GPUs-PyTorch-4, DGX_H100-4_GPUs-PyTorch-DeepSeek-1, DGX_H100-4_GPUs-PyTorch-Others-2" |
|
PR_Github #58463 [ run ] triggered by Bot. Commit: |
|
PR_Github #58463 [ run ] completed with state |
6bcab6e to
9d0f18d
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #59124 [ run ] triggered by Bot. Commit: |
|
PR_Github #59124 [ run ] completed with state
|
9d0f18d to
3e22022
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #59212 [ run ] triggered by Bot. Commit: |
|
PR_Github #59212 [ run ] completed with state
|
|
/bot -h |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. 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. |
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #59311 [ run ] triggered by Bot. Commit: |
|
PR_Github #59311 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_B200-PyTorch-2, DGX_B200-PyTorch-3, DGX_B200-2_GPUs-PyTorch-1, DGX_B200-8_GPUs-PyTorch-2, DGX_B200-8_GPUs-PyTorch-3, DGX_H100-4_GPUs-PyTorch-DeepSeek-1, GB200-8_GPUs-2_Nodes-PyTorch-1, GB200-8_GPUs-2_Nodes-PyTorch-2" |
|
PR_Github #59369 [ run ] triggered by Bot. Commit: |
|
PR_Github #59369 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-2, DGX_B200-8_GPUs-PyTorch-3, DGX_H100-4_GPUs-PyTorch-DeepSeek-1" |
…tent bounce/GIL/test fixes Squashed set of the v2 (Python) KV-cache transceiver work. Performance: - Skip the idle per-step transfer-status consensus on the generation worker (a worker that never sends KV has nothing to reconcile), and batch the remaining cancelled/failed/completed consensus into a single allgather. This removes the per-step host collectives that were inflating the timed gen-worker device step, recovering C++ (v1) parity. The skip gates only on monotonic never-reset flags whose value flips in lockstep across all TP/PP ranks, so no rank can skip a collective while a peer posts it. - Release the GIL in the NixlTransferAgent connection-setup bindings so the blocking connection handshake no longer starves other threads during warmup. Bounce / data-plane: - Fix latent correctness/leak/hang bugs in the KV-cache bounce path. - Make the bounce min_blocks gate an environment knob instead of a public arg. Tests: - Fix a pre-existing single-process cache-transceiver mock deadlock under attention-DP + pipeline parallelism (per-group barriers instead of one global barrier) and a flaky zero-timeout transfer-agent poll assertion; drop a dangling fast-path consensus test. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
3e22022 to
373c096
Compare
|
/bot skip --comment "Verified that the failed pipeline is unrelated to this change, so no retry is needed to save resources." |
|
PR_Github #59623 [ skip ] triggered by Bot. Commit: |
|
PR_Github #59623 [ skip ] completed with state |
Shixiaowei02
changed the title
6 participants
Description
This pull request introduces several improvements and bug fixes to the bounce (coalesced transfer) path in the disaggregated KV cache implementation. The main themes are enhanced robustness (especially around error handling and resource cleanup), improved configurability via environment variables, and stricter correctness checks for coalesced transfers. Additionally, some refactoring and logging enhancements are included to aid debugging and testing.
Bounce Path Robustness & Error Handling:
orphan_reservationmethod to the bounce allocator interface and implementations, ensuring that partially completed or cancelled transfers properly quarantine (rather than leak or prematurely release) bounce regions. This is called on cancellation, teardown, or failed transfer initiation. [1] [2] [3] [4] [5]_deliver_kv_to_agentso that if building a coalesced send request fails, the receiver is notified and the region is not leaked.Configurability & Tuning:
TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS) to override the minimum block count required for the bounce path, allowing advanced/test tuning without code changes. The config logic now reads and clamps this value safely. [1] [2] [3]Correctness & Safety of Coalesced Transfers:
Refactoring & Minor Improvements:
agentBindings.cpp.gil_scoped_releasecall guards to all pybind11 bindings forNixlTransferAgentto avoid deadlocks and improve concurrency.These changes collectively make the bounce path more robust, configurable, and correct, while improving observability and testability.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.Summary by CodeRabbit
New Features
Bug Fixes
Tests