[CUDA] Pack logical TMEM buffers into shared tcgen05.alloc arenas - #2831
Merged
Conversation
LowerSharedTmem gave every T.alloc_tmem buffer its own tcgen05.alloc, rounded up to a power of two. Small buffers therefore each burned a 32-column allocation: a 384-column accumulator next to three 4-column block-scale buffers needs 396 columns but asked the hardware for 512 + 32 + 32 + 32 = 608 of the 512 a CTA has, so that kernel could not be expressed at all. Kernels worked around it by declaring one wide buffer and slicing scale factors out of it by hand. Plan the allocations instead. PlanTmemArenas packs logical buffers into shared allocations, and each buffer's column offset inside the one it shares is folded into every address it forms -- both the encoded coordinate of a BufferLoad, which tcgen05.ld/st consume, and the (base, offset) argument pairs of the tl.ptx_tcgen05_* intrinsics. Buffers at least 32 columns wide keep the 32-column alignment tcgen05.alloc would have given them alone; narrower ones use the 4-column granularity of tcgen05.cp.32x128b.warpx4, the same boundary CUTLASS and DeepGEMM place these operands on. A buffer joins an existing allocation only when that strictly reduces the total column count, so packing can never grow a kernel's TMEM footprint and a kernel whose buffers already fit lowers exactly as before. The mxfp8 block-scaled GEMM drops from 320 to 288 columns; flash attention and single-accumulator GEMMs are untouched. Buffers with an explicit T.deallocate_tmem keep their own allocation, since releasing one early would release whatever shares it. A base address that reaches an operation this pass cannot offset is rejected rather than silently addressing the start of the allocation. A block that allocates more than 512 columns is now reported at compile time instead of failing inside tcgen05.alloc on device.
|
👋 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
📝 WalkthroughWalkthroughTMEM lowering now packs compatible buffers into shared physical arenas. It tracks column offsets, preserves explicit deallocation lifetimes, emits arena-based allocation sequences, rebases supported accesses, and validates capacity. CUDA tests cover packing, addressing, allocation order, and roundtrip behavior. ChangesTMEM arena packing and lowering
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant TMEMBufferDiscovery
participant ArenaPlanner
participant TMEMLowering
participant TCGEN05Operations
TMEMBufferDiscovery->>ArenaPlanner: Provide ordered buffers and column sizes
ArenaPlanner->>TMEMLowering: Provide arenas and per-buffer offsets
TMEMLowering->>TCGEN05Operations: Emit rebased base and offset operands
TCGEN05Operations-->>TMEMLowering: Execute packed buffer accesses
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Rachmanino
changed the title
tcgen05.alloc arenas
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/cuda/transform/lower_shared_tmem.cc (1)
515-525: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
std::stable_sortfor the allocation order.
std::sortis not stable, so arenas with equalnum_cols_allocatedcan be emitted in an order that depends on the standard library implementation.PlanTmemArenasalready usesstd::stable_sortto keep the plan reproducible. Match that here so the emittedtcgen05.allocsequence stays deterministic.♻️ Proposed change
- std::sort(init_mtmem_calls_.begin(), init_mtmem_calls_.end(), - compare_by_num_cols_desc); + std::stable_sort(init_mtmem_calls_.begin(), init_mtmem_calls_.end(), + compare_by_num_cols_desc);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuda/transform/lower_shared_tmem.cc` around lines 515 - 525, Replace std::sort with std::stable_sort in the allocation ordering around compare_by_num_cols_desc, preserving the descending num_cols ordering while retaining the original order for equal-sized allocations and matching PlanTmemArenas’ deterministic behavior.
🤖 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/cuda/transform/lower_shared_tmem.cc`:
- Around line 315-338: Update GetNumB32ColsRequired to retain the analyzer
bounds for both shape dimensions as int64_t, and validate each extent’s lower
bound and upper bound before converting to int; reject non-constant or
non-positive extents and values exceeding the TMEM row/column limits. Only
narrow to int after validation, preserving the existing row and rounded 32-bit
column calculations.
---
Nitpick comments:
In `@src/cuda/transform/lower_shared_tmem.cc`:
- Around line 515-525: Replace std::sort with std::stable_sort in the allocation
ordering around compare_by_num_cols_desc, preserving the descending num_cols
ordering while retaining the original order for equal-sized allocations and
matching PlanTmemArenas’ deterministic behavior.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 17aaa9bd-245e-4504-be22-403faf0c7d5a
📒 Files selected for processing (3)
src/cuda/transform/lower_shared_tmem.cctesting/python/language/test_tilelang_language_tmem_copy.pytesting/python/transform/test_tilelang_transform_tmem_arena_packing.py
Comment on lines
+315
to
+338
| int GetNumB32ColsRequired(const Buffer &buffer) const { | ||
| ICHECK_EQ(buffer->shape.size(), 2U); | ||
|
|
||
| auto analyzer = std::make_shared<arith::Analyzer>(); | ||
| arith::ConstIntBound phy_col_bounds = | ||
| analyzer->const_int_bound(buffer->shape[1]); | ||
| int num_value_cols_required = phy_col_bounds->max_value; | ||
| arith::Analyzer analyzer; | ||
| int num_rows_required = | ||
| analyzer.const_int_bound(buffer->shape[0])->max_value; | ||
| ICHECK(num_rows_required <= kTmemNumDatapaths) | ||
| << "The number of rows required for tmem buffer " << buffer->name | ||
| << " is " << num_rows_required << ", which exceeds the maximum of " | ||
| << kTmemNumDatapaths << " rows"; | ||
|
|
||
| int num_value_cols_required = | ||
| analyzer.const_int_bound(buffer->shape[1])->max_value; | ||
| // Layout column coordinates count values of buffer->dtype; PTX TMEM | ||
| // allocation counts 32-bit columns. Round up so a final partially | ||
| // occupied b32 column is included. | ||
| int num_cols_required = | ||
| (num_value_cols_required * GetValueBitWidth(buffer->dtype) + 31) / 32; | ||
| ICHECK(num_cols_required <= 512) | ||
| ICHECK(num_cols_required <= kTmemNumColumns) | ||
| << "The number of columns required for tmem buffer " << buffer->name | ||
| << " is " << num_cols_required | ||
| << ", which exceeds the maximum of 512 columns"; | ||
| << " is " << num_cols_required << ", which exceeds the maximum of " | ||
| << kTmemNumColumns << " columns"; | ||
| return num_cols_required; | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Keep the const-int bounds in int64_t and check the lower bound.
ConstIntBound::max_value is int64_t. A TMEM buffer whose shape is not a constant yields the analyzer's positive-infinity sentinel. The narrowing conversion to int can then produce a small or negative value, both ICHECKs pass, and the arena is sized from a wrong column count. Validate in int64_t and reject a non-constant or non-positive extent before narrowing.
🛡️ Proposed hardening
- arith::Analyzer analyzer;
- int num_rows_required =
- analyzer.const_int_bound(buffer->shape[0])->max_value;
- ICHECK(num_rows_required <= kTmemNumDatapaths)
+ arith::Analyzer analyzer;
+ int64_t num_rows_required =
+ analyzer.const_int_bound(buffer->shape[0])->max_value;
+ ICHECK(num_rows_required >= 1 && num_rows_required <= kTmemNumDatapaths)
<< "The number of rows required for tmem buffer " << buffer->name
<< " is " << num_rows_required << ", which exceeds the maximum of "
<< kTmemNumDatapaths << " rows";
- int num_value_cols_required =
- analyzer.const_int_bound(buffer->shape[1])->max_value;
+ int64_t num_value_cols_required =
+ analyzer.const_int_bound(buffer->shape[1])->max_value;
+ ICHECK_GE(num_value_cols_required, 1)
+ << "TMEM buffer " << buffer->name
+ << " must have a constant column extent, but got " << buffer->shape[1];
// Layout column coordinates count values of buffer->dtype; PTX TMEM
// allocation counts 32-bit columns. Round up so a final partially
// occupied b32 column is included.
- int num_cols_required =
+ int64_t num_cols_required =
(num_value_cols_required * GetValueBitWidth(buffer->dtype) + 31) / 32;
ICHECK(num_cols_required <= kTmemNumColumns)
<< "The number of columns required for tmem buffer " << buffer->name
<< " is " << num_cols_required << ", which exceeds the maximum of "
<< kTmemNumColumns << " columns";
- return num_cols_required;
+ return static_cast<int>(num_cols_required);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/cuda/transform/lower_shared_tmem.cc` around lines 315 - 338, Update
GetNumB32ColsRequired to retain the analyzer bounds for both shape dimensions as
int64_t, and validate each extent’s lower bound and upper bound before
converting to int; reject non-constant or non-positive extents and values
exceeding the TMEM row/column limits. Only narrow to int after validation,
preserving the existing row and rounded 32-bit column calculations.
Pins the plan the packer produces, not just that it produces one: - 396 columns across four buffers become one 512-column allocation at offsets 0 / 384 / 388 / 392, matching DeepGEMM's hand-written layout for this kernel column for column - block-scale factors share an allocation while the accumulator keeps its own, because packing all three would cost more than it saves - equal-cost cases keep separate allocations, so a kernel that already fits is untouched - allocations are issued widest first, as PTX requires - packing never allocates more than separate allocations would, over six shape combinations - an explicit T.deallocate_tmem keeps its own allocation - a dynamic pipeline-stage coordinate still carries the column offset - a block over the 512-column budget is reported
Covers the tcgen05.st/ld address path for a buffer packed at a nonzero column offset, which no existing kernel exercises: the block-scale buffers that share an allocation today are never read back into registers. The wide buffer's roundtrip is the assertion that matters. The narrow one would still round-trip if its store and its load dropped the column offset together -- but its store would land on the wide buffer's first columns, so only the wide readback catches it. Verified by mutation: dropping the offset in LowerSharedTmem fails this test.
Gate every TMEM test the same way, matching the decorators the tcgen05 language and kernel tests already carry. These three lower IR for a hardcoded sm_100 target and need no device, so this does trade running them on any machine for consistency with the rest of the TMEM suite.
Rachmanino
force-pushed
the
feat/tmem-arena-packing
branch
from
July 31, 2026 11:42
72c32b3 to
50276e5
Compare
Collaborator
|
LGTM cc @LeiWang1999 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
#2827
Problem
LowerSharedTmemgives everyT.alloc_tmembuffer its owntcgen05.alloc, rounded up to a power of two with a minimum of 32 columns. Narrow buffers therefore each burn a full 32-column allocation even when they need 4.That is not just waste — it makes some kernels inexpressible. A 384-column accumulator (three 128-column pipeline stages) next to three 4-column block-scale buffers needs 396 columns, but as four separate allocations it asks the hardware for
512 + 32 + 32 + 32 = 608of the 512 columns a CTA has. Today nothing reports this; the kernel just compiles and misbehaves on device.The workaround has been to declare one wide TMEM buffer and slice scale factors out of it by hand, passing column offsets manually — which is what a compiler should be doing.
CUTLASS solves the same problem with a bump allocator over one allocation (
cutlass::detail::find_tmem_tensor_col_offset, used bysm100_blockscaled_mma_warpspecialized.hppto placeaccumulator → tCtSFA → tCtSFB), and DeepGEMM hand-writes the equivalent layout.Approach
Plan the allocations before emitting them.
PlanTmemArenaspacks logical TMEM buffers into shared physical allocations, and each buffer's column offset inside the allocation it shares is folded into every address it forms.Two addressing paths carry TMEM addresses today, and both are handled:
BufferLoadon a TMEM buffer — used bytcgen05.ld/tcgen05.st— gains the offset in the encoded coordinate.(base address Var, offset)argument pairs of thetl.ptx_tcgen05_*intrinsics (_mma_ss,_mma_ts,_mma_blockscaled_ss,cp_warpx4), which all codegen to*(uint32_t*)base + offset, get the offset added to the second half of the pair.A base address that reaches anything else is rejected with a diagnostic rather than silently addressing the start of the allocation.
Packing policy. A buffer joins an existing allocation only when that strictly reduces the total number of allocated columns. Each placement therefore costs at most the buffer's own standalone allocation, so:
Widest-first placement keeps wide buffers on 32-column boundaries.
Alignment. A buffer at least 32 columns wide starts on a 32-column boundary — the same alignment
tcgen05.allocwould have given it alone — so sharing an allocation cannot leave an accumulator or a TMEM-resident A operand less aligned than it is today. Narrower buffers use the 4-column granularity oftcgen05.cp.32x128b.warpx4, the boundary CUTLASS and DeepGEMM place these very operands on.Lifetimes. A buffer with an explicit
T.deallocate_tmemkeeps its own allocation: releasing it early would release whatever shares it. Every other buffer in a block is live for the whole block, so this needs no liveness analysis to be correct. It does leave value on the table — buffers with disjoint live ranges could reuse the same columns outright instead of only filling each other's power-of-two padding — which is marked as a TODO on the planner.New diagnostic. A block that allocates more than 512 columns is now reported at compile time (unless it manages a lifetime by hand, where the peak is not the sum).
Effect on existing kernels
Measured by capturing
device_kernel.cufor every TMEM example under baseline and under this change (separateTILELANG_CACHE_DIRs, so no cache hit crosses versions) and diffing:gemm_mxfp8_blockscaled_1d1dgrouped_gemm_mxfp8_blockscaled_1d1dfp8_fp4_gemm_1d1d_sm100mha_fwd_bshd,gqa_fwd_bshdgemm_tcgen5mma{,_ws,_ws_persistent,_ws_clc}All of them still pass their own correctness checks.
The whole difference in the generated CUDA for the block-scaled GEMM is:
One fewer address word, one fewer alloc/dealloc pair, and the MMA loads one base address instead of two — so this removes a shared-memory load rather than adding one.
Wall-clock is not a usable signal on the machine I measured on (shared with another tenant): three interleaved rounds of the block-scaled GEMM gave baseline 3032 / 2931 / 1303 TFLOP/s against 2824 / 3031 / 1324, i.e. ±7% within a round and 2.3x across rounds. The code-level facts above are the claim; the timings only confirm nothing catastrophic.
Tests
New
testing/python/transform/test_tilelang_transform_tmem_arena_packing.py:tcgen05.allocfrom being larger)T.deallocate_tmemkeeps its own allocationNew GPU test in
test_tilelang_language_tmem_copy.py: a store/load roundtrip through two buffers sharing one allocation. The wide buffer's roundtrip is the assertion that matters — the narrow one would still round-trip if its store and load dropped the offset together, but its store would land on the wide buffer's columns.Existing coverage, on a B200:
testing/python/transform/: 380 passed, 7 skippedtcgen05GEMM / sliced-operand / TMEM-copy / atom-MMA language tests plus the bf16-TS and int8 TMEM kernel tests: 58 passed, 1 skippedmha_fwd_bshd,gqa_fwd_bshd,fp8_fp4_gemm_1d1d_sm100, and the threegemm_sm100variantsThe new tests were validated by mutation: dropping the column offset in
LowerSharedTmemfails the GPU roundtrip and three of the planning tests.Notes
LowerSharedTmemrather than in a separate pass, since buffer collection, column sizing and allocation emission are all already there. It is a self-contained function and can be split out if a second consumer appears.tl.ptx_tcgen05_*intrinsic passes a TMEM operand's base address immediately followed by its offset. That convention is now documented at the point where it is depended on, and a base address reaching any other context is a hard error.