Repository navigation
Conversation
|
Important Review skippedReview was skipped as selected files did not have any reviewable changes.
 |
| Layer / File(s) | Summary |
|---|---|
Relaxed store helper and barrier ordering cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu |
Barrier::st_flag_relaxed performs a relaxed system-scope store. Barrier::sync uses it for future CTA slots and release-stores the current block’s flag. |
Priority:
Estimated code review effort: 2 (Simple) | ~10 minutes
Change: Bug fix
Suggested reviewers: bowenfu, farazkh80
Merge Risk: 
b37b7
-
No merge-blocking issue was identified in the barrier-ordering change; proceed with normal validation.
🚥 Pre-merge checks | ✅ 4 | ❌ 1
-
❌ Failed checks (1 warning)
Check name Status Explanation Resolution Docstring Coverage ⚠️ WarningDocstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 1 files. Write docstrings for the functions missing them to satisfy the coverage threshold. ✅ Passed checks (4 passed)
-
Check name Status Explanation Title check ✅ Passed The title follows the repository format and clearly describes the barrier flag ordering fix. Description check ✅ Passed The description explains the bug and fix, provides detailed test coverage and results, and completes the PR checklist. Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request. Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
-
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
b37b7a0 to
c811e46
Compare
yilin-void
left a comment
yilin-void
left a comment
There was a problem hiding this comment.
LGTM, thanks for the fix!
|
/bot run |
|
PR_Github #81177 [ run ] triggered by Bot. Commit: |
|
PR_Github #81177 [ run ] completed with state
|
…lot last Barrier::sync writes the new flag to all kBarrierFlagCount slots, starting with the block's own slot, and the peer block waits only until its slot is != prev_flag. The flag cycles 0/1/2, so when the next two-shot launch has a larger grid, a block that is not launched now can find its slot still holding the flag from two barriers ago, which is also != prev_flag, and pass before the peer has written its data. Write the slots of the blocks that are not launched first, relaxed, and the block's own slot last with st.release.sys, so that the one release store orders all of them. Two-shot at T=5 then T=64 (hidden 7168, bf16) on 4xB200, no injected delay: 196/200 calls wrong before, 0/200 after; 2 GPUs (T=3 -> 512) 99/100 -> 0/100. The fix also drops one release store per slot: two-shot latency at hidden 7168 is unchanged at T=4096 and up to 66% lower at T=64 on 8 GPUs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zixin Huang <zixinh@nvidia.com>
c811e46 to
e8bcab1
Compare
|
/bot run |
|
PR_Github #81375 [ run ] triggered by Bot. Commit: |
|
PR_Github #81375 [ run ] completed with state
|
Dev Engineer Review
Barrier::syncnow writes flags for CTAs not launched in the current grid with relaxed system-scope stores, then writes the current block’s slot with a release store. This preserves ordering for the slot that peer CTAs poll while reducing release stores. The change targets stale-flag reads when a larger-grid call follows a smaller-grid call. The reported correctness and latency results are author-reported; no independent test results or review findings were supplied.QA Engineer Review
No test changes.
Per-File QA Perspective
cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu: Verify two-shot AllReduce when a larger grid follows a smaller grid, including the flag ordering and output correctness. Also check small-grid latency and confirm latency remains stable at larger grids.Description
The two-shot AllReduce-fusion kernel can return wrong results, with no error, when a two-shot call is followed by one that launches more CTAs.
Barrier::syncwrites the new flag to allkBarrierFlagCountslots of the peer, so that CTAs that are not launched now also see the current flag. It writes the block's own slot first. The peer block waits only until its slot is!= prev_flag. The flag cycles 0/1/2, so in a later launch with a larger grid, a block can find its slot still holding the flag from two barriers ago. That value is also!= prev_flag, so the block passes the barrier before the peer has written its data, and it reads stale data.The fix writes the slots of the CTAs that are not launched now first, with relaxed stores, and the block's own slot last, with
st.release.sys. The peer acquires that slot, so the one release store orders all of the earlier stores. One file changes:cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu(+13/-3).A side effect is that each block now issues one release store per barrier instead of one per slot, which makes the two-shot kernel faster for small grids.
Test Coverage
The kernel was built standalone from
main(687af96) before and after the change, on 8x B200. The workspace comes fromCustomAllReduceHelper.allocate_allreduce_fusion_workspace. PatternkAllReduce, bf16, two-shot forced. Each trial runs two back-to-back two-shot calls with different token counts, with no injected delay, and compares both outputs exactly to the expected sum (the inputs are small integers, so the sum is exact in bf16).TRTLLM_ENABLE_PDL=0The failure needs a small two-shot grid followed by a larger one. It did not show up without an injected delay on 8 GPUs, or when both calls have more than 128 tokens. After the change, no configuration gave a wrong output.
Two-shot latency, bf16, hidden 7168, CUDA graph of 500 calls, slowest rank, best of 3 (us/call):
Through the production op: in the
nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29container (itsallReduceFusionKernels.{cu,h}are identical tomain),allreduce_fusion_opwas replaced throughLD_PRELOADby a build of this file before and after the change. The same 5 -> 64 sequence, run throughtensorrt_llm._torch.distributed.AllReduce(strategy=AllReduceStrategy.TWOSHOT)on 4x B200, 50 trials each:fusion_opNONERESIDUAL_RMS_NORM(checksresidual_out)Existing tests, with the same before / after builds:
tests/unittest/_torch/multi_gpu/test_allreduce.pyandtest_ar_residual_norm.pyon 8 GPUs: 46 passed, 10 skipped for stock, before and after. The skips are the 8test_row_linear_residual_norm_fusioncases that are skipped in-tree and the 2 cases that need exactly 4 GPUs.test_minimax_allreduce_rms_qk: 2/2 passed before and after.None of the existing tests changes the token count between two-shot calls, so they do not catch this bug.
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.@yilin-void @hyukn could you take a look? This changes the two-shot
Barrier::syncfrom #3064. cc @kaiyux since #18834 discusses the two-shot barrier and grid size.🤖 Generated with Claude Code