[CUDA] Adopt multi-staged buffers in examples - #2836
Merged
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
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:
📝 WalkthroughWalkthroughThe PR generalizes attention kernels from fixed double buffers to stage-indexed storage and synchronization. It consolidates MLA query and output buffers. It also updates layout remapping and CUDA descriptor generation for static and dynamic shared-memory slice offsets, adds sliced WGMMA coverage, and adjusts three benchmark unroll directives. ChangesStaged attention kernels
Static slice descriptor lowering
SM120 benchmark loop directives
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Producer
participant StageBarriers
participant AttentionConsumer
participant OutputBuffer
Producer->>StageBarriers: signal stage readiness after staged KV copy
AttentionConsumer->>StageBarriers: wait for selected stage
AttentionConsumer->>AttentionConsumer: compute attention and accumulate output
AttentionConsumer->>StageBarriers: release the processed stage
AttentionConsumer->>OutputBuffer: write the combined output
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b7bdea6b5c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+440
to
+442
| for (size_t i = 0; i < indices.size(); ++i) { | ||
| multi_dim_indices.Set( | ||
| i, analyzer_->Simplify(indices[i] + multi_dim_indices[i])); |
There was a problem hiding this comment.
When an address_of/tl.access_ptr already has nonzero multidimensional indices and a PTX intrinsic supplies an additional element offset, these lines add the offset's coordinates independently without propagating carries. For example, shape [M, 16], indices [r, 15], and offset 1 are mapped as [r, 16] instead of [r + 1, 0]; passing the out-of-range coordinate through layout->Forward can select the wrong shared-memory address and corrupt ptx_ldmatrix results. Normalize the combined linear position, as the previous implementation did, before applying the layout.
Useful? React with 👍 / 👎.
| << "The original buffer shape must match the layout input shape, but " | ||
| "got " | ||
| << old_shape << " and " << input_shape; | ||
| ICHECK(relax::CanProveShapeEqual(new_shape, output_shape, analyzer_)) |
There was a problem hiding this comment.
When a shared layout has replicate_extent > 1, makeBufferWithLayout intentionally prepends that replicate axis to new_buffer->shape, so new_shape cannot equal layout->OutputShape(). Any remapped access_ptr or address_of for such a buffer now fails this ICHECK during compilation even though this representation was previously supported; the validation and index construction need to account for the extra leading replicate dimension.
Useful? React with 👍 / 👎.
Collaborator
Author
|
@regression-perf |
![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: 4
🧹 Nitpick comments (5)
examples/minference/example_vertical_slash_sparse_attn.py (1)
245-279: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument or assert the two-stage assumption of this prefetch schedule.
The stage index now uses
% num_stages, so the code reads as generic. The schedule is not generic: it primes exactly one block, issues one prefetch per iteration, and callsComputewithcount=1. That leaves exactly one outstandingcp.asyncgroup, which is only correct fornum_stages == 2. If someone raisesnum_stages, the extra buffers stay unused and the wait count stops matching the pipeline depth.Add an assertion next to the
num_stagesdefinition so the constraint fails fast.♻️ Proposed guard
num_stages = 2 + assert num_stages == 2, "the vertical-column prefetch schedule keeps one cp.async group in flight"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/minference/example_vertical_slash_sparse_attn.py` around lines 245 - 279, Add a fail-fast assertion beside the num_stages definition in the surrounding sparse-attention setup, requiring exactly two stages to match the prefetch schedule and Compute count=1 pipeline. Leave the existing stage-indexing and prefetch logic unchanged.tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py (1)
427-512: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate static/dynamic slice-offset resolution logic.
The dynamic-slice, static-slice, and non-sliced branches for
B_base_ptrrepeat the same pattern used forA_base_ptrright below, and the identical block is duplicated again intcgen05_macro_generator.pyandwgmma_sp_macro_generator.py(6 occurrences total). Extract a shared helper, for exampleresolve_slice_base_ptr(base_ptr, slice_byte_offset) -> (base_ptr, is_sliced), and reuse it in all six locations to avoid divergent fixes over time.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py` around lines 427 - 512, Extract the repeated slice-offset handling into a shared `resolve_slice_base_ptr(base_ptr, slice_byte_offset)` helper that returns the adjusted pointer and `is_sliced` flag, preserving dynamic, nonzero static, and zero-offset behavior. Replace the duplicated branches in `init_wgmma_b_desc`, `init_wgmma_a_desc`, and the corresponding four locations in the tcgen05 and wgmma_sp macro generators with this helper.tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py (1)
805-872: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate static/dynamic slice-offset resolution logic.
The dynamic-slice, static-slice, and non-sliced branches for
B_base_ptrandA_base_ptrrepeat the identical eight-line pattern for both descriptors in this file. The same pattern also appears inwgmma_macro_generator.pyandwgmma_sp_macro_generator.py. Extract a shared helper, for exampleresolve_slice_base_ptr(base_ptr, slice_byte_offset) -> (base_ptr, is_sliced), and reuse it across all six call sites. This reduces the risk that a future fix to this logic is applied in some sites but missed in others.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py` around lines 805 - 872, Extract the repeated static, dynamic, and non-sliced offset handling into a shared resolve_slice_base_ptr helper returning the adjusted base pointer and is_sliced flag. Replace the duplicated logic in init_tcgen05_b_desc, init_tcgen05_a_desc, and the corresponding four call sites in wgmma_macro_generator.py and wgmma_sp_macro_generator.py, preserving existing pointer adjustment and dynamic descriptor-offset behavior.tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py (2)
353-362: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate static/dynamic slice-offset resolution logic.
Same pattern as Lines 196-210 in this file, and as
wgmma_macro_generator.py/tcgen05_macro_generator.py. Reuse the same extracted helper here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py` around lines 353 - 362, Extract the repeated B slice-offset resolution in the current generator into the shared helper already used by the analogous logic around lines 196-210 and in wgmma_macro_generator.py/tcgen05_macro_generator.py. Replace the inline isinstance/static-offset branching around B_base_ptr with that helper while preserving dynamic offsets, nonzero static pointer adjustment, and the b_is_sliced result.
196-210: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate static/dynamic slice-offset resolution logic.
The A and B slice-offset branches here repeat the same static/dynamic/zero pattern used in
wgmma_macro_generator.pyandtcgen05_macro_generator.py. Extract a shared helper and reuse it here and at Lines 353-362 in this file.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py` around lines 196 - 210, Extract the repeated slice-offset resolution pattern into a shared helper, covering dynamic offsets, nonzero static offsets, and zero offsets while returning the resolved sliced state and base pointer. Replace the duplicated A/B logic in the current block and the corresponding handling around lines 353-362, and reuse the existing helper from wgmma_macro_generator.py and tcgen05_macro_generator.py where applicable.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/deepseek_mla/example_mla_decode_ws.py`:
- Around line 161-165: Add an assertion near the split producer’s num_stages = 2
definition in examples/deepseek_mla/example_mla_decode_ws.py:161-165 requiring
(seqlen_kv // num_split) to be divisible by block_N * num_stages, preventing
trailing stages beyond NI; the sibling non-split producer at
examples/deepseek_mla/example_mla_decode_ws.py:341-345 requires no direct change
because this shared guard protects both MLA decode producers.
In `@examples/deepseek_v32/README.md`:
- Line 157: Update the documented KV_shared allocation snippet in README.md to
use the combined feature width D + D_tail, matching the staged KV buffer
allocation in sparse_mla_fwd_pipelined.py. Keep the existing num_stages and BI
dimensions unchanged.
In `@src/transform/lower_tile_op.cc`:
- Around line 409-445: Update RemapAccessIndices to validate new_shape against
the buffer shape produced by makeBufferWithLayout, including shared-buffer
replicate_extent, rather than layout->OutputShape(). Preserve the existing
shape-equality validation and index remapping behavior for non-replicated
buffers.
In `@tilelang/language/loop.py`:
- Around line 285-293: Update the validation around the explicit and
unroll-factor handling in the loop unroll logic to also reject an existing
"pragma_unroll_factor" in annotations when explicit is true, including when
explicit is enabled by the explicit argument. Preserve the current
mutual-exclusion error for the unroll_factor argument and ensure conflicting
annotations are not emitted.
---
Nitpick comments:
In `@examples/minference/example_vertical_slash_sparse_attn.py`:
- Around line 245-279: Add a fail-fast assertion beside the num_stages
definition in the surrounding sparse-attention setup, requiring exactly two
stages to match the prefetch schedule and Compute count=1 pipeline. Leave the
existing stage-indexing and prefetch logic unchanged.
In `@tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py`:
- Around line 805-872: Extract the repeated static, dynamic, and non-sliced
offset handling into a shared resolve_slice_base_ptr helper returning the
adjusted base pointer and is_sliced flag. Replace the duplicated logic in
init_tcgen05_b_desc, init_tcgen05_a_desc, and the corresponding four call sites
in wgmma_macro_generator.py and wgmma_sp_macro_generator.py, preserving existing
pointer adjustment and dynamic descriptor-offset behavior.
In `@tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py`:
- Around line 427-512: Extract the repeated slice-offset handling into a shared
`resolve_slice_base_ptr(base_ptr, slice_byte_offset)` helper that returns the
adjusted pointer and `is_sliced` flag, preserving dynamic, nonzero static, and
zero-offset behavior. Replace the duplicated branches in `init_wgmma_b_desc`,
`init_wgmma_a_desc`, and the corresponding four locations in the tcgen05 and
wgmma_sp macro generators with this helper.
In `@tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py`:
- Around line 353-362: Extract the repeated B slice-offset resolution in the
current generator into the shared helper already used by the analogous logic
around lines 196-210 and in wgmma_macro_generator.py/tcgen05_macro_generator.py.
Replace the inline isinstance/static-offset branching around B_base_ptr with
that helper while preserving dynamic offsets, nonzero static pointer adjustment,
and the b_is_sliced result.
- Around line 196-210: Extract the repeated slice-offset resolution pattern into
a shared helper, covering dynamic offsets, nonzero static offsets, and zero
offsets while returning the resolved sliced state and base pointer. Replace the
duplicated A/B logic in the current block and the corresponding handling around
lines 353-362, and reuse the existing helper from wgmma_macro_generator.py and
tcgen05_macro_generator.py where applicable.
🪄 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: d1b39319-1029-40ec-b40c-7545902fc99c
📒 Files selected for processing (22)
examples/deepseek_mla/example_mla_decode_ws.pyexamples/deepseek_v32/README.mdexamples/deepseek_v32/sparse_mla_bwd.pyexamples/deepseek_v32/sparse_mla_fwd.pyexamples/deepseek_v32/sparse_mla_fwd_pipelined.pyexamples/deepseek_v32/sparse_mla_fwd_seesaw.pyexamples/flash_attention_sm100/gqa_fwd_bshd.pyexamples/flash_attention_sm100/mha_fwd_bshd.pyexamples/minference/example_vertical_slash_sparse_attn.pymaint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.pysrc/backend/common/op/reduce.hsrc/transform/loop_partition.ccsrc/transform/lower_tile_op.ccsrc/transform/unroll_loop.cctesting/python/language/test_tilelang_language_wgmma_gemm.pytesting/python/transform/test_tilelang_transform_unroll_loop.pytilelang/cuda/intrinsics/macro/tcgen05_macro_generator.pytilelang/cuda/intrinsics/macro/wgmma_macro_generator.pytilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.pytilelang/cuda/pipeline.pytilelang/language/loop.pytilelang/language/tir/ir.py
💤 Files with no reviewable changes (1)
- src/transform/loop_partition.cc
| # ... load KV into buffer 1 | ||
| T.cp_async_barrier_noinc(bar_k_1_ready[0]) | ||
| num_stages = 2 | ||
| KV_shared = T.alloc_shared([num_stages, BI, D], dtype) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Align the documented KV buffer shape with the kernel.
sparse_mla_fwd_pipelined.py line 83 allocates the staged KV buffer with the combined feature width D + D_tail. The snippet shows only D.
📝 Proposed fix
num_stages = 2
-KV_shared = T.alloc_shared([num_stages, BI, D], dtype)
+KV_shared = T.alloc_shared([num_stages, BI, D + D_tail], dtype)
bar_k_ready = T.alloc_barrier(arrive_count=[128] * num_stages)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/deepseek_v32/README.md` at line 157, Update the documented KV_shared
allocation snippet in README.md to use the combined feature width D + D_tail,
matching the staged KV buffer allocation in sparse_mla_fwd_pipelined.py. Keep
the existing num_stages and BI dimensions unchanged.
Comment on lines
+409
to
445
| Array<PrimExpr> RemapAccessIndices(const Array<PrimExpr> &indices, | ||
| const Array<PrimExpr> &old_shape, | ||
| const Array<PrimExpr> &new_shape, | ||
| const Layout &layout, | ||
| const Optional<PrimExpr> &offset) { | ||
| ICHECK_EQ(indices.size(), old_shape.size()) | ||
| << "The access rank must match the original buffer rank, but got " | ||
| << indices << " and " << old_shape; | ||
| ICHECK(!old_shape.empty()) | ||
| << "Layout-remapped access pointers do not support scalar buffers"; | ||
| const Array<PrimExpr> input_shape = layout->InputShape(); | ||
| const Array<PrimExpr> output_shape = layout->OutputShape(); | ||
| ICHECK(relax::CanProveShapeEqual(old_shape, input_shape, analyzer_)) | ||
| << "The original buffer shape must match the layout input shape, but " | ||
| "got " | ||
| << old_shape << " and " << input_shape; | ||
| ICHECK(relax::CanProveShapeEqual(new_shape, output_shape, analyzer_)) | ||
| << "The remapped buffer shape must match the layout output shape, but " | ||
| "got " | ||
| << new_shape << " and " << output_shape; | ||
|
|
||
| // Delinearize only the additional offset. Keeping the original indices in | ||
| // multidimensional form preserves slice-local expressions for the layout. | ||
| PrimExpr remaining_offset = | ||
| analyzer_->Simplify(offset.value_or(make_zero(indices[0].dtype()))); | ||
| Array<PrimExpr> multi_dim_indices; | ||
| for (int i = static_cast<int>(old_shape.size()) - 1; i >= 0; --i) { | ||
| multi_dim_indices.insert(multi_dim_indices.begin(), | ||
| floormod(remaining_offset, old_shape[i])); | ||
| remaining_offset = floordiv(remaining_offset, old_shape[i]); | ||
| } | ||
| for (size_t i = 0; i < indices.size(); ++i) { | ||
| multi_dim_indices.Set( | ||
| i, analyzer_->Simplify(indices[i] + multi_dim_indices[i])); | ||
| } | ||
| return layout->Forward(multi_dim_indices); | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for replicate_extent handling and any tests/examples exercising it together with access-ptr remapping.
rg -n "replicate_extent" src/transform/lower_tile_op.cc
rg -n "IsSharedBuffer" src -g '*.cc' -g '*.h'
rg -n "kLayoutMap" src -g '*.cc' -g '*.h' -C2Repository: tile-ai/tilelang
Length of output: 13113
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== lower_tile_op.cc relevant sections =="
sed -n '1,180p' src/transform/lower_tile_op.cc
echo "--"
sed -n '300,470p' src/transform/lower_tile_op.cc
echo "--"
sed -n '470,560p' src/transform/lower_tile_op.cc
echo "== replicate_extent usages in repo =="
rg -n "replicate_extent|MakeBufferWithLayoutBuffer|makeBufferWithLayout" src test tests examples 2>/dev/null || true
echo "== buffer access remapping functions names =="
rg -n "RemapAccessIndices|HandleAccessPtrAndOffset|address_of|tvm_access_ptr|tl::access_ptr|access_ptr" src/transform/lower_tile_op.cc src -g '*.cc' -g '*.h' -C1Repository: tile-ai/tilelang
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Find kLayoutMap writers and layout_remap usages =="
rg -n "layout_remap_|Set\(attr::kLayoutMap|kLayoutMap|LowerTileOpPass::" src test tests examples 2>/dev/null || true
echo "== Layout inference annotations context =="
sed -n '700,760p' src/transform/layout_inference.cc
sed -n '1260,1310p' src/transform/layout_inference.cc
echo "== Access-ptr lowering transformation context =="
sed -n '1,130p' src/transform/lower_access_ptr.ccRepository: tile-ai/tilelang
Length of output: 11098
Avoid crashing tvm_access_ptr remapping on shared buffers with replicate_extent.
makeBufferWithLayout adds the shared-buffer replicate_extent into the new buffer shape, but RemapAccessIndices still checks new_shape against layout->OutputShape(). A shared buffer that needs replicate_extent > 1 can reach this check through tvm_address_ptr/tl::access_ptr lowering and shared-LDMA/stores, so the ICHECK would abort. Check against the makeBufferWithLayout shape produced for this buffer instead of the raw layout output shape.
🤖 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/transform/lower_tile_op.cc` around lines 409 - 445, Update
RemapAccessIndices to validate new_shape against the buffer shape produced by
makeBufferWithLayout, including shared-buffer replicate_extent, rather than
layout->OutputShape(). Preserve the existing shape-equality validation and index
remapping behavior for non-replicated buffers.
Comment on lines
285
to
293
| if explicit: | ||
| annotations["pragma_unroll_explicit"] = True | ||
| else: | ||
| explicit = annotations.get("pragma_unroll_explicit", False) | ||
|
|
||
| if unroll_factor is not None: | ||
| # check pragma_unroll_explicit must be False | ||
| if annotations.get("pragma_unroll_explicit", True): | ||
| raise ValueError("pragma_unroll_explicit must be True when unroll_factor is not None") | ||
| if explicit: | ||
| raise ValueError("T.unroll's explicit and unroll_factor params are mutually exclusive.") | ||
| annotations.update({"pragma_unroll_factor": unroll_factor}) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject a factor supplied through annotations when explicit mode is enabled.
T.unroll(..., explicit=True, annotations={"pragma_unroll_factor": 1}) passes this validation. It emits both annotations. The unroll pass then treats the loop as explicitly expandable instead of preserving the factor-based loop.
Proposed fix
- if unroll_factor is not None:
- if explicit:
- raise ValueError("T.unroll's explicit and unroll_factor params are mutually exclusive.")
+ if explicit and ("pragma_unroll_factor" in annotations or unroll_factor is not None):
+ raise ValueError("T.unroll's explicit and unroll_factor params are mutually exclusive.")
+
+ if unroll_factor is not None:
annotations.update({"pragma_unroll_factor": unroll_factor})📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if explicit: | |
| annotations["pragma_unroll_explicit"] = True | |
| else: | |
| explicit = annotations.get("pragma_unroll_explicit", False) | |
| if unroll_factor is not None: | |
| # check pragma_unroll_explicit must be False | |
| if annotations.get("pragma_unroll_explicit", True): | |
| raise ValueError("pragma_unroll_explicit must be True when unroll_factor is not None") | |
| if explicit: | |
| raise ValueError("T.unroll's explicit and unroll_factor params are mutually exclusive.") | |
| annotations.update({"pragma_unroll_factor": unroll_factor}) | |
| if explicit: | |
| annotations["pragma_unroll_explicit"] = True | |
| else: | |
| explicit = annotations.get("pragma_unroll_explicit", False) | |
| if explicit and ("pragma_unroll_factor" in annotations or unroll_factor is not None): | |
| raise ValueError("T.unroll's explicit and unroll_factor params are mutually exclusive.") | |
| if unroll_factor is not None: | |
| annotations.update({"pragma_unroll_factor": unroll_factor}) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tilelang/language/loop.py` around lines 285 - 293, Update the validation
around the explicit and unroll-factor handling in the loop unroll logic to also
reject an existing "pragma_unroll_factor" in annotations when explicit is true,
including when explicit is enabled by the explicit argument. Preserve the
current mutual-exclusion error for the unroll_factor argument and ensure
conflicting annotations are not emitted.
Performance Regression Test ReportTriggered by: @Yongqi-Zhuo Results
Artifacts
|
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
testing/python/layout/test_tilelang_cute.py (1)
1042-1042: 🚀 Performance & Scalability | 🔵 TrivialRun CUDA performance regression tests before merge.
The assertion validates generated TMA code and runtime correctness, but it does not detect throughput regressions from looped TMA loads. Compare this path with the pre-change baseline.
🤖 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 `@testing/python/layout/test_tilelang_cute.py` at line 1042, Update the TMA coverage around the generated-source assertion to include a CUDA performance regression check for looped TMA loads, comparing the measured throughput against the pre-change baseline before merge. Preserve the existing tma_load generation and runtime-correctness assertions.
🤖 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.
Nitpick comments:
In `@testing/python/layout/test_tilelang_cute.py`:
- Line 1042: Update the TMA coverage around the generated-source assertion to
include a CUDA performance regression check for looped TMA loads, comparing the
measured throughput against the pre-change baseline before merge. Preserve the
existing tma_load generation and runtime-correctness assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: cbe03a26-c442-4d26-a639-0526ac954100
📒 Files selected for processing (1)
testing/python/layout/test_tilelang_cute.py
Yongqi-Zhuo
force-pushed
the
sliced-kernels
branch
from
August 3, 2026 15:20
3eba453 to
fd263f4
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tilelang/tools/pass_visualizer/core.py (1)
126-130: 🚀 Performance & Scalability | 🔵 TrivialRun CUDA performance regression tests for the new pass order.
UnrollLoopand the addedSimplifynow run beforePipelinePlanning. Test explicit, disabled, dynamic, and factor-based unroll cases, including warp-specialized kernels.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/tools/pass_visualizer/core.py` around lines 126 - 130, Update the pass sequence around the stages list to retain UnrollLoop and the preceding Simplify before PipelinePlanning, then run CUDA performance regression tests covering explicit, disabled, dynamic, and factor-based unroll configurations, including warp-specialized kernels.
🤖 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.
Nitpick comments:
In `@tilelang/tools/pass_visualizer/core.py`:
- Around line 126-130: Update the pass sequence around the stages list to retain
UnrollLoop and the preceding Simplify before PipelinePlanning, then run CUDA
performance regression tests covering explicit, disabled, dynamic, and
factor-based unroll configurations, including warp-specialized kernels.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: faa685b8-e943-4c33-a2c7-8e16d9bb7e68
📒 Files selected for processing (24)
examples/deepseek_mla/example_mla_decode_ws.pyexamples/deepseek_v32/README.mdexamples/deepseek_v32/sparse_mla_bwd.pyexamples/deepseek_v32/sparse_mla_fwd.pyexamples/deepseek_v32/sparse_mla_fwd_pipelined.pyexamples/deepseek_v32/sparse_mla_fwd_seesaw.pyexamples/flash_attention_sm100/gqa_fwd_bshd.pyexamples/flash_attention_sm100/mha_fwd_bshd.pyexamples/minference/example_vertical_slash_sparse_attn.pymaint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.pysrc/backend/common/op/reduce.hsrc/transform/loop_partition.ccsrc/transform/lower_tile_op.ccsrc/transform/unroll_loop.cctesting/python/language/test_tilelang_language_wgmma_gemm.pytesting/python/layout/test_tilelang_cute.pytesting/python/transform/test_tilelang_transform_unroll_loop.pytilelang/cuda/intrinsics/macro/tcgen05_macro_generator.pytilelang/cuda/intrinsics/macro/wgmma_macro_generator.pytilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.pytilelang/cuda/pipeline.pytilelang/language/loop.pytilelang/language/tir/ir.pytilelang/tools/pass_visualizer/core.py
💤 Files with no reviewable changes (1)
- src/transform/loop_partition.cc
🚧 Files skipped from review as they are similar to previous changes (22)
- maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py
- tilelang/cuda/pipeline.py
- tilelang/language/loop.py
- tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py
- testing/python/layout/test_tilelang_cute.py
- tilelang/language/tir/ir.py
- tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py
- examples/deepseek_v32/sparse_mla_bwd.py
- tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py
- testing/python/transform/test_tilelang_transform_unroll_loop.py
- examples/flash_attention_sm100/mha_fwd_bshd.py
- src/transform/unroll_loop.cc
- examples/minference/example_vertical_slash_sparse_attn.py
- src/transform/lower_tile_op.cc
- examples/deepseek_v32/README.md
- examples/deepseek_v32/sparse_mla_fwd.py
- testing/python/language/test_tilelang_language_wgmma_gemm.py
- examples/deepseek_mla/example_mla_decode_ws.py
- src/backend/common/op/reduce.h
- examples/flash_attention_sm100/gqa_fwd_bshd.py
- examples/deepseek_v32/sparse_mla_fwd_seesaw.py
- examples/deepseek_v32/sparse_mla_fwd_pipelined.py
Yongqi-Zhuo
force-pushed
the
sliced-kernels
branch
from
August 3, 2026 16:35
fd263f4 to
50c2afe
Compare
LeiWang1999
previously approved these changes
Aug 3, 2026
Collaborator
Author
|
@regression-perf |
Performance Regression Test ReportTriggered by: @Yongqi-Zhuo Results
Artifacts
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
examples/deepseek_mla/example_mla_decode_ws.py (1)
162-165: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRestore the
NIdivisibility guard before using both stages.The
NIcalculations at Line 74 and Line 255 useT.ceildiv(..., block_N), but the loops at Line 161 and Line 341 always execute both stage values. WhenNIis odd, the final stage computes tileNIat Line 165 and Line 345. This reads beyondKVandK_pe, and the consumers process data that was not loaded.Add the guard near Line 34 so both kernels require complete two-stage tiles.
Proposed fix
num_stages = 2 + assert (seqlen_kv // num_split) % (block_N * num_stages) == 0This is the same unresolved finding from the previous review; the current code still lacks the guard.
Also applies to: 342-345
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/deepseek_mla/example_mla_decode_ws.py` around lines 162 - 165, Restore the NI divisibility guard near the shared kernel setup so both decode kernels reject configurations where NI is not divisible by 2. Ensure the guard covers the stage loops in the kernels containing the stage and kv_indices calculations, preventing the final incomplete stage from reading or consuming out-of-range KV and K_pe data.
🧹 Nitpick comments (1)
maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py (1)
162-162: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftRun the SM120 performance regression before merging.
The default
T.unrollbehavior preservespragma_unroll_explicit=False. Compare correctness and throughput before and after this change on the target SM120 GPU. Include two regression runs to detect code-generation or performance variance.Also applies to: 212-212, 250-250
🤖 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 `@maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py` at line 162, Run the SM120 benchmark regression for the loops using T.unroll in the relevant GEMM paths, comparing correctness and throughput before and after the change with pragma_unroll_explicit=False preserved. Execute two runs on the target SM120 GPU for each version to identify code-generation or performance variance, and verify no correctness or throughput regression before merging.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/deepseek_mla/example_mla_decode_ws.py`:
- Around line 91-92: Ensure the stage loop around T.unroll in the MLA decode
path is explicitly unrolled by passing explicit=True, or otherwise guard the
CUDA backend against an unexpanded constant loop while preserving the loop’s
kUnrolled behavior.
---
Duplicate comments:
In `@examples/deepseek_mla/example_mla_decode_ws.py`:
- Around line 162-165: Restore the NI divisibility guard near the shared kernel
setup so both decode kernels reject configurations where NI is not divisible by
2. Ensure the guard covers the stage loops in the kernels containing the stage
and kv_indices calculations, preventing the final incomplete stage from reading
or consuming out-of-range KV and K_pe data.
---
Nitpick comments:
In `@maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py`:
- Line 162: Run the SM120 benchmark regression for the loops using T.unroll in
the relevant GEMM paths, comparing correctness and throughput before and after
the change with pragma_unroll_explicit=False preserved. Execute two runs on the
target SM120 GPU for each version to identify code-generation or performance
variance, and verify no correctness or throughput regression before merging.
🪄 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: 24212b76-a8a0-46af-ab33-6533872d8a1c
📒 Files selected for processing (4)
examples/deepseek_mla/example_mla_decode_ws.pyexamples/deepseek_v32/sparse_mla_fwd_pipelined.pyexamples/deepseek_v32/sparse_mla_fwd_seesaw.pymaint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py
🚧 Files skipped from review as they are similar to previous changes (2)
- examples/deepseek_v32/sparse_mla_fwd_seesaw.py
- examples/deepseek_v32/sparse_mla_fwd_pipelined.py
Comment on lines
+91
to
+92
| for stage in T.unroll(num_stages): | ||
| T.barrier_wait(bar_k_ready[stage], (i_i & 1)) |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="examples/deepseek_mla/example_mla_decode_ws.py"
test "$(rg -o 'T\.unroll\(num_stages\)' "$file" | wc -l)" -eq 6
rg -n -C 6 \
'def unroll|pragma_unroll_explicit|explicit_unroll|UnrollLoop' \
tilelang src
rg -n -C 8 \
'compute_gmma_descriptor|slice_byte_offset|increase_descriptor_offset' \
tilelangRepository: tile-ai/tilelang
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate file outline/stat =="
wc -l "examples/deepseek_mla/example_mla_decode_ws.py"
ast-grep outline "examples/deepseek_mla/example_mla_decode_ws.py" --view compact || true
echo "== relevant unroll loops and barrier/operand context =="
nl -ba "examples/deepseek_mla/example_mla_decode_ws.py" | sed -n '70,175p;255,355p'
echo "== focused TIL loop definitions =="
nl -ba "tilelang/language/tir/ir.py" | sed -n '1,130p'
nl -ba "tilelang/language/loop.py" | sed -n '220,305p'
echo "== UnrollLoop pass config relevant lines =="
nl -ba "src/transform/unroll_loop.cc" | sed -n '40,230p;350,385p'
nl -ba "src/transform/loop_partition.cc" | sed -n '235,275p'
echo "== deterministic model of explicit annotation placement =="
python3 - <<'PY'
from pathlib import Path
p = Path('tilelang/language/tir/ir.py')
text = p.read_text()
start = text.index('def unroll(')
end = text.index('def thread_binding(')
body = text[start:end]
print("unroll body:")
print(body)
print("explicit_true_unroll_in_body=", 'return _ir.unroll(start=start, stop=stop, explicit=explicit, annotations=annotations)' in body)
PYRepository: tile-ai/tilelang
Length of output: 486
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate file relevant context =="
awk 'NR>=70 && NR<=175 {printf "%5d %s\n", NR, $0} NR>=255 && NR<=355 {printf "%5d %s\n", NR, $0}' examples/deepseek_mla/example_mla_decode_ws.py
echo "== unroll API definitions =="
awk '/^def unroll\(/,/^def thread_binding\(/' tilelang/language/tir/ir.py
awk '/^def unroll\(/,/^def .*load_array\(/' tilelang/language/loop.py | head -n 120
echo "== unroll pass config relevant lines =="
awk 'NR>=40 && NR<=230 {printf "%5d %s\n", NR, $0} NR>=350 && NR<=385 {printf "%5d %s\n", NR, $0}' src/transform/unroll_loop.cc
awk 'NR>=235 && NR<=275 {printf "%5d %s\n", NR, $0}' src/transform/loop_partition.ccRepository: tile-ai/tilelang
Length of output: 28887
Keep the stage loop explicitly unrolled.
T.unroll(num_stages) emits pragma_unroll_explicit: False, so UnrollLoop may only leave the loop marked as kUnrolled instead of expanding it. Guard the resulting CUDA backend against an unexpanded constant loop, or use T.unroll(num_stages, explicit=True).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/deepseek_mla/example_mla_decode_ws.py` around lines 91 - 92, Ensure
the stage loop around T.unroll in the MLA decode path is explicitly unrolled by
passing explicit=True, or otherwise guard the CUDA backend against an unexpanded
constant loop while preserving the loop’s kUnrolled behavior.
Collaborator
Author
|
@regression-perf |
Performance Regression Test ReportTriggered by: @Yongqi-Zhuo Results
Artifacts
|
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
2 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.
Thanks to my series of PRs (#2380, #2452, #2785), you can write arbitrary slicing expressions to express the multiple stages of a multi-version buffer. So you no longer need to write brittle code that declares multiple versions of a buffer for multiple times, just to pass it to the exact same T.copy, T.gemm, etc., due to lack of support for slicing. Use slicing and indexing. That's all.
Summary
lower_tile_op.cc.Testing
C++ style / lint notes
src/transform/lower_tile_op.cc.docs/developer_guide/cpp_style.md.