[CUDA] Extend the GEMM FMA fallback to SM75 - #2811
Merged
Merged
Conversation
On SM75, SelectInst routed GEMMs to the plain MMA path even when their dtype combination had no native Turing mma.sync atom. In particular, fp32, fp64, and bf16 selected SM80-only atoms: bf16 failed an nvcc static_assert, while fp32 and fp64 compiled into runtime traps that invalidated the CUDA context. Add AllowTuringMma alongside the Volta check introduced in tile-ai#2339 and select the existing cuda.fma implementation when SM75 has no native atom. Keep the check dtype-only because the SM75 emitter supports the same operand scopes and transposes as the generic MMA path. Valid f16 combinations remain on native MMA rather than being demoted to the slower and less accurate serial FMA fallback. Add compile-only dispatch checks pinned to an explicit sm_75 target, so they run on any CUDA runner, verifying that fp32, fp64, and bf16 use the fallback while f16 remains on MMA. Add execution checks gated on compute capability 7.5.
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
Contributor
📝 WalkthroughWalkthroughTuring GEMM instruction selection now validates datatype compatibility before using ChangesTuring GEMM Fallback
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant GEMMKernel
participant GemmSelectInst
participant TuringEligibility
participant CUDAImplementation
GEMMKernel->>GemmSelectInst: select GEMM instruction for sm_75
GemmSelectInst->>TuringEligibility: check operand and accumulator dtypes
TuringEligibility-->>GemmSelectInst: eligible or unsupported
GemmSelectInst->>CUDAImplementation: use mma.sync or cuda.fma
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/cuda/op/gemm.cc (1)
132-142: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a descriptive dtype variable name.
Proposed refactor
- DataType a = op.a_->dtype; - if (a != op.b_->dtype) { + DataType a_dtype = op.a_->dtype; + if (a_dtype != op.b_->dtype) { return false; } - if (a == DataType::Float(16)) { + if (a_dtype == DataType::Float(16)) { return op.c_->dtype == DataType::Float(16) || op.c_->dtype == DataType::Float(32); } - if ((a.is_int() || a.is_uint()) && (a.bits() == 8 || a.bits() == 4)) { + if ((a_dtype.is_int() || a_dtype.is_uint()) && + (a_dtype.bits() == 8 || a_dtype.bits() == 4)) {As per path instructions, “Parameters and local variables should use descriptive lower_snake names; avoid ambiguous
Tfor API parameters.”🤖 Prompt for 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. In `@src/cuda/op/gemm.cc` around lines 132 - 142, Rename the local dtype variable `a` in the GEMM dtype validation logic to a descriptive lower_snake_case name, such as one representing the input dtype, and update all comparisons and checks in that block to use the new name.Source: Path instructions
🤖 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.
Nitpick comments:
In `@src/cuda/op/gemm.cc`:
- Around line 132-142: Rename the local dtype variable `a` in the GEMM dtype
validation logic to a descriptive lower_snake_case name, such as one
representing the input dtype, and update all comparisons and checks in that
block to use the new name.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 95b20a2e-8464-466a-9186-45f742bc85f0
📒 Files selected for processing (3)
src/cuda/op/gemm.cctesting/python/cuda/test_cuda_mma_sm75_dispatch.pytesting/python/kernel/test_tilelang_kernel_sm75_gemm_fma.py
cklxx
approved these changes
Jul 30, 2026
Contributor
|
good job |
Contributor
Author
thanks! |
LeiWang1999
approved these changes
Jul 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On SM75 (Turing), a
T.gemmwhose dtype combination has no native Turingmma.syncatom crashes instead of computing. FP32 and FP64 compile but trap when the kernel launches, leaving the CUDA context in an error state:BF16 fails during nvcc compilation:
The FP32, FP64, and BF16 cases in
test_tilelang_kernel_gemm.pyall hit these failures on an SM75 GPU. The device-side asserts from FP32 and FP64 also cause later tests in the same process to fail.Root Cause
On Turing,
SelectInstrouted GEMMs to the plain MMA path even when their dtype combination had no native SM75 atom. FP32, FP64, and BF16 therefore selected SM80-only atoms.#2339 added the
cuda.fmacorrectness fallback for GEMMs without a usable tensor-core MMA path, but only connected it to Volta.Fix
Add
AllowTuringMmaalongside the Volta check and select the existingcuda.fmaimplementation when SM75 has no native atom for the requested dtype combination.Unlike the Volta check, the Turing check is intentionally dtype-only. The SM75 emitter supports the same operand scopes and transpose combinations as the generic MMA path, so scope-based routing would incorrectly demote valid FP16
sr/rscases to the fallback. Besides being slower, the fallback accumulates serially, and FP16 accumulation loses accuracy over long K.Tested
test_tilelang_kernel_sm75_gemm_fma.pyfollowing the SM70 pattern from [TileOP] Add SM70 GEMM FMA fallback #2339. Compile-only checks pinned to an explicitsm_75target, so they run on any CUDA runner, assert that FP32, FP64, and BF16 lower through the fallback while FP16 stays on MMA. Execution checks gated on compute capability 7.5 compare the results against PyTorch.SelectInstchange makes the fallback tests fail with the errors above, while the FP16 test still passes.sm_75, Turing) with CUDA 12.4:test_tilelang_kernel_gemm.pyandtest_tilelang_kernel_gemm_sm75.pypass, with FP16, INT8, and INT4, including thesr/rsfragment cases, keeping their native MMA lowering.Summary
cuda.fmafallback for unsupported FP32, FP64, and BF16 combinations.C++ style / lint notes
docs/developer_guide/cpp_style.mdor public C++ APIs.