Multi utcmma opt

[yanqinz2]

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-* (see label list).

Affected area

Summary

Fix a series of potential issue during a throughout audit of tcgen.05/config

API and compatibility impact

Testing

Summary by CodeRabbit

• Bug Fixes
  • Improved GEMM reliability for multi-tile and block-scaled workloads.
  • Corrected handling of partial scale-factor data and tensor-memory sizing.
  • Added validation for unsupported matrix dimensions and hardware access ranges, with clearer error messages.
  • Improved synchronization during matrix multiplication output processing to prevent incomplete or inconsistent results.
  • Enhanced support for 64-row tile configurations across standard and grouped matrix multiplication workloads.
  • Improved reliability across additional matrix multiplication and mixture-of-experts workloads.

[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: ddba281f-c9d3-4ef8-9cb9-32f9ba072018

📥 Commits

Reviewing files that changed from the base of the PR and between 02ad618 and dc37844.

📒 Files selected for processing (16)

• python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py
• python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py
• python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py
• python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py
• python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py
• python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py

---

📝 Walkthrough

Walkthrough

Frost GEMM compilation now validates MMA N dimensions and block-scale TMEM spans. Multi-GEMM sizing and scale-factor atom counting are corrected. GEMM epilogues update elected-operation guards, select 64-row TMEM load shapes, and wait for loads before synchronization.

Changes

Frost GEMM runtime changes

Layer / File(s)	Summary
Compiler sizing and MMA validation python/cudnn/gemm/frost/compiler.py, python/cudnn/gemm/frost/kernel_registry.py	Multi-GEMM TMEM sizing, partial scale-factor atom counting, SFA/SFB TMEM span validation, and MMA N-dimension checks are updated. JIT compilation and template rejection invoke the dimension checks.
1-CTA barrier, TMA, and epilogue ordering python/cudnn/gemm/frost/kernel_templates/*1ctamma.py, python/cudnn/gemm/frost/kernel_templates/*1ctamma_static.py	Barrier and TMA operations use narrower elect_one guards. 64-row CTA tiles use SHAPE_16X32BX2 with zero offset. Epilogues wait for TMEM loads before fences and accumulator or deallocation signaling.
2-CTA barrier, TMA, and epilogue ordering python/cudnn/gemm/frost/kernel_templates/*2ctamma.py, python/cudnn/gemm/frost/kernel_templates/*2ctamma_static.py	Barrier and TMA control flow is reorganized around elected operations. Epilogues wait for pending TMEM loads before fences, accumulator signaling, and TMEM deallocation synchronization.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• NVIDIA/cudnn-frontend#504: This PR changes the same Frost GEMM elect_one control flow in the kernel templates.

Suggested labels: mod-cutedsl, cat-perf-bug

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 2 | ❌ 3

❌ Failed checks (2 warnings, 1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description includes a brief summary but omits the Why and Related issues sections, leaves the affected area and testing incomplete, and does not state compatibility impact.	Complete all template sections, select the affected area, state compatibility impact, and provide exact testing commands and results.
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title relates to the changes but is vague and uses unclear terminology, so it does not clearly identify the primary tcgen.05 fixes.	Use a specific title such as "Fix tcgen.05 configuration and multi-CTA MMA issues".

✅ Passed checks (2 passed)

Check name	Status	Explanation
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 PR with unit tests

---

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