[BugFix][CUDA] Fall back for non-MMA GEMM K shapes

[JayceSu98]

Fixes #2715

Summary

CUDA GEMM lowering selected the shared-memory layout and the MMA emitter independently. For K=20 and K=28, the swizzled layout fails its stride check; bypassing that check only moves the failure to the MMA atom divisibility check.

The layout and emitter now use the same MMA-shape check. Supported shapes keep the existing MMA path; other valid K extents use the padded shared-memory layout and the existing CUDA FMA emitter. Both choices are made in TileLang's GEMM lowering, so callers do not need to pad the mathematical K extent.

Changes

• Add a shared GEMM-shape eligibility check before selecting tensor-core MMA.
• Route non-MMA K extents through the existing CUDA FMA emitter.
• Use the padded shared-memory layout when the bank-swizzle stride contract is not satisfied.
• Add numerical CUDA regressions for K=16, K=20, and K=28.

Review Notes

• K=16 remains on the existing optimized path.
• K=20 and K=28 use the fallback path.
• The change does not pad or reinterpret the mathematical K extent.

Validation

• testing/python/issue/test_tilelang_issue_gemm_non_mma_k_fallback.py: 3 passed on NVIDIA A100.
• testing/python/issue/test_tilelang_issue_gemm_non_mma_k_fallback.py: 3 passed on NVIDIA H100.
• Numerical output compared against PyTorch GEMM for all three K shapes.
• clang-format and git diff --check passed.

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

📝 Walkthrough

Walkthrough

GEMM layout selection now falls back to padded layouts for unsupported strides, while instruction selection falls back from MMA to CUDA FMA for unsupported K shapes. A CUDA regression test covers FP16 GEMM cases with K values 16, 20, and 28.

Changes

GEMM shape fallback

Layer / File(s)	Summary
GEMM dispatch and layout fallback src/cuda/op/gemm.cc, src/layout/gemm_layouts.cc	MMA eligibility validates operand widths and K divisibility, unsupported shapes use CUDA FMA, and invalid swizzle strides use the padded GEMM layout.
Regression validation testing/python/issue/test_tilelang_issue_2715.py	A CUDA-parameterized FP16 GEMM test validates numerical results for K dimensions 16, 20, and 28.

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

Possibly related PRs

• tile-ai/tilelang#2339: Updates CUDA GEMM dispatch for unsupported shapes.
• tile-ai/tilelang#2390: Addresses CUDA GEMM handling of non-divisible K dimensions.
• tile-ai/tilelang#2811: Adds related MMA eligibility gating in src/cuda/op/gemm.cc.

Suggested reviewers: leiwang1999, yongqi-zhuo

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes address issue #2715 by selecting padded layouts for invalid strides and avoiding MMA selection for unsupported K shapes.
Out of Scope Changes check	✅ Passed	All changes, including MMA eligibility checks and regression tests, directly support the linked issue and stated PR objectives.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the CUDA fallback added for GEMM K shapes that are incompatible with MMA.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@testing/python/issue/test_tilelang_issue_2715.py`:
- Around line 29-40: Extend test_gemm_falls_back_for_non_swizzle_stride to
inspect the lowered CUDA source or instruction output for each parameterized k
value. Assert that k=16 uses the optimized CUDA MMA/GEMM instruction path, while
k=20 and k=28 use the scalar FMA fallback, preserving the existing numerical
correctness assertion.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 4a6753e2-f5ad-458d-a650-2e23913931eb

📥 Commits

Reviewing files that changed from the base of the PR and between e01c498 and a2949c0.

📒 Files selected for processing (3)

• src/cuda/op/gemm.cc
• src/layout/gemm_layouts.cc
• testing/python/issue/test_tilelang_issue_2715.py