[CUDA] Support SM120 block-scaled GEMM fragments and odd warp atom grids

[sepcnt]

Fixes #3256.
Support fragment A and row-major scale fragments in SM120 block-scaled GEMM, and fix compact scales for odd per-warp atom grids.

Summary

• Added SM120 block-scaled GEMM support for fragment-resident A operands with shared-memory B operands.
• Added fragment-scale loading for row-major scale layouts.
• Added support for odd per-warp MMA atom grids, including the 8x1 partition, and updated compact-scale row selection.
• Added regression coverage for fragment operands and scales, transposed A, compact layouts, and multiple warp policies.

C++ style / lint notes

• The PR changes C++ code but does not change docs/developer_guide/cpp_style.md.
• CI includes the C++ API Style Audit (warning only) step.
• No audit findings were supplied. Warning-only suggestions are separate from correctness, build, and test issues.

Testing

• Test execution results were not supplied.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

Note

Currently processing new changes in this PR. This may take a few minutes, please wait...

⚙️ Run configuration

Configuration used: Repository: tile-ai/tilelang/.coderabbit.yaml

Review profile: CHILL

Plan: Advanced

Run ID: da8bdfc1-f571-478e-b9cc-49d06e92b4b2

📥 Commits

Reviewing files that changed from the base of the PR and between 8d984b9 and 2dc510b.

📒 Files selected for processing (1)

• src/cuda/op/gemm_blockscaled.cc

 ________________________________________________________________
< Fully armed and operationally intelligent code reviewer bunny. >
 ----------------------------------------------------------------
  \
   \   (\__/)
       (•ㅅ•)
       / 　 づ

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: tile-ai/tilelang/.coderabbit.yaml

Review profile: CHILL

Plan: Advanced

Run ID: 6c682149-332b-4e8f-882d-97ecde06f300

📥 Commits

Reviewing files that changed from the base of the PR and between 7d73309 and 8d984b9.

📒 Files selected for processing (1)

• testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py

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

---

📝 Walkthrough

Walkthrough

Changes

SM120 block-scaled GEMM support

Layer / File(s)	Summary
Operand contract and fragment-A lowering src/cuda/op/gemm_blockscaled.cc, tilelang/cuda/op/gemm/gemm_mma_sm120.py, testing/python/language/test_tilelang_language_gemm_blockscaled.py	SM120 selection and lowering now accept fragment A with shared B. Fragment A requires a full region and row-major scales.
Fragment layouts and odd warp grids tilelang/cuda/intrinsics/macro/mma_sm120_macro_generator.py	The emitter adds fragment and scale load layouts, fragment-A addressing, and selector offset handling for odd warp atom grids.
NVFP4 CUDA coverage testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py	Tests cover fragment operands and scales, transposed A, warp policies, odd grids, compact layouts, and invalid compact fragment configurations.

Priority: ➖ Normal

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant NVFP4Test
  participant SelectBlockScaledGemmInst
  participant GemmMMASm120BlockScaled
  participant TensorCoreIntrinEmitterSM120
  NVFP4Test->>SelectBlockScaledGemmInst: request SM120 blockscaled GEMM
  SelectBlockScaledGemmInst->>GemmMMASm120BlockScaled: accept fragment A and shared B
  GemmMMASm120BlockScaled->>TensorCoreIntrinEmitterSM120: infer fragment and scale layouts
  TensorCoreIntrinEmitterSM120->>GemmMMASm120BlockScaled: provide load layouts and fragment stride
  GemmMMASm120BlockScaled->>NVFP4Test: lower and execute GEMM

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 11.76% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 34 functions across 6 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely summarizes the main changes: SM120 block-scaled GEMM fragment support and odd warp atom grid support.
Linked Issues check	✅ Passed	The changes meet the coding requirements in #3256. SM120 instruction selection accepts fragment A with shared B and fragment C. GemmMMASm120BlockScaled supports fragment A and row-major fragment sca…
Out of Scope Changes check	✅ Passed	The changed files directly implement #3256. Operand validation, scale layouts, SM120 instruction selection, lowering, and emitter changes support register-resident A and scales, row-major scales, and …

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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