Fix to #652: SM100 Grouped GEMM: read the TMEM accumulator before releasing the overlapping stage by GarlGuo · Pull Request #654 · NVIDIA/cudnn-frontend

GarlGuo · 2026-08-18T22:06:12Z

Before submitting

I agree to license this contribution under the terms of LICENSE.txt.
I ran pre-commit run and committed any formatting changes.
I added GitHub labels: one cat-*, one or more mod-*, and one orig-*.

Affected area

FE OSS kernels or CuTeDSL

Summary

Moves the TMEM→register cute.copy above the fence/release block in the epilogue of three
SM100 grouped-GEMM kernels — dsrelu, srelu, and quant.

Related issues

Fixes #652

Not tested:

Runtime evidence covers the dSReLU kernel only; srelu and quant carry the identical
source pattern but were not executed.
The Rubin-family quant_rubin kernel has the same defect and is deliberately not touched as no sm_107 hardware was available to validate it on.

Summary by CodeRabbit

Bug Fixes
- Improved grouped GEMM quantization and activation processing by preserving accumulator data before pipeline release.
- Helps ensure reliable results when overlapping accumulator operations are used.

…ing stage The SM100 dSReLU/SReLU/quant grouped-GEMM epilogues call acc_pipeline.consumer_release() -- publishing the overlapping accumulator stage as empty -- before issuing the TMEM->register cute.copy of a subtile that still lies inside the overlap region, so the MMA producer waiting on that mbarrier may overwrite accumulator columns the epilogue has not read yet. cute.arch.fence_view_async_tmem_load() cannot cover that load: both of its candidate lowerings order only prior operations (PTX ISA 9.3 sections 9.7.17.8.5 and 9.7.17.11.1), and the load has not been issued when the fence executes. Section 9.7.17.6.4.4 specifies tcgen05.ld -> tcgen05.wait::ld -> fence -> mbarrier.arrive for exactly this producer/consumer pair. Eight sibling kernels in the same package already order it that way; wgrad is the same idiom with the same pipeline object. Moves the TMEM read above the fence/release in the three SM100 kernels. Pure statement reorder: the release condition is unchanged, and the loop body is the same multiset of statements before and after. Fixes NVIDIA#652

coderabbitai · 2026-08-18T22:06:35Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4404a5da-0b3f-49ae-b3db-08bab2f0ac92

📥 Commits

Reviewing files that changed from the base of the PR and between 041ac2e and 91ced06.

📒 Files selected for processing (3)

python/cudnn/gemm/cutedsl/grouped/dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py
python/cudnn/gemm/cutedsl/grouped/quant/grouped_gemm_quant.py
python/cudnn/gemm/cutedsl/grouped/srelu/moe_blockscaled_grouped_gemm_srelu_quant.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Walkthrough

Walkthrough

The three SM100 grouped GEMM epilogues now read each TMEM accumulator subtile before releasing the overlapping accumulator stage and advancing pipeline state.

Changes

Grouped GEMM epilogue synchronization

Layer / File(s)	Summary
TMEM load before accumulator release `python/cudnn/gemm/cutedsl/grouped/dsrelu/...`, `python/cudnn/gemm/cutedsl/grouped/quant/grouped_gemm_quant.py`, `python/cudnn/gemm/cutedsl/grouped/srelu/...`	The epilogues copy the accumulator subtile before the overlapping-accumulator fence, conditional consumer release, and pipeline-state advance.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 91ced

The change reorders the accumulator read before stage release in three grouped-GEMM kernels. No actionable merge-blocking risk remains; the unexecuted SReLU and quant paths should receive normal owner follow-up.

Suggested labels: orig-nv-eng, cat-bugfix

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the fix for issue `#652` and the required TMEM accumulator read ordering.
Description check	✅ Passed	The description covers the affected area, summary, linked issue, and testing limitations, but omits explicit Why, API impact, and Testing sections.
Linked Issues check	✅ Passed	The three file changes implement issue `#652` by moving each TMEM-to-register copy before the accumulator-stage release logic.
Out of Scope Changes check	✅ Passed	All changes are limited to the three SM100 grouped-GEMM epilogues identified by issue `#652`, with no unrelated modifications.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

GarlGuo · 2026-08-18T22:12:43Z

Rubin kernel has the same pattern

cudnn-frontend/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py

Line 2265 in 041ac2e

tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]

but unfortunately we don't have a sm107-compatible device to verify the fix.

Anerudhan · 2026-08-19T19:39:13Z

Thanks @GarlGuo for the PR.
We will review this and update the repo with this PR.

GarlGuo · 2026-08-20T00:47:38Z

@saltyminty @Anerudhan can we merge the PR if it looks good? Megatron (the upstream) is affected by this bug and we need to use Megatron for training soon.

…the overlapping stage (#673) Apply the same fix as #654 to the Rubin (SM107) kernel moe_blockscaled_grouped_gemm_quant_rubin.py, which has the identical pattern: the overlapping accumulator stage was released before the TMEM->register copy of that stage's subtile, allowing the MMA warp to overwrite the accumulator while the epilogue was still reading it. Move the t2r copy ahead of the early-release block so the accumulator is read before consumer_release. The other Rubin kernels (glu, glu_hadamard_quant, dglu) already perform the copy before the release, and wgrad_rubin has overlapping_accum disabled, so only this file is affected. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

…ing stage (NVIDIA#654) The SM100 dSReLU/SReLU/quant grouped-GEMM epilogues call acc_pipeline.consumer_release() -- publishing the overlapping accumulator stage as empty -- before issuing the TMEM->register cute.copy of a subtile that still lies inside the overlap region, so the MMA producer waiting on that mbarrier may overwrite accumulator columns the epilogue has not read yet. cute.arch.fence_view_async_tmem_load() cannot cover that load: both of its candidate lowerings order only prior operations (PTX ISA 9.3 sections 9.7.17.8.5 and 9.7.17.11.1), and the load has not been issued when the fence executes. Section 9.7.17.6.4.4 specifies tcgen05.ld -> tcgen05.wait::ld -> fence -> mbarrier.arrive for exactly this producer/consumer pair. Eight sibling kernels in the same package already order it that way; wgrad is the same idiom with the same pipeline object. Moves the TMEM read above the fence/release in the three SM100 kernels. Pure statement reorder: the release condition is unchanged, and the loop body is the same multiset of statements before and after. Fixes NVIDIA#652 Co-authored-by: Wentao Guo <wg0420@princeton.edu>

coderabbitai Bot mentioned this pull request Aug 19, 2026

Rubin perf investigate #667

Merged

3 tasks

Anerudhan requested a review from yanqinz2 August 19, 2026 19:35

Anerudhan assigned GarlGuo Aug 19, 2026

Anerudhan added orig-external Reported or requested by an external user, customer, or community contributor. mod-cutedsl CuTeDSL kernels, generated kernels, examples, or related integration work. cat-bugfix labels Aug 19, 2026

Anerudhan added this to the Frontend 1.28.0 milestone Aug 19, 2026

saltyminty approved these changes Aug 19, 2026

View reviewed changes

Anerudhan self-requested a review August 19, 2026 21:19

Anerudhan mentioned this pull request Aug 19, 2026

Fix SM107 (Rubin) grouped GEMM quant: read TMEM accumulator before releasing the overlapping stage #673

Merged

Anerudhan approved these changes Aug 20, 2026

View reviewed changes

Anerudhan merged commit 287bda2 into NVIDIA:develop Aug 20, 2026
1 check passed

Anerudhan mentioned this pull request Sep 2, 2026

1.28.0 rc #851

Merged

Navigation Menu

Uh oh!

Fix to #652: SM100 Grouped GEMM: read the TMEM accumulator before releasing the overlapping stage - #654

Fix to #652: SM100 Grouped GEMM: read the TMEM accumulator before releasing the overlapping stage#654
Anerudhan merged 1 commit into
NVIDIA:developfrom
GarlGuo:pr-652

GarlGuo commented Aug 18, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Aug 18, 2026 •

edited

Loading

Walkthrough

Changes

Uh oh!

GarlGuo commented Aug 18, 2026

Uh oh!

Anerudhan commented Aug 19, 2026

Uh oh!

GarlGuo commented Aug 20, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

GarlGuo commented Aug 18, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Before submitting

Affected area

Summary

Related issues

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

GarlGuo commented Aug 18, 2026

Uh oh!

Anerudhan commented Aug 19, 2026

Uh oh!

GarlGuo commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

GarlGuo commented Aug 18, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 18, 2026 •

edited

Loading

GarlGuo commented Aug 20, 2026 •

edited

Loading