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

[Anerudhan]

Summary

Follow-up to #654, which fixed the SM100 grouped GEMM kernels to read the TMEM accumulator before releasing the overlapping accumulator stage. As noted in #654 (comment), the Rubin (SM107) kernel moe_blockscaled_grouped_gemm_quant_rubin.py has the identical pattern.

With overlapping_accum enabled, the epilogue released the accumulator stage (acc_pipeline.consumer_release) before performing the TMEM→register copy of that stage's subtile. Once released, the MMA warp can start overwriting the accumulator while the epilogue is still reading it, corrupting results. This PR applies the same reorder as #654: move the t2r copy ahead of the fence + early-release block.

Scope

Checked all Rubin kernels containing the iter_acc_early_release_in_epilogue early-release pattern:

Kernel	Status
grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py	Fixed here (release preceded copy)
grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py	Already correct (copy before release)
grouped/glu_hadamard_quant/moe_blockscaled_grouped_gemm_glu_hadamard_quant_rubin.py	Already correct (copy before release)
grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py	Already correct (copy before release)
grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py	Not affected (overlapping_accum = False)

Summary by CodeRabbit

• Bug Fixes
  • Improved the reliability of grouped quantized matrix multiplication by ensuring accumulated results are captured before pipeline resources are released.

[coderabbitai]

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: 0e8ab7b1-1788-4ed7-9ec7-b60baa2e770a

📥 Commits

Reviewing files that changed from the base of the PR and between 3041838 and d14c960.

📒 Files selected for processing (1)

• python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py

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

---

📝 Walkthrough

Walkthrough

The quantized Rubin grouped GEMM epilogue now copies the TMEM accumulator fragment before the overlapping-accumulator fence and pipeline-stage release.

Changes

Quantized Rubin GEMM epilogue

Layer / File(s)	Summary
Accumulator copy ordering python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py	Moves the TMEM accumulator fragment copy before the overlapping-accumulator fence and early pipeline-stage release.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Merge Risk: ⚪ Minimal · up to d14c9

This is a localized ordering fix in the Rubin grouped GEMM path, and no actionable merge-blocking risk remains beyond normal checks and review.

Possibly related PRs

• NVIDIA/cudnn-frontend#667: Related epilogue synchronization changes move accumulator loads before release and fence operations.
• NVIDIA/cudnn-frontend#654: Related quant_rubin grouped GEMM changes apply the same accumulator-read ordering.

Suggested reviewers: sraman-rgb

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description explains the bug and scope, but it omits several template sections, including affected area, testing, compatibility impact, and checklist status.	Add the required template sections and checklist results, including affected area, API and compatibility impact, exact testing commands and results, and related issues.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title clearly identifies the Rubin grouped GEMM quant race fix and the required read-before-release ordering.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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