[None][fix] Fail loudly on asymmetric warmup batch under attention-DP - #14126
Merged
lfr-0531 merged 1 commit intoMay 14, 2026
Merged
Conversation
hyukn
requested review from
niukuo and
yuanjingx87
and removed request for
a team
May 14, 2026 07:47
hyukn
marked this pull request as draft
May 14, 2026 07:48
hyukn
force-pushed
the
fix/warmup-deadlock-attention-dp
branch
from
May 14, 2026 07:51
52452c8 to
bfdcda7
Compare
…under attention-DP Under attention-DP, each TP rank's KV cache capacity can differ at runtime, causing _create_warmup_request to return None on some ranks while others proceed into forward() with collective ops — deadlocking the job. Instead of silently skipping warmup (which hides the root cause), allgather batch validity across TP ranks and raise a RuntimeError with per-rank curr_max_num_tokens diagnostics when asymmetry is detected. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
hyukn
force-pushed
the
fix/warmup-deadlock-attention-dp
branch
from
May 14, 2026 07:59
bfdcda7 to
10e3e32
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
hyukn
marked this pull request as ready for review
May 14, 2026 08:08
hyukn
requested review from
joyang-nv and
qiaoxj07
and removed request for
a team
May 14, 2026 08:08
hyukn
changed the title
Collaborator
|
PR_Github #48335 [ run ] triggered by Bot. Commit: |
qiaoxj07
approved these changes
May 14, 2026
qiaoxj07
added a commit
to qiaoxj07/TensorRT-LLM
that referenced
this pull request
May 14, 2026
…ream PR NVIDIA#14126 NVIDIA#14126 ("Fail loudly on asymmetric warmup batch under attention-DP") addresses the same rank-divergent warmup-config failure mode this PR previously patched in `_get_max_shape_warmup_requests`. Its philosophy (raise RuntimeError with diagnostics, ask user to bump `--kv_cache_free_gpu_mem_fraction`) is more conservative than the silent fix-and-proceed we had here, and keeps the warmup contract explicit. To avoid duplicating the upstream effort and to reduce review surface in this PR, drop the local synchronization. The remaining fixes (autotuner pre/post barriers and AutoTuner.choose_one collective vote) are independent and stay. Note: this revert is safe only if PR NVIDIA#14126 (or an equivalent) lands before merging this branch into mainline. If neither lands, the rank-divergent warmup-config deadlock remains observable on DSv4-Pro ADP runs and would need to be re-introduced. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
5 tasks
Collaborator
|
PR_Github #48335 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Under attention-DP, each TP rank's KV cache capacity can differ at runtime, causing
_create_warmup_requestto returnNoneon some ranks while others proceed intoforward()with collective ops — deadlocking the job.Instead of silently skipping warmup (which hides the root cause), this PR adds
_assert_all_tp_ranks_have_warmup_batchthat allgathers batch validity across TP ranks and raisesRuntimeErrorwith per-rankcurr_max_num_tokensdiagnostics when asymmetry is detected.Test Coverage
pytest tests/unittest/_torch/passes (no asymmetry in test env)--kv_cache_free_gpu_mem_fraction)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.