Add warp specialization benchmarks and improve efficiency - #10
Merged
silentCoder-dev merged 8 commits intoApr 9, 2026
Merged
Conversation
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds BERT token unpad/pad utilities, two TileLang-based attention benchmark suites (FlashMLA and GQA backward) with reference/benchmark harnesses, and compiler updates: shared-memory boundary handling and removal of a Z3 memory-limit constraint. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (2)
examples/auto_schedule/gqa_bwd_benchmark.py (1)
607-613: Wasteful postprocessing: zeros allocated for dK/dV but immediately discarded.In the split branch,
mod_postis called with dummy zero tensors for dK/dV (which allocates memory), but only dQ is used from the result. SincedKanddVfrom line 609-610 are already in fp16 and summed directly on line 613, consider either:
- Creating a simpler postprocess kernel that only handles dQ
- Or avoiding the dummy allocations
♻️ Possible optimization
- dq, _, _ = mod_post(dq, torch.zeros_like(k, dtype=torch.float32), torch.zeros_like(v, dtype=torch.float32))+ # Only postprocess dQ since dK/dV are already fp16 from split kernel+ mod_post_dq = flashattn_bwd_postprocess(total_q, total_kv, H, HEAD_KV, D_HEAD_QK, D_HEAD_V)+ dq, _, _ = mod_post_dq(dq, torch.zeros(1, device=q.device, dtype=torch.float32).expand_as(k), + torch.zeros(1, device=q.device, dtype=torch.float32).expand_as(v))Alternatively, consider factoring out a dQ-only postprocess kernel.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_schedule/gqa_bwd_benchmark.py` around lines 607 - 613, The postprocess allocates zero tensors for dK/dV by calling flashattn_bwd_postprocess (mod_post) with dummy tensors even though only dQ is used and dk/dv are computed separately; fix by adding a dQ-only postprocess path (e.g., implement flashattn_bwd_postprocess_dq or add an optional argument to flashattn_bwd_postprocess to skip dK/dV outputs) and call that from the split branch instead of creating torch.zeros_like(k)/torch.zeros_like(v), or modify mod_post to accept None for dK/dV and avoid allocating them; update the call in the split branch where mod_post(dq, ...) is invoked so only dQ work is performed and leave dk/dv summation (dk = dk.sum(0); dv = dv.sum(0)) unchanged.examples/auto_schedule/bert_padding.py (1)
201-201: Unused variabledimfrom commented-out code.The variable
dimis assigned but never used since the implementation was changed to useindex_put_first_axis. Consider removing it.♻️ Proposed fix
def pad_input(hidden_states, indices, batch, seqlen): """ Arguments: hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. indices: (total_nnz) Return: hidden_states: (batch, seqlen, ...) """ - dim = hidden_states.shape[-1] # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) # output[indices] = hidden_states output = index_put_first_axis(hidden_states, indices, batch * seqlen) return rearrange(output, "(b s) ... -> b s ...", b=batch)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_schedule/bert_padding.py` at line 201, Remove the now-unused variable assignment dim = hidden_states.shape[-1] since the code was refactored to use index_put_first_axis; locate the assignment (variable name dim) near the handling of hidden_states and delete that line so there are no unused locals left in the scope where index_put_first_axis is used.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/auto_schedule/bert_padding.py`:
- Line 179: The line computing cu_seqlens uses an incorrect dtype reference
torch.torch.int32 which will crash; update the F.pad(torch.cumsum(...,
dtype=torch.torch.int32), ...) call to use the correct dtype torch.int32 (i.e.,
change torch.torch.int32 to torch.int32) for the torch.cumsum on
seqlens_in_batch so cu_seqlens is created with a valid integer dtype.
- Line 107: The dtype argument incorrectly uses torch.torch.int32 causing an
AttributeError; update the call that computes cu_seqlens (the line using F.pad
and torch.cumsum with seqlens_in_batch) to use a valid dtype symbol like
torch.int32 (i.e., dtype=torch.int32) instead of torch.torch.int32 so the call
to torch.cumsum and F.pad succeeds at runtime.
In `@examples/auto_schedule/flashmla_benchmark.py`:
- Around line 603-605: The code writes the CUDA source to a relative path
"../debug/flashmla_{label}.cu" which is resolved from the caller's CWD; change
this to compute the dump directory relative to this source file using
pathlib.Path(__file__).resolve().parent / "debug" and build the file path as
debug_dir / f"flashmla_{label}.cu", create the debug_dir with
debug_dir.mkdir(parents=True, exist_ok=True) before writing, and ensure you add
"from pathlib import Path" near the top of the file; keep using
kernel.get_kernel_source() for the file contents.
- Around line 580-601: Validate input shapes in main before calling the kernel
functions (flashattn_auto, flashattn_manual, flashattn_warp_specialize):
explicitly check and raise a clear error for unsupported cases such as kv_heads
!= 1, dim not divisible by 2 (since code uses dim // 2), heads < BLOCK_H or
heads % BLOCK_H != 0 (avoid kv_group_num // BLOCK_H == 0 or dropped tail
blocks), and any other preconditions your kernels require (e.g., BLOCK_H and
BLOCK_N constraints); perform these checks at the start of main (or immediately
before kernel = kernel_fn(...)) and raise ValueError with a descriptive message
so failures occur fast and reproducibly rather than inside JIT/low-level
asserts.
- Around line 610-621: The code currently benchmarks and appends results even
when profiler.assert_allclose fails; change flow so that after catching the
exception from profiler.assert_allclose you skip benchmarking and do not append
to results: either raise/continue immediately in the except block or guard the
benchmarking calls (profiler.do_bench, computing tflops and
results.append((label, latency, tflops, correct))) behind an if correct: check.
Ensure you reference profiler.assert_allclose, profiler.do_bench, total_flops,
and results so the fix is applied at the same spot.
In `@examples/auto_schedule/gqa_bwd_benchmark.py`:
- Around line 789-793: The CLI flag --causal is being overridden by the
hardcoded assignment args.causal = True, so either remove the
add_argument("--causal", ...) if causal must always be true, or delete the line
args.causal = True so the parser-provided value (args.causal) takes effect;
locate the parser.add_argument("--causal", ...) and the args.causal = True
assignment and remove the one that doesn't match the intended behavior.
In `@src/transform/merge_shared_memory_allocations.cc`:
- Around line 569-592: The current loop removes every Evaluate(0) via the
is_noop lambda, which can delete intentional empty TIR statements; instead,
change the filter so it only removes Evaluate(0) nodes that are known remnants
of stripped boundary markers (e.g., by checking a specific boundary
attribute/flag preserved on those nodes) rather than any EvaluateNode with
is_zero(value). Update the filtering around seq->seq to consult that marker (or
a helper that recognizes explicit boundary remnants) when deciding to skip a
stmt, and keep the unwrap_decl_buffers helper unchanged.
- Around line 556-559: In MergeSharedMemoryAllocations(), the current branch
that checks op->attr_key == attr::kAutoScheduleSharedMemoryBoundary returns
StmtExprMutator::VisitStmt(op->body), which removes the
kAutoScheduleSharedMemoryBoundary annotations prematurely; instead preserve the
attribute during the first (dynamic) rewrite so the later static pass can see
boundaries — i.e., stop stripping the attr in the VisitStmt branch (or only
strip when a final-pass flag is set), ensure the traversal returns/mutates the
node without deleting op->attr_key, and perform actual removal of
kAutoScheduleSharedMemoryBoundary only after both dynamic and static
shared-memory rewrite passes complete.
In `@tilelang/transform/z3_scheduler.py`:
- Line 323: Re-enable the shared-memory budget constraint by restoring the
solver.add(...) call that enforces z3.Sum([buffer_vars[i] * buffer_sizes[i] for
i in range(len(buffer_sizes))]) <= memory_limit so buffer_vars are constrained
by the actual memory_limit; specifically, uncomment or re-add the
solver.add(...) using z3.Sum and ensure buffer_vars, buffer_sizes and
memory_limit are Z3-compatible (wrap numeric Python values with z3.IntVal or
cast buffer_sizes/memory_limit to Z3 Ints if needed) so the Sum expression and
comparison are valid in the solver.
---
Nitpick comments:
In `@examples/auto_schedule/bert_padding.py`:
- Line 201: Remove the now-unused variable assignment dim =
hidden_states.shape[-1] since the code was refactored to use
index_put_first_axis; locate the assignment (variable name dim) near the
handling of hidden_states and delete that line so there are no unused locals
left in the scope where index_put_first_axis is used.
In `@examples/auto_schedule/gqa_bwd_benchmark.py`:
- Around line 607-613: The postprocess allocates zero tensors for dK/dV by
calling flashattn_bwd_postprocess (mod_post) with dummy tensors even though only
dQ is used and dk/dv are computed separately; fix by adding a dQ-only
postprocess path (e.g., implement flashattn_bwd_postprocess_dq or add an
optional argument to flashattn_bwd_postprocess to skip dK/dV outputs) and call
that from the split branch instead of creating
torch.zeros_like(k)/torch.zeros_like(v), or modify mod_post to accept None for
dK/dV and avoid allocating them; update the call in the split branch where
mod_post(dq, ...) is invoked so only dQ work is performed and leave dk/dv
summation (dk = dk.sum(0); dv = dv.sum(0)) unchanged.
🪄 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: 8d2f94c0-fc7f-49cf-87a5-c204425d6775
📒 Files selected for processing (5)
examples/auto_schedule/bert_padding.pyexamples/auto_schedule/flashmla_benchmark.pyexamples/auto_schedule/gqa_bwd_benchmark.pysrc/transform/merge_shared_memory_allocations.cctilelang/transform/z3_scheduler.py
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
examples/auto_schedule/gqa_bwd_benchmark.py (1)
807-810:⚠️ Potential issue | 🟡 Minor
--causalCLI flag is still being overridden.Line 810 (
args.causal = True) makes Line 807 ineffective. Either remove the flag or remove the override so CLI input actually works.Suggested fix
parser.add_argument("--causal", action="store_true", help="Causal flag") parser.add_argument("--groups", type=int, default=16, help="groups") args = parser.parse_args() -args.causal = True main(args.batch, args.h, args.n_ctx, args.d_head_qk, args.d_head_v, args.groups, args.causal)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_schedule/gqa_bwd_benchmark.py` around lines 807 - 810, The CLI boolean flag defined with parser.add_argument("--causal", action="store_true", ...) is being overwritten by the line args.causal = True which makes the flag ineffective; remove the explicit override (delete the args.causal = True line) so args.causal reflects the user's CLI input, or alternatively if you want causal to default to True without requiring the flag, change the add_argument call to set default=True and still remove the args.causal = True override; locate these in the parser.add_argument and args.causal usage to apply the change.
🧹 Nitpick comments (2)
examples/auto_schedule/gqa_bwd_benchmark.py (2)
571-571: Drop unusedseqlens_kfromctx.save_for_backward(RUF059).Line 583 unpacks
seqlens_kbut never uses it; storing it also increases autograd context size unnecessarily.Suggested fix
- ctx.save_for_backward(q_unpad, k_unpad, v_unpad, o_unpad, lse, seqlens_q, seqlens_k, cu_seqlens_q, cu_seqlens_k)+ ctx.save_for_backward(q_unpad, k_unpad, v_unpad, o_unpad, lse, seqlens_q, cu_seqlens_q, cu_seqlens_k) ... - q, k, v, o, lse_clone, seqlens_q, seqlens_k, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors+ q, k, v, o, lse_clone, seqlens_q, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensorsAlso applies to: 583-583
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_schedule/gqa_bwd_benchmark.py` at line 571, The backward save currently stores seqlens_k via ctx.save_for_backward(q_unpad, k_unpad, v_unpad, o_unpad, lse, seqlens_q, seqlens_k, cu_seqlens_q, cu_seqlens_k) but seqlens_k is never used in the backward pass; remove seqlens_k from the ctx.save_for_backward call and from the corresponding unpack (where saved tensors are unpacked) so the autograd context doesn't hold that unused tensor; update any tuple unpacking that expects seqlens_k (e.g., in the backward function that destructures saved tensors) to match the new saved list.
556-576:seqlens_k/max_seqlen_kare accepted but not meaningfully used for kernel launch sizing.The wrapper stores both Q/K sequence metadata, but backward kernel sizing uses
ctx.max_seqlen_qonly. If this benchmark is self-attention-only, enforce that invariant explicitly; otherwise wiremax_seqlen_kinto K/V-tiled launch dimensions.Suggested guard for current self-attention assumption
def forward(ctx, q, k, v, seqlens_q, seqlens_k, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal, groups=1): + if not torch.equal(seqlens_q, seqlens_k):+ raise ValueError("This benchmark currently assumes self-attention (seqlens_q must equal seqlens_k).") BATCH, N_CTX, H, D_HEAD_QK = q.shapeAlso applies to: 598-609, 632-633
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_schedule/gqa_bwd_benchmark.py` around lines 556 - 576, The forward wrapper accepts seqlens_k/max_seqlen_k but only uses max_seqlen_q for kernel sizing—either enforce the self-attention invariant or actually use K/V sizes: in forward (function forward) either assert seqlens_k equals seqlens_q (or raise) to document self-attention-only behavior, or change the kernel launch and calls to flashattn_fwd to compute/accept K/V tiling using max_seqlen_k/total_kv and pass max_seqlen_k into the kernel; also store ctx.max_seqlen_k (and any K-specific indices/values) so backward sizing code (the corresponding backward kernel logic referenced around the later backward wrapper) can use the correct K/V dimensions. Ensure consistency where ctx.max_seqlen_q is used so both q and k maxima are available or the self-attention assertion is present.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/auto_schedule/gqa_bwd_benchmark.py`:
- Around line 187-189: Rename the ambiguous lambda parameter `l` in
make_dq_layout to a non-ambiguous name (e.g., `b`) to satisfy the Ruff E741
lint; update the T.Layout call in function make_dq_layout so the lambda
signature becomes (b, h, d) and the returned list uses `b` (i.e., [h, b, d])
while keeping the T.Layout and make_dq_layout identifiers unchanged.
- Line 151: The identifier named "O" is ambiguous and triggers Ruff E741; rename
this parameter/variable to a clear, unambiguous name (e.g., "output",
"output_tensor", or "out_tensor") wherever it's declared as "O: T.Tensor(shape,
dtype), # type: ignore" and update all usages/references in the same scope
(function/class/method) to the new name to ensure consistency and pass linting.
---
Duplicate comments:
In `@examples/auto_schedule/gqa_bwd_benchmark.py`:
- Around line 807-810: The CLI boolean flag defined with
parser.add_argument("--causal", action="store_true", ...) is being overwritten
by the line args.causal = True which makes the flag ineffective; remove the
explicit override (delete the args.causal = True line) so args.causal reflects
the user's CLI input, or alternatively if you want causal to default to True
without requiring the flag, change the add_argument call to set default=True and
still remove the args.causal = True override; locate these in the
parser.add_argument and args.causal usage to apply the change.
---
Nitpick comments:
In `@examples/auto_schedule/gqa_bwd_benchmark.py`:
- Line 571: The backward save currently stores seqlens_k via
ctx.save_for_backward(q_unpad, k_unpad, v_unpad, o_unpad, lse, seqlens_q,
seqlens_k, cu_seqlens_q, cu_seqlens_k) but seqlens_k is never used in the
backward pass; remove seqlens_k from the ctx.save_for_backward call and from the
corresponding unpack (where saved tensors are unpacked) so the autograd context
doesn't hold that unused tensor; update any tuple unpacking that expects
seqlens_k (e.g., in the backward function that destructures saved tensors) to
match the new saved list.
- Around line 556-576: The forward wrapper accepts seqlens_k/max_seqlen_k but
only uses max_seqlen_q for kernel sizing—either enforce the self-attention
invariant or actually use K/V sizes: in forward (function forward) either assert
seqlens_k equals seqlens_q (or raise) to document self-attention-only behavior,
or change the kernel launch and calls to flashattn_fwd to compute/accept K/V
tiling using max_seqlen_k/total_kv and pass max_seqlen_k into the kernel; also
store ctx.max_seqlen_k (and any K-specific indices/values) so backward sizing
code (the corresponding backward kernel logic referenced around the later
backward wrapper) can use the correct K/V dimensions. Ensure consistency where
ctx.max_seqlen_q is used so both q and k maxima are available or the
self-attention assertion is present.
🪄 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: 6108a8bb-3c81-453f-9c39-5fec3d70ff66
📒 Files selected for processing (2)
examples/auto_schedule/gqa_bwd_benchmark.pysrc/transform/merge_shared_memory_allocations.cc
✅ Files skipped from review due to trivial changes (1)
- src/transform/merge_shared_memory_allocations.cc
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
silentCoder-dev
merged commit Apr 9, 2026
753fa7e
into
silentCoder-dev:auto-schedule
1 of 2 checks passed
Uh oh!
There was an error while loading. Please reload this page.
This was referenced Apr 9, 2026
Merged
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 freeto 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.
Summary by CodeRabbit
New Features
New Features
Improvements
Improvements