[TRTLLM-11492][fix] Fix benchmark disagg deadlock by eliminating blocking fill loop - #12208
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #38922 [ run ] triggered by Bot. Commit: |
|
PR_Github #38922 [ run ] completed with state
|
94c4723 to
f9b82d4
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #39122 [ run ] triggered by Bot. Commit: |
|
PR_Github #39122 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #39289 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughIntroduces a benchmark gating mechanism for disaggregated broadcasting in PyExecutor. A new Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes 🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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.
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/pyexecutor/py_executor.py`:
- Around line 1798-1805: The gate currently counts all
scheduled_batch.generation_requests (local_gen_count) including ADP-injected
dummy requests, which lets total_gen_count
(sum(self.dist.tp_allgather(local_gen_count)) when self.enable_attention_dp)
reach benchmark_req_queues_size prematurely; change the local count to exclude
requests where request.is_attention_dp_dummy is true (compute
local_real_gen_count = number of generation_requests with is_attention_dp_dummy
== False) and use that in the allgather instead of local_gen_count, keeping the
existing branches that use self.enable_attention_dp and self.dist.tp_allgather;
additionally, add an ADP regression that creates a mix of real and dummy
generation requests to verify can_forward (or the method enclosing this logic)
only opens after the required number of real requests complete KV transfer.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 744888f9-916d-4891-8b07-02c9ebb2c257
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_benchmark_disagg.py
|
PR_Github #39289 [ run ] completed with state
|
Tabrizian
left a comment
Tabrizian
left a comment
There was a problem hiding this comment.
Did you also try the e2e benchmarks to make sure it still works?
|
/bot run --disable-fail-fast |
|
PR_Github #39344 [ run ] triggered by Bot. Commit: |
|
PR_Github #39344 [ run ] completed with state
|
f9b82d4 to
f76c75c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #39501 [ run ] triggered by Bot. Commit: |
|
PR_Github #39501 [ run ] completed with state
|
f76c75c to
5ce5ce2
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #39760 [ run ] triggered by Bot. Commit: |
|
PR_Github #39760 [ run ] completed with state
|
|
PR_Github #42566 [ run ] completed with state
|
Design doc covering the deadlock fix in benchmark disaggregated serving mode: problem statement, root cause analysis (structural starvation in fill loop), comparison with prior PRs (NVIDIA#12091, NVIDIA#12206), non-blocking gate design, ADP dummy request handling, and test coverage. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Made-with: Cursor
…wn handling The previous dummy skip guard incorrectly skipped ADP dummy insertion during taper-down (when ranks empty at different rates due to e.g. speculative decoding acceptance variance). The guard checked is_benchmark_disagg which is True for the entire benchmark run, but dummies should only be skipped during the fill phase. Add _benchmark_fill_phase_active flag that starts True and is cleared when the can_forward gate opens. _should_skip_dummy_for_benchmark_disagg now checks this flag instead, so: - Fill phase: skip dummies (gate closed, dummies would get stuck) - After fill: normal dummy lifecycle (handles taper-down correctly) This also removes the fill_phase_complete/some_ranks_permanently_empty logic, which is no longer needed since the flag cleanly separates the two phases. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Made-with: Cursor
|
/bot help |
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 --disable-fail-fast |
|
PR_Github #42755 [ run ] triggered by Bot. Commit: |
|
PR_Github #42755 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #42811 [ run ] triggered by Bot. Commit: |
|
PR_Github #42811 [ run ] completed with state |
Design doc covering the deadlock fix in benchmark disaggregated serving mode: problem statement, root cause analysis (structural starvation in fill loop), comparison with prior PRs (NVIDIA#12091, NVIDIA#12206), non-blocking gate design, ADP dummy request handling, and test coverage. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Made-with: Cursor
…fix plan Restructure benchmark-disagg-nonblocking-gate/ into benchmark-disagg-fill/ with v0-v4 lineage. Add root-cause analysis for the post-PR-NVIDIA#12208 hang (nvbug 6071070) on wideep_kimi-k2-thinking-fp4_8k1k_ctx8_gen1_dep32_bs256 and document the step-1 gate rewrite and step-2 external orchestrator plans. Preserves PR NVIDIA#12208 design docs under 01-history-nonblocking-gate/. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
…stuck-state finding Document the v2.2 follow-up to PR NVIDIA#13347 and a separate latent issue discovered while validating it locally: * 07-fill-phase-flow-control.md (new): v2.2 production fix. Reintroduces the per-iteration `tp_size` admission cap that PR NVIDIA#12091 had as part of the blocking fill loop, and that PR NVIDIA#12208 implicitly removed when it eliminated that loop. Without it, on `dep32 × bs256 == con8192` the GEN server is OOM-killed during burst admission before any KV transfer can drain. Includes safety analysis, test coverage notes, and a tunable-cap discussion. * 08-fill-phase-stuck-state-finding.md (new): a separate latent deadlock observed during local 8-GPU validation, where the state-based gate stays closed because some active requests remain in INIT after CTX has finished prefilling. Explicitly framed as independent of v2.1 / v2.2 (with a per-patch coverage table proving none of those patches touch the deadlock's mechanism), with repro instructions and recommended next steps. * README.md: lineage table gains a v2.2 row; current-state section now enumerates four issues with v2.1/v2.2 attribution; documents table and reading order updated; "Relationship between the steps" notes v2.1 and v2.2 both fall under Step 1. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Made-with: Cursor
…rst-OOM The state-based fill gate fixed the "wrong predicate" failure mode in the wide-EP Kimi-K2-Thinking gen-only test, but uncovered a latent fill-phase memory-pressure issue: with the gate now correctly waiting for all `benchmark_req_queues_size` requests to leave INIT, GEN admits the entire preloaded waiting queue (`tp_size * max_batch_size` requests) in a single iteration. This spikes peak KV-cache reservations + NIXL recv-buffer reservations and OOM-kills the GEN server before any KV transfer can drain, manifesting in CI as `STEP ... DUE to SIGNAL Killed` on `dep32 x bs256 == con8192`. PR NVIDIA#12091 originally throttled per-iteration admission to `tp_size` as part of its blocking fill loop. PR NVIDIA#12208 removed that throttle along with the loop, replacing it with a non-blocking gate. The throttle was the only implicit memory-pressure regulator on the GEN side, and its removal was the latent regression that the state-based gate (commit f98e2fd) just made visible. Reintroduce the per-iteration cap as a deterministic predicate inside `_pop_from_waiting_queue`, scoped to the benchmark fill phase only: if (self.is_benchmark_disagg and self._benchmark_fill_phase_active and not self.is_warmup): max_new_requests = min(max_new_requests, self.dist.tp_size) The cap is `tp_size` so admission averages one request per rank per iteration, regardless of cluster size. With the 100ms outer-loop sleep, filling an 8192-deep queue takes ~25s, which is the same wall-clock fill time as PR NVIDIA#12091's per-rank rate (independent of dep size). The cap is only active under three conjunctive conditions: * `is_benchmark_disagg`: normal serving never has a queue this large * `_benchmark_fill_phase_active`: cleared once the gate fires, so taper-down and steady-state admission are unaffected * `not is_warmup`: warmup admits synthetic requests; must not slow Tests: TestBenchmarkFillAdmissionFlowControl (2 parameterized) verifies both the cap during fill and the lack-of-cap once the fill phase ends. Total benchmark-disagg unit tests: 45 passing. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Made-with: Cursor
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
During the initial benchmark fill, admit at most tp_size new requests per executor iteration. This restores the conservative one-request-per-rank fill rate from PR NVIDIA#12091 without reintroducing the blocking fill loop removed by PR NVIDIA#12208. The cap only applies to non-warmup benchmark fill. Once the gate opens, normal admission resumes. Tests cover tp_size 1, 4, and 32, plus the lifecycle path that calls the real waiting-queue admission helper. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
6 participants
Summary by CodeRabbit
Refactor
Tests
Description
Fix a deadlock in benchmark disaggregated serving mode and improve code quality.
Problem: In benchmark disagg mode, the GEN executor's
_prepare_and_schedule_batchcontained a blocking fill loop that held control until a batch of requests was fetched. During this blocking wait, the executor could not service KV transfers, check timeouts, detect errors, or handle control requests.The severity depends on the transport backend:
_check_disagg_gen_transfer_statusnever runs inside the fill loop, GEN never receives, CTX's sends block, and CTX cannot free KV cache blocks — causing a hard deadlock when CTX capacity <batch_size.Solution: Replace the blocking fill loop with a non-blocking
can_forwardgate in the executor loops. Each main-loop iteration now performs a full processing cycle — fetch, service transfers, check timeouts, handle errors — then checks readiness via_is_benchmark_disagg_fill_complete. This eliminates the deadlock for all backends and ensures the executor remains responsive during the fill phase.Key changes
Core fix:
_prepare_and_schedule_batch, making it fully non-blocking (single fetch per call)_is_benchmark_disagg_fill_completehelper that counts generation-ready requests (excludes ADP dummy requests) and uses TP allgather for ADP, local counting otherwise_check_benchmark_disagg_gatehelper that consolidates the gate logic shared by both_executor_loopand_executor_loop_overlapcan_forwardgate to_executor_loop(previously only existed in_executor_loop_overlap)ADP dummy request fix:
_pad_attention_dp_dummy_requestinto focused helpers:_count_schedulable_active_requests(counts requests past KV transfer) and_should_skip_dummy_for_benchmark_disagg(encapsulates skip decision)_benchmark_fill_phase_activeflag to distinguish the fill phase (dummies skipped — gate blocks collectives, stuck dummies waste KV slots) from taper-down (dummies allowed — normal add-forward-terminate lifecycle handles ranks emptying at different rates due to e.g. speculative decoding acceptance variance)Truein benchmark disagg mode and is cleared whencan_forwardgate opensis_attention_dp_dummyrequests from the generation count in the fill-complete check, preventing premature gate openingCode quality:
is_benchmark_disaggattribute for consistent condition checks (replaces repeatedbenchmark_req_queues_size > 0 and kv_cache_transceiver is not None)RuntimeErrorfor precondition violation in_is_benchmark_disagg_fill_completetuple[bool, bool]return type annotation on_check_benchmark_disagg_gatelogger.debugfor fill-progress messages (avoids log spam at 0.1s intervals)This PR is co-worked with Cursor agent.
Test Coverage
Added
tests/unittest/_torch/executor/test_benchmark_disagg.pywith 40 tests across 8 test classes:TestFillCompleteNonADPTestFillCompleteADPTestFillCompleteADPDummyExclusionTestCanForwardGatingTestCheckBenchmarkDisaggGateTestPadAttentionDpDummyBenchmarkDisaggTestIncrementalFillScenarioTestPrepareAndScheduleBatchNoBlock_prepare_and_schedule_batchcalls fetch exactly once per invocation (non-blocking)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.