[CUDA] Reject conflicting SM120 scale fragment layouts

[LeiWang1999]

Summary

Passing the same fragment as SFA and SFB to an SM120 block-scaled GEMM could silently overwrite one operand's layout requirement and produce incorrect scale reads. Layout inference now rejects incompatible requirements with guidance to use separate scale fragments or shared-memory scales.

Fixes #3283.

Changes

• Compare the full fragment layouts before assigning a layout to an existing buffer entry.
• Preserve reuse when the layout requirements agree; report a clear error when they conflict.
• Add 18 CUDA lowering regression cases covering all three warp policies, shared/fragment A, aliased scale fragments, separate scale fragments, and aliased shared-memory scales.

Validation

• cmake --build build -j8
• ./format.sh
• python -m pytest testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py testing/python/language/test_tilelang_language_gemm_blockscaled.py -q --tb=short — 138 passed, 21 skipped.
• Before the fix, all six conflicting-fragment cases failed to raise; after the fix, all 18 new cases passed.

The new tests lower explicitly to sm_120a without device compilation. Numerical execution on SM120 hardware was not tested.

Summary

• infer_layout() now checks scale-load layouts when SFA and SFB use the same fragment buffer. It raises a ValueError when the layouts differ and directs callers to use separate fragment buffers or shared-memory scales.
• Added SM120 lowering tests across warp policies and A/scale allocation variants. Aliased fragment scales are expected to fail; other combinations are expected to lower and emit the block-scaled MMA call.

Validation

The supplied summary does not report test or build results for this change. Numerical execution on SM120 hardware is not reported.

[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.

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: 91bc3aeb-2880-4655-8093-8c781c4a1f32

📥 Commits

Reviewing files that changed from the base of the PR and between 7a5f446 and 9904e7d.

📒 Files selected for processing (2)

• testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py
• tilelang/cuda/op/gemm/gemm_mma_sm120.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

SM120 block-scaled GEMM layout inference now detects incompatible layouts when SFA and SFB share a fragment buffer. A parameterized lowering test covers aliasing, allocation choices, and warp policies.

Changes

SM120 scale-buffer alias validation

Layer / File(s)	Summary
Scale layout check and regression coverage tilelang/cuda/op/gemm/gemm_mma_sm120.py, testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py	infer_layout raises a layout-mismatch ValueError when aliased fragment scales require different layouts. The test checks this case and verifies that the other parameter combinations lower and emit the block-scaled MMA call.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Medium

Suggested reviewers: sepcnt

Merge Risk: ⚪ Minimal · up to 9904e

No merge-blocking issue is established; the scale-layout conflict check can proceed through normal validation.

Security Architecture Review

Security architecture risk: 🔵 Low · up to 9904e

The change prevents incompatible scale-fragment layouts from being accepted. It does not appear to add an entrypoint or change deployment behavior, although numerical execution on SM120 hardware was not verified.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The directly evidenced effect is rejection of conflicting fragment-scale aliases during SM120 GEMM lowering; the test change does not itself create a production entrypoint.

Trust Boundaries and Controls

• observed — The added control compares a buffer's existing fragment layout with the operand-specific requirement and rejects disagreement before reassignment.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.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
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: rejecting conflicting SM120 scale fragment layouts during layout inference.
Linked Issues check	✅ Passed	Issue #3283 requires SM120 lowering to reject conflicting fragment-layout requirements when SFA and SFB alias, and requires lowering-level regression coverage. At the reviewed head, `GemmMMASm120Block…
Out of Scope Changes check	✅ Passed	The reviewed changes are limited to SM120 block-scaled GEMM layout validation and regression tests for scale-buffer aliasing. These changes directly support issue #3283. No unrelated product behavior …

• 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.}