Repository navigation
[https://nvbugs/6892151][fix] Add __syncthreads() after the shared-memory init in both kernels; make the… - #20104
trtllm-agent wants to merge 2 commits into
Conversation
…e and one-warp large-batch path GroupRMSNormBaseKernel and GroupRMSNormKernelLargeBatch initialize shared memory in some warps while lane 0 of other warps writes partial sums to the same slots, with no barrier in between, so warp contributions are lost and outputs are inflated. Add a __syncthreads() after the initialization; the sm < 900 __syncwarp() that only ordered lanes within a warp is now redundant. When every input fits in one warp (<= 256 fp16/bf16 or <= 128 fp32 elements), the large batch kernel launches a single warp, so the warp 1 branch that converts input 1's sum of squares into its rsqrt never runs and input 1 is scaled by its raw sum of squares. Let warp 0 initialize every row and also reduce input 1 in that case, leaving the launch configuration and the kernel-selection heuristic unchanged. Add one-warp large-batch shapes and a repeated-call test for all three ops. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
|
Important Review skippedReview was skipped as selected files did not have any reviewable changes.
 |
| Layer / File(s) | Summary |
|---|---|
Base kernel shared-memory synchronization cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu |
The base kernel adds a block-wide barrier after shared-memory initialization and removes a conditional warp synchronization before shared-memory writes. |
Large-batch initialization, reduction, and tests cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu, tests/unittest/_torch/modules/test_group_rmn_norm.py |
The large-batch kernel initializes every shared-memory row and handles input 1 reduction when only one warp is launched. Tests add the [256, 256] hidden-dimension case and compare operator outputs with reference outputs over 100 calls for each of four configurations. |
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Bug fix
Suggested reviewers: bowenfu
Merge Risk: 
ff681
-
The targeted test covers the reported one-warp regression, and no actionable merge-blocking risk remains after normal checks.
Pre-merge checks | 
-
❌ Failed checks (1 warning)
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. Write docstrings for the functions missing them to satisfy the coverage threshold. ✅ Passed checks (4 passed)
-
Check name Status Explanation Title check The title identifies the NVBugs issue, fix type, and primary shared-memory synchronization change. It is related and concise, although the displayed ending is truncated. Description check The description explains the root causes, fixes, test coverage, reproduction status, and related bug. It does not use the template's exact "Description" and "Test Coverage" headings and omits the chec… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
ℹ️ Recent review info
-
⚙️ Run configuration
-
- Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
0b2caa1d-3ef7-42bb-b89b-4aa6469c9c3f
📒 Files selected for processing (2)
-
cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cutests/unittest/_torch/modules/test_group_rmn_norm.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
ℹ️ Recent review info
-
⚙️ Run configuration
-
- Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
0b2caa1d-3ef7-42bb-b89b-4aa6469c9c3f
📒 Files selected for processing (2)
-
cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cutests/unittest/_torch/modules/test_group_rmn_norm.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
-
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Stamp on behalf of runtime devs, delegating proper review to @NVIDIA/trt-llm-kernels-devs; please ping me if you think this is not accurate
|
Bug engineer ownership @karljang, as the engineer assigned to the linked bug, you are responsible for driving this bot PR to resolution. Please coordinate review, address feedback and required code changes, get CI passing, and verify the original failure before driving it to merge. If the PR is obsolete or superseded, close it with a reason. Keep this PR updated with blockers and the next step. |
|
/bot run --disable-fail-fast |
|
PR_Github #80196 [ run ] triggered by Bot. Commit: |
|
PR_Github #80196 [ run ] completed with state
|
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
|
/bot run |
|
PR_Github #81930 [ run ] triggered by Bot. Commit: |
|
PR_Github #81930 [ run ] completed with state
|
|
/bot run |
|
PR_Github #82448 [ run ] triggered by Bot. Commit: |
|
PR_Github #82448 [ run ] completed with state |
|
/bot run |
Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com>
|
/bot run |
|
PR_Github #82766 [ run ] triggered by Bot. Commit: |
|
PR_Github #82770 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #82766 [ run ] completed with state |
|
PR_Github #82771 [ run ] triggered by Bot. Commit: |
|
PR_Github #82770 [ reuse-pipeline ] completed with state |
|
/bot run |
|
PR_Github #82795 [ run ] triggered by Bot. Commit: |
|
PR_Github #82771 [ run ] completed with state |
|
PR_Github #82795 [ run ] completed with state
|
|
/bot run |
|
PR_Github #84430 [ run ] triggered by Bot. Commit: |
Summary
warp_idx==1step never runs, so input 1 is never normalised (each output is x·sum(x²)·w).__syncthreads()after the shared-memory init in both kernels; make the large-batch init loop step byblockDim.yso one warp clears both rows; whenblockDim.y==1, warp 0 also reduces input 1 and writes its rsqrt. Launch configuration, heuristic and kernel selection are unchanged. The tests add the[256,256]one-warp shape and a 100-call repeated-calls test.bash .repair-bot/run_repro.shTest plan
Links
Reproduction comparison
Not run: the NVBug report states the failure ToT reproduced (=== REP_END )
Signature: === REP_END
Dev Engineer Review
The kernels add a block-wide barrier after shared-memory initialization. The large-batch kernel now initializes all shared-memory rows and normalizes input 1 when it launches one warp. The reported launch configuration, heuristic, and kernel selection remain unchanged.
QA Engineer Review
tests/unittest/_torch/modules/test_group_rmn_norm.pyadds the[256, 256]large-batch case andtest_group_rms_norm_repeated_calls. The repeated-call test exercises the base, large-batch, and heuristic operators 100 times across four configurations and compares results with an RMSNorm reference. This unit-test file does not require an integration test-list entry. Coverage verdict: sufficient for the reported defects. Test execution results were not supplied.Per-File QA Perspective
cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu: Verify that the new synchronization prevents shared-memory initialization races and that the single-warp path normalizes input 1.tests/unittest/_torch/modules/test_group_rmn_norm.py: Covers the single-warp shape and repeated calls across three operators. It is a unit-test file, so integration-list membership does not apply.