Repository navigation
[CUDA] Support tcgen5mma gemm ts - #1866
Merged
Merged
Conversation
… with 2 val, to be consistent with mma_ts's behaviour
…bf16_gemm_tcgen5_ts.py
|
👋 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
|
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 TCGEN5 TS (A from TMEM, B from SMEM) GEMM support: store metadata/templates and device store helpers, direction-aware tmem copy lowering, a TCGEN5 MMA TS macro/emitter and GEMM TS layout/lowering, dtype-routing, plus example and tests for chained GEMM workflows. Changes
Sequence Diagram(s)sequenceDiagram
participant Host as Host
participant Kernel as Kernel (TCGEN5 TS)
participant TMEM as TensorMem (TMEM)
participant SMEM as SharedMem (SMEM)
participant Acc as Accumulator
Host->>Kernel: launch chained GEMM (A: TMEM, B: SMEM)
rect rgba(100,150,255,0.5)
Kernel->>TMEM: tcgen05.ld (load A fragments)
TMEM-->>Kernel: A fragments
Kernel->>SMEM: load B fragments
SMEM-->>Kernel: B fragments
Kernel->>Acc: ptx_tcgen05_mma_ts (A × B -> accum)
end
rect rgba(150,200,100,0.5)
Kernel->>TMEM: tcgen05.st (store accum / chain buffer to TMEM)
TMEM-->>Kernel: (optional reload for next stage)
end
rect rgba(200,150,100,0.5)
Kernel->>Acc: second ptx_tcgen05_mma_ts (second stage)
Kernel->>Host: write final output
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 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)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). 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 (5)
tilelang/intrinsics/tcgen05_macro_generator.py (2)
405-431: Duplicatedaccess_ptr_fromhelper — extract to a shared method.This is a verbatim copy of the
access_ptr_fromclosure defined intcgen05mma(lines 237–263). Consider extracting it as a private method onTensorCoreIntrinEmitter(e.g.,_access_ptr_from) to eliminate the duplication.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/tcgen05_macro_generator.py` around lines 405 - 431, The access_ptr_from helper is duplicated; extract it from the closure into a single private method on the TensorCoreIntrinEmitter class (e.g., def _access_ptr_from(self, buffer_or_load_or_region, access_type="r"): ...) and replace the local closures named access_ptr_from in both tcgen05_macro_generator.py and the tcgen05mma closure with calls to TensorCoreIntrinEmitter._access_ptr_from to remove duplication while preserving the same behavior for Buffer, BufferLoad and BufferRegion handling and the same offset/stride calculations.
365-365: Unused unpacked variablesenable_wsandenable_2cta.The TS path doesn't use
enable_ws(unlike the SS path which passes it toptx_tcgen05_mma_ss). Prefix with underscores to signal intent and silence the linter.Proposed fix
- atom_m, atom_n, atom_k, enable_ws, enable_2cta = (int(x) for x in meta) + atom_m, atom_n, atom_k, _enable_ws, _enable_2cta = (int(x) for x in meta)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/tcgen05_macro_generator.py` at line 365, The unpacking line sets atom_m, atom_n, atom_k, enable_ws, enable_2cta but enable_ws and enable_2cta are unused in the TS path; change their names to begin with underscores (e.g., _enable_ws, _enable_2cta) in the unpacking expression so the linter knows they are intentionally unused and to avoid warnings — update the tuple unpacking where this occurs (the line assigning atom_m, atom_n, atom_k, enable_ws, enable_2cta) and any subsequent references if any exist.examples/gemm_sm100/gemm_tcgen5mma_ts.py (1)
84-84: B1 row offset hardcoded to 0 — only correct whenblock_N1 == N1.
T.copy(B1[0, k * block_K], B1_shared)assumes the entire N1 dimension fits in a single block. This is fine for this example (N1 == block_N1 == 128), but a comment noting this constraint would help future readers.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/gemm_sm100/gemm_tcgen5mma_ts.py` at line 84, T.copy currently uses a hardcoded row 0 for B1 which only works when the entire N1 fits in one block; replace the literal 0 with the computed row block offset (e.g., row_offset = n1_block * block_N1 or whatever outer N1-block loop/index is used) and call T.copy(B1[row_offset, k * block_K], B1_shared) so the code works when block_N1 < N1; update any nearby loop variable name references (the N1 block index) to match the actual identifier in this file.src/tl_templates/cuda/tcgen_05_st.h (1)
1-332: Consider generating these templates programmatically.This file is ~1300 lines of highly repetitive inline assembly. Each shape variant follows an identical pattern with only the instruction mnemonic and register count differing. A code generator (e.g., a Python script or C++ preprocessor macros) would reduce maintenance burden and eliminate the risk of copy-paste errors across variants.
That said, this approach is consistent with how the existing
tcgen05_ldtemplates are likely structured, so deferring this is fine.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/tl_templates/cuda/tcgen_05_st.h` around lines 1 - 332, The file contains ~1300 lines of repetitive inline-asm in the class templates tmem_st_32dp32bNx<false>::copy and tmem_st_32dp32bNx<true>::copy; replace the manual duplication by generating the instruction variants programmatically (e.g., a small generator script or compile-time macro) so you emit the same asm patterns with varying mnemonic suffixes and register counts instead of hand-copying each N-case; target the copy< int N > function bodies (and the two specializations for Unpack16 true/false) and produce code that loops/generates the correct asm string and operand lists for N in {1,2,4,8,16,32,64,128} to eliminate the repetitive blocks and reduce maintenance risk.testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py (1)
85-93: Stage 1 implicitly requiresbN1 == N1(single tile over N1).
B1[0, k * bK]always starts from row 0, meaning the entire N1 dimension must fit in one tile. This is a necessary constraint for the chained pattern (stage 2 MMA TS consumes the full intermediate result in one shot), but it's not enforced. Consider adding an assertion for safety.🛡️ Suggested guard
def chained_gemm( M, N1, N2, K, bM, bN1, bN2, bK, in_dtype, out_dtype, accum_dtype, threads, ): """Chained GEMM: SS → tcgen05.ld → cast → tcgen05.st → MMA TS → output.""" + assert bN1 == N1, "Chained GEMM requires bN1 == N1 (full intermediate tile)"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py` around lines 85 - 93, Stage 1 assumes the entire N1 dimension is in a single tile because B1 is indexed as B1[0, k * bK]; add an explicit runtime or compile-time assertion to enforce bN1 == N1 before the pipelined loop (or at kernel entry). Locate the Stage 1 code that uses B1[0, k * bK] and S_tmem (the pipelined T.Pipelined loop and T.gemm call) and insert a check that raises/throws or fails the build when bN1 != N1 to prevent incorrect execution when N1 is tiled.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@examples/gemm_sm100/gemm_tcgen5mma_ts.py`:
- Line 84: T.copy currently uses a hardcoded row 0 for B1 which only works when
the entire N1 fits in one block; replace the literal 0 with the computed row
block offset (e.g., row_offset = n1_block * block_N1 or whatever outer N1-block
loop/index is used) and call T.copy(B1[row_offset, k * block_K], B1_shared) so
the code works when block_N1 < N1; update any nearby loop variable name
references (the N1 block index) to match the actual identifier in this file.
In `@src/tl_templates/cuda/tcgen_05_st.h`:
- Around line 1-332: The file contains ~1300 lines of repetitive inline-asm in
the class templates tmem_st_32dp32bNx<false>::copy and
tmem_st_32dp32bNx<true>::copy; replace the manual duplication by generating the
instruction variants programmatically (e.g., a small generator script or
compile-time macro) so you emit the same asm patterns with varying mnemonic
suffixes and register counts instead of hand-copying each N-case; target the
copy< int N > function bodies (and the two specializations for Unpack16
true/false) and produce code that loops/generates the correct asm string and
operand lists for N in {1,2,4,8,16,32,64,128} to eliminate the repetitive blocks
and reduce maintenance risk.
In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py`:
- Around line 85-93: Stage 1 assumes the entire N1 dimension is in a single tile
because B1 is indexed as B1[0, k * bK]; add an explicit runtime or compile-time
assertion to enforce bN1 == N1 before the pipelined loop (or at kernel entry).
Locate the Stage 1 code that uses B1[0, k * bK] and S_tmem (the pipelined
T.Pipelined loop and T.gemm call) and insert a check that raises/throws or fails
the build when bN1 != N1 to prevent incorrect execution when N1 is tiled.
In `@tilelang/intrinsics/tcgen05_macro_generator.py`:
- Around line 405-431: The access_ptr_from helper is duplicated; extract it from
the closure into a single private method on the TensorCoreIntrinEmitter class
(e.g., def _access_ptr_from(self, buffer_or_load_or_region, access_type="r"):
...) and replace the local closures named access_ptr_from in both
tcgen05_macro_generator.py and the tcgen05mma closure with calls to
TensorCoreIntrinEmitter._access_ptr_from to remove duplication while preserving
the same behavior for Buffer, BufferLoad and BufferRegion handling and the same
offset/stride calculations.
- Line 365: The unpacking line sets atom_m, atom_n, atom_k, enable_ws,
enable_2cta but enable_ws and enable_2cta are unused in the TS path; change
their names to begin with underscores (e.g., _enable_ws, _enable_2cta) in the
unpacking expression so the linter knows they are intentionally unused and to
avoid warnings — update the tuple unpacking where this occurs (the line
assigning atom_m, atom_n, atom_k, enable_ws, enable_2cta) and any subsequent
references if any exist.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
tilelang/tileop/gemm/gemm_tcgen05.py (2)
101-104: Ruff TRY003: long inline exception message.The static analysis tool flags the multi-word f-string message inside
ValueError. Consider moving it to a custom exception or a named constant if the project's lint config enforces TRY003.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/tileop/gemm/gemm_tcgen05.py` around lines 101 - 104, The inline multi-word f-string used in the ValueError inside gemm_tcgen05.py is flagged by TRY003; replace the long inline message with either a named constant or a custom exception and reference it from the check in the method that calls is_gemm_ss()/is_gemm_ts() — e.g., define a module-level constant (e.g., MSG_TCGEN5MMA_INVALID_SCOPE) or a small exception class (e.g., TCGEN5MMAError) and raise that with the formatted details using A.scope() and B.scope() in the raise site instead of embedding the long f-string directly in the ValueError.
144-156:_gemm_ss_cond/_gemm_ssnames are misleading for the TS variant.Both inner
prim_funcclosures and the capturedA_shared/B_sharedvariable names imply SS-only semantics, yet these closures are returned verbatim for the TS path as well. The actual SS vs TS dispatch is delegated tomma_emitter.tcgen05mma, so the current code works, but the naming makes it harder to reason about the TS lowering path.♻️ Suggested rename
- `@T.prim_func` - def _gemm_ss_cond() -> None: + `@T.prim_func` + def _gemm_cond() -> None: if thread_var // 32 == thread_bounds.min // warp_size: - mma_emitter.tcgen05mma(A_shared, B_shared, C_local, mbarptr, clear_accum) + mma_emitter.tcgen05mma(A_shared, B_shared, C_local, mbarptr, clear_accum) - `@T.prim_func` - def _gemm_ss() -> None: + `@T.prim_func` + def _gemm() -> None: mma_emitter.tcgen05mma(A_shared, B_shared, C_local, mbarptr, clear_accum) return ( - _Simplify(_gemm_ss, inline_let=True) + _Simplify(_gemm, inline_let=True) if analyzer.can_prove(thread_bounds.extent == warp_size) - else _Simplify(_gemm_ss_cond, inline_let=True) + else _Simplify(_gemm_cond, inline_let=True) )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/tileop/gemm/gemm_tcgen05.py` around lines 144 - 156, Rename the misleading SS-specific symbols to neutral names so the TS path is clearer: change the inner prim_funcs _gemm_ss_cond and _gemm_ss to something generic like _gemm_inner_cond and _gemm_inner, and rename captured variables A_shared and B_shared to neutral names (e.g., A_tile / B_tile or A_buf / B_buf); update all places these symbols are referenced (including the return expression and any captures used by mma_emitter.tcgen05mma, and keep the same thread_bounds / warp_size conditional logic) so behavior is unchanged but names no longer imply SS-only semantics.tilelang/intrinsics/tcgen05_macro_generator.py (1)
444-470: Extractaccess_ptr_fromas a class method — exact duplicate of lines 258–284Both
tcgen05mmaandtcgen05mma_tsdefine an identical local helper. Promote it to a private class method (e.g._access_ptr_from) and replace both copies withself._access_ptr_from(...). As a bonus, theTypeErrorfixup (Ruff TRY004) only needs to be applied once.♻️ Proposed refactor — add one class method, remove both inner defs
+ `@staticmethod` + def _access_ptr_from(buffer_or_load_or_region, access_type: str = "r"): + if isinstance(buffer_or_load_or_region, Buffer): + return buffer_or_load_or_region.access_ptr(access_type) + elif isinstance(buffer_or_load_or_region, BufferLoad): + buffer_load = buffer_or_load_or_region + offset, stride = 0, 1 + buffer = buffer_load.buffer + for i, shape in enumerate(reversed(buffer.shape)): + indice = buffer_load.indices[len(buffer_load.indices) - i - 1] + if isinstance(indice, (tvm.tir.IntImm, tvm.tir.PrimExpr)): + offset += indice * stride + elif isinstance(indice, tvm.tir.Ramp): + offset += indice.base * stride + else: + raise TypeError(f"Unsupported index type: {type(indice)}") + stride *= shape + return buffer.access_ptr(access_type, offset=offset) + elif isinstance(buffer_or_load_or_region, BufferRegion): + buffer_region = buffer_or_load_or_region + buffer = buffer_region.buffer + offset, stride = 0, 1 + for i, shape in enumerate(reversed(buffer.shape)): + offset += buffer_region.region[len(buffer_region.region) - i - 1].min * stride + stride *= shape + return buffer.access_ptr(access_type, offset=offset) + else: + raise TypeError(f"Unsupported buffer type: {type(buffer_or_load_or_region)}")Then in both
_warp_mmaand_warp_mma_ts:- def access_ptr_from(buffer_or_load_or_region, access_type: str = "r"): - ... # (remove entire local def) - A_ptr = access_ptr_from(A_buf, "r") - B_ptr = access_ptr_from(B_buf, "r") + A_ptr = self._access_ptr_from(A_buf, "r") + B_ptr = self._access_ptr_from(B_buf, "r")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/tcgen05_macro_generator.py` around lines 444 - 470, Extract the duplicated local function access_ptr_from into a private class method named _access_ptr_from and replace the local defs inside tcgen05mma and tcgen05mma_ts with calls to self._access_ptr_from(...); specifically, remove the inline access_ptr_from in both _warp_mma and _warp_mma_ts, add a single method _access_ptr_from(self, buffer_or_load_or_region, access_type: str = "r") implementing the existing logic (including the same offset/stride computation and error handling), and update callers in _warp_mma and _warp_mma_ts to call self._access_ptr_from(...) so the TypeError/Try004 fixup is centralized.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py`:
- Line 85: Replace the ambiguous Unicode multiplication sign in the inline
comments so they use an ASCII symbol; locate the comment "Stage 1: SS GEMM —
S_tmem = A × B1^T (fp32 accumulator)" (and the similar comment at the other
occurrence) and change the multiplication sign '×' to a plain '*' or 'x' so the
comment reads e.g. "S_tmem = A * B1^T (fp32 accumulator)".
- Around line 86-93: The chained_gemm implementation currently loads B1 using a
hard-coded row index B1[0, k * bK] which is only correct when there is no tiling
over N1; add a defensive runtime check at the start of chained_gemm (or right
before the loop) that verifies bN1 == N1 (or else raise/assert) and/or update
the B1 load to use the correct tile index (e.g., B1[bn1, k * bK] or the
corresponding loop variable) so the code does not silently read row 0 when bN1 <
N1; reference the chained_gemm function and the existing assert_chained_gemm to
mirror its guard.
In `@tilelang/intrinsics/tcgen05_macro_generator.py`:
- Around line 473-478: The branch in tcgen05mma_ts currently only handles
BufferRegion and Buffer and raises an unhelpful ValueError when given a
BufferLoad; update the A_buf handling to also accept BufferLoad (extract the
underlying buffer data via A_buf.buffer.data, matching access_ptr_from behavior)
and change the raised error to a TypeError with a clear message for unsupported
types; locate the logic around A_buf in tcgen05_macro_generator.py
(tcgen05mma_ts) and add an elif isinstance(A_buf, BufferLoad): a_tmem_data =
A_buf.buffer.data branch, then replace the final raise ValueError(...) with
raise TypeError(...) describing the unsupported A_buf type.
---
Nitpick comments:
In `@tilelang/intrinsics/tcgen05_macro_generator.py`:
- Around line 444-470: Extract the duplicated local function access_ptr_from
into a private class method named _access_ptr_from and replace the local defs
inside tcgen05mma and tcgen05mma_ts with calls to self._access_ptr_from(...);
specifically, remove the inline access_ptr_from in both _warp_mma and
_warp_mma_ts, add a single method _access_ptr_from(self,
buffer_or_load_or_region, access_type: str = "r") implementing the existing
logic (including the same offset/stride computation and error handling), and
update callers in _warp_mma and _warp_mma_ts to call self._access_ptr_from(...)
so the TypeError/Try004 fixup is centralized.
In `@tilelang/tileop/gemm/gemm_tcgen05.py`:
- Around line 101-104: The inline multi-word f-string used in the ValueError
inside gemm_tcgen05.py is flagged by TRY003; replace the long inline message
with either a named constant or a custom exception and reference it from the
check in the method that calls is_gemm_ss()/is_gemm_ts() — e.g., define a
module-level constant (e.g., MSG_TCGEN5MMA_INVALID_SCOPE) or a small exception
class (e.g., TCGEN5MMAError) and raise that with the formatted details using
A.scope() and B.scope() in the raise site instead of embedding the long f-string
directly in the ValueError.
- Around line 144-156: Rename the misleading SS-specific symbols to neutral
names so the TS path is clearer: change the inner prim_funcs _gemm_ss_cond and
_gemm_ss to something generic like _gemm_inner_cond and _gemm_inner, and rename
captured variables A_shared and B_shared to neutral names (e.g., A_tile / B_tile
or A_buf / B_buf); update all places these symbols are referenced (including the
return expression and any captures used by mma_emitter.tcgen05mma, and keep the
same thread_bounds / warp_size conditional logic) so behavior is unchanged but
names no longer imply SS-only semantics.
Comment on lines
+86
to
+93
| for k in T.Pipelined(T.ceildiv(K, bK), num_stages=1): | ||
| T.copy(A[by * bM, k * bK], A_shared) | ||
| T.copy(B1[0, k * bK], B1_shared) | ||
| T.gemm( | ||
| A_shared, B1_shared, S_tmem, | ||
| transpose_B=True, mbar=mbar1, wg_wait=-1, clear_accum=k == 0, | ||
| ) | ||
| T.mbarrier_wait_parity(mbar1, k % 2) |
Contributor
There was a problem hiding this comment.
B1 is always loaded from row 0 — chained_gemm itself carries no guard for bN1 == N1.
Line 88 hard-codes B1[0, k * bK], which is correct only when bN1 == N1 (no tiling over the N1 dimension). The assertion bN1 == N1 exists in assert_chained_gemm (line 138) but not inside chained_gemm itself, so a direct call to chained_gemm with bN1 < N1 would silently produce an incorrect result.
Consider adding the constraint as a runtime guard (or at least a comment) inside chained_gemm:
🛡️ Suggested defensive check
def chained_gemm(
M, N1, N2, K, bM, bN1, bN2, bK,
in_dtype, out_dtype, accum_dtype, threads,
):
"""Chained GEMM: SS → tcgen05.ld → cast → tcgen05.st → MMA TS → output."""
+ assert bN1 == N1, (
+ f"chained_gemm requires bN1 == N1 (B1 is not tiled over N1), got bN1={bN1}, N1={N1}"
+ )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py` around
lines 86 - 93, The chained_gemm implementation currently loads B1 using a
hard-coded row index B1[0, k * bK] which is only correct when there is no tiling
over N1; add a defensive runtime check at the start of chained_gemm (or right
before the loop) that verifies bN1 == N1 (or else raise/assert) and/or update
the B1 load to use the correct tile index (e.g., B1[bn1, k * bK] or the
corresponding loop variable) so the code does not silently read row 0 when bN1 <
N1; reference the chained_gemm function and the existing assert_chained_gemm to
mirror its guard.
Comment on lines
+473
to
+478
| if isinstance(A_buf, BufferRegion): | ||
| a_tmem_data = A_buf.buffer.data | ||
| elif isinstance(A_buf, Buffer): | ||
| a_tmem_data = A_buf.data | ||
| else: | ||
| raise ValueError(f"Unsupported A_buf type for TS variant: {type(A_buf)}") |
Contributor
There was a problem hiding this comment.
BufferLoad A_buf falls through to a misleading ValueError
access_ptr_from (used for B) explicitly handles BufferLoad, so callers can and do pass BufferLoad objects. If is_tensor_memory() checks the .buffer scope on a BufferLoad, tcgen05mma_ts would be entered with such a value and fail here with an unhelpful message. Also, a wrong-type argument should raise TypeError, not ValueError (Ruff TRY004).
🐛 Proposed fix
- if isinstance(A_buf, BufferRegion):
- a_tmem_data = A_buf.buffer.data
- elif isinstance(A_buf, Buffer):
- a_tmem_data = A_buf.data
- else:
- raise ValueError(f"Unsupported A_buf type for TS variant: {type(A_buf)}")
+ if isinstance(A_buf, BufferRegion):
+ a_tmem_data = A_buf.buffer.data
+ elif isinstance(A_buf, Buffer):
+ a_tmem_data = A_buf.data
+ elif isinstance(A_buf, BufferLoad):
+ a_tmem_data = A_buf.buffer.data
+ else:
+ raise TypeError(f"Unsupported A_buf type for TS variant: {type(A_buf)}")🧰 Tools
🪛 Ruff (0.15.1)
[warning] 478-478: Prefer TypeError exception for invalid type
(TRY004)
[warning] 478-478: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tilelang/intrinsics/tcgen05_macro_generator.py` around lines 473 - 478, The
branch in tcgen05mma_ts currently only handles BufferRegion and Buffer and
raises an unhelpful ValueError when given a BufferLoad; update the A_buf
handling to also accept BufferLoad (extract the underlying buffer data via
A_buf.buffer.data, matching access_ptr_from behavior) and change the raised
error to a TypeError with a clear message for unsupported types; locate the
logic around A_buf in tcgen05_macro_generator.py (tcgen05mma_ts) and add an elif
isinstance(A_buf, BufferLoad): a_tmem_data = A_buf.buffer.data branch, then
replace the final raise ValueError(...) with raise TypeError(...) describing the
unsupported A_buf type.
Hale423
force-pushed
the
feat/tcgen5mma-gemm-ts
branch
from
February 22, 2026 03:55
9e7c8c6 to
e1e8723
Compare
Member
|
Thanks for your contributions! would you mind taking a look? @Rachmanino |
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (4)
tilelang/tileop/gemm/gemm_tcgen05.py (1)
144-157: TS path reuses_gemm_ss/_gemm_ss_condnaming — consider renaming for clarity.Now that
lower()supports both SS and TS (the dispatch to TS happens insidemma_emitter.tcgen05mma), the local function names_gemm_ssand_gemm_ss_condare misleading. A more neutral name (e.g._gemm_mma/_gemm_mma_cond) would better reflect that these handle both variants.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/tileop/gemm/gemm_tcgen05.py` around lines 144 - 157, Rename the misleading local functions _gemm_ss and _gemm_ss_cond to neutral names reflecting they handle both SS and TS (for example _gemm_mma and _gemm_mma_cond); update the function definitions (def _gemm_ss_cond -> def _gemm_mma_cond, def _gemm_ss -> def _gemm_mma), update the conditional return to reference the new names, and ensure any local references or calls to _gemm_ss/_gemm_ss_cond within this scope (including the analyzer.can_prove branch and the mma_emitter invocation) are changed accordingly so the behavior is unchanged but the names match their generalized role.tilelang/intrinsics/tcgen05_macro_generator.py (2)
444-470:access_ptr_fromis duplicated verbatim from the SS path (lines 258–284).This helper is an exact copy of the one inside
tcgen05mma. Extract it to a shared instance method (or a module-level helper) to eliminate the duplication.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/tcgen05_macro_generator.py` around lines 444 - 470, The access_ptr_from helper is duplicated (the function access_ptr_from in this file and the identical one inside tcgen05mma); extract it into a single shared function (e.g., a module-level helper named access_ptr_from) and replace both copies with calls to that shared helper; update any local references to Buffer, BufferLoad, BufferRegion, and calls that previously used the nested function so behavior remains identical.
405-425: B descriptor parameter computation is duplicated from the SS path.Lines 407–423 replicate lines 212–232 in
tcgen05mmanearly identically. Consider extracting this into a shared helper (e.g.,_compute_b_descriptor_params) to reduce duplication and make future fixes apply to both paths.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/tcgen05_macro_generator.py` around lines 405 - 425, The B-descriptor parameter computation in tcgen05_macro_generator.py (the block computing b_leading_byte_offset, b_stride_byte_offset and bk_atom_size) is duplicated from tcgen05mma; extract that logic into a shared helper function (e.g., _compute_b_descriptor_params) that accepts the inputs used in the block (b_is_k_major, b_swizzle_mode, elems_in_bytes, n_dim, k_dim, b_swizzle_atom_elems, micro_size_k, etc.) and returns the computed values (b_leading_byte_offset, b_stride_byte_offset, bk_atom_size); then replace the duplicated code in both tcgen05_macro_generator.py and tcgen05mma to call _compute_b_descriptor_params and remove the duplicate blocks so both paths use the single helper.testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py (1)
157-178: Consider adding a test case whereMis not a multiple of 128.All current test cases use
bM=128withMbeing a multiple of 128. Since TCGEN5MMA atoms can operate onatom_m ∈ {32, 64, 128}, a test with smallerbM(e.g., 64 or 32) could improve coverage of the TMEM offset geometry for differentinterleavefactors in the TS path.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py` around lines 157 - 178, Add a new CUDA test function (decorated with `@tilelang.testing.requires_cuda` and `@tilelang.testing.requires_cuda_compute_version`(10)) that calls assert_chained_gemm with an M value that is not a multiple of 128 (e.g., use M=64) to exercise smaller TCGEN5MMA atom_m cases; reference the existing test functions test_chained_gemm_128/test_chained_gemm_256 for placement and reuse the same argument pattern but replace the first M dimension with 64 so assert_chained_gemm is invoked with a non-128 M to validate TMEM offset/interleave geometry.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py`:
- Line 85: The comment "Stage 1: SS GEMM — S_tmem = A × B1^T (fp32 accumulator)"
(and the similar comment on line 100) uses the ambiguous Unicode multiplication
symbol ×; replace it with a clear ASCII operator like 'x' or '*' (e.g., "S_tmem
= A x B1^T" or "S_tmem = A * B1^T") to avoid encoding/visibility issues and
update both comment occurrences.
- Around line 86-93: B1 is always loaded from row 0 which breaks when N is
blocked; in the pipelined gemm loop replace the fixed row index in the load
(B1[0, k * bK]) with the correct block row index (e.g. B1[bn * bN1, k * bK] or
the appropriate bn variable used for N1-blocking) and update any boundary check
so partial final blocks are handled rather than relying on an external bN1 == N1
assumption; adjust the load into B1_shared and any related guards so B1 indexing
uses the block index (bn) consistently with A’s by index and respects tails.
In `@tilelang/intrinsics/tcgen05_macro_generator.py`:
- Around line 472-478: In tcgen05mma_ts, handle the case where A_buf is a
BufferLoad (in addition to BufferRegion and Buffer) by extracting the underlying
buffer's data (e.g., set a_tmem_data = A_buf.buffer.data or A_buf.buffer.data
depending on BufferLoad shape) instead of falling through to the exception; also
change the raised exception from ValueError to TypeError for unsupported A_buf
types to match type-mismatch semantics (update the branch that currently checks
isinstance(A_buf, BufferRegion)/Buffer and the final raise to reflect these
changes).
---
Nitpick comments:
In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py`:
- Around line 157-178: Add a new CUDA test function (decorated with
`@tilelang.testing.requires_cuda` and
`@tilelang.testing.requires_cuda_compute_version`(10)) that calls
assert_chained_gemm with an M value that is not a multiple of 128 (e.g., use
M=64) to exercise smaller TCGEN5MMA atom_m cases; reference the existing test
functions test_chained_gemm_128/test_chained_gemm_256 for placement and reuse
the same argument pattern but replace the first M dimension with 64 so
assert_chained_gemm is invoked with a non-128 M to validate TMEM
offset/interleave geometry.
In `@tilelang/intrinsics/tcgen05_macro_generator.py`:
- Around line 444-470: The access_ptr_from helper is duplicated (the function
access_ptr_from in this file and the identical one inside tcgen05mma); extract
it into a single shared function (e.g., a module-level helper named
access_ptr_from) and replace both copies with calls to that shared helper;
update any local references to Buffer, BufferLoad, BufferRegion, and calls that
previously used the nested function so behavior remains identical.
- Around line 405-425: The B-descriptor parameter computation in
tcgen05_macro_generator.py (the block computing b_leading_byte_offset,
b_stride_byte_offset and bk_atom_size) is duplicated from tcgen05mma; extract
that logic into a shared helper function (e.g., _compute_b_descriptor_params)
that accepts the inputs used in the block (b_is_k_major, b_swizzle_mode,
elems_in_bytes, n_dim, k_dim, b_swizzle_atom_elems, micro_size_k, etc.) and
returns the computed values (b_leading_byte_offset, b_stride_byte_offset,
bk_atom_size); then replace the duplicated code in both
tcgen05_macro_generator.py and tcgen05mma to call _compute_b_descriptor_params
and remove the duplicate blocks so both paths use the single helper.
In `@tilelang/tileop/gemm/gemm_tcgen05.py`:
- Around line 144-157: Rename the misleading local functions _gemm_ss and
_gemm_ss_cond to neutral names reflecting they handle both SS and TS (for
example _gemm_mma and _gemm_mma_cond); update the function definitions (def
_gemm_ss_cond -> def _gemm_mma_cond, def _gemm_ss -> def _gemm_mma), update the
conditional return to reference the new names, and ensure any local references
or calls to _gemm_ss/_gemm_ss_cond within this scope (including the
analyzer.can_prove branch and the mma_emitter invocation) are changed
accordingly so the behavior is unchanged but the names match their generalized
role.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds missing Blackwell (SM100) TileLang support for storing from registers back to Tensor Memory (tcgen05.st) and introduces the TCGEN5MMA TS variant (A from TMEM, B from shared), enabling chained GEMM pipelines (SS → cast → ST → TS) used by FlashAttention-style kernels.
Changes:
- Add PTX inline-asm templates for
tcgen05.stand wire them into SM100 copy helpers. - Extend TCGEN5MMA lowering/macros to dispatch and emit the TS (TMEM-Shared) MMA variant.
- Add end-to-end bf16 chained GEMM tests and a runnable SM100 example.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tilelang/tileop/gemm/gemm_tcgen05.py | Adds TS variant layout inference + allows TS in lowering path. |
| tilelang/tileop/gemm/gemm_base.py | Adds is_gemm_ts() and adjusts in_dtype derivation for TS. |
| tilelang/intrinsics/tcgen05_macro_generator.py | Adds tcgen05mma_ts emission and TS dispatch in tcgen05mma. |
| testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py | New bf16 tests covering SS baseline and chained SS→TS GEMM with source assertions. |
| src/tl_templates/cuda/tcgen_05_st.h | New tcgen05.st inline PTX templates for multiple shapes/modifiers. |
| src/tl_templates/cuda/copy_sm100.h | Adds tcgen05_st_* wrappers + recursive splitting core for stores. |
| src/op/gemm_py.cc | Adjusts dtype selection for allowing TCGEN5MMA when A is from TMEM (TS). |
| src/op/copy.cc | Extends TMEM copy lowering to support the store direction and modifier behavior. |
| src/layout/tcgen05_layout.h | Declares new getTcgen05MetaSt_* metadata helpers for stores. |
| src/layout/tcgen05_layout.cc | Implements getTcgen05MetaSt_* metadata for store instructions. |
| examples/gemm_sm100/gemm_tcgen5mma_ts.py | New example demonstrating chained GEMM using tcgen05.st + MMA TS. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+453
to
+456
| if isinstance(indice, (tvm.tir.IntImm, tvm.tir.PrimExpr)): | ||
| offset += indice * stride | ||
| elif isinstance(indice, tvm.tir.Ramp): | ||
| offset += indice.base * stride |
There was a problem hiding this comment.
In access_ptr_from, the check isinstance(indice, (tvm.tir.IntImm, tvm.tir.PrimExpr)) makes the tvm.tir.Ramp branch unreachable because Ramp is a PrimExpr. This will incorrectly treat Ramp indices as scalar expressions when computing the offset. Handle Ramp before the generic PrimExpr case, or narrow the generic check to scalar index types only.
Suggested change
| if isinstance(indice, (tvm.tir.IntImm, tvm.tir.PrimExpr)): | |
| offset += indice * stride | |
| elif isinstance(indice, tvm.tir.Ramp): | |
| offset += indice.base * stride | |
| if isinstance(indice, tvm.tir.Ramp): | |
| offset += indice.base * stride | |
| elif isinstance(indice, (tvm.tir.IntImm, tvm.tir.PrimExpr)): | |
| offset += indice * stride |
Member
|
@Hale423 Please run |
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/tl_templates/cuda/tcgen_05_st.h (1)
10-1300: Consider code-generating these templates to reduce ~1300 lines of copy-paste risk.This file contains 8 nearly identical class specializations with mechanical variation (shape string, register count formula, N upper bound, optional
unpack::16b). Each is a longif constexprchain hand-enumerating operand placeholders up to%128. A script-generated or macro-based approach (as is common for PTX wrapper headers) would eliminate the risk of off-by-one placeholder errors in the larger N cases and simplify future maintenance (e.g., adding new shapes).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/tl_templates/cuda/tcgen_05_st.h` around lines 10 - 1300, The file repeats long hand-rolled if constexpr chains in tmem_st_32dp32bNx::copy, tmem_st_16dp64bNx::copy, tmem_st_16dp128bNx::copy, tmem_st_16dp256bNx::copy and their Unpack16 specializations with manual "%1..%128" placeholders and duplicated operand lists; replace these manual chains with a generator or macro that builds the assembly mnemonic, the brace-list of operand placeholders, and the corresponding "r"(src_ptr[i]) operand list programmatically (including the Unpack16 variant of the mnemonic and per-class N upper bound), or add a small constexpr helper that formats the asm template and fills operands based on template parameters (used by all classes and the composite wrappers tmem_st_32dp64bNx/tmem_st_32dp128bNx/tmem_st_32dp256bNx), ensuring src_ptr indexing and dst_addr offsets follow the same formulas currently used and preserving the static_assert constraints.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@src/tl_templates/cuda/tcgen_05_st.h`:
- Around line 10-1300: The file repeats long hand-rolled if constexpr chains in
tmem_st_32dp32bNx::copy, tmem_st_16dp64bNx::copy, tmem_st_16dp128bNx::copy,
tmem_st_16dp256bNx::copy and their Unpack16 specializations with manual
"%1..%128" placeholders and duplicated operand lists; replace these manual
chains with a generator or macro that builds the assembly mnemonic, the
brace-list of operand placeholders, and the corresponding "r"(src_ptr[i])
operand list programmatically (including the Unpack16 variant of the mnemonic
and per-class N upper bound), or add a small constexpr helper that formats the
asm template and fills operands based on template parameters (used by all
classes and the composite wrappers
tmem_st_32dp64bNx/tmem_st_32dp128bNx/tmem_st_32dp256bNx), ensuring src_ptr
indexing and dst_addr offsets follow the same formulas currently used and
preserving the static_assert constraints.
Hale423
force-pushed
the
feat/tcgen5mma-gemm-ts
branch
from
February 24, 2026 01:37
2284392 to
bfd10cd
Compare
Contributor
Author
|
Hi @LeiWang1999 @Rachmanino , looks like the format is refreshed, please kindly take a review, thanks. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py (1)
100-104:⚠️ Potential issue | 🟡 Minor
B1[0, ...]still assumes full N1 tile insidechained_gemm.Only
assert_chained_gemmguardsbN1 == N1; direct calls tochained_gemmcan still produce wrong results. Add a local guard (or tileB1).🛡️ Suggested guard
def chained_gemm( M, N1, N2, K, bM, bN1, bN2, bK, in_dtype, out_dtype, accum_dtype, threads, ): """Chained GEMM: SS → tcgen05.ld → cast → tcgen05.st → MMA TS → output.""" + assert bN1 == N1, ( + f"chained_gemm requires bN1 == N1 (B1 not tiled over N1), got bN1={bN1}, N1={N1}" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py` around lines 100 - 104, The test's Stage 1 loop calls chained_gemm and uses B1[0, ...] assuming the full N1 tile; add a local guard or tile so bN1 == N1 before calling chained_gemm: either assert bN1 == N1 locally or copy/slice the B1 tile into a N1-sized buffer (e.g., create B1_tile = B1[0, k * bK : k * bK + bK] or similar) and pass that into chained_gemm; reference the chained_gemm invocation and the existing assert_chained_gemm helper to ensure the same safety is applied for direct calls that bypass assert_chained_gemm and prevent incorrect results when bN1 != N1.tilelang/intrinsics/tcgen05_macro_generator.py (1)
478-484:⚠️ Potential issue | 🟡 MinorSupport
BufferLoadfor TMEM A in TS path.
is_tensor_memoryacceptsBufferLoad, so TS dispatch can reach this branch and currently fails with aValueError. Add aBufferLoadcase (and preferTypeErrorfor unsupported types).🔧 Suggested fix
- if isinstance(A_buf, BufferRegion): - a_tmem_data = A_buf.buffer.data - elif isinstance(A_buf, Buffer): - a_tmem_data = A_buf.data - else: - raise ValueError(f"Unsupported A_buf type for TS variant: {type(A_buf)}") + if isinstance(A_buf, BufferRegion): + a_tmem_data = A_buf.buffer.data + elif isinstance(A_buf, Buffer): + a_tmem_data = A_buf.data + elif isinstance(A_buf, BufferLoad): + a_tmem_data = A_buf.buffer.data + else: + raise TypeError(f"Unsupported A_buf type for TS variant: {type(A_buf)}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tilelang/intrinsics/tcgen05_macro_generator.py` around lines 478 - 484, The TMEM A resolution branch currently handles BufferRegion and Buffer but not BufferLoad, causing a ValueError when a BufferLoad reaches this TS path; add an isinstance(A_buf, BufferLoad) branch that resolves a_tmem_data the same way as BufferRegion (e.g., A_buf.buffer.data or the appropriate attribute on BufferLoad), and change the final exception to raise a TypeError(f"Unsupported A_buf type for TS variant: {type(A_buf)}") instead of ValueError; update the logic around the a_tmem_data assignment so A_buf, BufferRegion, Buffer, and BufferLoad are all handled consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/gemm_sm100/gemm_tcgen5mma_ts.py`:
- Around line 96-104: The loop reads B1[0, k * block_K] unconditionally which is
unsafe when the last tile has block_N1 < N1; update the inner pipeline (the k
loop with T.copy(B1[0, k * block_K], B1_shared) and subsequent T.gemm usage) to
guard accesses: either tile B1 so the copy uses a per-tile row index instead of
constant 0, or add a bounds check (e.g., conditional on the target column/row
being < N1 or block_N1) before T.copy and T.gemm so you don't read past
available columns; adjust the related symbols B1, block_N1, N1, the
T.copy(B1...) call and the T.gemm that consumes B1_shared accordingly.
---
Duplicate comments:
In `@testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py`:
- Around line 100-104: The test's Stage 1 loop calls chained_gemm and uses B1[0,
...] assuming the full N1 tile; add a local guard or tile so bN1 == N1 before
calling chained_gemm: either assert bN1 == N1 locally or copy/slice the B1 tile
into a N1-sized buffer (e.g., create B1_tile = B1[0, k * bK : k * bK + bK] or
similar) and pass that into chained_gemm; reference the chained_gemm invocation
and the existing assert_chained_gemm helper to ensure the same safety is applied
for direct calls that bypass assert_chained_gemm and prevent incorrect results
when bN1 != N1.
In `@tilelang/intrinsics/tcgen05_macro_generator.py`:
- Around line 478-484: The TMEM A resolution branch currently handles
BufferRegion and Buffer but not BufferLoad, causing a ValueError when a
BufferLoad reaches this TS path; add an isinstance(A_buf, BufferLoad) branch
that resolves a_tmem_data the same way as BufferRegion (e.g., A_buf.buffer.data
or the appropriate attribute on BufferLoad), and change the final exception to
raise a TypeError(f"Unsupported A_buf type for TS variant: {type(A_buf)}")
instead of ValueError; update the logic around the a_tmem_data assignment so
A_buf, BufferRegion, Buffer, and BufferLoad are all handled consistently.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
examples/gemm_sm100/gemm_tcgen5mma_ts.pysrc/op/copy.ccsrc/op/gemm_py.ccsrc/tl_templates/cuda/copy_sm100.hsrc/tl_templates/cuda/tcgen_05_st.htesting/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.pytilelang/intrinsics/tcgen05_macro_generator.pytilelang/tileop/gemm/gemm_tcgen05.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/op/gemm_py.cc
Comment on lines
+96
to
+104
| for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=1): | ||
| T.copy(A[by * block_M, k * block_K], A_shared) | ||
| T.copy(B1[0, k * block_K], B1_shared) | ||
| T.gemm(A_shared, B1_shared, S_tmem, transpose_A=False, transpose_B=True, mbar=mbar1, wg_wait=-1, clear_accum=k == 0) | ||
| T.mbarrier_wait_parity(mbar1, k % 2) | ||
|
|
||
| T.copy(S_tmem, S_local) | ||
| T.copy(S_local, P_local) | ||
| T.copy(P_local, P_tmem) |
Contributor
There was a problem hiding this comment.
Guard against B1[0, ...] when block_N1 < N1.
B1 is always read from row 0, so this example is only correct for full-row tiles. Add a guard (or tile B1) to avoid silent misuse.
🛡️ Suggested guard
def chained_gemm(
M,
N1,
N2,
K,
block_M,
block_N1,
block_N2,
block_K,
in_dtype,
out_dtype,
accum_dtype,
threads,
):
"""Test 2: Full chained GEMM with tcgen05.st"""
+ assert block_N1 == N1, (
+ f"chained_gemm assumes block_N1 == N1 (B1 not tiled), got block_N1={block_N1}, N1={N1}"
+ )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/gemm_sm100/gemm_tcgen5mma_ts.py` around lines 96 - 104, The loop
reads B1[0, k * block_K] unconditionally which is unsafe when the last tile has
block_N1 < N1; update the inner pipeline (the k loop with T.copy(B1[0, k *
block_K], B1_shared) and subsequent T.gemm usage) to guard accesses: either tile
B1 so the copy uses a per-tile row index instead of constant 0, or add a bounds
check (e.g., conditional on the target column/row being < N1 or block_N1) before
T.copy and T.gemm so you don't read past available columns; adjust the related
symbols B1, block_N1, N1, the T.copy(B1...) call and the T.gemm that consumes
B1_shared accordingly.
Collaborator
|
Really appreciate your contribution! Also sry for the delay, i'll review this soon |
Rachmanino
reviewed
Feb 24, 2026
Rachmanino
reviewed
Feb 24, 2026
…getTcgen05MetaLd_*
2 tasks done
Contributor
Author
|
Seems like these two failing checks are not owned by this feature |
LeiWang1999
previously approved these changes
Feb 26, 2026
LeiWang1999
left a comment
LeiWang1999
left a comment
Member
There was a problem hiding this comment.
LGTM, but copilot's suggestion should be applied.
Added a TODO comment for future support of tcgen05.cp in conjunction with LowerTmemCopy.
This was referenced Mar 18, 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
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.
Fixes #1862
Summary
This PR implements two missing Blackwell (SM100) primitives in TileLang:
tcgen05.st(Register → Tensor Memory Store): Enables writing data from registers back into Tensor Memory (TMEM), which is the reverse of the existingtcgen05.ld. This is essential for chained GEMM patterns where intermediate results need type conversion (e.g., fp32→bf16) before being fed into the next MMA stage.tcgen05mma_ts(TMEM-Shared MMA variant): Enables GEMM where operand A is read directly from TMEM while operand B comes from shared memory. Combined withtcgen05.st, this unlocks the chained GEMM pattern:SS GEMM → ld → cast → st → TS GEMM, which is the core building block for Flash Attention on Blackwell.Key Changes
New files:
src/tl_templates/cuda/tcgen_05_st.h— PTX inline assembly templates for alltcgen05.stshape variants (32dp32b, 32dp64b, 32dp128b, 32dp256b) with optionalunpack::16bmodifiertesting/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py— Test cases covering SS baseline, chained GEMM (128×128, 256×128), and non-power-of-2 K dimension (K=192)examples/gemm_sm100/gemm_tcgen5mma_ts.py— Example demonstrating chained MMA: D = (A × B1) × B2Modified files:
src/op/copy.cc— ExtendedLowerTmemCopyto supporttcgen05.stdirection; critical fix: disabledunpack::16bfor bf16 stores so data is written packed (2 bf16 per 32-bit TMEM column), matching MMA TS's read expectationsrc/layout/tcgen05_layout.{cc,h}— AddedgetTcgen05MetaSt_*functions providing Fragment metadata for store operationssrc/tl_templates/cuda/copy_sm100.h— Addedtcgen05_st_corerecursive splitting template and wrapper functions for all store shape variantssrc/op/gemm_py.cc— UpdatedallowTcgen5Mmato correctly determineab_dtypefor TS variant (A from TMEM has accumulator dtype, B from shared has input dtype)tilelang/intrinsics/tcgen05_macro_generator.py— Addedtcgen05mma_tsmethod andmake_mma_store_layoutfor TMEM layout inferencetilelang/tileop/gemm/gemm_base.py— Addedis_gemm_ts()detection and correctin_dtypederivation for TS varianttilelang/tileop/gemm/gemm_tcgen05.py— Added TS variant support ininfer_layoutandlowerbf16 Pack/Unpack Fix
The most critical fix addresses a data corruption issue with bf16 in TMEM:
tcgen05.stwas generatingunpack::16b, which unpacks each uint32 register into 2 separate TMEM columns (128 columns total). ButMMA TSexpects bf16 data packed (2 values per column, 64 columns).unpack::16bis disabled (use_pack_unpack_modifier = false), so 64 uint32 registers (each containing 2 packed bf16) are written directly to 64 TMEM columns.effective_chunksstill halves correctly for the register count.Test Plan
tcgen05_st_corerecursive splitting (192 → 128+64)tcgen05mma_tsandtcgen05_stappear in outputSummary by CodeRabbit
New Features
Examples
Tests
Documentation
Bug Fixes