Add SM120 dense block-scaled MMA and essential support for SageAttention3#2253
Add SM120 dense block-scaled MMA and essential support for SageAttention3#2253sepcnt wants to merge 6 commits into
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:
📝 WalkthroughWalkthroughAdds a TileLang intrinsic ChangesBlackwell SM120 Intrinsics
Sequence DiagramsequenceDiagram
participant User as User Code
participant TLFrontend as TileLang Frontend
participant TLCodegen as CUDA Codegen
participant PTXUtil as PTX Utilities
participant DeviceTemplate as Device Template
User->>TLFrontend: call ptx_mma_blockscale(...)
TLFrontend->>TLCodegen: call_intrin tl.ptx_mma_blockscale
TLCodegen->>PTXUtil: GetMMABlockScaleInfo(...)
PTXUtil->>PTXUtil: validate & lookup config
PTXUtil->>TLCodegen: return MMABlockScaleInfo
TLCodegen->>TLCodegen: set need_mma_blockscale_instruction_h_
TLCodegen->>PTXUtil: PrintMMABlockScaleAssembly(...)
TLCodegen->>DeviceTemplate: include mma_blockscale.h
TLCodegen->>DeviceTemplate: call tl::mma_sync_blockscale<...>(...)
DeviceTemplate->>DeviceTemplate: select dispatcher & emit SM120 asm
🎯 4 (Complex) | ⏱️ ~45 minutes
🚥 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 |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/backend/cuda/codegen/codegen_cuda.cc (1)
4435-4443:⚠️ Potential issue | 🟠 Major | ⚡ Quick winApply packed int4/int1 shared-size normalization to
shared.warptoo.
shared.warpnow allocates via the shared-memory path, but the packed-size adjustment still only checksscope == "shared". That over-allocates shared memory for packed int4/int1 buffers and can push kernels over shared-memory limits.Suggested fix
- if ((alloc_dtype == DataType::Int(4) || alloc_dtype == DataType::UInt(4)) && - scope == "shared") { + if ((alloc_dtype == DataType::Int(4) || alloc_dtype == DataType::UInt(4)) && + (scope == "shared" || scope == "shared.warp")) { constant_size = (constant_size + 1) / 2; - } else if (alloc_dtype == DataType::Int(1) && scope == "shared") { + } else if (alloc_dtype == DataType::Int(1) && + (scope == "shared" || scope == "shared.warp")) { constant_size = constant_size / 32; }🤖 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/backend/cuda/codegen/codegen_cuda.cc` around lines 4435 - 4443, The shared-memory size normalization for packed types only runs when scope == "shared", causing over-allocation for scope == "shared.warp"; update the normalization logic around alloc_dtype/DataType::Int(4)/DataType::UInt(4) and DataType::Int(1) so it applies when scope is "shared" or "shared.warp" (same predicate used later where stream emits vid and constant_size), i.e., broaden the condition that adjusts constant_size to include "shared.warp" so packed int4/int1 buffers get the same division handling before emitting the declaration.
🤖 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.
Outside diff comments:
In `@src/backend/cuda/codegen/codegen_cuda.cc`:
- Around line 4435-4443: The shared-memory size normalization for packed types
only runs when scope == "shared", causing over-allocation for scope ==
"shared.warp"; update the normalization logic around
alloc_dtype/DataType::Int(4)/DataType::UInt(4) and DataType::Int(1) so it
applies when scope is "shared" or "shared.warp" (same predicate used later where
stream emits vid and constant_size), i.e., broaden the condition that adjusts
constant_size to include "shared.warp" so packed int4/int1 buffers get the same
division handling before emitting the declaration.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c341c0d4-8d51-4c8c-8f60-7f5a9dedefe1
📒 Files selected for processing (9)
src/backend/cuda/codegen/codegen_cuda.ccsrc/op/builtin.ccsrc/op/builtin.hsrc/tl_templates/cuda/instruction/mma.htilelang/language/__init__.pytilelang/language/ast/ir.pytilelang/language/builtin.pytilelang/language/tir/ir.pytilelang/language/tir/op.py
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 `@src/backend/cuda/codegen/codegen_cuda.cc`:
- Around line 4485-4486: The shared.warp allocation path is routed to the same
emitter as "shared" but the packed-size adjustment for int4/int1 only checks
scope == "shared", causing over-allocation; update the packed-size calculation
(the branch that computes constant_size for int4/int1) to include scope ==
"shared.warp" so that constant_size is reduced for packed types when scope is
"shared" or "shared.warp", keeping the allocation emission that writes
vid[constant_size] consistent with the sizing logic.
🪄 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: bf228a2c-d8d6-4963-b54a-bab9bd7c5652
📒 Files selected for processing (10)
src/backend/cuda/codegen/codegen_cuda.ccsrc/op/builtin.ccsrc/op/builtin.hsrc/tl_templates/cuda/instruction/mma.hsrc/tl_templates/cuda/instruction/mma_blockscale.htilelang/language/__init__.pytilelang/language/ast/ir.pytilelang/language/builtin.pytilelang/language/tir/ir.pytilelang/language/tir/op.py
✅ Files skipped from review due to trivial changes (1)
- tilelang/language/init.py
55ae476 to
70203e2
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
src/backend/cuda/codegen/ptx.cc (1)
1138-1153: ⚡ Quick winMake the operand numbering explicit.
This
templates << ... << arg_counter++ ...chain is already tripping static analysis and is harder than necessary to reason about. Building each placeholder group in separate steps keeps the numbering deterministic and easier to audit.Proposed cleanup
inline std::tuple<std::string, std::string, std::string> GetMMABlockScaleOperands() { std::stringstream templates, inputs, outputs; int arg_counter = 0; - templates << "{%" << arg_counter++ << ", %" << arg_counter++ << ", %" - << arg_counter++ << ", %" << arg_counter++ << "}, "; - templates << "{%" << arg_counter++ << ", %" << arg_counter++ << ", %" - << arg_counter++ << ", %" << arg_counter++ << "}, "; - templates << "{%" << arg_counter++ << ", %" << arg_counter++ << "}, "; - templates << "{%" << arg_counter++ << ", %" << arg_counter++ << ", %" - << arg_counter++ << ", %" << arg_counter++ << "}, "; - templates << "{%" << arg_counter++ << "}, {%" << arg_counter++ << ", %" - << arg_counter++ << "}, "; - templates << "{%" << arg_counter++ << "}, {%" << arg_counter++ << ", %" - << arg_counter++ << "}"; + auto append_group = [&](int count) { + templates << "{"; + for (int i = 0; i < count; ++i) { + if (i != 0) { + templates << ", "; + } + templates << "%" << arg_counter++; + } + templates << "}"; + }; + append_group(4); + templates << ", "; + append_group(4); + templates << ", "; + append_group(2); + templates << ", "; + append_group(4); + templates << ", "; + append_group(1); + templates << ", "; + append_group(2); + templates << ", "; + append_group(1); + templates << ", "; + append_group(2);🤖 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/backend/cuda/codegen/ptx.cc` around lines 1138 - 1153, The GetMMABlockScaleOperands function builds a template string using repeated inline arg_counter++ expressions which is hard to read and flags static analysis; replace the chained templates << ... << arg_counter++ sequences by constructing each placeholder group in separate steps: declare arg_counter, then for each group create a local string or use templates << with explicit placeholders using arg_counter, incrementing arg_counter in its own statement after each placeholder (e.g., set int idx = arg_counter++; templates << "{%" << idx << ", %" << (arg_counter++) << ", %" << (arg_counter++) << ", %" << (arg_counter++) << "}, "; etc.), so the numbering is explicit and deterministic while preserving the exact order produced by the current GetMMABlockScaleOperands implementation.
🤖 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 `@src/backend/cuda/codegen/codegen_cuda.cc`:
- Around line 2613-2668: The MMA blockscale branch forwards raw element offsets
into the pointer arithmetic registrations "(A_offset)" and "(B_offset)" but for
packed 4-bit dtypes (MXF4/NVF4) those logical element indices must be normalized
to byte offsets (divide by 2) like the ptx_mma()/GetBufferRef() paths do; modify
code around where replacer.register_rule("(A_offset)", a_offset) and
replacer.register_rule("(B_offset)", b_offset) to detect packed 4-bit types (use
the dtype_a_enum/dtype_b_enum values from DTypeFromString/DTypeEnumToString) and
replace a_offset/b_offset with an expression that divides the offset by 2 (or
otherwise converts element index to byte index) before registering the rules.
- Around line 3693-3716: The bf16 packing path in the tl::stg64_pack_b16
lowering in codegen_cuda.cc builds the uint2 payload using __pack_half2 which is
for FP16 and will mispack nv_bfloat16 inputs; update the tl::stg64_pack_b16
branch (the block that prints the uint2{__pack_half2(...), __pack_half2(...)}
payload via PrintExpr on op->args[1..4]) to detect when the operand element type
is bf16 (nv_bfloat16 / bfloat16_t) and emit the BF16-specific vector pack
intrinsic (e.g., __pack_nv_bfloat162 or the appropriate cuda_bf16.h helper)
instead of __pack_half2 for those operands, and ensure the bf16 helper header is
included or the intrinsic is available before emission so the resulting u32
lanes contain correct BF16 encodings.
In `@src/tl_templates/cuda/instruction/mma_blockscale.h`:
- Around line 5-8: The header conditionally excludes <type_traits> and <utility>
when __CUDACC_RTC__ is defined but the file (mma_blockscale.h) uses
std::remove_extent_t, std::extent_v, std::index_sequence,
std::make_index_sequence, std::is_same_v, std::is_void_v and
std::integral_constant unconditionally; update the includes so those headers are
available for the __CUDACC_RTC__ path—i.e., move or add `#include` <type_traits>
and `#include` <utility> outside the `#ifndef` __CUDACC_RTC__ guard (or add
equivalent includes within the RTC branch) so the symbols used by
functions/types in this header resolve during NVRTC/RTC compilation.
---
Nitpick comments:
In `@src/backend/cuda/codegen/ptx.cc`:
- Around line 1138-1153: The GetMMABlockScaleOperands function builds a template
string using repeated inline arg_counter++ expressions which is hard to read and
flags static analysis; replace the chained templates << ... << arg_counter++
sequences by constructing each placeholder group in separate steps: declare
arg_counter, then for each group create a local string or use templates << with
explicit placeholders using arg_counter, incrementing arg_counter in its own
statement after each placeholder (e.g., set int idx = arg_counter++; templates
<< "{%" << idx << ", %" << (arg_counter++) << ", %" << (arg_counter++) << ", %"
<< (arg_counter++) << "}, "; etc.), so the numbering is explicit and
deterministic while preserving the exact order produced by the current
GetMMABlockScaleOperands implementation.
🪄 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: 372a95d2-8b3c-4572-8e0b-abaac6f806a0
📒 Files selected for processing (12)
src/backend/cuda/codegen/codegen_cuda.ccsrc/backend/cuda/codegen/codegen_cuda.hsrc/backend/cuda/codegen/ptx.ccsrc/backend/cuda/codegen/ptx.hsrc/op/builtin.ccsrc/op/builtin.hsrc/tl_templates/cuda/instruction/mma_blockscale.htilelang/language/__init__.pytilelang/language/ast/ir.pytilelang/language/builtin.pytilelang/language/tir/ir.pytilelang/language/tir/op.py
✅ Files skipped from review due to trivial changes (2)
- tilelang/language/tir/ir.py
- tilelang/language/init.py
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/backend/cuda/codegen/ptx.cc (1)
1142-1153: ⚡ Quick winDon’t treat this as a correctness bug for C++17 (
<<is sequenced)
Given this project builds with C++17 (CMAKE_CXX_STANDARD 17), the chainedstd::stringstream <<insertions are sequenced in chain order (eachE1 << E2sequences side effects of the left operand before the right). So the%Nplaceholder numbering here should not permute, and the reported evaluation-order risk is unlikely to apply.
Optional: refactor the stream-building to avoidarg_counter++inline for readability, but it’s not required for correctness.🤖 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/backend/cuda/codegen/ptx.cc` around lines 1142 - 1153, The stream construction uses multiple inline arg_counter++ side-effects inside chained templates << operations (symbols: arg_counter, templates in ptx.cc) which is legal under C++17 but harms readability and risks future confusion; refactor by computing the placeholder numbers into local ints (e.g., n0 = arg_counter++; n1 = arg_counter++; ...) or by using a small helper function/loop to generate each "{%d, ...}" segment and then stream those variables into templates to remove inline post-increments and make the intent explicit.
🤖 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/language/ast/ir.py`:
- Around line 1883-1885: The export for ptx_mma_blockscale is currently exposing
the raw _tir_op.ptx_mma_blockscale and thus doesn't accept the dtype= calling
convention; wrap it with the same helper used for the other PTX MMA helpers by
assigning ptx_mma_blockscale = _dtype_forward(_tir_op.ptx_mma_blockscale) so it
behaves like ptx_mma and ptx_mma_sp and supports the dtype= argument.
---
Nitpick comments:
In `@src/backend/cuda/codegen/ptx.cc`:
- Around line 1142-1153: The stream construction uses multiple inline
arg_counter++ side-effects inside chained templates << operations (symbols:
arg_counter, templates in ptx.cc) which is legal under C++17 but harms
readability and risks future confusion; refactor by computing the placeholder
numbers into local ints (e.g., n0 = arg_counter++; n1 = arg_counter++; ...) or
by using a small helper function/loop to generate each "{%d, ...}" segment and
then stream those variables into templates to remove inline post-increments and
make the intent explicit.
🪄 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: 3a19daae-e26a-4e11-8205-e0bb4c02d8f5
📒 Files selected for processing (12)
src/backend/cuda/codegen/codegen_cuda.ccsrc/backend/cuda/codegen/codegen_cuda.hsrc/backend/cuda/codegen/ptx.ccsrc/backend/cuda/codegen/ptx.hsrc/op/builtin.ccsrc/op/builtin.hsrc/tl_templates/cuda/instruction/mma_blockscale.htilelang/language/__init__.pytilelang/language/ast/ir.pytilelang/language/builtin.pytilelang/language/tir/ir.pytilelang/language/tir/op.py
✅ Files skipped from review due to trivial changes (1)
- tilelang/language/init.py
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 `@src/backend/cuda/codegen/ptx.cc`:
- Around line 1399-1402: The registered placeholder keys are wrong:
replacer.register_rule currently uses "{.alayout}" and "{.blayout}" but the
assembly template expects .{alayout} and .{blayout}, so the tokens never match.
Update the calls to register_rule to register "{alayout}" -> "." + A_layout and
"{blayout}" -> "." + B_layout (keep the same values), ensuring
replacer.register_rule is called with "{alayout}" and "{blayout}" so the
template's .{alayout}/.{blayout} are replaced correctly; look for the
replacer.register_rule calls and the A_layout/B_layout variables in the PTX
lowering code.
🪄 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: dfa45c48-586f-4bef-9532-d87dbc6933b7
📒 Files selected for processing (12)
src/backend/cuda/codegen/codegen_cuda.ccsrc/backend/cuda/codegen/codegen_cuda.hsrc/backend/cuda/codegen/ptx.ccsrc/backend/cuda/codegen/ptx.hsrc/op/builtin.ccsrc/op/builtin.hsrc/tl_templates/cuda/instruction/mma_blockscale.htilelang/language/__init__.pytilelang/language/ast/ir.pytilelang/language/builtin.pytilelang/language/tir/ir.pytilelang/language/tir/op.py
b529b43 to
09f219b
Compare
276dc53 to
d497ab8
Compare
708c2eb to
c0a0590
Compare
|
Thanks for your contribution! A recent pr introduce support for tcgen5 fp4 blockscaled gemm, and I personally recommend following similar API standard:) |
| dtype: DType, | ||
| scope: str = "local.var", | ||
| *, | ||
| init: PrimExpr | int | float | None = None, |
There was a problem hiding this comment.
Is there any better ways to control the allocation position (better automatically analyzed)? We can discuss with @LeiWang1999
There was a problem hiding this comment.
better to provide a higher-level API, like existing T.tcgen5_gemm_blockscaled for tcgen5, rather than a thin wrapped ptx intrinsic.
There was a problem hiding this comment.
I personally think that it would be better to move these instructions to tl_templates/cuda/gemm_sm120.h or somewhere else
| from tilelang.carver.arch import driver | ||
|
|
||
| from examples.sage_attention_sm120.tir_helpers import ( | ||
| mma_m16n32k64_blockscale_f32, |
There was a problem hiding this comment.
As I mentioned before, compared to using these thin PTX intrin wrappers, consider supporting automatically lowering from tile level operation T.gemm and T.copy to mma atom and ldmatrix/stmatrix.
b12c2bc to
e1e312a
Compare
Expose a TileLang builtin for Blackwell fp4/fp6 block-scale MMA, add the CUDA codegen lowering, and provide small helper builtins used by warp-local bf16 stores.
Summary
This PR adds SM120 block-scale MMA support and the TileLang compiler/runtime features needed to implement SageAttention3 FP4 raw-core kernels in pure TileLang. It introduces a generic
T.ptx_mma_blockscaleintrinsic formxf8f6f4andmxf4nvf4, adds the SM120 SageAttention3 FP4 example kernel, and fixes reduce/storage-sync lowering needed by the new path and existing language tests.Changes
Add
tl.ptx_mma_blockscalebuiltin and Python frontend bindings.Add CUDA codegen support for lowering block-scale MMA calls through
tl::mma_sync_blockscale.Add
mma_blockscale.hwith SM120 block-scale MMA dispatcher specializations.Add PTX-side validation metadata for supported block-scale MMA configs:
mxf8f6f4,m16n8k32,scale_vec=1,ue8m0mxf4nvf4,m16n8k64,scale_vec=2,ue8m0mxf4nvf4,m16n8k64,scale_vec=4,ue4m3Add
PrintMMABlockScaleAssemblyhelpers and shared dtype/string mapping for block-scale MMA validation.Add
T.pack_b16for packing four b16 values into a global store.Add
examples/sage_attention_sm120with a pure TileLang SM120 SageAttention3 FP4 raw-core kernel and helper macros.Extend
T.ptx_ldmatrixwith optional matrix shape/dtype metadata and add CUDA lowering form8n16.x4.shared.b8x16.b4x16_p64.Add
T.lds32and CUDA shared-memory 32-bit load helpers for packed scale/FP4 data paths.Add strict
T.copy(..., prefer_instruction="ldsm"|"stsm")handling so matrix-copy requests fail clearly instead of silently falling back.Add
T.tma_copy(..., tma_element_dtype=...)and improve TMA descriptor projection, folded base offsets,elem_offsethandling, and transaction-byte accounting.Preserve explicit pack lanes when reshaping packed byte/sub-byte layouts.
Add
role_scoped=Trueallocation support so branch-local register arrays can stay inside warp-specialized producer/consumer roles.Improve CUDA codegen for scoped
pragma_unroll_factor, aligned narrow local allocations, and FP4 localaddress_of.Update
ReduceLowererbatch handling to clamp requested batch size to the physical per-thread output element count before divisibility checks.Use float32 temporary accumulators for fp16/bf16
sumandabssum, then cast back to the destination dtype.Avoid packed fp16/bf16 sum reducers where low-precision pairwise accumulation can exceed existing reduce-test tolerances.
Remove allocation-wide alias ranges from thread-storage-sync disjointness proofs so pointer/buffer conflict checks stay conservative.
Add handle byte offset support in TMA descriptor parsing for nvrtc
Summary by CodeRabbit