[https://nvbugs/6438586][fix] Make gen-only benchmark insufficient-KV fail-fast ADP-safe - #16253
Conversation
… fail-fast ADP-safe In gen-only benchmark mode, the "Insufficient KV cache" fail-fast in _prepare_and_schedule_batch fired on rank-local conditions (fitting_disagg_gen_init_requests / stuck INIT requests). Under attention DP, a subset of ranks could enter _handle_errors (which runs a collective response gather) while the remaining ranks sat in the fill gate's tp_allgather. The mismatched collectives desync the TP group: peers crash with "TypeError: '<' not supported between instances of 'list' and 'int'" in _is_benchmark_disagg_fill_complete, or hang until the HangDetector hard-kills all ranks via MPI_Abort. The check also mistook requests deferred by the disagg transfer admission controller (added in NVIDIA#15356) for stuck requests: deferred requests are waiting on in-flight KV transfers, not on KV capacity, so the fail-fast killed benchmarks whose fill was making progress. Fixes: - Skip the fail-fast when the admission controller is deferring requests behind active transfers (wait_for_disagg_gen_transfer_progress). - Reach an ADP consensus (any-rank allgather, run on every iteration regardless of local state) before failing, so all ranks enter _handle_errors together and stay collective-aligned. - In _handle_errors' fatal path, enter the waiting-queue response gather on every ADP rank (it was rank-0 only, same desync class). The feat/deepseek_v4 branch fixed the same desync in 95245c9 but the fix was never ported to main. Verified on Lyris (4x GB200, DeepSeek-V4-Flash ctx tp8 / gen dep8, gen-only benchmark, concurrency 512): with per-rank divergence the old code hangs and MPI_Aborts at benchmark start; with the fix all ranks fail together with the real error message, and a healthy config runs the benchmark to completion. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughBenchmark disaggregated generation now defers fail-fast until transfer progress stops and the request queue is saturated, coordinates stuck decisions across Attention-DP ranks, and preserves collective response-enqueue participation during fatal errors. Tests cover healthy, peer-stuck, and locally-stuck scenarios. ChangesAttention-DP fail-fast synchronization
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Scheduler
participant TPAllgather
participant ErrorHandler
participant ResponseQueue
Scheduler->>TPAllgather: Gather local stuck status
TPAllgather-->>Scheduler: Return rank consensus
Scheduler->>ErrorHandler: Handle synchronized failure
ErrorHandler->>ResponseQueue: Enqueue responses collectively
Suggested reviewers: 🚥 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/unittest/_torch/executor/test_benchmark_disagg.py (1)
1068-1141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage verdict: sufficient for the consensus decision, one follow-up gap.
The new tests cover the
local_stuckpredicate and theany(tp_allgather(...))consensus well (healthy no-op, peer-stuck kill-together, local-stuck vote). Since_handle_errorsis mocked here, the companion collective-safe change in_handle_errors(the ADP_enqueue_responses(waiting_responses)on an empty local list) is never exercised — no test asserts every rank enters that gather in lockstep. Recommend a follow-up test in this file that drives the fatal_handle_errorspath under ADP with an emptywaiting_responseson the non-rank-0 stub and asserts_enqueue_responsesparticipation, to lock in the fix against regressions.As per path instructions ("suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR").
🤖 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/unittest/_torch/executor/test_benchmark_disagg.py` around lines 1068 - 1141, Coverage for the ADP consensus decision is sufficient, but add a follow-up test in this file for the fatal _handle_errors path: configure a non-rank-0 ADP executor with empty waiting_responses, invoke the error handling flow, and assert _enqueue_responses is called with the empty list so every rank participates in the response gather.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/unittest/_torch/executor/test_benchmark_disagg.py`:
- Around line 1068-1141: Coverage for the ADP consensus decision is sufficient,
but add a follow-up test in this file for the fatal _handle_errors path:
configure a non-rank-0 ADP executor with empty waiting_responses, invoke the
error handling flow, and assert _enqueue_responses is called with the empty list
so every rank participates in the response gather.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2acfdbc2-f16d-4e04-9598-ffa51a8ce926
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_benchmark_disagg.py
|
PR_Github #58700 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
PR_Github #58704 [ kill ] triggered by Bot. Commit: |
|
PR_Github #58700 [ run ] completed with state |
|
PR_Github #58704 [ kill ] completed with state |
2 participants
Summary by CodeRabbit
Bug Fixes
Tests
Description
Fixes https://nvbugs/6438586 (and the crash mode of https://nvbugs/6438658).
In gen-only benchmark mode, the "Insufficient KV cache" fail-fast in
_prepare_and_schedule_batchfired on rank-local conditions (fitting_disagg_gen_init_requests/ stuck INIT requests). Under attention DP, a subset of ranks could enter_handle_errors(which runs a collective response gather) while the remaining ranks sat in the fill gate'stp_allgather. The mismatched collectives desync the TP group: peer ranks crash withTypeError: '<' not supported between instances of 'list' and 'int'in_is_benchmark_disagg_fill_complete(nvbug 6438586), or hang until the HangDetector hard-kills all ranks viaMPI_Abort.The check also mistook requests deferred by the disagg transfer admission controller (#15356) for stuck requests: deferred requests are waiting on in-flight KV transfers, not on KV capacity, so the fail-fast could kill benchmarks whose fill was still making progress (nvbug 6438658).
Changes:
wait_for_disagg_gen_transfer_progress).anyover atp_allgather, run on every iteration regardless of local state) before failing, so all ranks enter_handle_errorstogether and stay collective-aligned._handle_errors' fatal path, enter the waiting-queue response gather on every ADP rank (it was rank-0 only — same desync class).Note:
feat/deepseek_v4fixed the same desync in 95245c9 ("[TRTLLM-12403][fix] Fix deepseekv4 stall") but the fix was never ported to main, which is why 1.3.0rc15.post1 (dsv4) is unaffected while rc21 (main) regressed.Verified on Lyris (4x GB200, DeepSeek-V4-Flash ctx tp8 / gen dep8 ADP, gen-only benchmark, concurrency 512, TensorRT-LLM 1.3.0rc21 container):
MPI_Aborts all GEN ranks at benchmark start.Test Coverage
tests/unittest/_torch/executor/test_benchmark_disagg.py:test_admission_deferral_does_not_kill— deferral by the admission controller must not trigger the fail-fast.test_adp_consensus_runs_allgather_even_when_healthy— the consensus collective runs every iteration to keep ranks aligned.test_adp_consensus_kills_all_ranks_when_peer_is_stuck/test_adp_consensus_local_stuck_reported— consensus semantics.TestFillPhaseEndToEnd::test_full_lifecyclefor the new consensus collective.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.