[Feature] Flash Attention 1SM/2SM kernels with features supported#2360
Open
chengyupku wants to merge 51 commits into
Open
[Feature] Flash Attention 1SM/2SM kernels with features supported#2360chengyupku wants to merge 51 commits into
chengyupku wants to merge 51 commits into
Conversation
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…en support Adds SM100 single-SM attention kernel variants and supporting infra changes across reduce ops, PTX codegen, tcgen05 macros, shared-tmem lowering, and language builtins/allocate helpers. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Core TileLang fixes: - lower_tile_op.cc: workspace at outermost block scope (fixes SMEM aliasing under warp specialization — workspace no longer overlaps live K/V buffers) - layout.cc: handle non-zero dom->min in Layout/Fragment constructors (enables TMEM column-slice with arbitrary offsets) - inject_fence_proxy.cc: classify initialize_tcgen05_descriptor as kNone (removes spurious fence_proxy_async before GEMM descriptor init) New intrinsic infrastructure: - T.exp2_poly: degree-3 minimax polynomial exp2 on CUDA cores (builtin.h/cc, codegen_cuda.cc, math_intrinsics.py). Works correctly but pure polynomial is slower than SFU — needs SFU+FMA interleaving to be beneficial. Codegen infrastructure for device function outlining: - codegen_cuda.h: outlined_fns_stream_ + shared_state tracking members - codegen_cuda.cc: kWarpSpecializationScope handler (placeholder for full __noinline__ device function emission) Kernel (attention_kernel_1sm.py kq2): - Removed ballot-skip (SMEM reduce overhead exceeds rescale savings) - Moved O rescale from math WGs into mma_load WG (overlaps with softmax) - Added scores_scale SMEM communication (mbar_scale0/1 barriers) - Added disable_thread_storage_sync pass config (removes unnecessary __syncthreads from K-loop, verified correct via explicit mbar waits) Test scripts for __noinline__ device function outlining experiments: - test_noinline_outline.py: NVRTC compilation pipeline test - test_outline_nvrtc.py: source transformation + compilation experiment Performance: Maverick (B=8,H=40,KV=8,S=4096,D=128) Before: 703 TFLOPS (77% of .so) After: 761 TFLOPS (83.4% of .so) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* `VisitStmt_(IfThenElseNode)`: lift the `!current_loop_var_name_.empty()`
guard so we also outline top-level `if(tid_x)` chains (each branch
becomes its own `__device__ __noinline__` function with no
`int __loop_extent` parameter). This lets warp-role bodies that contain
their own for-loop reach the outlined function with their init code
intact — previously the kernel-frame init (e.g. `scores_max=-INF`)
was unreachable from the outlined fn and the math warps started from
zero, producing `inf` output.
* `EmitOutlinedDeviceFunction` / ForNode-based outlining: filter
`local_allocs_` to just the buffer vars the branch actually touches
(new `BufferVarTouchCollector` TIR visitor). Cuts per-fn local
declarations from 36 → 4–14 in the FA kq_stages=2 kernel. ptxas was
already DCE-ing the dead `= {}` initializers so no immediate perf
delta, but the generated source is now readable and the next step
(chunked rescale + correction-warp split) doesn't have to wade
through 30 unused fragments per fn.
* `PrintExtraAttrs`: prefer `__launch_bounds__` over `__maxnreg__`
whenever the requested max_nreg fits 64K/threads. The legacy path
always swapped to `__maxnreg__` when `T.set_max_nreg` was used, which
is mutually exclusive with `__launch_bounds__` and forced ptxas into
a uniform per-thread cap — incompatible with FA4's per-warpgroup
register donation (math 184 / correction 64 / other 80). Now we only
emit `__maxnreg__` when the requested count truly exceeds the
launch_bounds default.
* `T.annotate_min_blocks_per_sm(0)`: allow `n=0` to emit
`__launch_bounds__(N, 0)` (no `minBlocks` preference), so ptxas can
size around the per-warp `setmaxnreg` requests instead of the
64K/N average.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…FLOPS)
Folds per-warp init and math epilogue into the for-loop body so the
codegen's outlining path can see them. Each warp role's branch becomes
a `__device__ __noinline__` function whose registers stay alive across
iterations:
* math0/math1: `scores_max=-INF` / `logsum=0` init at `k==0`,
`T.copy(logsum, logsum_shared)` epilogue at `k==loop_range-1`
* mma_load: TMEM zero + Q/K[0]/V[0] TMA + wait at `k==0`; post-loop
epilogue (normalize O via TMEM, TMA store) stays outside the loop
because it works through SMEM/TMEM rather than private registers.
This sidesteps the SMEM-aliasing regression we'd otherwise hit if we
merged init+loop+epilogue under one top-level `if(tid)` chain: the
shared-memory liveness analyzer attributes touches at the alloc level,
so co-locating everything under one branch defeats the
`O_shared`↔`K_shared` aliasing and pushes SMEM from 192 KB to 264 KB
(> 228 KB B200 limit). Keeping the epilogue outside the for-loop
preserves the analyzer's separate lifetime points.
Driver: `main()` now exposes `--kq_stages` (default 2, the fast path —
the slower kq_stages=1 path was the previous default and gave 600
instead of 770 TFLOPS on the user's H=64 config). Enable
`tl.outline_warp_spec_branches` by default in PASS_CFG so the kernel
actually gets the device-fn split out of the box.
`test_noinline_outline.py`: pin to `/usr/local/cuda-13.1`
(`/usr/local/cuda` -> 12.8 on this box, which doesn't assemble SM100a
`max.ftz.f32` 3-op / `tcgen05.*`), and use `-gencode
arch=compute_100a,code=sm_100a` instead of `-arch=sm_100a` (the latter
lowers to compute_100 PTX which ptxas rejects).
H=64 KV=64 S=4096 D=128 bench: 600 → 770 TFLOPS (~84% of avo 912 ref).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two paired bugs prevented T.copy(O_tmem[:, START:END], O_chunk) from working when START != 0 — i.e. for FA4-style chunked rescale where the mma_load warp walks O_tmem in 16/32-col slices instead of the full [128, 128] tile. * tcgen05_layout.cc:expandTcgen05Layout — the per-warpgroup chunk math computed `FloorDiv(iter_col, num_cols_each_wg) * WARPGROUP_SIZE` to pick the right warpgroup, but `iter_col` is the *absolute* physical column, not relative to the slice. A T.copy at cols [32, 64) with num_cols_each_wg=32 then produced thread indices in [128, 256), shifting the access onto a warpgroup that doesn't exist (mma_load has a single 128-thread warpgroup). Switch the chunk math to `iter_col - col_dom->min` so it's relative to the slice base, not TMEM origin. * copy.cc:LowerTmem — the BufferLoad/Store address used `loop_vars[i]->dom->min` for the tmem-side index, but MakeIterVars always emits loop_vars with min=0; the slice's actual starting column lives only in `op.src_range[i]->min` (or dst_range for a store). Result: every chunked T.copy would issue the tcgen05 instruction with column offset 0 regardless of the slice, silently writing every chunk into TMEM cols [0..N) and corrupting the output. Read the slice base from the right range array. With both fixes, chunked-rescale `T.copy(O_tmem[:, c*16:c*16+16], buf)` compiles and produces correct results. Numerical check: max_abs=0.002 on H=8 S=1024 D=128. Bench at H=64 S=4096 D=128: 771 → 778 TFLOPS (small win from cleaner register footprint in mma_load; the bigger gain is that this unblocks A1 — chunked rescale is the precondition for fitting mma_load below 128 regs/thread at 512 threads). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The mma_load warp's in-loop rescale used to read/write the entire [128, dim] O fragment per iter, costing O_reg[128] + O0_reg[128] = 256 regs/thread (unspillable, because they're tcgen05 operands). Switch to the FA4 reference's correction_warp_fn pattern: walk O_tmem in 32-col slices via the just-added slice notation T.copy(O_tmem[:, c*32:...]), holding only O_chunk[32] (32 regs/thread) at a time. The k==0 TMEM zero-init also moves to the chunk to avoid resurrecting O0_reg[128] in the outlined fn. H=64 S=4096 D=128: 771 → 778 TFLOPS (~1%). The bigger payoff is structural — chunked rescale is the precondition for fitting mma_load below the 128-reg cap at 512 threads. (Math still uses S0_reg[128] + P0_cast[128] = 192 regs and trips the same wall on its outlined fn, so 512 threads still needs a math-side chunking pass. Holding at 384 until that lands.) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* copy.cc:LowerTmem — when emitting the tcgen05_ld/st address for a sliced T.copy(P_tmem[:, START:END], P_chunk) where P_tmem is bf16, the BufferLoad's column index must be in 32-bit cells, not in bf16 elements (TMEM packs 2 bf16 per cell). The previous fix passed the logical column directly, producing +32 on a chunk starting at bf16 col 32 when the correct cell offset is +16 — writes to the wrong half of the next chunk, NaN output. Divide by elements_per_cell = 32 / dtype.bits and assert alignment. * codegen_cuda.cc:PrintExtraAttrs — always emit __launch_bounds__ for multi-thread launches, never __maxnreg__. The previous logic switched to __maxnreg__(N) when T.set_max_nreg requested N greater than 64K/threads, which is the exact case where __maxnreg__ is unusable: it caps every thread at N regardless of warpgroup, so __maxnreg__(176) at 512 threads asks ptxas for 90K regs and is rejected. The FA4 register-donation idiom needs per-warpgroup variance via runtime setmaxnreg.inc/dec — ptxas reads those instructions and sizes per-warpgroup allocations to fit 64K, but only when launch_bounds (not __maxnreg__) is in effect. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replace the full-tile P_cast[128, 128] write with 4 chunks of P_chunk[128, 32]. P_cast was unspillable (tcgen05_st source), holding 64 packed bf16 regs/thread alongside S_reg[128]'s 128 regs and forcing math to peak at 192 regs/thread. The chunked write keeps P's unspillable footprint at 16 regs. H=64 S=4096 D=128: 778 → 782 TFLOPS. 512 still blocked though — the tcgen05_ld_32dp32bNx<128> for S_reg pins 128 regs at the load instruction; with surrounding live state ptxas needs 146+ regs, over the 128 cap. Fitting 512 needs chunked S loads (4 partial tcgen05_ld<32> with per-chunk partial reductions), a bigger softmax restructure. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replace math's full S_reg[128,128] load with a 2-pass chunked algorithm — load 32 cols at a time into S_chunk[128,32], reduce max incrementally across 4 chunks, then re-load and exp+sum+cast+write P chunk by chunk. Math's unspillable register footprint drops from 128 (S_reg pinned by tcgen05_ld_32dp32bNx<128>) to 32 (S_chunk), so the outlined math fn fits below the 128-reg cap at 512 threads. Pass 2 MUST run in REVERSE chunk order (3 → 2 → 1 → 0). P_tmem aliases S_tmem at col_offset=64 (S cells [64,128) overlap P cells [0,64)); writing P chunk 0 to cells [64,80) clobbers data chunk 2 later needs to read from cells [64,96). Reverse order keeps every read ahead of any write that touches its cells: c=3: read S[96:128], write P[96:128] (= S[112:128]). No conflict. c=2: read S[64:96], write P[64:96] (= S[96:112)). No conflict. c=1: read S[32:64], write P[32:64] (= S[80:96)). No conflict. c=0: read S[0:32), write P[0:32) (= S[64:80)). No conflict. Also chunks the post-loop mma_load O-normalize epilogue (full O_reg was the last unspillable 128-reg block in its own outlined fn). H=64 S=4096 D=128 bench: 782 (384 threads) → 701 (512). The drop comes from 2x S TMEM reads + no register donation yet (math stays at the launch_bounds default 128 because ptxas's static allocator won't let setmaxnreg.inc<176> over-subscribe the 64K reg file at 512 threads — 8*176 + 8*80 = 64K exactly, but ptxas's per-warpgroup accounting refuses it). The 2-pass overhead needs to be amortized by other wins (persistent block, deeper K/V pipeline) for 512 to beat 384. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- src/transform/inject_tcgen05_fence.cc: extend IsPlainBarrierArrive to
recognize ptx_arrive_barrier_lane0. Without this, the fence pass
silently skipped every lane-elect arrive, breaking any cross-WG TMEM
handoff that uses mbarrier_arrive(lane0=True).
- examples/flash_attention_sm100/attention_kernel_1sm.py: add per-WG
T.set_max_nreg in the kq_stages=2 path (mma_load WG -> 232, donor WG
-> 24); ~8-9% delta (701 -> 753 TFLOPS at s=16384 h=16). Adds sister
main_kq2_split prim_func for the 4-role split (correction WG split
off mma_load); currently broken (max_abs=0.3-0.7) due to cross-WG
TMEM coherency: ptxas drops FENCE.VIEW.ASYNC.T before UTCHMMA.
- examples/flash_attention_sm100/test_{single,cross}_wg_tmem.py:
minimal repros for the cross-WG TMEM handoff; both currently fail at
TileLang layout inference ("Failed to find a suitable instruction
for tcgen05.st") and need Changes 3/4 (partial-row + warp-granular
copy.cc) before they can run.
- task.md: refresh status (best stable = 753 TFLOPS at s=16384 h=16,
~83% of 912 target) and call out the 4-role split / cross-WG TMEM
fence as the top priority.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The warp index guarding tcgen05.alloc was hardcoded to 12. Expose it as a tmem_alloc_warp annotation so kernels can specify the allocating warp explicitly. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Remove archived stubs, debug repros, reference CUDAs, and benchmark scripts from flash_attention_sm100/ (17 files) - Rename attention_kernel_1sm.py → gqa_fwd_bshd_pipelined_1sm.py - Rename attention_kernel_2sm.py → gqa_fwd_bshd_pipelined_2sm.py - Remove benchmark wrappers, ref.py, and test from grouped_gemm/ (5 files) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
… PrintExtraAttrs Runtime cluster launch via cudaLaunchKernelEx already handles cluster_dims; the compile-time __cluster_dims__ attribute was redundant. MaxNRegExtractor only fed a degenerate __maxnreg__ fallback for threads==1 launches that never occurs in practice. Revert PrintExtraAttrs to the original logic: emit __launch_bounds__(threads, min_blocks_per_sm) via LaunchConfigExtractor. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…face
Barrier codegen:
- Replace inline PTX emission with Barrier::arrive()/wait()/arrive_and_expect_tx()
- barrier.h dispatches release/acquire ordering via #if __CUDA_ARCH__ >= 900
- Revert ptx.cc TMA barrier instructions to original (TMA engine guarantees ordering)
- Remove ptx_arrive_barrier_lane0 op — use explicit `if tid % 32 == 0` at DSL level
Dead code removal:
- Remove tcgen05_ld_x16, tcgen05_st_x16, tcgen05_mma_1sm_ts_128x64_bmn_x2{,_contig}
from codegen, builtin.cc/h, builtin.py, and inject_tcgen05_fence.cc
- Remove use_x16, skip_expect_tx, expect_tx_bytes from copy.cc
- Remove StoreFenceMarksTcgen05Use from inject_tcgen05_fence.cc
TMA interface unification:
- create_tma_descriptor() accepts Buffer directly (extracts .data Var)
- Revert lower_hopper_intrin.cc GetDescriptorBaseVar (no longer needed)
- 2SM kernel: tma_load_2cta_2d → tma_load_2sm (unified calling convention)
- FA kernels: pass Buffer directly instead of T.access_ptr(buf, "r")
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…otation - lower_shared_tmem.cc / lower_shared_barrier.cc: infer use_2cta from function's cluster_dims attribute (product > 1) instead of SBlock annotation - Remove T.use_2cta_tmem() annotation entirely — its responsibilities are now: - use_2cta: auto-inferred from cluster_dims - alloc_warp: parameter on T.alloc_tmem(alloc_warp=N) - init_thread: removed (default shuffle_elect is sufficient) - compact_shared_state: removed (no perf impact) - Remove T.annotate_min_blocks_per_sm(1) from 2SM kernel (default value) - Remove 10 dead Python DSL functions from builtin.py Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…, delete task.md - Remove cudaFuncSetAttribute(MaxRegistersPerThread) from wrapper.py — __launch_bounds__ + in-body setmaxnreg.inc/dec is sufficient for warp-specialized register donation - Revert cmake/FindPipCUDAToolkit.cmake to upstream (manual nvcc search was a workaround for a specific build environment) - Delete task.md (development notes, not needed in repo) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- lower_tile_op.cc: only place workspace at outermost scope when warp-specialization is enabled (prevents merge pass aliasing with concurrently TMA-written buffers). Non-WS kernels keep innermost scope. - Remove dead ops: ptx_arrive_barrier_lane0, mbarrier_wait_parity_lane0, tma_load_2cta_2d, pack_bf16_pair (registration + codegen handlers) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- barrier.h: revert to original CUTLASS typedef (release/acquire not needed) - cluster.h: revert arrive/wait to plain aligned (keep cluster_id_x) - tcgen_05.h: keep memory clobbers, revert commit instruction to shared::cluster - lower_hopper_intrin.cc: fully revert (shuffle_elect + direct Var access work) - copy.cc: revert emit_arrive logic to original (after-TMA arrive works) - Remove tcgen05_mma_1sm_ss_128x128_commit (dead: no DSL or kernel uses it) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR extends TileLang with comprehensive SM100 TCGEN05 and warp-specialization support across the compiler stack, language frontend, and examples. It introduces new builtin intrinsics for TCGEN05/TMEM/barrier operations, updates CUDA codegen to support compact shared-state management and device-function outlining, adds transform passes for 2CTA/alias TMEM handling, expands CUDA templates with nomask MMA variants and SM100-specific helpers, and provides three new example kernels (1SM and 2SM FlashAttention, and grouped GEMM variants with FP8 support). The frontend API gains optional barrier initialization-thread control, TMEM aliasing and warp-scoped allocation, and new math/cluster-identification intrinsics. Existing FlashAttention kernels are extended with new variants and CLI controls. ChangesSM100 Compiler and Kernel Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
|
👋 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! 🚀 |
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- gemm_tcgen05.py: use upstream's _shared_layout_continuity + infer_shared_layout(self.B, ...) while keeping matrix="A" for TS-MMA A operand layout - builtin.py: accept upstream's removal of _IS_HIP_AVAILABLE and _get_tl_op (no longer used) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
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: 12
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
examples/flash_attention_sm100/gqa_fwd_bshd.py (1)
1363-1375:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winReject unknown
variantvalues instead of silently runningwasp.Anything not matched above falls through to
flashattn_wasp(...), so a typo still runs a kernel while the banner prints the bad variant string. That makes benchmark output easy to misread.Suggested fix
- else: + elif variant == "wasp": kernel = flashattn_wasp( batch, heads, seq_len, dim, is_causal, groups=groups, block_M=128, block_N=64, threads=256, num_stages=2, ) + else: + raise ValueError(f"Unsupported variant: {variant}")🤖 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/flash_attention_sm100/gqa_fwd_bshd.py` around lines 1363 - 1375, The code currently falls through to calling flashattn_wasp for any unrecognized variant, so update the variant dispatch logic to explicitly validate the variant variable and reject unknown values instead of defaulting to wasp; locate the branch that calls flashattn_wasp and replace the implicit fallback with an explicit conditional that only calls flashattn_wasp when variant == "wasp" (or is in an allowed set like {"wasp", ...}), and otherwise raise a clear exception (e.g., ValueError) or log an error and exit to prevent silent runs with typos in variant.tilelang/language/builtin.py (1)
468-507:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
lane0is currently a no-op inmbarrier_wait_parity.
lane0is accepted by the public API but never affects emitted IR, so callers get different behavior than requested.Suggested fix
def mbarrier_wait_parity(mbarrier: BarrierType, parity: int | Var, *, lane0: bool = False): @@ """ mbarrier = _mbar_to_buffer_load(mbarrier) - return tirx.call_intrin("handle", tirx.op.Op.get("tl.mbarrier_wait_parity"), mbarrier, parity) + ann = {"lane0_only": True} if lane0 else None + return tirx.call_intrin( + "handle", + tirx.op.Op.get("tl.mbarrier_wait_parity"), + mbarrier, + parity, + annotations=ann, + )🤖 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/builtin.py` around lines 468 - 507, The public API parameter lane0 in mbarrier_wait_parity is ignored; update mbarrier_wait_parity so the lane0 flag is encoded into the emitted IR by passing it into the tirx.call_intrin invocation (alongside mbarrier and parity) so the intrinsic sees the lane0 setting. Locate mbarrier_wait_parity, keep the existing mbarrier = _mbar_to_buffer_load(mbarrier) and then add lane0 (converted to an appropriate IR/int value) as an extra argument to tirx.call_intrin (the call producing Op.get("tl.mbarrier_wait_parity")) so the intrinsic receives and can act on the lane0 flag. Ensure callers retain the same signature.
🧹 Nitpick comments (4)
tilelang/language/gemm_op.py (1)
245-245: 💤 Low valueType hint could be more precise.
The default value
Nonesuggests the type should beBarrierType | None = Nonefor consistency with other optional parameters in this file (e.g., line 158, 291).Suggested fix
- mbar: BarrierType = None, + mbar: BarrierType | None = None,🤖 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/gemm_op.py` at line 245, The parameter mbar currently declares "mbar: BarrierType = None" but uses None as a default so its type hint should allow None; update the signature for the mbar parameter to be nullable (e.g., BarrierType | None or Optional[BarrierType]) to match other optional parameters in this module (ensure any imports for Optional/typing are added if needed) — locate the mbar parameter declaration in the gemm op function/class (the mbar parameter in gemm_op.py) and change its annotation accordingly.tilelang/language/allocate.py (1)
223-223: Remove or document unused_tmem_alloc_counter.
tilelang/language/allocate.pydefines_tmem_alloc_counter = [0](line 223), and there are no other references to_tmem_alloc_counteranywhere else in the codebase (only this occurrence was found). Remove it if it’s not needed, or add a TODO/comment explaining the intended future use.🤖 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/allocate.py` at line 223, The symbol `_tmem_alloc_counter` is defined but unused; either delete the definition `_tmem_alloc_counter = [0]` from allocate.py to remove dead code, or annotate it with a clear TODO comment explaining intended future use (e.g., what it will count, when it will be referenced) so it’s not mistaken for accidental leftover; reference the exact symbol `_tmem_alloc_counter` when making the change and ensure tests/imports still pass after removal or that the comment explains why the placeholder remains.examples/grouped_gemm/example_grouped_gemm_fwd_sm90_fp8_to_bf16.py (1)
24-55: ⚡ Quick winConsider extracting shared helper functions to avoid duplication.
construct_inputs_bf16_fp8andtorch_gmm_bf16_fp8_refare identical copies fromexample_grouped_gemm_fwd_sm100_fp8_to_bf16.py. Consider moving these to a shared utility module (e.g., the existingexample_grouped_gemm_fwd.pypattern) to reduce maintenance burden.🤖 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/grouped_gemm/example_grouped_gemm_fwd_sm90_fp8_to_bf16.py` around lines 24 - 55, The two functions construct_inputs_bf16_fp8 and torch_gmm_bf16_fp8_ref are duplicated from another example; extract them into a shared helper module (e.g., the existing example_grouped_gemm_fwd utility) and have this file import those helpers instead of defining them inline. Concretely: move the implementations of construct_inputs_bf16_fp8 and torch_gmm_bf16_fp8_ref into the shared module, export them with the same names, then in this example remove the duplicated definitions and add an import for construct_inputs_bf16_fp8 and torch_gmm_bf16_fp8_ref so both example files use the single shared implementation.tilelang/language/math_intrinsics.py (1)
432-439: ⚡ Quick winPrefer the registered intrinsic op over raw extern for
exp2_approx.Using
tl.tcgen05_exp2f_approxviacall_intrinkeeps lowering behavior consistent with the rest of the TCGEN05 surface.Suggested refactor
def exp2_approx(x: PrimExpr) -> PrimExpr: @@ """ x = tirx.convert(x) - return tirx.call_pure_extern(x.dtype, "tl::tcgen05_exp2f_approx", x) + return tirx.call_intrin(x.dtype, tirx.op.Op.get("tl.tcgen05_exp2f_approx"), x)🤖 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/math_intrinsics.py` around lines 432 - 439, Replace the raw extern call in exp2_approx with the registered intrinsic path: keep the tirx.convert(x) conversion, then call the registered intrinsic via tirx.call_intrin instead of tirx.call_pure_extern so lowering follows the TCGEN05 intrinsic registration (replace the call to tirx.call_pure_extern(x.dtype, "tl::tcgen05_exp2f_approx", x) with a tirx.call_intrin variant that invokes the registered TCGEN05 exp2 approximate intrinsic while preserving the return dtype and argument x).
🤖 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/flash_attention_sm100/gqa_fwd_bshd_pipelined_1sm.py`:
- Around line 1042-1099: The split-epilogue always emits fixed 32x128 TMA stores
(tl::attention_1sm_epilogue_tma_store_32x128) for both halves which can write
past Output on the final partial tile; guard those calls or assert full-tile
precondition. Fix by checking the per-page row bounds before each store: compute
remaining_rows = seq_len - (bx * kq2_total_M) and for the second half use
remaining_rows2 = seq_len - (bx * kq2_total_M + kq2_block_M), then only invoke
the T.call_extern for O0_shared[cw * 32, 0] when cw*32 < remaining_rows and for
O1_shared[cw * 32, 0] when cw*32 < remaining_rows2; alternatively, add an
upfront validation that seq_len is a multiple of kq2_total_M (full 256-row
tiles) and early-fail if not. Ensure checks reference bx, kq2_total_M,
kq2_block_M, seq_len, O0_shared/O1_shared and the
tl::attention_1sm_epilogue_tma_store_32x128 calls.
- Around line 254-255: The code only checks dim == block_N but many parts are
hard-coded for a 128x128 tile (functions/externs like qk_mma_128x128_dsl,
softmax_warp_dsl with 128-element buffers/loops, and any *_d128 helpers), so
kernels can be generated with incompatible shapes; update the validation in
attention_kernel_1sm to require block_M == 128, block_N == 128 and dim == 128
(or otherwise make the helpers generic), and either: (a) enforce the fixed
128x128 contract by raising a ValueError if block_M != 128 or block_N != 128 or
dim != 128, referencing attention_kernel_1sm and the qk_mma_128x128_dsl /
softmax_warp_dsl / *_d128 helpers; or (b) refactor those helpers to accept
parameterized tile sizes and replace hard-coded 128 usages so the existing dim
== block_N check is sufficient.
In `@examples/flash_attention_sm100/gqa_fwd_bshd_pipelined_2sm.py`:
- Around line 445-451: The d128 online-softmax path only updates rmax_local when
acc_scale_log2 < -8.0, causing rmax_local to remain stale for small increases;
change the logic in the d128 branch (the block that computes acc_scale_log2, rs,
and updates rmax_local for kb==0 vs else) so rmax_local is always assigned to
new_max every iteration (like the d256 path) and compute rs =
T.tcgen05_exp2f_approx(acc_scale_log2) only when acc_scale_log2 < -8.0 (clamping
underflow then); in short, always advance rmax_local = new_max and only apply
the underflow-rescale to rs when needed.
- Around line 1692-1695: The correction path must compute PV barrier parity from
the committed V index (as produced by d256_issue_pv_dsl which sets v_global = 2
* tkb + 1) instead of deriving it from tkb; update the wait logic that currently
uses expressions like (prev // kv_stages) & 1 and (last_tkb // kv_stages) & 1 to
compute parity from prev_v_idx and last_v_idx (e.g., (prev_v_idx // kv_stages) &
1 and (last_v_idx // kv_stages) & 1) so the mb_pv wait matches the producer’s
v_global parity and prevents reading O0_tmem before PV commit; apply the same
replacement in the other occurrences noted around the later blocks.
In `@examples/flash_attention_sm100/gqa_fwd_bshd.py`:
- Around line 587-590: The K/V stage mbarriers (mb_k_loaded, mb_k_consumed,
mb_v_loaded, mb_v_consumed) are allocated with count 1 but
T.mbarrier_arrive(...) is invoked from every thread in the producer/MMA warps,
advancing the barrier warp-size times and allowing premature slot release; fix
by allocating these barriers with a warp-sized count (e.g., 32) or alternately
guard the T.mbarrier_arrive calls so only lane 0 calls them—update the
alloc_barrier calls that create mb_k_loaded, mb_k_consumed, mb_v_loaded and
mb_v_consumed or add a lane check around T.mbarrier_arrive invocations to ensure
one arrival per warp.
- Around line 1247-1259: The epilogue reads logsum_shared from the
correction/MMA half immediately after the softmax half wrote it (writes to
logsum_shared via T.copy when tid < 128) without synchronization, risking stale
reads; fix by inserting a CTA/shared barrier (or equivalent warp-group barrier)
after the T.copy(logsum, logsum_shared) write and before any reads of
logsum_shared (i.e., before the branch where tid >= 128 that performs
T.copy(logsum_shared, logsum) and the normalization loop on O_reg), or
alternatively move the final normalization (the loop dividing O_reg by logsum
and subsequent T.copy to O_shared/Output) into the same half that owns logsum to
avoid cross-half reads; locate uses of logsum_shared, T.copy, tid, O_reg,
O_shared, and Output to apply the barrier or move the code.
- Around line 176-187: The function flashattn_ts_v2 declares a num_stages
parameter but the pipeline is hard-coded to num_stages=1 at the T.Pipelined(...)
call, so either wire the argument through or remove it from the API; update the
T.Pipelined invocation to use the function parameter (num_stages=num_stages) so
callers can control staging, or if staging must be fixed, delete the num_stages
parameter from flashattn_ts_v2 and its callers to avoid a misleading API.
In `@src/cuda/codegen/codegen_cuda.cc`:
- Around line 6840-6852: The outlined-branch construction in the loop (symbols:
ExtractLeadingSetMaxNReg, can_emit_directly, EmitOutlinedDeviceFunction and
vectors fn_names, branch_prologues, outlined_bodies, emit_direct) currently
omits captured scalar temps, causing helpers to miss outer Bind/local.var state;
update this path to collect and forward scalar capture metadata
(scalar_var_infos_ and local_scalar_var_infos_) for each outlined_body whether
emitted directly or via EmitOutlinedDeviceFunction, and ensure the
EmitOutlinedDeviceFunction call and any direct-emission dispatch include these
scalar var info structures (or a wrapper) so outlined helpers receive the same
scalar/thread-local context that kDeviceFuncScope already threads through.
- Around line 5424-5429: The code currently records all shared allocations as
"uint*" when outline_warp_spec_enabled_ is true, which breaks helper signatures
for non-uint32 storage; change the tmem_var_infos_ push to use the actual
emitted storage type: call PrintType (or reuse the emitted storage type string
produced for alloc_dtype) to form "<storage_type>*" instead of hard-coding
"uint*", and push that string (alongside vid and buffer) into tmem_var_infos_ so
outlined helper signatures match the real shared-buffer pointer type (use the
existing alloc_dtype/PrintType usage to build the correct pointer type).
- Around line 5194-5197: The code currently strips the
"tl_compact_tmem_base_alias" attribute without recording which buffer it marks,
causing compact_tmem_base_var_ (used later to emit __tl_compact_tbase) to
default to the first compacted singleton; update the handler for the
tl_compact_tmem_base_alias AttrStmt (the branch that calls PrintStmt(op->body))
to capture and store the actual buffer variable being aliased instead of just
printing the body — e.g., extract the buffer Var/IterVar name from op->body (or
from op->node/attr value) and assign it to compact_tmem_base_var_ so the later
selection logic (the code around compact_tmem_base_var_ usage and the fallback
at the 5417-5420 area) uses the recorded target rather than the first seen
allocation. Ensure the same recorded mapping is consulted where
__tl_compact_tbase is generated so the alias targets the correct buffer.
In `@src/cuda/transform/inject_fence_proxy.cc`:
- Around line 222-224: The code dereferences call->op without checking that call
is non-null; update the proxy classification logic (the branch that compares
call->op to initialize_tcgen05_descriptor()) to first verify the evaluated node
is a CallNode (i.e., that call is non-null) and return ProxyEvent::kNone if it
is not; locate the logic that gets eval->value.as<CallNode>() (used by
ClassifyCallProxyEvent) and add the null/instance check before using call->op to
prevent crashes.
In `@src/tl_templates/cuda/reduce.h`:
- Around line 190-201: The template NamedBarrier is still instantiated with its
default barrier_id_ (1); update the SM90+ reduction lowering in
src/cuda/op/reduce.cc so AllReduce passes the per-WG barrier ID into the
template (i.e., instantiate tl::NamedBarrier with the runtime/per-WG barrier_id
as the second template parameter instead of using the default), ensure the
per-WG ID is derived from the WG/thread_offset as the codegen intends, and keep
the pair-allocation rule (AllReduce reads barrier_id and uses barrier_id+1)
while bounding IDs within the available hardware range so concurrent WGs do not
collide on the same {barrier_id, barrier_id+1} pair.
---
Outside diff comments:
In `@examples/flash_attention_sm100/gqa_fwd_bshd.py`:
- Around line 1363-1375: The code currently falls through to calling
flashattn_wasp for any unrecognized variant, so update the variant dispatch
logic to explicitly validate the variant variable and reject unknown values
instead of defaulting to wasp; locate the branch that calls flashattn_wasp and
replace the implicit fallback with an explicit conditional that only calls
flashattn_wasp when variant == "wasp" (or is in an allowed set like {"wasp",
...}), and otherwise raise a clear exception (e.g., ValueError) or log an error
and exit to prevent silent runs with typos in variant.
In `@tilelang/language/builtin.py`:
- Around line 468-507: The public API parameter lane0 in mbarrier_wait_parity is
ignored; update mbarrier_wait_parity so the lane0 flag is encoded into the
emitted IR by passing it into the tirx.call_intrin invocation (alongside
mbarrier and parity) so the intrinsic sees the lane0 setting. Locate
mbarrier_wait_parity, keep the existing mbarrier =
_mbar_to_buffer_load(mbarrier) and then add lane0 (converted to an appropriate
IR/int value) as an extra argument to tirx.call_intrin (the call producing
Op.get("tl.mbarrier_wait_parity")) so the intrinsic receives and can act on the
lane0 flag. Ensure callers retain the same signature.
---
Nitpick comments:
In `@examples/grouped_gemm/example_grouped_gemm_fwd_sm90_fp8_to_bf16.py`:
- Around line 24-55: The two functions construct_inputs_bf16_fp8 and
torch_gmm_bf16_fp8_ref are duplicated from another example; extract them into a
shared helper module (e.g., the existing example_grouped_gemm_fwd utility) and
have this file import those helpers instead of defining them inline. Concretely:
move the implementations of construct_inputs_bf16_fp8 and torch_gmm_bf16_fp8_ref
into the shared module, export them with the same names, then in this example
remove the duplicated definitions and add an import for
construct_inputs_bf16_fp8 and torch_gmm_bf16_fp8_ref so both example files use
the single shared implementation.
In `@tilelang/language/allocate.py`:
- Line 223: The symbol `_tmem_alloc_counter` is defined but unused; either
delete the definition `_tmem_alloc_counter = [0]` from allocate.py to remove
dead code, or annotate it with a clear TODO comment explaining intended future
use (e.g., what it will count, when it will be referenced) so it’s not mistaken
for accidental leftover; reference the exact symbol `_tmem_alloc_counter` when
making the change and ensure tests/imports still pass after removal or that the
comment explains why the placeholder remains.
In `@tilelang/language/gemm_op.py`:
- Line 245: The parameter mbar currently declares "mbar: BarrierType = None" but
uses None as a default so its type hint should allow None; update the signature
for the mbar parameter to be nullable (e.g., BarrierType | None or
Optional[BarrierType]) to match other optional parameters in this module (ensure
any imports for Optional/typing are added if needed) — locate the mbar parameter
declaration in the gemm op function/class (the mbar parameter in gemm_op.py) and
change its annotation accordingly.
In `@tilelang/language/math_intrinsics.py`:
- Around line 432-439: Replace the raw extern call in exp2_approx with the
registered intrinsic path: keep the tirx.convert(x) conversion, then call the
registered intrinsic via tirx.call_intrin instead of tirx.call_pure_extern so
lowering follows the TCGEN05 intrinsic registration (replace the call to
tirx.call_pure_extern(x.dtype, "tl::tcgen05_exp2f_approx", x) with a
tirx.call_intrin variant that invokes the registered TCGEN05 exp2 approximate
intrinsic while preserving the return dtype and argument x).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8d324abd-8175-4e97-8702-c4287033f37a
📒 Files selected for processing (39)
3rdparty/tvmexamples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.pyexamples/flash_attention_sm100/gqa_fwd_bshd.pyexamples/flash_attention_sm100/gqa_fwd_bshd_pipelined_1sm.pyexamples/flash_attention_sm100/gqa_fwd_bshd_pipelined_2sm.pyexamples/grouped_gemm/example_grouped_gemm_fwd_sm100.pyexamples/grouped_gemm/example_grouped_gemm_fwd_sm100_fp8.pyexamples/grouped_gemm/example_grouped_gemm_fwd_sm100_fp8_to_bf16.pyexamples/grouped_gemm/example_grouped_gemm_fwd_sm90_fp8_to_bf16.pysrc/cuda/codegen/codegen_cuda.ccsrc/cuda/codegen/codegen_cuda.hsrc/cuda/codegen/codegen_cutedsl.ccsrc/cuda/op/copy.ccsrc/cuda/op/copy_analysis.ccsrc/cuda/op/finalize_reducer.ccsrc/cuda/transform/inject_fence_proxy.ccsrc/cuda/transform/inject_tcgen05_fence.ccsrc/cuda/transform/lower_shared_barrier.ccsrc/cuda/transform/lower_shared_tmem.ccsrc/layout/tcgen05_layout.ccsrc/op/builtin.ccsrc/op/builtin.hsrc/tl_templates/cuda/cluster.hsrc/tl_templates/cuda/copy_sm100.hsrc/tl_templates/cuda/instruction/tcgen05mma.hsrc/tl_templates/cuda/reduce.hsrc/tl_templates/cuda/tcgen_05.hsrc/transform/lower_tile_op.cctilelang/cuda/intrinsics/macro/tcgen05_macro_generator.pytilelang/cuda/op/gemm/gemm_tcgen05.pytilelang/language/__init__.pytilelang/language/allocate.pytilelang/language/annotations.pytilelang/language/builtin.pytilelang/language/gemm_op.pytilelang/language/kernel.pytilelang/language/math_intrinsics.pytilelang/language/tir/op.pytilelang/transform/pass_config.py
Comment on lines
+254
to
+255
| if dim != block_N: | ||
| raise ValueError(f"attention_kernel_1sm split DSL requires dim == block_N; got dim={dim}, block_N={block_N}") |
Contributor
There was a problem hiding this comment.
Validate the fixed 128x128 contract here.
The implementation below is still specialized to block_M == 128, block_N == 128, and dim == 128 (qk_mma_128x128_dsl, softmax_warp_dsl's 128-element buffers/loops, and the *_d128 extern helpers all hard-code that shape). The current check only enforces dim == block_N, so calls like attention_kernel_1sm(..., dim=256) or block_M != 128 can compile a kernel whose barrier/layout logic no longer matches the actual tile shape.
Suggested fix
- if dim != block_N:
- raise ValueError(f"attention_kernel_1sm split DSL requires dim == block_N; got dim={dim}, block_N={block_N}")
+ if dim != 128 or block_M != 128 or block_N != 128:
+ raise ValueError(
+ "attention_kernel_1sm currently supports block_M=128, block_N=128, dim=128 only"
+ )🤖 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/flash_attention_sm100/gqa_fwd_bshd_pipelined_1sm.py` around lines
254 - 255, The code only checks dim == block_N but many parts are hard-coded for
a 128x128 tile (functions/externs like qk_mma_128x128_dsl, softmax_warp_dsl with
128-element buffers/loops, and any *_d128 helpers), so kernels can be generated
with incompatible shapes; update the validation in attention_kernel_1sm to
require block_M == 128, block_N == 128 and dim == 128 (or otherwise make the
helpers generic), and either: (a) enforce the fixed 128x128 contract by raising
a ValueError if block_M != 128 or block_N != 128 or dim != 128, referencing
attention_kernel_1sm and the qk_mma_128x128_dsl / softmax_warp_dsl / *_d128
helpers; or (b) refactor those helpers to accept parameterized tile sizes and
replace hard-coded 128 usages so the existing dim == block_N check is
sufficient.
Comment on lines
+1042
to
+1099
| # ---- Split epilogue TMA store ---- | ||
| if tid >= 448 and tid < 480: | ||
| output_desc = T.create_tma_descriptor( | ||
| 9, | ||
| 4, | ||
| Output, | ||
| dim, | ||
| heads, | ||
| seq_len, | ||
| batch, | ||
| 2, | ||
| dim * 2, | ||
| heads * dim * 2, | ||
| seq_len * heads * dim * 2, | ||
| 64, | ||
| 1, | ||
| 32, | ||
| 1, | ||
| 1, | ||
| 1, | ||
| 1, | ||
| 1, | ||
| 0, | ||
| 3, | ||
| 2, | ||
| 0, | ||
| ) | ||
| with T.device_func(): | ||
| if (T.get_thread_binding() & 31) == 0: | ||
| T.tcgen05_wait_barrier(mbar_epi0, 0) | ||
| T.fence_proxy_async() | ||
| for cw in T.unroll(4): | ||
| T.call_extern( | ||
| "void", | ||
| "tl::attention_1sm_epilogue_tma_store_32x128", | ||
| output_desc, | ||
| T.access_ptr(O0_shared[cw * 32, 0], "r"), | ||
| bx * kq2_total_M + cw * 32, | ||
| by, | ||
| bz, | ||
| ) | ||
| T.tma_store_arrive() | ||
|
|
||
| T.tcgen05_wait_barrier(mbar_epi1, 0) | ||
| T.fence_proxy_async() | ||
| for cw in T.unroll(4): | ||
| T.call_extern( | ||
| "void", | ||
| "tl::attention_1sm_epilogue_tma_store_32x128", | ||
| output_desc, | ||
| T.access_ptr(O1_shared[cw * 32, 0], "r"), | ||
| bx * kq2_total_M + kq2_block_M + cw * 32, | ||
| by, | ||
| bz, | ||
| ) | ||
| T.tma_store_arrive() | ||
| T.tma_store_wait(0) | ||
|
|
Contributor
There was a problem hiding this comment.
Guard the split epilogue stores on the final partial tile.
This kernel launches ceildiv(seq_len, 256) CTAs, but the epilogue always issues fixed 32x128 TMA stores for both halves of the 256-row tile. On the tail CTA, bx * kq2_total_M + ... can exceed seq_len, so these calls can write past the end of Output. If the kernel only supports full 256-row tiles, validate that up front; otherwise add per-page bounds checks before each store.
🤖 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/flash_attention_sm100/gqa_fwd_bshd_pipelined_1sm.py` around lines
1042 - 1099, The split-epilogue always emits fixed 32x128 TMA stores
(tl::attention_1sm_epilogue_tma_store_32x128) for both halves which can write
past Output on the final partial tile; guard those calls or assert full-tile
precondition. Fix by checking the per-page row bounds before each store: compute
remaining_rows = seq_len - (bx * kq2_total_M) and for the second half use
remaining_rows2 = seq_len - (bx * kq2_total_M + kq2_block_M), then only invoke
the T.call_extern for O0_shared[cw * 32, 0] when cw*32 < remaining_rows and for
O1_shared[cw * 32, 0] when cw*32 < remaining_rows2; alternatively, add an
upfront validation that seq_len is a multiple of kq2_total_M (full 256-row
tiles) and early-fail if not. Ensure checks reference bx, kq2_total_M,
kq2_block_M, seq_len, O0_shared/O1_shared and the
tl::attention_1sm_epilogue_tma_store_32x128 calls.
Comment on lines
+445
to
+451
| if kb == 0: | ||
| rmax_local = new_max | ||
| else: | ||
| acc_scale_log2 = (rmax_local - new_max) * softmax_scale_log2 | ||
| if acc_scale_log2 < -8.0: | ||
| rs = T.tcgen05_exp2f_approx(acc_scale_log2) | ||
| rmax_local = new_max |
Contributor
There was a problem hiding this comment.
The d128 online-softmax rescale only updates on huge max jumps.
Line 449 gates both rs and rmax_local behind acc_scale_log2 < -8.0. In the normal case where the running max increases by a small amount, rs stays 1 and rmax_local never advances, so later blocks are normalized against a stale max. The d256 path on Lines 1516-1518 updates the running max every iteration; the d128 path needs the same behavior, with any underflow clamp applied after computing the new max.
Also applies to: 1516-1518
🤖 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/flash_attention_sm100/gqa_fwd_bshd_pipelined_2sm.py` around lines
445 - 451, The d128 online-softmax path only updates rmax_local when
acc_scale_log2 < -8.0, causing rmax_local to remain stale for small increases;
change the logic in the d128 branch (the block that computes acc_scale_log2, rs,
and updates rmax_local for kb==0 vs else) so rmax_local is always assigned to
new_max every iteration (like the d256 path) and compute rs =
T.tcgen05_exp2f_approx(acc_scale_log2) only when acc_scale_log2 < -8.0 (clamping
underflow then); in short, always advance rmax_local = new_max and only apply
the underflow-rescale to rs when needed.
Comment on lines
+1692
to
+1695
| v_global = 2 * tkb + 1 | ||
| v_stage = v_global % kv_stages | ||
| v_phase = (v_global // kv_stages) & 1 | ||
| p_phase = (tkb // 2) & 1 |
Contributor
There was a problem hiding this comment.
Use the committed V index when waiting on mb_pv in the d256 correction path.
d256_issue_pv_dsl drives mb_pv with v_global = 2 * tkb + 1, so the barrier phase advances on v_global // kv_stages. The correction path currently waits with (prev // kv_stages) & 1 and (last_tkb // kv_stages) & 1, which diverges from the producer's generation once tkb crosses a stage group. That can block on the wrong parity or let correction read O0_tmem before the final PV commit completes. Compute the parity from prev_v_idx / last_v_idx instead.
Also applies to: 1963-1969, 1983-1996
🤖 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/flash_attention_sm100/gqa_fwd_bshd_pipelined_2sm.py` around lines
1692 - 1695, The correction path must compute PV barrier parity from the
committed V index (as produced by d256_issue_pv_dsl which sets v_global = 2 *
tkb + 1) instead of deriving it from tkb; update the wait logic that currently
uses expressions like (prev // kv_stages) & 1 and (last_tkb // kv_stages) & 1 to
compute parity from prev_v_idx and last_v_idx (e.g., (prev_v_idx // kv_stages) &
1 and (last_v_idx // kv_stages) & 1) so the mb_pv wait matches the producer’s
v_global parity and prevents reading O0_tmem before PV commit; apply the same
replacement in the other occurrences noted around the later blocks.
Comment on lines
+176
to
+187
| @tilelang.jit(out_idx=[3], pass_configs=PASS_CFG) | ||
| def flashattn_ts_v2( | ||
| batch, | ||
| heads, | ||
| seq_len, | ||
| dim, | ||
| is_causal, | ||
| groups=1, | ||
| block_M=128, | ||
| block_N=128, | ||
| num_stages=2, | ||
| ): |
Contributor
There was a problem hiding this comment.
num_stages is currently a no-op in flashattn_ts_v2.
The public signature exposes num_stages, but Line 248 hard-codes T.Pipelined(..., num_stages=1), so callers cannot actually tune staging. Either thread the argument through or drop it from the API to avoid a misleading contract.
Also applies to: 248-249
🤖 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/flash_attention_sm100/gqa_fwd_bshd.py` around lines 176 - 187, The
function flashattn_ts_v2 declares a num_stages parameter but the pipeline is
hard-coded to num_stages=1 at the T.Pipelined(...) call, so either wire the
argument through or remove it from the API; update the T.Pipelined invocation to
use the function parameter (num_stages=num_stages) so callers can control
staging, or if staging must be fixed, delete the num_stages parameter from
flashattn_ts_v2 and its callers to avoid a misleading API.
Comment on lines
+5194
to
+5197
| if (op->attr_key == "tl_compact_tmem_base_alias") { | ||
| PrintStmt(op->body); | ||
| return; | ||
| } |
Contributor
There was a problem hiding this comment.
tl_compact_tmem_base_alias never selects the alias target.
This attribute is stripped without recording which buffer it marks, so compact_tmem_base_var_ still falls back to the first compacted shared uint32[1] allocation seen in Lines 5417-5420. If another compacted singleton is allocated earlier, __tl_compact_tbase aliases the wrong buffer and the rewritten TMEM base accesses use the wrong address.
Also applies to: 5417-5420
🤖 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/codegen/codegen_cuda.cc` around lines 5194 - 5197, The code
currently strips the "tl_compact_tmem_base_alias" attribute without recording
which buffer it marks, causing compact_tmem_base_var_ (used later to emit
__tl_compact_tbase) to default to the first compacted singleton; update the
handler for the tl_compact_tmem_base_alias AttrStmt (the branch that calls
PrintStmt(op->body)) to capture and store the actual buffer variable being
aliased instead of just printing the body — e.g., extract the buffer Var/IterVar
name from op->body (or from op->node/attr value) and assign it to
compact_tmem_base_var_ so the later selection logic (the code around
compact_tmem_base_var_ usage and the fallback at the 5417-5420 area) uses the
recorded target rather than the first seen allocation. Ensure the same recorded
mapping is consulted where __tl_compact_tbase is generated so the alias targets
the correct buffer.
Comment on lines
+5424
to
+5429
| if (outline_warp_spec_enabled_) { | ||
| // Track all shared variables for device function params | ||
| std::ostringstream type_os; | ||
| PrintType(alloc_dtype, type_os); | ||
| tmem_var_infos_.push_back({vid, "uint*", buffer}); | ||
| } |
Contributor
There was a problem hiding this comment.
Use the real shared-buffer pointer type in outlined helper signatures.
Every shared allocation is recorded as uint* here, regardless of its actual element type. Once an outlined branch touches a half_t*, bfloat16_t*, or any other non-uint32 shared buffer, the generated helper signature no longer matches the call site and NVCC rejects the kernel. Reuse the emitted storage type instead of hard-coding uint*.
🤖 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/codegen/codegen_cuda.cc` around lines 5424 - 5429, The code
currently records all shared allocations as "uint*" when
outline_warp_spec_enabled_ is true, which breaks helper signatures for
non-uint32 storage; change the tmem_var_infos_ push to use the actual emitted
storage type: call PrintType (or reuse the emitted storage type string produced
for alloc_dtype) to form "<storage_type>*" instead of hard-coding "uint*", and
push that string (alongside vid and buffer) into tmem_var_infos_ so outlined
helper signatures match the real shared-buffer pointer type (use the existing
alloc_dtype/PrintType usage to build the correct pointer type).
Comment on lines
+6840
to
+6852
| std::vector<std::string> fn_names; | ||
| std::vector<std::vector<Stmt>> branch_prologues; | ||
| std::vector<Stmt> outlined_bodies; | ||
| std::vector<bool> emit_direct; | ||
| for (const auto &branch : branches) { | ||
| std::vector<Stmt> prologue; | ||
| Stmt outlined_body = ExtractLeadingSetMaxNReg(branch.body, &prologue); | ||
| bool direct = can_emit_directly(outlined_body); | ||
| emit_direct.push_back(direct); | ||
| fn_names.push_back(direct ? std::string() | ||
| : EmitOutlinedDeviceFunction(outlined_body)); | ||
| branch_prologues.push_back(std::move(prologue)); | ||
| outlined_bodies.push_back(outlined_body); |
Contributor
There was a problem hiding this comment.
Outlined warp-branch calls drop captured scalar temps.
This dispatch path only forwards barriers, shared/TMEM buffers, kernel params, and persistent locals. A branch that reads an outer Bind temp or local.var never passes it here, even though the kDeviceFuncScope path already threads scalar_var_infos_ and local_scalar_var_infos_ through the outlined call. The resulting helper will miss required state as soon as outlining is enabled for a branch that closes over scalars.
Also applies to: 6896-6918
🤖 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/codegen/codegen_cuda.cc` around lines 6840 - 6852, The
outlined-branch construction in the loop (symbols: ExtractLeadingSetMaxNReg,
can_emit_directly, EmitOutlinedDeviceFunction and vectors fn_names,
branch_prologues, outlined_bodies, emit_direct) currently omits captured scalar
temps, causing helpers to miss outer Bind/local.var state; update this path to
collect and forward scalar capture metadata (scalar_var_infos_ and
local_scalar_var_infos_) for each outlined_body whether emitted directly or via
EmitOutlinedDeviceFunction, and ensure the EmitOutlinedDeviceFunction call and
any direct-emission dispatch include these scalar var info structures (or a
wrapper) so outlined helpers receive the same scalar/thread-local context that
kDeviceFuncScope already threads through.
Comment on lines
+222
to
+224
| if (call->op.same_as(initialize_tcgen05_descriptor())) { | ||
| return ProxyEvent::kNone; | ||
| } |
Contributor
There was a problem hiding this comment.
Guard call before dereferencing in proxy classification.
At Line 222, call->op is accessed without a null check. ClassifyCallProxyEvent is called with eval->value.as<CallNode>(), so non-call Evaluate values can crash this pass.
Suggested fix
- if (call->op.same_as(initialize_tcgen05_descriptor())) {
+ if (call != nullptr && call->op.same_as(initialize_tcgen05_descriptor())) {
return ProxyEvent::kNone;
}📝 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 (call->op.same_as(initialize_tcgen05_descriptor())) { | |
| return ProxyEvent::kNone; | |
| } | |
| if (call != nullptr && call->op.same_as(initialize_tcgen05_descriptor())) { | |
| return ProxyEvent::kNone; | |
| } |
🤖 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/inject_fence_proxy.cc` around lines 222 - 224, The code
dereferences call->op without checking that call is non-null; update the proxy
classification logic (the branch that compares call->op to
initialize_tcgen05_descriptor()) to first verify the evaluated node is a
CallNode (i.e., that call is non-null) and return ProxyEvent::kNone if it is
not; locate the logic that gets eval->value.as<CallNode>() (used by
ClassifyCallProxyEvent) and add the null/instance check before using call->op to
prevent crashes.
Comment on lines
+190
to
201
| // barrier_id selects which of the 16 hardware named barriers to use. | ||
| // IDs 1 and 2 are kReduce_0 / kReduce_1 in thread_sync_types.h. When two | ||
| // WGs in the same CTA do concurrent inter-thread reductions, they must use | ||
| // DIFFERENT barrier IDs or the bar.sync waits on the wrong cohort and the | ||
| // kernel deadlocks. The codegen passes a per-WG unique ID derived from the | ||
| // thread_offset of the reduce. | ||
| // AllReduce reads `barrier_id` and uses it for the first intra-reduce sync, | ||
| // plus `barrier_id + 1` for the second; codegen must allocate IDs in pairs. | ||
| template <int all_threads, int barrier_id_ = 1> struct NamedBarrier { | ||
| static constexpr int barrier_id = barrier_id_; | ||
| template <int phase = barrier_id_> static TL_DEVICE void sync() { | ||
| asm volatile("bar.sync %0, %1;" : : "r"(phase), "r"(all_threads)); |
Contributor
There was a problem hiding this comment.
The per-WG barrier ID is still not wired into SM90+ reductions.
src/cuda/op/reduce.cc:30-54 still lowers every SM>=90 reduce to tl::NamedBarrier<all_threads>, so this template keeps instantiating with the default barrier_id_=1. AllReduce now consumes barrier_id and barrier_id + 1, which means concurrent WGs still collide on the same {1,2} pair and the deadlock this PR is trying to avoid remains reachable.
🤖 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/tl_templates/cuda/reduce.h` around lines 190 - 201, The template
NamedBarrier is still instantiated with its default barrier_id_ (1); update the
SM90+ reduction lowering in src/cuda/op/reduce.cc so AllReduce passes the per-WG
barrier ID into the template (i.e., instantiate tl::NamedBarrier with the
runtime/per-WG barrier_id as the second template parameter instead of using the
default), ensure the per-WG ID is derived from the WG/thread_offset as the
codegen intends, and keep the pair-allocation rule (AllReduce reads barrier_id
and uses barrier_id+1) while bounding IDs within the available hardware range so
concurrent WGs do not collide on the same {barrier_id, barrier_id+1} pair.
Performance Regression Test ReportTriggered by: @chengyupku Results
Artifacts
|
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/language/builtin.py (1)
468-506:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUnused
lane0parameter has no effect.The
lane0parameter was added to the signature but is never used in the function body. The intrinsic call only forwardsmbarrierandparity. Callers passinglane0=Truewill observe no behavior change.Compare with
tcgen05_commit_1smwhich correctly uses thelane0parameter to set an annotation:def tcgen05_commit_1sm(mbar, *, lane0: bool = False): ... ann = {"lane0_only": True} if lane0 else {} return tirx.call_intrin(..., annotations=ann)Either forward
lane0to the intrinsic (e.g., via annotation) or remove the parameter.🤖 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/builtin.py` around lines 468 - 506, The lane0 parameter on mbarrier_wait_parity is unused; update mbarrier_wait_parity to propagate lane0 into the tirx.call_intrin annotations like tcgen05_commit_1sm does (e.g., build ann = {"lane0_only": True} if lane0 else {} and pass annotations=ann to tirx.call_intrin), or if lane0 is unnecessary, remove the lane0 parameter from the mbarrier_wait_parity signature and any callers; locate the function mbarrier_wait_parity and apply one of these two changes to ensure lane0 behavior is either implemented or the parameter eliminated.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tilelang/language/builtin.py`:
- Around line 468-506: The lane0 parameter on mbarrier_wait_parity is unused;
update mbarrier_wait_parity to propagate lane0 into the tirx.call_intrin
annotations like tcgen05_commit_1sm does (e.g., build ann = {"lane0_only": True}
if lane0 else {} and pass annotations=ann to tirx.call_intrin), or if lane0 is
unnecessary, remove the lane0 parameter from the mbarrier_wait_parity signature
and any callers; locate the function mbarrier_wait_parity and apply one of these
two changes to ensure lane0 behavior is either implemented or the parameter
eliminated.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8d9d20f4-0b9d-4973-aace-407dd2800db7
📒 Files selected for processing (1)
tilelang/language/builtin.py
…xpect_tx codegen The CUTLASS Barrier methods already accept uint32_t — explicit casts broke test assertions that check for `.wait(0)` in generated code. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
10 tasks
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
1 participant
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.
Summary
Add warp-specialized FlashAttention forward kernels targeting SM100 (Blackwell), along with the necessary codegen, DSL, and runtime infrastructure.
New Kernels
gqa_fwd_bshd_pipelined_1sm.py— 1SM 4-role warp-specialized FA (softmax×2, correction, MMA+producer), 512 threads, GQA/MHA, causal supportgqa_fwd_bshd_pipelined_2sm.py— 2CTA cluster FA (D128 + D256), persistent tile loopgqa_fwd_bshd.py— Additional variants (ts_v2, fa4_pipe, fa4_ws)Codegen & Infrastructure
__device__ __noinline__per-role functions with independent register allocation)alloc_tmem(alias=parent, col_offset=N))cluster_dims > 1max3,fmax2_ftz,exp2_poly,exp2_approxcreate_tma_descriptoraccepts Buffer directlytma_load_2smfor 1SM/2SM TMA loadsBug Fixes
Performance (B200, BF16, non-causal, batch=8, q_heads=40, kv_heads=8, dim=128)
Test plan
Summary by CodeRabbit
New Features
Improvements