Skip to content

Add SM120 dense block-scaled MMA and essential support for SageAttention3#2253

Open
sepcnt wants to merge 6 commits into
tile-ai:mainfrom
sepcnt:mxf4nvf4
Open

Add SM120 dense block-scaled MMA and essential support for SageAttention3#2253
sepcnt wants to merge 6 commits into
tile-ai:mainfrom
sepcnt:mxf4nvf4

Conversation

@sepcnt

@sepcnt sepcnt commented May 23, 2026

Copy link
Copy Markdown
Contributor

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_blockscale intrinsic for mxf8f6f4 and mxf4nvf4, 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_blockscale builtin and Python frontend bindings.

  • Add CUDA codegen support for lowering block-scale MMA calls through tl::mma_sync_blockscale.

  • Add mma_blockscale.h with SM120 block-scale MMA dispatcher specializations.

  • Add PTX-side validation metadata for supported block-scale MMA configs:

    • mxf8f6f4, m16n8k32, scale_vec=1, ue8m0
    • mxf4nvf4, m16n8k64, scale_vec=2, ue8m0
    • mxf4nvf4, m16n8k64, scale_vec=4, ue4m3
  • Add PrintMMABlockScaleAssembly helpers and shared dtype/string mapping for block-scale MMA validation.

  • Add T.pack_b16 for packing four b16 values into a global store.

  • Add examples/sage_attention_sm120 with a pure TileLang SM120 SageAttention3 FP4 raw-core kernel and helper macros.

  • Extend T.ptx_ldmatrix with optional matrix shape/dtype metadata and add CUDA lowering for m8n16.x4.shared.b8x16.b4x16_p64.

  • Add T.lds32 and 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_offset handling, and transaction-byte accounting.

  • Preserve explicit pack lanes when reshaping packed byte/sub-byte layouts.

  • Add role_scoped=True allocation 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 local address_of.

  • Update ReduceLowerer batch handling to clamp requested batch size to the physical per-thread output element count before divisibility checks.

  • Use float32 temporary accumulators for fp16/bf16 sum and abssum, 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

  • New Features
    • Added SM120 block‑scaled MMA intrinsic with scale‑vector support and runtime dispatch.
    • Frontend/TIR/TileLang APIs exposed for the new intrinsic and a pack_b16(v0, v1) helper.
  • Chores
    • CUDA backend and codegen updated to emit the block‑scale MMA assembly and wire the intrinsic.
    • Backend builtin registered so the intrinsic is available to the compiler stack.

Review Change Stack

@github-actions

Copy link
Copy Markdown

👋 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

coderabbitai Bot commented May 23, 2026

Copy link
Copy Markdown
Contributor

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Adds a TileLang intrinsic ptx_mma_blockscale (SM120 block-scaled MMA) with PTX config/assembly helpers, CUDA device template mma_sync_blockscale, codegen emission, and a frontend pack_b16 helper and exports.

Changes

Blackwell SM120 Intrinsics

Layer / File(s) Summary
Op registration
src/op/builtin.h, src/op/builtin.cc
Declares and registers tl.ptx_mma_blockscale (17 inputs, opaque effect).
PTX codegen utilities
src/backend/cuda/codegen/ptx.h, src/backend/cuda/codegen/ptx.cc
Adds DTypeToString, MMABlockScaleInfo, GetMMABlockScaleInfo, GetMMABlockScaleOperands, and PrintMMABlockScaleAssembly to validate configs and emit mma.sync.aligned ... .block_scale PTX assembly.
SM120 device template
src/tl_templates/cuda/instruction/mma_blockscale.h
Implements device-side dispatch, SM120 inline-assembly macros, concrete impls for supported dtype/scale combinations, and tl::mma_sync_blockscale API.
CUDA codegen integration
src/backend/cuda/codegen/codegen_cuda.h, src/backend/cuda/codegen/codegen_cuda.cc
Adds need_mma_blockscale_instruction_h_ flag; Finish() conditionally includes mma_blockscale.h; visitor lowers tl.ptx_mma_blockscale calls to tl::mma_sync_blockscale<...>(...).
TileLang frontend and helpers
tilelang/language/builtin.py, tilelang/language/__init__.py, tilelang/language/ast/ir.py, tilelang/language/tir/ir.py, tilelang/language/tir/op.py
Adds pack_b16 helper, exposes and forwards ptx_mma_blockscale through TIR/AST layers, and re-exports helper in package __init__.

Sequence Diagram

sequenceDiagram
  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
Loading

🎯 4 (Complex) | ⏱️ ~45 minutes

"🐰 I packed two b16s into a word with care,
PTX strings hummed, and templates took the air,
From TileLang call down to SM120's lair,
Block-scale fma leapt fast and fair,
My whiskers twitch — math dancing there."

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The PR title directly matches the main objective: adding SM120 dense block-scaled MMA support. It accurately summarizes the primary change across multiple files (builtin ops, CUDA codegen, PTX support, and Python bindings).
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ 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

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Apply packed int4/int1 shared-size normalization to shared.warp too.

shared.warp now allocates via the shared-memory path, but the packed-size adjustment still only checks scope == "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

📥 Commits

Reviewing files that changed from the base of the PR and between bc310e0 and cc55167.

📒 Files selected for processing (9)
  • src/backend/cuda/codegen/codegen_cuda.cc
  • src/op/builtin.cc
  • src/op/builtin.h
  • src/tl_templates/cuda/instruction/mma.h
  • tilelang/language/__init__.py
  • tilelang/language/ast/ir.py
  • tilelang/language/builtin.py
  • tilelang/language/tir/ir.py
  • tilelang/language/tir/op.py

@sepcnt sepcnt changed the title Add MXF4NVF4 block-scale MMA builtin [WIP] Add block-scale MMA builtin May 23, 2026
@sepcnt sepcnt changed the title [WIP] Add block-scale MMA builtin Add SM120 block-scaled MMA builtin May 23, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between cc55167 and 47dd240.

📒 Files selected for processing (10)
  • src/backend/cuda/codegen/codegen_cuda.cc
  • src/op/builtin.cc
  • src/op/builtin.h
  • src/tl_templates/cuda/instruction/mma.h
  • src/tl_templates/cuda/instruction/mma_blockscale.h
  • tilelang/language/__init__.py
  • tilelang/language/ast/ir.py
  • tilelang/language/builtin.py
  • tilelang/language/tir/ir.py
  • tilelang/language/tir/op.py
✅ Files skipped from review due to trivial changes (1)
  • tilelang/language/init.py

Comment thread src/backend/cuda/codegen/codegen_cuda.cc Outdated
@sepcnt sepcnt marked this pull request as draft May 23, 2026 17:38
@sepcnt sepcnt marked this pull request as ready for review May 23, 2026 18:37
@sepcnt sepcnt force-pushed the mxf4nvf4 branch 2 times, most recently from 55ae476 to 70203e2 Compare May 23, 2026 18:44
@sepcnt sepcnt changed the title Add SM120 block-scaled MMA builtin Add SM120 dense block-scaled MMA support May 23, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

🧹 Nitpick comments (1)
src/backend/cuda/codegen/ptx.cc (1)

1138-1153: ⚡ Quick win

Make 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

📥 Commits

Reviewing files that changed from the base of the PR and between 47dd240 and 70203e2.

📒 Files selected for processing (12)
  • src/backend/cuda/codegen/codegen_cuda.cc
  • src/backend/cuda/codegen/codegen_cuda.h
  • src/backend/cuda/codegen/ptx.cc
  • src/backend/cuda/codegen/ptx.h
  • src/op/builtin.cc
  • src/op/builtin.h
  • src/tl_templates/cuda/instruction/mma_blockscale.h
  • tilelang/language/__init__.py
  • tilelang/language/ast/ir.py
  • tilelang/language/builtin.py
  • tilelang/language/tir/ir.py
  • tilelang/language/tir/op.py
✅ Files skipped from review due to trivial changes (2)
  • tilelang/language/tir/ir.py
  • tilelang/language/init.py

Comment thread src/cuda/codegen/codegen_cuda.cc Outdated
Comment thread src/backend/cuda/codegen/codegen_cuda.cc Outdated
Comment thread src/tl_templates/cuda/instruction/mma_blockscaled.h

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
src/backend/cuda/codegen/ptx.cc (1)

1142-1153: ⚡ Quick win

Don’t treat this as a correctness bug for C++17 (<< is sequenced)
Given this project builds with C++17 (CMAKE_CXX_STANDARD 17), the chained std::stringstream << insertions are sequenced in chain order (each E1 << E2 sequences side effects of the left operand before the right). So the %N placeholder numbering here should not permute, and the reported evaluation-order risk is unlikely to apply.
Optional: refactor the stream-building to avoid arg_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

📥 Commits

Reviewing files that changed from the base of the PR and between 70203e2 and e4d73c9.

📒 Files selected for processing (12)
  • src/backend/cuda/codegen/codegen_cuda.cc
  • src/backend/cuda/codegen/codegen_cuda.h
  • src/backend/cuda/codegen/ptx.cc
  • src/backend/cuda/codegen/ptx.h
  • src/op/builtin.cc
  • src/op/builtin.h
  • src/tl_templates/cuda/instruction/mma_blockscale.h
  • tilelang/language/__init__.py
  • tilelang/language/ast/ir.py
  • tilelang/language/builtin.py
  • tilelang/language/tir/ir.py
  • tilelang/language/tir/op.py
✅ Files skipped from review due to trivial changes (1)
  • tilelang/language/init.py

Comment thread tilelang/language/ast/ir.py Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between e4d73c9 and dd693f0.

📒 Files selected for processing (12)
  • src/backend/cuda/codegen/codegen_cuda.cc
  • src/backend/cuda/codegen/codegen_cuda.h
  • src/backend/cuda/codegen/ptx.cc
  • src/backend/cuda/codegen/ptx.h
  • src/op/builtin.cc
  • src/op/builtin.h
  • src/tl_templates/cuda/instruction/mma_blockscale.h
  • tilelang/language/__init__.py
  • tilelang/language/ast/ir.py
  • tilelang/language/builtin.py
  • tilelang/language/tir/ir.py
  • tilelang/language/tir/op.py

Comment thread src/backend/cuda/codegen/ptx.cc Outdated
@sepcnt sepcnt force-pushed the mxf4nvf4 branch 2 times, most recently from b529b43 to 09f219b Compare May 29, 2026 15:34
@sepcnt sepcnt force-pushed the mxf4nvf4 branch 4 times, most recently from 276dc53 to d497ab8 Compare June 5, 2026 03:28
@sepcnt sepcnt changed the title Add SM120 dense block-scaled MMA support Add SM120 dense block-scaled MMA and essential support for SageAttention3 Jun 5, 2026
@LeiWang1999 LeiWang1999 self-requested a review June 5, 2026 05:28
@sepcnt sepcnt force-pushed the mxf4nvf4 branch 2 times, most recently from 708c2eb to c0a0590 Compare June 10, 2026 14:08
@Rachmanino

Copy link
Copy Markdown
Collaborator

Thanks for your contribution! A recent pr introduce support for tcgen5 fp4 blockscaled gemm, and I personally recommend following similar API standard:)

Comment thread tilelang/language/loop.py Outdated
@sepcnt sepcnt requested a review from Rachmanino June 12, 2026 08:04
dtype: DType,
scope: str = "local.var",
*,
init: PrimExpr | int | float | None = None,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any better ways to control the allocation position (better automatically analyzed)? We can discuss with @LeiWang1999

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

better to provide a higher-level API, like existing T.tcgen5_gemm_blockscaled for tcgen5, rather than a thin wrapped ptx intrinsic.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@sepcnt sepcnt force-pushed the mxf4nvf4 branch 2 times, most recently from b12c2bc to e1e312a Compare June 17, 2026 08:23
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants