Repository navigation
[CUDA] Add native SM75 MMA GEMM support for FP16, INT8 and INT4 - #2198
Conversation
|
👋 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! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds complete SM75/Turing support for TileLang CUDA MMA GEMM: routing GEMM dispatch to select Turing-specific emitter, implementing SM75-aware MMA shape policies with optimized configurations for FP16, INT8, and INT4, adding inline-PTX hardware MMA implementations, and validating end-to-end correctness with kernel and unit tests. ChangesSM75 Turing MMA Support
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (5)
src/tl_templates/cuda/instruction/mma.h (1)
165-168: 💤 Low valueInconsistent namespace qualification.
Every other
TL_DEFINE_MMA_DISPATCHERrow in this file references CuTe structs withcute::(e.g.,cute::SM80_16x8x32_S32S4S4S32_TNon lines 171/175). The two new INT4 rows omit the prefix because the structs are defined locally intl::detail. It works thanks to ADL/current-namespace lookup, but adding an explicittl::detail::(or moving the struct definitions undercute::alongside the SM75 includes if you intend them as CuTe extensions) makes the divergence obvious at the dispatch site.🤖 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/tl_templates/cuda/instruction/mma.h` around lines 165 - 168, The two TL_DEFINE_MMA_DISPATCHER entries for INT4/UInt4 (SM75_8x8x32_S32S4S4S32_TN and SM75_8x8x32_S32U4U4S32_TN) are missing the namespace qualifier used elsewhere; update those dispatcher invocations to consistently reference the structs with an explicit namespace (either prefix with tl::detail:: if the structs live in tl::detail, or move/rename the struct definitions into the cute:: namespace and use cute::) so the dispatch lines match the other entries (e.g., the same pattern used for cute::SM80_16x8x32_S32S4S4S32_TN).tilelang/cuda/intrinsics/macro/mma_macro_generator.py (2)
236-247: 💤 Low valueDocument why
local_size_c == 2reuses the FP64 store map.The added
or local_size_c == 2clause silently extends the FP64 8×8 store layout to Turing INT8/INT4 (which also produce 2 outputs per lane over an 8×8 tile). Functionally correct because the per-lane layout form8n8k16/m8n8k32matchesm8n8k4, but the condition reads as a coincidental proxy. A brief comment (or a more explicit predicate likeself.M_DIM == 8 and self.n_dim == 8) will make the intent clear to readers who only knowmma_store_index_map_fp64by name.🤖 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 `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py` around lines 236 - 247, The condition in get_store_index_map that checks "or local_size_c == 2" silently reuses mma_store_index_map_fp64 for cases where two outputs per lane occur (e.g., Turing INT8/INT4) but is unclear; update get_store_index_map to either add a brief clarifying comment explaining why local_size_c == 2 implies the FP64 8x8 store layout is valid, or replace the heuristic with an explicit predicate (for example check self.M_DIM == 8 and self.n_dim == 8 or another clear structural condition) so the intent is explicit; reference get_store_index_map, local_size_c, mma_store_index_map_fp64 and DataType(self.accum_dtype).bits when making the change.
190-202: 💤 Low valueConsider folding the Turing fp16 branch into the integer branch.
Both branches set the same
n_dim,micro_size_y,warp_rows,warp_cols. The only meaningful difference is theassert m_dim == 8in the first branch, which is invariant given the constructor already pinsM_DIMfor Turing 4-bit/8-bit dtypes. Optional consolidation:♻️ Possible consolidation
- if k_dim == 4 or (self.is_turing and k_dim in (16, 32)): - # fp64 or Turing integer path: m_dim must be 8, n_dim 8 - assert m_dim == 8, f"For fp64/Turing integer MMA, m_dim must be 8, got {m_dim}" - self.n_dim = 8 - self.micro_size_y = 8 - self.warp_rows = warp_row_tiles // m_dim - self.warp_cols = warp_col_tiles // 8 - elif k_dim == 8 and self.is_turing: - # Turing fp16 path - self.n_dim = 8 - self.micro_size_y = 8 - self.warp_rows = warp_row_tiles // m_dim - self.warp_cols = warp_col_tiles // 8 + if k_dim == 4 or (self.is_turing and k_dim in (8, 16, 32)): + # fp64 (m8n8k4) or Turing path (m16n8k8 / m8n8k16 / m8n8k32): + # all use 8-wide N tiles. + if k_dim == 4 or (self.is_turing and k_dim in (16, 32)): + assert m_dim == 8, f"For fp64/Turing integer MMA, m_dim must be 8, got {m_dim}" + self.n_dim = 8 + self.micro_size_y = 8 + self.warp_rows = warp_row_tiles // m_dim + self.warp_cols = warp_col_tiles // 8🤖 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 `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py` around lines 190 - 202, Combine the two branches that assign identical values into a single conditional to remove duplication: detect the union condition (k_dim == 4 or (self.is_turing and k_dim in (16, 32)) or (k_dim == 8 and self.is_turing)) and then set self.n_dim = 8, self.micro_size_y = 8, self.warp_rows = warp_row_tiles // m_dim, and self.warp_cols = warp_col_tiles // 8 in one place; preserve the assert m_dim == 8 only for the fp64/Turing-integer case by guarding it (e.g., if k_dim == 4 or (self.is_turing and k_dim in (16,32)): assert m_dim == 8 ...) so you keep the invariant while removing duplicated assignments involving m_dim, k_dim, self.is_turing, warp_row_tiles, and warp_col_tiles in mma_macro_generator.py.tilelang/cache/kernel_cache.py (1)
138-153: 💤 Low valueSort
dirnamesto make the codegen stamp deterministic across runs and machines.
os.walkdoesn't guarantee subdirectory traversal order—POSIXreaddirreturns entries in filesystem-dependent order, so two checkouts with identical contents can produce differenttilelang_codegendigests (cache misses across machines/CI runners, or after acp -rthat changes inode allocation). You already iteratesorted(filenames); mirror that by mutatingdirnamesin place:♻️ Proposed fix
for root, suffixes in roots: - for dirpath, _, filenames in os.walk(root): + for dirpath, dirnames, filenames in os.walk(root): + dirnames.sort() for filename in sorted(filenames): if suffixes is not None and not filename.endswith(suffixes): continue🤖 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 `@tilelang/cache/kernel_cache.py` around lines 138 - 153, The os.walk loop currently ignores and doesn't sort subdirectories (it uses "for dirpath, _, filenames in os.walk(root)"), causing non-deterministic traversal; update the loop to capture dirnames ("for dirpath, dirnames, filenames in os.walk(root)") and call dirnames.sort() immediately so subdirectory order is deterministic before iterating filenames (keep existing sorted(filenames) logic and the rest of the file hashing logic in the same block).tilelang/cuda/op/gemm/gemm_mma.py (1)
25-26: 💤 Low valueHoist
target_is_turingimport to module top-level.Other
tilelang.utils.*symbols are imported at module scope (line 8). Placing this import inside the method is inconsistent and slightly slower per invocation; move it to the top of the file unless there is a circular-import concern.♻️ Proposed refactor
from tilelang.utils.language import is_shared, is_fragment, is_full_region +from tilelang.utils.target import target_is_turing from tilelang import tvm as tvmwarp_col_tiles = int(self.N // n_warp) - from tilelang.utils.target import target_is_turing - emitter = TensorCoreIntrinEmitter(🤖 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 `@tilelang/cuda/op/gemm/gemm_mma.py` around lines 25 - 26, The import of target_is_turing is currently inside a function/method in gemm_mma.py; hoist the import to the module top alongside the other tilelang.utils imports (where other tilelang.utils.* symbols are imported) so it's imported once at import-time rather than per-call, updating the import list to include target_is_turing; if moving it causes a circular import, keep the local import but add a comment explaining the circular dependency and why it must remain local. Ensure references in any function/method that used the in-function import (e.g., the function that previously contained "from tilelang.utils.target import target_is_turing") continue to use the now-top-level symbol.
🤖 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 `@examples/gemm/example_gemm_int8_turing.py`:
- Around line 61-65: Add an explicit SM75 guard before invoking
matmul_int8_turing: call torch.cuda.get_device_capability() (or
torch.cuda.get_device_capability(0)) right after printing the device info and
verify it equals (7, 5); if not, raise a clear exception or sys.exit with a
message saying this example is Turing/SM75-only. Update the block around where
kernel = matmul_int8_turing(M, N, K) so the guard runs first and prevents
attempting to compile/run on non-SM75 GPUs.
- Around line 47-48: The test helper function ref_program uses the private
operator torch._int_mm which may break across PyTorch versions; replace it with
the public API torchao.kernel.intmm.int_matmul when available (import and call
int_matmul(A, B.t())), and add a conditional fallback in ref_program that
detects absence of torchao and either uses a safe alternative implementation or
raises a clear error message; update references to ref_program accordingly so
the code uses the public int_matmul when present and falls back gracefully if
not.
In `@src/tl_templates/cuda/instruction/mma.h`:
- Around line 57-105: The inline PTX in fma for SM75_8x8x32_S32S4S4S32_TN and
SM75_8x8x32_S32U4U4S32_TN hardcodes the saturating variant ".satfinite" but the
dispatcher/metadata declares Saturate=false; update the asm strings in these two
fma methods to the non-saturating mnemonics (remove ".satfinite" so they use
"mma.sync.aligned.m8n8k32.row.col.s32.s4.s4.s32" and
"mma.sync.aligned.m8n8k32.row.col.s32.u4.u4.s32" respectively) to match the
dispatcher, and ensure any dispatcher declarations for these types also state
Saturate=false or provide a separate _SATURATE variant; if saturation was
intentional, instead add a clear comment on the struct explaining why it
diverges from the dispatcher.
---
Nitpick comments:
In `@src/tl_templates/cuda/instruction/mma.h`:
- Around line 165-168: The two TL_DEFINE_MMA_DISPATCHER entries for INT4/UInt4
(SM75_8x8x32_S32S4S4S32_TN and SM75_8x8x32_S32U4U4S32_TN) are missing the
namespace qualifier used elsewhere; update those dispatcher invocations to
consistently reference the structs with an explicit namespace (either prefix
with tl::detail:: if the structs live in tl::detail, or move/rename the struct
definitions into the cute:: namespace and use cute::) so the dispatch lines
match the other entries (e.g., the same pattern used for
cute::SM80_16x8x32_S32S4S4S32_TN).
In `@tilelang/cache/kernel_cache.py`:
- Around line 138-153: The os.walk loop currently ignores and doesn't sort
subdirectories (it uses "for dirpath, _, filenames in os.walk(root)"), causing
non-deterministic traversal; update the loop to capture dirnames ("for dirpath,
dirnames, filenames in os.walk(root)") and call dirnames.sort() immediately so
subdirectory order is deterministic before iterating filenames (keep existing
sorted(filenames) logic and the rest of the file hashing logic in the same
block).
In `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py`:
- Around line 236-247: The condition in get_store_index_map that checks "or
local_size_c == 2" silently reuses mma_store_index_map_fp64 for cases where two
outputs per lane occur (e.g., Turing INT8/INT4) but is unclear; update
get_store_index_map to either add a brief clarifying comment explaining why
local_size_c == 2 implies the FP64 8x8 store layout is valid, or replace the
heuristic with an explicit predicate (for example check self.M_DIM == 8 and
self.n_dim == 8 or another clear structural condition) so the intent is
explicit; reference get_store_index_map, local_size_c, mma_store_index_map_fp64
and DataType(self.accum_dtype).bits when making the change.
- Around line 190-202: Combine the two branches that assign identical values
into a single conditional to remove duplication: detect the union condition
(k_dim == 4 or (self.is_turing and k_dim in (16, 32)) or (k_dim == 8 and
self.is_turing)) and then set self.n_dim = 8, self.micro_size_y = 8,
self.warp_rows = warp_row_tiles // m_dim, and self.warp_cols = warp_col_tiles //
8 in one place; preserve the assert m_dim == 8 only for the fp64/Turing-integer
case by guarding it (e.g., if k_dim == 4 or (self.is_turing and k_dim in
(16,32)): assert m_dim == 8 ...) so you keep the invariant while removing
duplicated assignments involving m_dim, k_dim, self.is_turing, warp_row_tiles,
and warp_col_tiles in mma_macro_generator.py.
In `@tilelang/cuda/op/gemm/gemm_mma.py`:
- Around line 25-26: The import of target_is_turing is currently inside a
function/method in gemm_mma.py; hoist the import to the module top alongside the
other tilelang.utils imports (where other tilelang.utils.* symbols are imported)
so it's imported once at import-time rather than per-call, updating the import
list to include target_is_turing; if moving it causes a circular import, keep
the local import but add a comment explaining the circular dependency and why it
must remain local. Ensure references in any function/method that used the
in-function import (e.g., the function that previously contained "from
tilelang.utils.target import target_is_turing") continue to use the
now-top-level symbol.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 356792ee-db78-486f-9626-956be3e7b6a1
📒 Files selected for processing (9)
CMakeLists.txtexamples/gemm/example_gemm_int8_turing.pysrc/tl_templates/cuda/instruction/mma.htesting/python/cuda/test_cuda_mma_sm75_dispatch.pytesting/python/kernel/test_tilelang_kernel_gemm_sm75.pytilelang/cache/kernel_cache.pytilelang/cuda/intrinsics/macro/mma_macro_generator.pytilelang/cuda/op/gemm/__init__.pytilelang/cuda/op/gemm/gemm_mma.py
| def ref_program(A, B): | ||
| return torch._int_mm(A, B.t()) |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Is torch._int_mm a stable public PyTorch API, and what is the recommended public alternative for INT8xINT8->INT32 GEMM on CUDA?
💡 Result:
No. torch._int_mm is an underscore-prefixed, non-public/internal PyTorch operator and is not part of PyTorch’s documented stable public API surface; PyTorch maintainers/dev discussions explicitly treat it as not a supported stable/documented API (e.g., requests/perf discussions note that “_int_mm and _scaled_mm are not part of a stable and documented API”) [1]. It can also have backend- or layout-specific behavior/performance issues (e.g., slower for row-major RHS) [2]. Recommended public alternative (public API) for INT8xINT8 -> INT32 GEMM on CUDA: - Use torchao’s integer matmul entry point instead of torch._int_mm. torchao implements “safe” integer GEMM with CUDA/cuBLAS/Triton paths and an explicit fallback, exposed via torchao.kernel.intmm.int_matmul (which returns int32 accumulations) [3]. This is a public, supported-by-project API (unlike torch._int_mm, which is internal). Concretely, the public alternative is torchao.kernel.intmm.int_matmul (or the torchao.quantization path that calls into torchao’s int-matmul). torchao’s int_matmul dispatches to a Triton-based kernel when available; otherwise it falls back to safe_int_mm, which for CUDA uses torch.ops.aten.mm.default with int32 accumulation as the fallback path [3]. If you specifically need a low-level CUDA INT8 GEMM kernel from PyTorch itself (rather than a third-party library), PyTorch does not currently expose a documented stable public INT8xINT8->INT32 GEMM API in the sources above; the internal operator torch._int_mm is the mechanism used by PyTorch/Inductor for int matmul lowering [4], but it remains not-stable/public per the discussion cited above [1].
Citations:
- 1:
_scaled_mmerrors and_int_mmis slow with row-major rhs matrix pytorch/pytorch#164491 - 2: _int_mm is slow with row-major rhs matrix pytorch/pytorch#165230
- 3: https://github.com/pytorch/ao/blob/main/torchao/kernel/intmm.py
- 4: https://github.com/pytorch/pytorch/blob/main/torch/_inductor/kernel/mm.py
🏁 Script executed:
# Check if the file exists and examine the context around lines 47-48
if [ -f "examples/gemm/example_gemm_int8_turing.py" ]; then
echo "=== File exists, checking lines around 47-48 ==="
cat -n "examples/gemm/example_gemm_int8_turing.py" | sed -n '40,55p'
else
echo "File not found, searching for it..."
find . -name "example_gemm_int8_turing.py" -type f
fiRepository: tile-ai/tilelang
Length of output: 612
Replace private torch._int_mm with a supported public API or conditional fallback.
torch._int_mm is an internal PyTorch operator that is not part of the stable public API and can change across PyTorch versions. The recommended public alternative is torchao.kernel.intmm.int_matmul (from the torchao library). If adding a dependency is not feasible, add a conditional fallback to prevent breakage:
Proposed change
def ref_program(A, B):
- return torch._int_mm(A, B.t())
+ # Prefer public torchao API; fall back if unavailable
+ try:
+ from torchao.kernel.intmm import int_matmul
+ return int_matmul(A, B.t())
+ except ImportError:
+ # Fallback: torch.matmul with int32 (may differ numerically from _int_mm)
+ return torch.matmul(A.to(torch.int32), B.t().to(torch.int32))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def ref_program(A, B): | |
| return torch._int_mm(A, B.t()) | |
| def ref_program(A, B): | |
| # Prefer public torchao API; fall back if unavailable | |
| try: | |
| from torchao.kernel.intmm import int_matmul | |
| return int_matmul(A, B.t()) | |
| except ImportError: | |
| # Fallback: torch.matmul with int32 (may differ numerically from _int_mm) | |
| return torch.matmul(A.to(torch.int32), B.t().to(torch.int32)) |
🤖 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 `@examples/gemm/example_gemm_int8_turing.py` around lines 47 - 48, The test
helper function ref_program uses the private operator torch._int_mm which may
break across PyTorch versions; replace it with the public API
torchao.kernel.intmm.int_matmul when available (import and call int_matmul(A,
B.t())), and add a conditional fallback in ref_program that detects absence of
torchao and either uses a safe alternative implementation or raises a clear
error message; update references to ref_program accordingly so the code uses the
public int_matmul when present and falls back gracefully if not.
| M, N, K = args.m, args.n, args.k | ||
| print(f"INT8 GEMM for Turing, M={M}, N={N}, K={K}") | ||
| print(f"Device: {torch.cuda.get_device_name()} capability={torch.cuda.get_device_capability()}") | ||
|
|
||
| kernel = matmul_int8_turing(M, N, K) |
There was a problem hiding this comment.
Add an explicit SM75 capability guard for clearer failures.
This example is Turing-specific, but it currently runs on any CUDA GPU. A check near Line 63 avoids confusing compile/runtime behavior on non-SM75 devices.
Proposed change
M, N, K = args.m, args.n, args.k
print(f"INT8 GEMM for Turing, M={M}, N={N}, K={K}")
- print(f"Device: {torch.cuda.get_device_name()} capability={torch.cuda.get_device_capability()}")
+ capability = torch.cuda.get_device_capability()
+ print(f"Device: {torch.cuda.get_device_name()} capability={capability}")
+ if capability != (7, 5):
+ raise RuntimeError(
+ f"This example targets Turing SM75, but got capability={capability}."
+ )🤖 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 `@examples/gemm/example_gemm_int8_turing.py` around lines 61 - 65, Add an
explicit SM75 guard before invoking matmul_int8_turing: call
torch.cuda.get_device_capability() (or torch.cuda.get_device_capability(0))
right after printing the device info and verify it equals (7, 5); if not, raise
a clear exception or sys.exit with a message saying this example is
Turing/SM75-only. Update the block around where kernel = matmul_int8_turing(M,
N, K) so the guard runs first and prevents attempting to compile/run on non-SM75
GPUs.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py`:
- Around line 236-238: stmatrix() currently hardcodes mma_store_index_map while
get_store_index_map()/make_mma_store_layout() switch to mma_store_index_map_fp64
when DataType(self.accum_dtype).bits == 64 or local_size_c == 2, causing
inconsistent lane maps; update stmatrix() to use the same selection logic as
get_store_index_map() (i.e. when DataType(self.accum_dtype).bits == 64 or
local_size_c == 2 use IndexMap.from_func(mma_store_index_map_fp64, ...)
otherwise use the existing mma_store_index_map) or simply call/get the index map
via get_store_index_map() so both layout inference and the store macro use the
identical lane map.
- Line 50: The new "uint4" dtype entry in mma_macro_generator.py is advertised
even though make_mma_load_layout() still rejects dtype_bits == 4, causing
gemm_sr/gemm_rs/gemm_rr layout inference to fail; fix by either implementing the
4-bit load transforms inside make_mma_load_layout() to accept dtype_bits == 4
(add the same load transform logic used for 8/16-bit types but adapted for 4-bit
packing) or revert/gate the "uint4" mapping so it is only used for the
shared/shared fragment-layout path (restrict the "uint4" entry or add a
conditional where the dtype-to-register-type map is consulted so that "uint4" is
not returned for fragment paths); reference make_mma_load_layout(), the "uint4"
mapping in mma_macro_generator.py, and gemm_sr/gemm_rs/gemm_rr when making the
change.
- Around line 186-194: The subclass
TensorCoreIntrinEmitterWithLadderTransform._initialize_micro_size currently has
a parameter mismatch versus the base class: change its signature to match the
base caller (accept two args: m_dim and k_dim, e.g. def
_initialize_micro_size(self, m_dim: int = 16, k_dim: int = 16)) and update any
internal uses that referenced an n_dim parameter to instead derive the column
dimension from the correct variable (either the passed k_dim or self.N_DIM as
appropriate for ladder transform logic); ensure all assertions and conditionals
(e.g. the use_m8n8/use_n8 branches and the assert on m_dim) reference the
corrected k_dim/name so fragment shape inference uses the actual k_dim passed by
the base class.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 07738dd0-16bd-499e-b2ea-e7bb6f93f81a
📒 Files selected for processing (4)
src/tl_templates/cuda/instruction/mma.htilelang/cache/kernel_cache.pytilelang/cuda/intrinsics/macro/mma_macro_generator.pytilelang/cuda/op/gemm/gemm_mma.py
🚧 Files skipped from review as they are similar to previous changes (2)
- src/tl_templates/cuda/instruction/mma.h
- tilelang/cache/kernel_cache.py
LeiWang1999
left a comment
LeiWang1999
left a comment
There was a problem hiding this comment.
we can introduce mma_sm75_macro_generator instead of is_turing?
|
Thanks for your contribution @Tokimorphling ! I left some comments. |
…tter to remove Turing-specific logic
There was a problem hiding this comment.
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 `@tilelang/cuda/intrinsics/macro/mma_sm75_macro_generator.py`:
- Around line 12-23: The _initialize_k_dim method currently allows k_dim to be
reduced below SM75 integer MMA requirements for small integer chunks; detect
when a_dtype is integer (e.g., a_dtype.bits in {4,8}) and if the computed k_dim
would be smaller than the minimum SM75 integer k (choose and document the
required minimum, e.g., 16), raise a clear ValueError (or assert) instructing
the caller to increase chunk or change dtype; locate and modify
_initialize_k_dim (referencing self.k_dim, self.chunk, and a_dtype) to perform
this validation after computing k_dim so downstream micro-tile selection logic
(the integer micro-tile path) never sees an unsupported small k_dim.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 28063ae8-87ae-4ed5-a8fd-e3114b751fb5
📒 Files selected for processing (4)
testing/python/cuda/test_cuda_mma_sm75_dispatch.pytilelang/cuda/intrinsics/macro/mma_macro_generator.pytilelang/cuda/intrinsics/macro/mma_sm75_macro_generator.pytilelang/cuda/op/gemm/gemm_mma.py
| def _initialize_k_dim(self, a_dtype=T.float16): | ||
| if isinstance(a_dtype, str): | ||
| a_dtype = DataType(a_dtype) | ||
| if a_dtype.bits == 4: | ||
| self.k_dim = min(32, self.chunk) | ||
| elif a_dtype.bits == 8: | ||
| self.k_dim = min(16, self.chunk) | ||
| elif a_dtype.bits == 16: | ||
| self.k_dim = min(8, self.chunk) | ||
| else: | ||
| self.k_dim = min(256 // a_dtype.bits, self.chunk) | ||
|
|
There was a problem hiding this comment.
Guard unsupported small-chunk integer SM75 MMA configs.
Line [16]-[19] can reduce integer k_dim below SM75 MMA requirements. For example, int8 with chunk=8 drives k_dim=8, but the integer micro-tile path (Line [50]-[58]) remains m8n8-style, while Line [32]-[33] selects m16n8k8. That mismatch can generate invalid instruction selection.
🔧 Proposed fix
def _initialize_k_dim(self, a_dtype=T.float16):
if isinstance(a_dtype, str):
a_dtype = DataType(a_dtype)
if a_dtype.bits == 4:
- self.k_dim = min(32, self.chunk)
+ if self.chunk < 32:
+ raise ValueError(f"SM75 int4 MMA requires chunk >= 32, got {self.chunk}")
+ self.k_dim = 32
elif a_dtype.bits == 8:
- self.k_dim = min(16, self.chunk)
+ if self.chunk < 16:
+ raise ValueError(f"SM75 int8 MMA requires chunk >= 16, got {self.chunk}")
+ self.k_dim = 16
elif a_dtype.bits == 16:
- self.k_dim = min(8, self.chunk)
+ if self.chunk < 8:
+ raise ValueError(f"SM75 fp16 MMA requires chunk >= 8, got {self.chunk}")
+ self.k_dim = 8
else:
self.k_dim = min(256 // a_dtype.bits, self.chunk)Also applies to: 31-39, 41-59
🤖 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 `@tilelang/cuda/intrinsics/macro/mma_sm75_macro_generator.py` around lines 12 -
23, The _initialize_k_dim method currently allows k_dim to be reduced below SM75
integer MMA requirements for small integer chunks; detect when a_dtype is
integer (e.g., a_dtype.bits in {4,8}) and if the computed k_dim would be smaller
than the minimum SM75 integer k (choose and document the required minimum, e.g.,
16), raise a clear ValueError (or assert) instructing the caller to increase
chunk or change dtype; locate and modify _initialize_k_dim (referencing
self.k_dim, self.chunk, and a_dtype) to perform this validation after computing
k_dim so downstream micro-tile selection logic (the integer micro-tile path)
never sees an unsupported small k_dim.
- gemm_sm90.h: B_type fallback incorrectly used A_type_cute instead of B_type_cute (copy-paste typo), affecting mixed-precision GEMM on SM90. - gemm.cc: Remove TargetIsTuring from kNPerWarp=16 condition. PR #2173 originally set this to 16 because SM75 used Volta-style 16x16x4 instructions (n_dim=16). After #2198, SM75 now uses native m16n8k8 instructions (n_dim=8) via GemmMMASm75, so kNPerWarp should be 8.
The SM75 GEMM emitter derives MMA shapes from input dtype bits alone, so it already emits the (16,8,8) shape for fp16->fp16 accumulation (identical to the registered fp16->fp32 case) and the (8,8,16) shape for uint8 inputs (identical to signed int8). But the C++ dispatcher table only covered the three configs tested by #2198 (fp16->fp32, int8, int4), so requesting fp16->fp16-accum or uint8 GEMM on Turing failed at nvcc with an opaque `tl::mma_sync: unsupported configuration` static_assert. Add the two missing tl::detail SM75 atoms and register their dispatchers, following the precedent set by #2198's hand-written int4/uint4 structs (CuTe ships only fp16->fp32 and signed-int8 SM75 atoms). The uint8 atom uses the non-saturating PTX form to match the existing signed-int8 path. Verified on a TITAN RTX (SM75): emitted PTX is mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 and mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32, and end-to-end GEMM numerics match torch.
2 participants
Summary
This PR adds native SM75/Turing MMA GEMM support for FP16, INT8, and INT4.
It addresses the same SM75 instruction-selection issue reported in #1529 and #1992, where
T.gemm()could lower to SM80-only MMA instructions on Turing GPUs. In addition, this PR adds native INT8 and INT4 SM75 MMA paths.Changes include:
m16n8k8dispatch.m8n8k16dispatch.m8n8k32dispatch.Related to #1529 and #1992.
Validation
Validated on RTX 2080 Ti / SM75.
INT8
Before
192a118967a376ee58eb8db6bdf07fccb329fb61, SM75 INT8 GEMM is lowered through the SM70 MMA path and fails to compile:After
192a118967a376ee58eb8db6bdf07fccb329fb61:INT4
Before
a6769c3dc5267f8fbf6139cb2da551d009904964, SM75 INT4 GEMM is lowered to the SM80-onlym16n8k64path:After
a6769c3dc5267f8fbf6139cb2da551d009904964:Tests
PYTHONPATH=/home/asuka/codes/tilelang \ LD_LIBRARY_PATH=/home/asuka/codes/tilelang/build/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} \ TL_TEMPLATE_PATH=/home/asuka/codes/tilelang/src \ python -m pytest testing/python/kernel/test_tilelang_kernel_gemm_sm75.py -qSummary by CodeRabbit
New Features
Tests