FA4 consumer Blackwell (sm_120) integration: forward + backward + dispatcher fixes - #1
Closed
thad0ctor wants to merge 224 commits into
Closed
FA4 consumer Blackwell (sm_120) integration: forward + backward + dispatcher fixes#1thad0ctor wants to merge 224 commits into
thad0ctor wants to merge 224 commits into
Conversation
`FlashAttentionForwardSm80.__call__` sets
self.use_tma_O = self.arch >= Arch.sm_90
but the base class kernel never constructs a TMA atom for O — it
passes None as `tma_atom_O` to `self.epilogue` (line ~1069). The check
exists because the file once intended to support a Hopper-style TMA-O
path that was never wired up here.
On Hopper / SM_100 hardware this is dead code because those archs use
their own forward classes (`FlashAttentionForwardSm90` /
`FlashAttentionForwardSm100`) with their own `__call__`. But
`FlashAttentionForwardSm120` inherits from this class, and
`FlashAttentionForwardBase.__init__` reads `self.arch` from the DSL,
which is `Arch.sm_120` on consumer Blackwell. The epilogue then takes
the TMA-output branch and crashes inside
`quack.copy_utils.tma_get_copy_fn` -> `cpasync.tma_partition` with
`AttributeError: 'NoneType' object has no attribute '_trait'`.
Static `arch = 80` on `FlashAttentionForwardSm120` was intended to
prevent this but is overwritten by `__init__`.
Force `use_tma_O = False` here; SM90 and SM100 are unaffected because
they have their own `__call__`.
Reproduced on RTX 5090 (SM_120, cuTeDSL 4.4.2, torch 2.10.0+cu128).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`FlashAttentionForwardSm80.epilogue` chose the rmem->smem store atom
via `get_smem_store_atom(self.arch.major*10 + self.arch.minor, ...)`,
which returns
- `CopyUniversalOp` for arch < 90 (or non-16-bit data), and
- `StMatrix8x8x16bOp(num_matrices=4)` (Hopper `stmatrix`) for
arch >= 90 on 16-bit data.
`stmatrix` is hardware-paired with WGMMA's output register layout. The
SM80 base class uses `mma.sync.aligned.m16n8k16` whose output register
layout is *not* what `stmatrix` consumes. With WGMMA-output the atom
permutes bytes from a fixed pattern of threads/registers; feeding it
SM80-MMA-output silently scrambles values across nearby register
lanes during the store.
On native SM_80 hardware this branch never fires because the DSL arch
is sm_80 < sm_90. The bug only surfaces when this class is reused on
SM_120 via `FlashAttentionForwardSm120`, where `self.arch` is read
from the DSL as `sm_120` and the >= 90 branch picks `stmatrix`.
Symptom: the kernel completes without error and returns the correct
output shape and a roughly correct output norm (each scrambled value
is replaced by a same-magnitude neighbour), but element-wise diffs vs
fp32 SDPA are 0.5-1.2 (non-causal) and 3.4-3.9 (causal), versus
SDPA-bf16's own ~0.003 and ~0.008. Determinism still holds and error
scales linearly with input magnitude — the precision/permutation
signature, not a logic bug.
Fix: pass a fixed `80` to `get_smem_store_atom` here so the SM80 base
class always takes the universal-copy path, matching its actual MMA
output layout.
Verified on RTX 5090 (SM_120, cuTeDSL 4.4.2, torch 2.10.0+cu128):
240/240 correctness configs pass against fp32 SDPA reference across
{fp16, bf16} x {causal, full} x B in {1,2} x S in {128..4096} x
{MHA, GQA, MQA} x D in {64, 128}, with max abs diff matching
SDPA-bf16's own (~0.012 worst case, ~0.0027 mean).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`pack_gqa.compute_ptr` calls
utils.elem_pointer(tensor, ((h_idx, m_idx),))
with a tensor whose layout is supposed to keep a composite
`(qhead_per_kvhead, seqlen_q)` first mode (created by
`pack_gqa_layout`). The slice `mO[None, 0]` that lands in
`compute_ptr` is meant to preserve that compositeness so the rank-2
coord matches.
On SM_120 with `cuTeDSL==4.4.2` the slice collapses the composite
mode into a rank-1 layout. `cute.crd2idx` then refuses the rank-2
coord and raises at trace time with
unable to compute crd2idx with
'!cute.layout<"(?):(?{i64 div=8})">'
and '!cute.coord<"((?,?))">'
resulting in a `ValueError: Operation creation failed` before the
kernel can run. Every default-policy GQA / MQA shape on consumer
Blackwell hits this.
The non-packed GQA path is numerically identical (pack_gqa is a
perf-only optimization for the GQA Q-load / O-store), so flipping
the auto-default to False on SM_120 makes GQA / MQA work out of the
box while a deeper cuTeDSL fix is investigated. Explicit
`pack_gqa=True` from the caller is still honoured.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add FlashAttentionForwardSm120Tma class that uses TMA (cp.async.bulk) for Q/K/V loads with 1 DMA warp + 4 MMA warps, enabling producer/consumer overlap via PipelineTmaAsync with mbarrier synchronization. Key design: - TMA-compatible SMEM swizzle: Swizzle(B, 4, 3) instead of (B, 3, 3) - KV double-buffering (kv_stages=2), 160 threads (5 warps), 99KB SMEM - All pipeline operations inlined in the mainloop (not delegated to a separate @cute.jit method), which avoids CuTe DSL compiler hangs when pipeline states flow through method boundaries - is_first=False with pre-reset softmax state eliminates the need for a compile-time is_first flag in the single-loop mainloop - Dispatch: TMA default for SM120 non-paged, non-varlen. Falls back to CpAsync for paged KV and varlen (TMA addressing constraints). Validated on SM121a (DGX Spark): - 8/8 configs pass: non-causal + causal, B=1/2, Sq=64/128/256, Sk=128/256/512, H=4/8, D=128 - All diffs 0.002-0.008 vs reference Contributed by Second Nature Computing (https://joinsecondnature.com) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Override self.arch = Arch.sm_80 after parent __init__ to prevent base class code paths from seeing the runtime arch (12.x) and enabling SM90+ features. The parent __init__ overwrites the class-level arch=80 attribute with the actual GPU arch. This was found by @2imi9 in Dao-AILab#2420 for the CpAsync kernel — same bug applies here. Add can_implement() check before TMA dispatch in interface.py so that configs exceeding SM120's 99KB SMEM (e.g. hdim=192 with kv_stages=2) fall back to the CpAsync kernel instead of failing at instantiation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com>
The base FlashAttentionForwardSm80.__call__ and FlashAttentionForwardSm100.__call__ both keep `stream` as the final parameter, with a comment: "Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI)". cute.compile binds arguments positionally against the compile_args list in interface.py, which ends with `current_stream`. The TMA kernel had `stream` at position 7 (right after softmax_scale). On this branch alone the kernel still works as advertised, but the mismatch breaks when composed with PRs that append further positional arguments to the compile path — most visibly when combined with Dao-AILab#2348's paged-KV plumbing or Dao-AILab#2439's dropout seeds, where the extra positions push `current_stream` onto a parameter that no longer exists or has the wrong type. Aligning with the base-class convention is mechanical and preserves correctness in isolation: Validation on SM121a (DGX Spark GB10), causal ∈ {False, True}, dtype ∈ {bf16, fp16}, B=1 S=256 H=4 D=64: causal=False bf16: max_diff=0.0020 PASS causal=False fp16: max_diff=0.0002 PASS causal=True bf16: max_diff=0.0078 PASS causal=True fp16: max_diff=0.0010 PASS Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com>
Block-sparse attention processes only the KV blocks specified by block_sparse_tensors rather than the full KV sequence. Two block types are supported: mask_blocks (partially masked, apply mask_mod per element) and full_blocks (fully unmasked, skip masking entirely). Design follows the same mma_one_n_block callback pattern as SM90/SM100. The SM80 base class gets a new mma_one_n_block_bs method (load K, load V, wait, GEMM QK, score_mod, mask, softmax, GEMM PV) and a corresponding run_block_sparse_mainloop_sm80 utility in block_sparse_utils.py that iterates mask blocks then full blocks, mirroring consume_block_sparse_loads. Key implementation details: - run_block_sparse_mainloop_sm80: iterate mask_blocks first (highest n), then full_blocks. First full block always gets mask_seqlen=True since it may be at a higher n position than any mask block. - mma_one_n_block_bs: no pipeline overlap (block address unknown ahead of time), load K then V with separate cp_async_wait_group(1)/wait_group(0). - SM120 inherits SM80 base class and gets block sparsity for free. - FlashAttentionForwardSm120.__init__ forces self.arch = Arch.sm_80 to prevent the SM80 epilogue from using TMA-O (which would crash on SM121a since tma_atom_O is None in this kernel variant). - SM120: num_splits clamped to 1 in interface.py (no SplitKV support yet). - Block sparsity assert removed from SM120 interface path. Validated on SM121a (DGX Spark GB10): - test_block_sparsity.py: 4621 passed, 40 skipped - causal/non-causal, various head dims and sequence lengths Contributed by Second Nature Computing (https://joinsecondnature.com) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Replace inline blocksparse_tensors[0]/[2] index access with the get_total_block_count() utility from block_sparse_utils.py. This keeps variable naming consistent with the rest of the block-sparse codebase (which unpacks by name, not by index position). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds SM120 TMA forward kernel, SM80 forward extensions (paged‑KV, block‑sparse), pack‑GQA-aware backward/postprocess changes, pointer/seqlen/atomic updates, a benchmark tuning harness, and multiple SM120 regression tests. ChangesBenchmark harness & measurement
SM120 TMA forward kernel
SM80 forward extensions (paged‑KV, block‑sparse, pack‑GQA wiring)
SM120 dispatch, validation, and pack‑GQA policies
Backward & postprocess (pack‑GQA, atomics, store paths)
PackGQA, seqlen, mask, and utility fixes
SM120 regression tests and worktree shims
🎯 4 (Complex) | ⏱️ ~60 minutes
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
flash_attn/cute/block_sparse_utils.py (1)
708-758: ⚡ Quick winVarlen path not supported – direct 4D indexing duplicates helper logic.
run_block_sparse_mainloop_sm80hard-codes the non-varlen 4D indexing pattern (lines 750-758) instead of usingget_curr_blocksparse_tensorswhich handles both varlen (2D) and non-varlen (4D) layouts. Other consumers likeconsume_block_sparse_loads(line 403) properly delegate to the helper.If varlen + block-sparse on SM80/SM120 is intended to be supported later, the function should accept
seqlen_infoand use the existing helper. If not supported, consider adding a compile-time assertion.♻️ Suggested refactor to use helper
`@cute.jit` def run_block_sparse_mainloop_sm80( blocksparse_tensors: BlockSparseTensors, batch_idx, head_idx, m_block, mma_one_n_block, mask_fn, mask_mod, fastdiv_mods, + seqlen_info: SeqlenInfoQK, qhead_per_kvhead: cutlass.Constexpr[int] = 1, q_subtile_factor: cutlass.Constexpr[int] = 1, ): ... - mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, *_ = blocksparse_tensors - m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) - - curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block_sparse] - curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None] - - if const_expr(full_block_cnt is not None): - curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse] - curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None] - else: - curr_full_block_cnt = Int32(0) - curr_full_block_idx = None + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, head_idx, m_block_sparse, blocksparse_tensors, seqlen_info + )🤖 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 `@flash_attn/cute/block_sparse_utils.py` around lines 708 - 758, run_block_sparse_mainloop_sm80 duplicates non-varlen 4D indexing for mask/full block tensors instead of using the shared helper; either update run_block_sparse_mainloop_sm80 to accept seqlen_info and call get_curr_blocksparse_tensors(...) (the same helper used by consume_block_sparse_loads) so it correctly handles both varlen (2D) and non-varlen (4D) layouts, or add a compile-time assertion in run_block_sparse_mainloop_sm80 that varlen layouts are not supported; locate the logic around curr_mask_block_cnt/curr_mask_block_idx and curr_full_block_cnt/curr_full_block_idx and replace it with the helper call (or assert) accordingly.flash_attn/cute/flash_fwd_sm120_tma.py (3)
566-566: 💤 Low valueRemove unused variable
n_block.This variable is computed but never used in the kernel. The loop at line 784 computes
cur_n_block = n_block_max - n_tile - 1directly.🧹 Suggested cleanup
n_block_min, n_block_max = block_info.get_n_block_min_max( seqlen, m_block, split_idx, num_splits ) - n_block = cutlass.max(n_block_max - 1, 0)🤖 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 `@flash_attn/cute/flash_fwd_sm120_tma.py` at line 566, Remove the unused local variable n_block (the expression n_block = cutlass.max(n_block_max - 1, 0)) since it is never referenced later; update the surrounding code to rely on the existing computation cur_n_block = n_block_max - n_tile - 1 (and any uses of n_block) so only n_block_max, n_tile, and cur_n_block remain; ensure no other code paths reference n_block before deleting the assignment and related dead-code.
23-44: 💤 Low valueRemove unused imports.
Several imports are flagged by static analysis as unused:
Constexpr(line 23) — used ascutlass.Constexprinstead of the bare namePackGQA(line 37)NamedBarrierFwd(line 38)FastDivmodDivisor(line 44)🧹 Suggested cleanup
-from cutlass import Constexpr, Float32, Int32, const_expr +from cutlass import Float32, Int32, const_expr-from flash_attn.cute.pack_gqa import PackGQA -from flash_attn.cute.named_barrier import NamedBarrierFwd-from cutlass.cute import FastDivmodDivisor🤖 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 `@flash_attn/cute/flash_fwd_sm120_tma.py` around lines 23 - 44, Remove the unused imports to clean up the module: delete the bare imports Constexpr, PackGQA, NamedBarrierFwd, and FastDivmodDivisor from the top import block; ensure code that currently refers to cutlass.Constexpr continues to use the qualified name (cutlass.Constexpr) instead of the removed bare Constexpr, and verify there are no references to PackGQA, NamedBarrierFwd, or FastDivmodDivisor elsewhere in this file (e.g., search for PackGQA, NamedBarrierFwd, FastDivmodDivisor) before removing them to avoid breaking references.
914-1003: 💤 Low value
mma_one_n_blockmethod is unused.This method is defined but never called. The kernel at lines 784–837 has the same logic inlined directly, with a comment at lines 780–783 explaining this is intentional to avoid CuTe DSL compiler hangs when pipeline states flow through method boundaries.
If this is dead code from development, consider removing it to reduce maintenance burden. If it's intended for future use, consider adding a comment indicating that.
🤖 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 `@flash_attn/cute/flash_fwd_sm120_tma.py` around lines 914 - 1003, The method mma_one_n_block is defined but never used (its logic is inlined inside the kernel to avoid CuTe DSL compiler hangs); either remove this dead function to reduce maintenance or keep it but add a clear comment above mma_one_n_block stating it is intentionally unused and kept for reference/future use (mention the kernel that inlines the logic and the compiler-hang rationale), so future readers know why it remains in the codebase.
🤖 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 `@flash_attn/cute/flash_fwd.py`:
- Around line 1384-1396: The call to apply_score_mod is passing seqlen as the
positional argument where softmax_scale belongs, causing an argument-order bug;
update the call in the const_expr(score_mod is not None) branch to pass
softmax_scale and seqlen by keyword (e.g., softmax_scale=softmax.softmax_scale,
seqlen=seqlen) or otherwise ensure softmax_scale is the 7th positional and
seqlen the 8th; reference the apply_score_mod invocation that currently uses
mma_params.thr_mma_qk, batch_idx, head_idx, m_block, acc_S, n_block, seqlen,
softmax_scale=... and mirror the correct ordering used in compute_one_n_block.
In `@flash_attn/cute/interface.py`:
- Around line 954-960: The local reassignment of is_varlen near the TMA kernel
selection overrides the earlier definition that includes seqused_q/seqused_k;
remove the redefinition (the lines that set is_varlen = cu_seqlens_q is not None
or cu_seqlens_k is not None) and use the previously computed is_varlen when
evaluating use_tma_sm120 (alongside page_table and use_block_sparsity) so
callers providing seqused_q/seqused_k are correctly treated as varlen.
---
Nitpick comments:
In `@flash_attn/cute/block_sparse_utils.py`:
- Around line 708-758: run_block_sparse_mainloop_sm80 duplicates non-varlen 4D
indexing for mask/full block tensors instead of using the shared helper; either
update run_block_sparse_mainloop_sm80 to accept seqlen_info and call
get_curr_blocksparse_tensors(...) (the same helper used by
consume_block_sparse_loads) so it correctly handles both varlen (2D) and
non-varlen (4D) layouts, or add a compile-time assertion in
run_block_sparse_mainloop_sm80 that varlen layouts are not supported; locate the
logic around curr_mask_block_cnt/curr_mask_block_idx and
curr_full_block_cnt/curr_full_block_idx and replace it with the helper call (or
assert) accordingly.
In `@flash_attn/cute/flash_fwd_sm120_tma.py`:
- Line 566: Remove the unused local variable n_block (the expression n_block =
cutlass.max(n_block_max - 1, 0)) since it is never referenced later; update the
surrounding code to rely on the existing computation cur_n_block = n_block_max -
n_tile - 1 (and any uses of n_block) so only n_block_max, n_tile, and
cur_n_block remain; ensure no other code paths reference n_block before deleting
the assignment and related dead-code.
- Around line 23-44: Remove the unused imports to clean up the module: delete
the bare imports Constexpr, PackGQA, NamedBarrierFwd, and FastDivmodDivisor from
the top import block; ensure code that currently refers to cutlass.Constexpr
continues to use the qualified name (cutlass.Constexpr) instead of the removed
bare Constexpr, and verify there are no references to PackGQA, NamedBarrierFwd,
or FastDivmodDivisor elsewhere in this file (e.g., search for PackGQA,
NamedBarrierFwd, FastDivmodDivisor) before removing them to avoid breaking
references.
- Around line 914-1003: The method mma_one_n_block is defined but never used
(its logic is inlined inside the kernel to avoid CuTe DSL compiler hangs);
either remove this dead function to reduce maintenance or keep it but add a
clear comment above mma_one_n_block stating it is intentionally unused and kept
for reference/future use (mention the kernel that inlines the logic and the
compiler-hang rationale), so future readers know why it remains in the codebase.
🪄 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: 9c771396-cc50-48ee-b2a8-55805d806e5e
📒 Files selected for processing (7)
flash_attn/cute/block_sparse_utils.pyflash_attn/cute/flash_bwd.pyflash_attn/cute/flash_fwd.pyflash_attn/cute/flash_fwd_sm120.pyflash_attn/cute/flash_fwd_sm120_tma.pyflash_attn/cute/interface.pyflash_attn/cute/utils.py
thad0ctor
added a commit
that referenced
this pull request
May 24, 2026
…shadowing 1. flash_fwd.py:1392 — apply_score_mod() was called with seqlen as the 7th positional argument, where the signature places softmax_scale. The original call then passed softmax_scale=... as a keyword, which would have raised 'multiple values for argument softmax_scale' under strict Python, or silently bound seqlen as the softmax_scale under CuTeDSL's relaxed semantics (producing wildly wrong attention scores in the score_mod path). Fix by passing both softmax_scale and seqlen by keyword to match the correct call pattern in compute_one_n_block at lines 1254-1266. 2. interface.py:955 — the SM120 dispatch block redefined is_varlen with a narrower check (cu_seqlens_q/cu_seqlens_k only), shadowing the outer-scope is_varlen defined at lines 626-631 which correctly includes seqused_q and seqused_k. A caller passing seqused_q/seqused_k without cu_seqlens would have been silently routed to the TMA kernel, which does not support varlen, producing wrong output or a crash. Remove the local redefinition so the outer is_varlen is used; add a comment so this doesn't get re-introduced. Both flagged by CodeRabbit on PR #1. Regression smokes pass: fwd: max abs diff 0.003906 vs SDPA, bwd: max abs diff 0.007812.
thad0ctor
added a commit
that referenced
this pull request
May 24, 2026
flash_fwd_sm120_tma.py: - Remove unused imports: Constexpr (the file uses cutlass.Constexpr via qualified name), PackGQA, NamedBarrierFwd, FastDivmodDivisor. - Remove unused local n_block at line 566 — the kernel's compute loop computes cur_n_block = n_block_max - n_tile - 1 directly and never references n_block. - Add comment above mma_one_n_block explaining why the method is kept despite being unused: its logic is inlined into kernel() (around lines 784-837) to avoid a CuTe DSL compiler hang when k_pipeline/v_pipeline consumer states flow through a method boundary. Kept as a reference template. block_sparse_utils.py: - run_block_sparse_mainloop_sm80 hard-codes non-varlen 4D indexing into the blocksparse_tensors NamedTuple. Add a docstring note documenting why this is intentionally narrow (the SM120 dispatcher in interface.py does not enable varlen + block-sparse together) and what would need to change to lift it (route through get_curr_blocksparse_tensors with seqlen_info threaded in). Regression smokes pass: fwd max abs diff 0.003906, bwd 0.007812 vs SDPA.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/cute/test_flash_attn_bwd_sm120_postprocess.py`:
- Around line 71-73: The inline continuation comment for the variable S in the
test (around the assignment "S = 1024") is misaligned and triggers Flake8
E114/E116; update the comment so its indentation matches the wrapped statement's
indentation (align the comment with the start of the "S = 1024" line or the
wrapped continuation indent) to satisfy linting rules, touching the
test_flash_attn_bwd_sm120_postprocess.py section containing S = 1024 and the
following comment.
In `@tests/cute/test_paged_kv_sm120.py`:
- Around line 36-38: Remove the unused imports causing Flake8 F401 in
tests/cute/test_paged_kv_sm120.py by deleting the import lines for math and os
(leave the necessary import for sys if used elsewhere); ensure only required
modules are imported so the top-of-file import block no longer includes math or
os.
🪄 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: b8475edd-852f-436c-8947-c425cf37be30
📒 Files selected for processing (10)
flash_attn/cute/block_sparse_utils.pyflash_attn/cute/flash_bwd_postprocess.pyflash_attn/cute/flash_fwd.pyflash_attn/cute/flash_fwd_sm120.pyflash_attn/cute/flash_fwd_sm120_tma.pyflash_attn/cute/interface.pyflash_attn/cute/pack_gqa.pyflash_attn/cute/seqlen_info.pytests/cute/test_flash_attn_bwd_sm120_postprocess.pytests/cute/test_paged_kv_sm120.py
thad0ctor
added a commit
that referenced
this pull request
May 24, 2026
CodeRabbit's second review on PR #1 flagged two lint issues introduced by the Phase 4-S bwd-postprocess test and the Phase 4-R paged-KV test: * tests/cute/test_flash_attn_bwd_sm120_postprocess.py:72 The wrapped comment on the second line of `S = 1024 # ...` was indented to align with the value (column 15), which trips E114 ("indentation is not a multiple of x (comment)") and E116 ("unexpected indentation (comment)"). Convert it to a normal 4-space-indented comment block above the assignment. * tests/cute/test_paged_kv_sm120.py:36-37 `import math` and `import os` were left over from an earlier draft of the regression test and are not referenced anywhere in the file (F401). Drop both. No functional change to either test. flake8 with --select=E114,E116,F401 is now clean on both files.
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)
flash_attn/cute/interface.py (1)
1105-1105:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPhase 5c tuned
num_stagesvalues are never applied to SM120 kernels.The
sm120_num_stagesvariable is computed at line 585 from the_SM120_TILE_LOOKUP(e.g.,(64, 1, 2048, 1): (64, 64, 2)setsnum_stages=2), but line 1105 hardcodesnum_stages=1instead of using the tuned value.This means the Phase 5c performance tuning for SM120 is ineffective.
🐛 Proposed fix
fa_fwd = FlashAttentionForwardSm120( dtype, head_dim, head_dim_v, qhead_per_kvhead, is_causal=causal, is_local=local, pack_gqa=pack_gqa, tile_m=tile_m, tile_n=tile_n, - num_stages=1, + num_stages=sm120_num_stages, num_threads=num_threads, Q_in_regs=False, score_mod=score_mod, mask_mod=mask_mod, has_aux_tensors=aux_tensors is not 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 `@flash_attn/cute/interface.py` at line 1105, The code hardcodes num_stages=1 where SM120 kernels should use the tuned sm120_num_stages computed from _SM120_TILE_LOOKUP; update the call/site that currently passes num_stages=1 to instead pass the variable sm120_num_stages (the value computed at line ~585) so Phase 5c SM120 tuning is applied, ensuring any downstream uses (function/class where num_stages is supplied) accept and propagate that variable rather than the literal 1.
🧹 Nitpick comments (1)
flash_attn/cute/interface.py (1)
899-901: 💤 Low valueComment is misleading — Phase 5c lookup does not run for SM80.
The comment claims
num_stagesis set by "Phase 5c per-shape lookup," but the_SM120_TILE_LOOKUPlogic (lines 535–592) is guarded byif arch // 10 == 12:, so it never executes for SM80. For SM80,sm120_num_stagesstays at the default value of 1.Consider updating the comment to reflect the actual behavior:
- # num_stages set by Phase 5c per-shape lookup above (defaults - # to 1; bumped to 2 for shapes where K/V pipelining wins). + # SM80 uses 1 stage; SM120 overrides this in its own branch. num_stages=sm120_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 `@flash_attn/cute/interface.py` around lines 899 - 901, The comment is misleading about Phase 5c affecting num_stages; update the comment near the num_stages=num_stages=sm120_num_stages assignment to say that the Phase 5c per-shape lookup (_SM120_TILE_LOOKUP) runs only when arch // 10 == 12 (SM120), so for SM80 sm120_num_stages remains the default 1; mention that num_stages is therefore conditionally adjusted only for SM120 and defaults to 1 on SM80 to avoid confusion.
🤖 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 `@flash_attn/cute/interface.py`:
- Line 1105: The code hardcodes num_stages=1 where SM120 kernels should use the
tuned sm120_num_stages computed from _SM120_TILE_LOOKUP; update the call/site
that currently passes num_stages=1 to instead pass the variable sm120_num_stages
(the value computed at line ~585) so Phase 5c SM120 tuning is applied, ensuring
any downstream uses (function/class where num_stages is supplied) accept and
propagate that variable rather than the literal 1.
---
Nitpick comments:
In `@flash_attn/cute/interface.py`:
- Around line 899-901: The comment is misleading about Phase 5c affecting
num_stages; update the comment near the num_stages=num_stages=sm120_num_stages
assignment to say that the Phase 5c per-shape lookup (_SM120_TILE_LOOKUP) runs
only when arch // 10 == 12 (SM120), so for SM80 sm120_num_stages remains the
default 1; mention that num_stages is therefore conditionally adjusted only for
SM120 and defaults to 1 on SM80 to avoid confusion.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a5f532f7-b808-4704-ac78-59bf83310c26
📒 Files selected for processing (3)
flash_attn/cute/interface.pytests/cute/test_flash_attn_bwd_sm120_postprocess.pytests/cute/test_paged_kv_sm120.py
💤 Files with no reviewable changes (1)
- tests/cute/test_paged_kv_sm120.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@flash_attn/cute/interface.py`:
- Around line 1104-1107: The forward JIT cache key (compile_key) currently omits
sm120_num_stages while the SM120 kernel selection uses it
(num_stages=sm120_num_stages), leading to cache collisions across shapes; update
the compile_key construction (the same place that builds the forward compile
cache key) to include sm120_num_stages so that compile_key distinguishes
variants by num_stages and prevents reusing a compiled kernel with the wrong
sm120_num_stages.
🪄 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: 8d4bc294-813d-427f-b7d8-eb453f985493
📒 Files selected for processing (1)
flash_attn/cute/interface.py
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)
flash_attn/cute/interface.py (1)
1028-1042:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
can_implementcheck uses hardcodednum_stages=1instead ofsm120_num_stages.The
can_implementcall at line 1033 passesnum_stages=1, but the kernel is instantiated at line 1053 withnum_stages=sm120_num_stages. Sincesm120_num_stagescan be 2 from the Phase 5c lookup table (e.g.,(64, 1, 2048, 1): (64, 64, 2)), the SMEM check validates for 1-stage but the kernel runs with 2 stages — potentially overflowing the 99 KB SMEM cap.🐛 Proposed fix
assert FlashAttentionForwardSm120.can_implement( dtype, head_dim, head_dim_v, tile_m, tile_n, - num_stages=1, num_threads=num_threads, is_causal=causal, + num_stages=sm120_num_stages, num_threads=num_threads, is_causal=causal, Q_in_regs=False, ), (🤖 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 `@flash_attn/cute/interface.py` around lines 1028 - 1042, The SM 12.0 validation wrongly hardcodes num_stages=1 when calling FlashAttentionForwardSm120.can_implement, which can mis-validate SMEM usage vs the actual kernel instantiation using sm120_num_stages; update the can_implement call to pass num_stages=sm120_num_stages (preserving the other parameters like dtype, head_dim, head_dim_v, tile_m, tile_n, num_threads, is_causal, Q_in_regs) so the SMEM/hang/divisibility checks reflect the real kernel configuration used later when creating the FlashAttentionForwardSm120 kernel.
🧹 Nitpick comments (2)
flash_attn/cute/flash_fwd.py (1)
1700-1701: 💤 Low valueDead variables:
smem_pipe_readandsmem_pipe_writeare unused.These variables are initialized but never read. Since
num_stages == 1is enforced by the assertion at line 1655, all SMEM accesses use hardcoded index0(e.g.,sK[None, None, 0]). Consider removing them.♻️ Proposed fix to remove dead variables
- smem_pipe_read = Int32(0) - smem_pipe_write = Int32(0) nb = n_block🤖 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 `@flash_attn/cute/flash_fwd.py` around lines 1700 - 1701, Remove the dead variables smem_pipe_read and smem_pipe_write which are initialized but never used; since the assertion enforcing num_stages == 1 (see assertion around num_stages) makes all SMEM accesses use the fixed index 0 (e.g., sK[None, None, 0]), delete the Int32(0) declarations for smem_pipe_read and smem_pipe_write and any related unused references so there is no unused state left in the flash_fwd.py top-level scope.flash_attn/cute/block_sparse_utils.py (1)
788-811: 💤 Low valueConsider explicitly passing
mask_mod=Nonefor full blocks for consistency.In
consume_block_sparse_loads(lines 470, 480), when transitioning from mask blocks to full blocks,mask_mod=Noneis explicitly passed. Here,mask_modis omitted for full blocks, relying on a default value inapply_mask. While this likely works ifapply_maskdefaultsmask_mod=None, explicit passing would be clearer and consistent with the SM90/SM100 consumer path.♻️ Proposed fix for explicit mask_mod=None
if const_expr(full_block_cnt is not None): if curr_full_block_cnt > 0: n_block = curr_full_block_idx[curr_full_block_cnt - 1] if curr_mask_block_cnt == 0: mma_one_n_block( n_block=n_block, - mask_fn=partial(mask_fn, mask_seqlen=True), + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True), is_first_n_block=True, ) else: mma_one_n_block( n_block=n_block, - mask_fn=partial(mask_fn, mask_seqlen=True), + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True), is_first_n_block=False, ) for j in cutlass.range(1, curr_full_block_cnt): n_block = curr_full_block_idx[curr_full_block_cnt - 1 - j] mma_one_n_block( n_block=n_block, - mask_fn=partial(mask_fn, mask_seqlen=False), + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False), is_first_n_block=False, )🤖 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 `@flash_attn/cute/block_sparse_utils.py` around lines 788 - 811, The full-block handling in the block-sparse consumer omits an explicit mask_mod, relying on apply_mask's default; change the two mma_one_n_block calls in the full-block section so they pass mask_mod=None (i.e., the first call with mask_fn=partial(mask_fn, mask_seqlen=True), is_first_n_block=True/False should also include mask_mod=None, and the subsequent calls using mask_seqlen=False should likewise include mask_mod=None) to match the consume_block_sparse_loads behavior and the SM90/SM100 path and make intent explicit.
🤖 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 `@flash_attn/cute/interface.py`:
- Around line 1028-1042: The SM 12.0 validation wrongly hardcodes num_stages=1
when calling FlashAttentionForwardSm120.can_implement, which can mis-validate
SMEM usage vs the actual kernel instantiation using sm120_num_stages; update the
can_implement call to pass num_stages=sm120_num_stages (preserving the other
parameters like dtype, head_dim, head_dim_v, tile_m, tile_n, num_threads,
is_causal, Q_in_regs) so the SMEM/hang/divisibility checks reflect the real
kernel configuration used later when creating the FlashAttentionForwardSm120
kernel.
---
Nitpick comments:
In `@flash_attn/cute/block_sparse_utils.py`:
- Around line 788-811: The full-block handling in the block-sparse consumer
omits an explicit mask_mod, relying on apply_mask's default; change the two
mma_one_n_block calls in the full-block section so they pass mask_mod=None
(i.e., the first call with mask_fn=partial(mask_fn, mask_seqlen=True),
is_first_n_block=True/False should also include mask_mod=None, and the
subsequent calls using mask_seqlen=False should likewise include mask_mod=None)
to match the consume_block_sparse_loads behavior and the SM90/SM100 path and
make intent explicit.
In `@flash_attn/cute/flash_fwd.py`:
- Around line 1700-1701: Remove the dead variables smem_pipe_read and
smem_pipe_write which are initialized but never used; since the assertion
enforcing num_stages == 1 (see assertion around num_stages) makes all SMEM
accesses use the fixed index 0 (e.g., sK[None, None, 0]), delete the Int32(0)
declarations for smem_pipe_read and smem_pipe_write and any related unused
references so there is no unused state left in the flash_fwd.py top-level scope.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4f343344-27a3-4116-b112-e6513be015c2
📒 Files selected for processing (9)
flash_attn/cute/block_sparse_utils.pyflash_attn/cute/flash_bwd_postprocess.pyflash_attn/cute/flash_fwd.pyflash_attn/cute/flash_fwd_sm120.pyflash_attn/cute/flash_fwd_sm120_tma.pyflash_attn/cute/interface.pyflash_attn/cute/pack_gqa.pyflash_attn/cute/seqlen_info.pyflash_attn/cute/utils.py
thad0ctor
changed the title
… + tile tuning Makes the three cherry-picked upstream SM120 PRs (Dao-AILab#2553, Dao-AILab#2349, Dao-AILab#2389) actually usable end-to-end on consumer Blackwell (RTX 5090, RTX PRO 6000 Blackwell). The upstream PRs alone leave SM120 forward dispatcher-buggy and backward broken; this commit adds the integration glue, backward support, real paged-KV + pack_gqa implementations, a subprocess-isolated per-shape tile lookup, and the test coverage to back it up. # Dispatcher fixes (SM120 forward + backward couldn't compile or run end-to-end # without these) - Initialize dQ_single_wg in the SM120 backward setup (was unbound) - Keep softmax_scale non-None for SM80/SM120 backward dK epilogue (inline log2 computation like SM90 does) - atomic_add_fp32: adopt the new keyword-only nvvm.atomicrmw signature in nvidia-cutlass-dsl >= 4.x - Drop unsupported is_split_kv kwarg on the SM120 forward path - Pass split_idx=0, num_splits=1, seqlen_info=seqlen at the SM80/SM120 get_total_block_count call site (arity mismatch fix) - Rename vec_size -> score_vec_size on the SM120 TMA forward (typo in upstream Dao-AILab#2349; was AttributeError on softcap/learnable_sink/score_mod) - Auto-downgrade pack_gqa=False when qhead_per_kvhead doesn't divide tile_m=128 (qwen2.5-7b's 7-way GQA otherwise fails cute.local_tile division) - Auto-downgrade pack_gqa=False when paged-KV is used (cross-feature interaction with PagedKVManager's K/V indexing) - Route head_dim > head_dim_v to the non-TMA SM120 path: bisection showed the hang lives in FlashAttentionForwardSm120Tma, not in the SM80-base mainloop as upstream diagnosed. Non-TMA can_implement accepts d > dv; the TMA path still rejects it so the dispatcher falls through. d > dv shapes now work (verified bitwise-identical to SDPA on the minimum repro). - Route SM120 through the shared _validate_head_dims helper (invalid head_dim was reaching the kernel and faulting with cudaErrorMisalignedAddress) - Clamp the cu_seqlens[batch_idx+1] read in SeqlenInfoQK.create so SM80/SM120 over-launched varlen tiles don't fault on a non-resident page - arch-gate FlashAttentionForwardBase.epilogue smem store atom: SM80/SM120 force the universal copy, SM90 keeps WGMMA-paired stmatrix (upstream PR Dao-AILab#2553's bc67a9c unconditionally forced 80, which silently switched SM90 forward through the universal-copy path) - Include sm120_num_stages in the forward compile cache key (different ns values with the same tile would otherwise share a key and the second call would reuse the first-compiled kernel) - Document why deterministic backward can't be lifted on SM120 (the SM80 base kernel itself lacks the dQ_semaphore code path; a feature gap shared with SM80) # SM120-specific kernel work - Real paged-KV forward via PagedKVManager on the SM80-base kernel, supported through head_dim <= 128. A paged-specific tile override (128, 128, ns=1) gates on page_table is not None and head_dim <= 128 so PagedKVManager's tile_n >= num_threads invariant holds. SMEM math fits: 48 KB at d=64, 72 KB at d=96, 96 KB at d=128 (cap 99 KB). - Real pack_gqa=True support: rewrite PackGQA.compute_ptr to compute the flat offset arithmetically from stride[0][0] and stride[0][1] rather than cute.crd2idx (which cuTeDSL 4.4-4.5 collapses through trailing slices). Call pack_gqa_layout in the SM80-base forward so packed Q is actually materialized (was missing — would have produced wrong output even after the crd2idx workaround). - Backward postprocess dQ smem-store atom: force universal copy on SM80/SM120 (same class of bug as the upstream Dao-AILab#2553 forward fix but in the dQ postprocess; left silent rmem->smem scrambling otherwise). Permanent regression test with a white-box source-inspection guard against reintroduction. - New D > 128 SM120 tile bracket (64, 64, ns=1) that fits the 99 KB SMEM cap for head_dim=256. # Forward tile selection (per-shape lookup) The SM120 forward dispatch now consults a tile + num_stages lookup keyed on (head_dim, qhead_per_kvhead, seqlen, causal). Shapes outside the lookup fall back to the head_dim-only brackets that match the pre-tuning defaults. The lookup was built from a subprocess-isolated sweep: each (cell, candidate) pair is measured in a fresh python process so JIT-cache pollution can't bias the rankings (a single-process sweep silently reuses compiled kernels across candidates with subtly different shapes). The top-3 candidates per cell get a reproducibility re-measurement; variance > 10% excludes a candidate. A candidate ships only when its mean TFLOPS beats the baseline tile by >= 2%; otherwise the cell falls back to baseline. # Test coverage added - tests/cute/test_paged_kv_sm120.py (38 cases): paged-KV correctness across page_size {16, 64, 256}, identity / permuted / shared page tables, GQA + MQA, d in {64, 96, 128}; expected NotImplementedError for d in {192, 256}; expected correctness (now, not rejection) for the paged + d > dv + varlen cross-feature combination. - tests/cute/test_flash_attn_bwd_sm120_postprocess.py (10 cases): backward dQ postprocess regression suite, combines numeric vs fp32-SDPA comparison with a white-box source-inspection guard against the buggy literal pattern. - tests/cute/test_flash_attn_sm120_dgtdv.py (11 cases): regression test for the Bug E d > dv non-TMA routing. 8 kernel-launch parametrizations plus 3 unit probes (TMA rejection, non-TMA acceptance, SMEM constraint). All kernel tests carry pytest-timeout(30) with --timeout-method=signal so a future TMA gate widening that re-introduces the GPU hang fails as a timeout instead of wedging the GPU. # What this is NOT - Real paged-KV at head_dim > 128: rejected at dispatch with a clear NotImplementedError. Lifting would require either a refactor of PagedKVManager (per-thread page-table fragment > 0 at tile_n < num_threads) or a separate kernel; the 99 KB SMEM cap precludes the simple (128, 128, ns=1) approach used for d <= 128. - Real fix for the TMA path d > dv hang: the kernel-level root cause needs cuda-gdb or instrumented bisection; the routing fix makes user-visible shapes correct today, but the TMA kernel itself is still latent-broken for d > dv. The can_implement gate ensures the TMA path is never selected for d > dv. - Deterministic backward on SM120: asserts off because the SM80 base kernel itself lacks the dQ_semaphore code path. Lift would need a feature port from SM90 into the SM80 base; out of scope here.
thad0ctor
force-pushed
the
sm120-integrate
branch
from
May 25, 2026 01:23
53b1ad5 to
cee3b54
Compare
Owner
Author
|
@coderabbitai review Force-pushed: squashed all thad0ctor commits into one (
Please do a full review on the squashed diff. |
|
✅ Actions performedFull review triggered. |
Owner
Author
|
@coderabbitai full review |
✅ Actions performedFull review triggered. |
Owner
Author
|
Independent verification of squashed tip ( Re-ran the full test plan from a worktree-isolated agent on a fresh branch off
GPU pinned by UUID to the RTX 5090 ( |
Owner
Author
|
@coderabbitai full review |
✅ Actions performedFull review triggered. |
Owner
Author
|
FA4 (this PR) vs FA2 2.8.3 benchmark — RTX 5090, sm_120, bf16, 160 paired cells + 8 FA4-only cells Head-to-head on the same RTX 5090, same torch 2.11.0+cu130. Fresh subprocess per cell, 3 warmup + 10 timed CUDA-event iterations, median latency, Dao FLOPs convention. All 168 cells returned Forward (40 paired cells):
Backward (40 paired cells):
FA4-only shapes (no FA2 reference):
These shapes either hung the GPU ( Full per-cell table now in the PR body under "Performance vs FlashAttention 2". |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/cute/test_flash_attn_bwd_sm120_postprocess.py`:
- Around line 151-156: The test's assertion in
tests/cute/test_flash_attn_bwd_sm120_postprocess.py is brittle because it checks
an exact multi-line string for "get_smem_store_atom(\n
self.arch,"; update the guard to use a regex that ignores whitespace (e.g., use
re.search with a pattern like r"get_smem_store_atom\s*\(\s*self\.arch\s*,") so
it fails if self.arch is still passed as the first positional argument but
survives formatting changes; locate the assertion referencing
flash_bwd_postprocess.FlashAttentionBackwardPostprocess and replace the string
containment check with a whitespace-agnostic regex match that ensures self.arch
is not the first positional parameter (consider also allowing store_atom_arch to
appear).
In `@tests/cute/test_paged_kv_sm120.py`:
- Around line 223-230: The tests use Python's built-in hash(page_table_pattern)
for seeding (in test_page_table_patterns and the other paged-case test that
calls _run_paged_case), which is randomized per process; replace that with a
deterministic hash function or mapping (e.g., compute a stable integer from
page_table_pattern using hashlib.sha256 or zlib.crc32 and then mask to 16 bits)
and pass that stable seed into _run_paged_case so runs are reproducible; update
both occurrences where seed=hash(page_table_pattern) & 0xFFFF to use the
deterministic conversion instead.
🪄 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: 40870c62-2a0d-4695-a0c5-d30a2d82106c
📒 Files selected for processing (13)
flash_attn/cute/block_sparse_utils.pyflash_attn/cute/flash_bwd.pyflash_attn/cute/flash_bwd_postprocess.pyflash_attn/cute/flash_fwd.pyflash_attn/cute/flash_fwd_sm120.pyflash_attn/cute/flash_fwd_sm120_tma.pyflash_attn/cute/interface.pyflash_attn/cute/pack_gqa.pyflash_attn/cute/seqlen_info.pyflash_attn/cute/utils.pytests/cute/test_flash_attn_bwd_sm120_postprocess.pytests/cute/test_flash_attn_sm120_dgtdv.pytests/cute/test_paged_kv_sm120.py
CodeRabbit (PR #1) flagged the exact-string membership check on `get_smem_store_atom(\n self.arch,` as brittle: any reformat of the source line would silently disarm the regression guard even if `self.arch` were still passed as the first positional arg. Switch the assertion to `re.search(r"get_smem_store_atom\(\s*self\.arch\s*,", src)` so the guard fails on the buggy pattern regardless of whitespace or line-wrap formatting. Closes-Upstream-Comment: #1
…lls), no dispatch win; swizzle+cp.async already done
…+occupancy-walled; spill is 8.5% of traffic (~2% upside), surgery not worth it
…head_dim-split needs 4x recompute, n=32 proxy shows +16% regress -> predicted net loss
…fault grid already exploits the K/V/Q/dO L2 locality) — backward confirmed at floor
The D256 Q-in-regs wide tile (128x64+Qregs+256t) was gated to non-varlen (cu_seqlens). varlen D256 forward (packed-sequence training) ran at the 64x64 fallback; the wide tile helps it too: RTX6000 A/B +7-11% (qwen3.5-122b qpkv16 1.041->1.150 vs FA2, qwen3.5-9b 1.038->1.124, qwen3.6-35b 1.039->1.116). Relaxed the cu_seqlens exclusion in sm120_d256_wide (seqused stays excluded, untested). Same SM80-base kernel, bit-identical retiling; correctness vs SDPA-varlen rel 1-5e-3, test_flash_attn_varlen.py d=256 1296 passed/0 failed.
Sliding-window (local) D256 wide tile was gated to non-varlen. Gemma packed sliding-window training benefits: RTX6000 A/B gemma4-31b qpkv2 w1024 +12% (0.978->1.095 vs FA2), e4b qpkv4 w512 +6.8% (0.885->0.946), e2b marginal. Relaxed cu_seqlens exclusion in sm120_local_d256_wide (seqused stays narrow). Correctness vs SDPA-windowed-varlen rel ~2.5e-3.
…sabled) SM120 hard-disabled SplitKV, so decode (seqlen_q<=8, large KV) launched only ~batch*num_head_kv CTAs streaming the whole KV cache: FA4 0.10-0.21x of FA2 at B=1 Sk>=16384. The split scheduler/range/combine/partial-buffers/per-split softmax all already existed (arch-neutral); the SM80-base kernel just hard-coded split off. Implemented split on flash_fwd.py (all const_expr(is_split_kv)-gated -> training byte-identical): split-aware O/LSE layout transpose, split_idx, BlockInfo split range, empty-split O=0/LSE=-inf guard, direct fp32 reg->gmem partial write, 3 epilogue call sites. Gated PDL griddepcontrol to arch>=90 in the combine (illegal on sm_80-compiled SM120). interface.py: removed the disable, force non-TMA for split, decode auto-trigger (seqlen_q<=8 -> request num_splits=0 before the pack_gqa disable; heuristic self-protects so prefill/large-batch is untouched). Decode now 0.32-0.65x of FA2 (~2-4x faster), correct vs SDPA (bottom-right) rel 2e-3..8e-3 across B/Sk/D/qpkv/seqlen_q. Training unaffected (early-trigger off for seqlen_q>8; wide-tile wins intact); pytest d256 seqlen4096 120 passed/0 failed. varlen/paged/seqused/MLA split left on the safe path.
…ode) ncu showed the SplitKV decode kernel compute-bound (81% SM, 19% DRAM) — the default 128x64 tile wastes the MMA on ~120 empty query rows (seqlen_q=1). A 16x64/32-thread tile for D128 seqlen_q<=8 decode cuts the waste: D128 q8 B1 Sk32768 0.406->0.646 vs FA2 (+59%), q4 0.535->0.816 (+53%). Cumulative decode (orig->SplitKV->+tile): 0.10-0.21 -> 0.40-0.54 -> 0.65-0.82x. D256 decode kept on the lookup path (no benefit from small tile). Training (seqlen_q>8) unaffected; correctness vs SDPA rel 5-8e-3.
…ted off) From-scratch memory-bound decode kernel: 1 CTA per (split,kv_head,batch) handles all qhead_per_kvhead query rows together (KV read once, no GQA redundancy), Q.K^T and P.V as GEMV (FMA+warp-shuffle, no wasted m16n8k16 MMA on empty query rows), cp.async-streamed K/V, fp32 partial O/LSE into the existing combine. Gated behind FLASH_ATTENTION_SM120_DECODE_KERNEL (default off); dispatch additive + early-returns (flag off byte-identical, zero regression). Re-validated (seqlen_q=1, causal=False): D256 decode 1.19-1.31x over baseline (D256 q4 B4 ~0.95x of FA2), tie on D128 (already tile-fixed), correct vs SDPA rel 9e-4..5.9e-3. ncu: D256 flips compute-bound (132us SM61%/DRAM30%) -> memory-leaning (90us SM18%/DRAM44%). Does NOT beat FA2 overall — still register-bound at 1 CTA/SM (DRAM 44%, not saturated); needs register reduction for >=2 CTA/SM. qpkv16 may fall back. Opt-in D256-decode win + foundation.
…ry-access-bound (uncoalesced, 843GB/s vs FA2 1.14TB/s), not warp-starved; 2 CTA/SM makes it 4x worse
…d (32/32B/sector), bottleneck is smem-latency at 1-CTA/SM; GEMV decode is ~2x slower than FA2 across the matrix (earlier 0.95x was cherry-picked); decode concluded
…nce path) The forced 128x128 tile for paged-KV D128 is ~1.4-1.9x slower than 64x64/128t (tile_n=128 + the paged cp.async load is inefficient). Paged D128 (head_dim<=128) now uses 64x64+128t, except qpkv5 (Hq40/Hkv8) which prefers 128x128 (kept). RTX6000 A/B: B1 Sq4096 q4 1.84x, c0 1.59x, B4 q8 1.38x, ps64 q1 1.84x. Paged-KV is the inference prefill path. Correctness vs SDPA rel ~1e-3; test_paged_kv_sm120 51/51 pass. D192/D256 paged unchanged.
… d96/D128-MHA leads verified marginal (agent baselines inflated); fp8/MLA gated-off (separate project)
Milestone 1 of the fp8 effort. The SM120 GEMV decode kernel is an FMA loop, not tensor-core, so fp8 KV needs no fp8 MMA atom and no V-transpose: K/V loads become e4m3 (Q stays bf16 — the serving case: quantized KV cache + live query), lifting cp.async from 8 to 16 elems per 16B load. k_descale folds into the QK score (commutes through softmax; LSE reflects it), v_descale folds into the O-normalization reciprocal. Both default 1.0 so the bf16 path is math-unchanged. Gated behind FLASH_ATTENTION_SM120_DECODE_KERNEL — default dispatch byte-identical (fp8 inputs without the flag still hit the q==k==v dtype assert). Smem regions sized max(kv_tile, fp32_reduction_scratch): bf16 tile dominates (legacy size), fp8 1-byte tile is padded so the recast-fp32 scratch still fits. Validated in-process on RTX PRO 6000: - correctness vs fp8-quantized ref (per-(b,kv-head) amax e4m3): worst rel-err 1.69e-3 (target <1e-2); all 12 decode shapes pass. - batch16 full grid: R2/R4 win 1.6-1.85x over both bf16-decode AND FA2. R8 regresses (0.7x bf16) — GEMV FMA count is TN*R regardless of dtype, so it's compute-bound there, not bandwidth-bound; left unguarded since fp8 R8 still needs this kernel to run and keeps the 2x KV-cache memory saving. Probe/bench/correctness scripts in agent_space/ (gitignored).
The fp8 KV-cache decode kernel regressed at GQA R8 (~0.7x its own bf16). Root cause was the fp32 cross-group reduction scratch: the single-level smem fan-out was red_acc = rpi*R*tpr*vec*4, which at fp8-R8/D128 is 64KB → 80KB total smem → ~1 CTA/SM occupancy. Replace it with a two-level reduction: (1) warp-shuffle butterfly merge across the gpw=32/tpr row-groups that share a lane within a warp, then (2) an smem fan-out across only nwarps (=4) instead of the rpi row-groups. Scratch is now sized by nwarps. fp8-R8 scratch drops 64→16KB (total 80→32KB). The bf16 path's K/V tile still dominates so its smem is byte-identical; the D256 gpw=1 edge runs 0 butterfly iterations and falls back to the original fan-out. Re-validated in-process on RTX PRO 6000: - R8 (h32/4) batch16 fp8/bf16: 0.70-0.73 → 0.79-0.85 (fp8 10-13% faster). Still <1.0 — the residual is the intrinsic fp8→fp32 conversion cost in the TN*R GEMV loop, which a scratch relayout can't remove. - R4/R2 batch16 unchanged at 1.54-1.85x. Correctness unchanged (worst fp8-vs- fp8ref rel-err 1.69e-3); the 25-case test_fp8_decode_sm120.py suite passes.
The fp8 KV-cache decode kernel is the ONLY sm_120 path that can consume an fp8
(e4m3/e5m2) K/V cache (fp8 prefill is a no-go; the standard forward asserts
q.dtype==k.dtype==v.dtype). Requiring FLASH_ATTENTION_SM120_DECODE_KERNEL=1 made
a quantized cache unusable without an env var. Auto-enable it whenever fp8 K/V +
bf16/fp16 Q is genuinely passed; the env flag remains the manual override for the
*bf16* decode kernel only.
For bf16/fp16 inputs the new predicates (fp8_kv_decode, want_fp8_decode) are
always False, so with the env flag unset the dispatch is byte-identical to before
(verified: a large pytest sweep's failures are pre-existing sm_120 feature gaps,
identical on clean HEAD). The fp8 path needs num_splits>=2 (the kernel's
ceil_div(seqlen_k, num_splits) tiler rejects 1), so it bumps a 1 to 2; combine
handles any count.
Add a guard: fp8 K/V in an unsupported configuration (seqlen_q>1, MHA, D192,
non-sm120, varlen/paged/local/mask/etc.) now raises NotImplementedError instead
of silently running the bf16 MMA over reinterpreted fp8 bytes.
New tests/cute/test_fp8_decode_sm120.py: 25 passed — fp8 decode correctness vs an
fp8-quantized reference, R in {2,4,8}, D128, Sk {4096,16384}, B {1,16}, causal,
parametrized over the env flag on/off plus an auto-enable assertion.
… FA2) The forward tile lookup for (head_dim=128, qpkv=4, S=2048, causal) used 64x96, ns=1, which on the RTX PRO 6000 is erratic vs FA2 — measured fa4/fa2 swinging 0.98-1.08 across seeds (sometimes a laggard). Switch to 128x64 with num_stages=2. Re-validated in-process (interleaved A/B, clock soak, CUDA-event median over 8 repeats x 3 seeds): the new config holds a tight 1.068-1.081x vs FA2 across all seeds, vs the old config's 0.983-1.077 spread. So this isn't a uniform speedup so much as removing a parity-prone config — the shape now reliably beats FA2 by ~7% instead of fluctuating around 1.0. Correctness rel-err vs the old tile 1.1e-3. (The only genuine forward dispatch lever found in a full re-sweep of the B2 grid; every other apparent gap was at floor or a cold-baseline measurement phantom.)
…A causal laggard The causal D128 MHA (qpkv1) forward trailed FA2 by ~6-7% across all seqlens (geo 0.943). Root cause: flash_fwd_sm120_tma.py declared LPT intent (lpt=self.is_causal or self.is_local in TileSchedulerArguments) but instantiated SingleTileScheduler, which silently ignores the lpt field — only SingleTileLPTScheduler consumes it. The longest-processing-time ordering meant to balance the causal triangle's load imbalance was never active. ncu confirmed pure tail-wave SM idle: causal SM% 60.7 vs FA2 68.5 at identical occupancy/grid, with masked-block skipping already correct and per-instruction stats identical to the ~parity dense path. Fix: route causal/local non-varlen to SingleTileLPTScheduler and pass a real seqlen_k (= mK_t.shape[0]; was hardcoded 0, which faults the LPT L2-swizzle sizing that divides by it — SingleTileScheduler ignored it). Dense and varlen paths unchanged. Output is bit-identical (only CTA->tile assignment changes); correctness re-verified vs SDPA, worst 2.97e-3 across causal MHA/GQA/local/D256. Re-validated with a per-shape-soak + round-wise-ratio harness (10 samples/cell). Causal D128 MHA fa4/fa2: S1024 0.93->1.04, S2048 0.96->1.05, S4096 0.94->0.99, S8192 0.94->0.96, S16384 0.94->0.95; geo 0.943 -> 0.996. S1024/S2048 now win. No real regression (dense provably unaffected drifts 0.5-2% = cross-run noise floor; causal GQA/D256 changes are within that). The fix is in the shared TMA-path scheduler, so it benefits any causal/local shape on that path.
…cy wall), kernel shot NO-GO
… shipping config The sm_120 perf campaign left ~15 FLASH_ATTENTION_SM120_* env-var escape hatches in the dispatch for A/B probing. Each gated an override around the tuned default; the env-unset path is the shipping/winning config. Collapse every hatch to that default and delete the override branches, dead os.environ reads, and now-unused mode variables (16 insertions, 140 deletions). No kernel files needed changes — they receive resolved decisions as parameters. Removed: QPKV5_S16384_QREGS, D256_QREGS128, D256_QPKV8/16/6_CAUSAL_QREGS, D256_WIDE, LOCAL_D256_WIDE, PACK_GQA_VALID_ROWS_FAST, QPKV5_S4096_NC_TMA, QPKV6_D256_HOOKS, QPKV6_D256_STATIC_CAUSAL_BLOCKS, QPKV5_HOOKS, FUSED_DKV, BWD_PACK_GQA_M_SPLITS, BWD_NONPACK_M_SPLITS, BWD_SKIP_FULL_CAUSAL_MASK. Kept (genuine gates): FLASH_ATTENTION_SM120_DECODE_KERNEL (experimental bf16 GEMV decode opt-in) and FLASH_ATTENTION_ARCH (arch override). Behavior-preserving: independently verified bit-identical output (max|edit-head| = 0.000e+00) across 6 shapes exercising the collapsed paths (d256 wide/causal, pack_gqa, qpkv5 hooks, local d256, lpt causal) vs HEAD with no env vars set, which is what ships. Backward dispatch identical over the tuned shape space.
…fset window bwd) Both surfaced during PR prep and confirmed in-process vs SDPA; neither is a test artifact. BUG 1 — learnable_sink + SplitKV (forward, silent wrong output). Each KV split folds exp(sink - max) into its own denominator/LSE, and the combine kernel does not de-duplicate it, so the sink is counted once per split. Worse, the sm_120 decode auto-split (seqlen_q<=8) did not check learnable_sink, so real decode workloads with an attention sink silently got wrong output (max-diff 2.3-3.5). Fix: force num_splits=1 whenever learnable_sink is set (single guard after every num_splits decision). Always correct; costs only the decode SplitKV speedup for sink models. Verified: sink + num_splits=3 now matches the reference (~7e-3). BUG 2 — negative-offset sliding-window backward (wrong dK/dV). For windows with a negative bound — window_size=(None,-X) or (-X,None) — the forward is correct but the backward column-range/local masking produces garbage dK/dV (nonzero grad at fully-masked positions, max-diff ~4). Seqlen-independent; finite non-negative windows are fine. Fix: _flash_attn_bwd raises NotImplementedError for these on sm_120 (loud refusal instead of silent wrong gradients); forward still works. Tests: skip deterministic and negative-offset-window (local_enum 2/3) backward on sm_120 with reasons. Slice on an sm_120 device: 99 passed / 93 skipped / 0 failed. README: document both limitations. (Proper fixes — apply sink once in combine; correct the offset-window bwd range — are left as follow-ups.)
Owner
Author
|
Superseded by a cleaned PR (rebased: experimental tuning env-flags collapsed to shipping defaults, internal campaign notes removed from the tree, two correctness bugs fixed — learnable_sink+SplitKV and negative-offset window backward). Opening the replacement now. |
thad0ctor
added a commit
that referenced
this pull request
Jun 4, 2026
…+SplitKV, neg-offset window bwd) Response to CodeRabbit review on PR #1/#2 plus the two proper bug fixes that replace the conservative guards. All changes validated in-process vs SDPA on sm_120 (RTX 5090 functional; RTX 6000 for the memory-heavy fp8 case). Correctness bug fixes (replace the guards added earlier this branch): - learnable_sink + SplitKV: the sink is now folded into the LSE only in split 0 (compute_sink_val suppresses it to -inf in splits >0 via a runtime bias), so the combine counts it exactly once. softmax.finalize hardened for empty SplitKV splits (row_max==-inf -> LSE=sink, no +inf poison). The interface num_splits=1 guard is now scoped to non-sm120 only (SM90/SM100 sink kernels were not given the split-0 gating; kept conservative there). Verified: sink + num_splits in {1,2,3} all match attention_ref (rel <=8e-3). - Negative-offset sliding-window backward (window_size (None,-X)/(-X,None)): the bwd dK/dV m-block-range prune produced empty/inverted ranges for one-sided open windows, leaving fully-masked key blocks with garbage gradients. Now the prune is skipped (runtime, per side) when a bound is negative; correctness comes from the per-block mask. Removed the interface NotImplementedError + the test skip. Verified: dK/dV rel-err 0.002-0.018 (was ~4); positive/symmetric windows keep the fast prune. CodeRabbit fixes applied (verified real on current HEAD): - flash_fwd_combine.py: don't enable PDL on sm_120 (arch_int 120 is >=90 but griddepcontrol is unavailable on the sm_80-compiled sm120 target). [Critical] - seqlen_info.py: clamp every per-batch read (mSeqUsed*, mCuTotalMBlocks, mCuBlockIdxOffsets are [num_batch]) so SM80/SM120 over-launched tiles stay in-allocation. [Critical] - flash_fwd.py: paged-KV unmasked-start clamp mirrors the dense path; add the PackGQA Q loader to the paged mainloop (path still gated off in interface). - flash_bwd.py: gate reuse_qk_dov_smem to single-stage pipelines. - pack_gqa.py: drop the dead zero_oob_rows param (+ 5 call sites); OOB rows are inert (not stored fwd; P==0 bwd). - interface.py: block fake-mode fp8-KV decode (would compile the wrong kernel). - test_flash_attn_sm120_local.py: non-destructive module teardown. - bench_master_bwd.py: phase_repro idempotency + tuned tile in the resume key. CodeRabbit items verified STALE/already-fixed or false-positive and skipped: score_mod arg order, unmasked-loop n_block_min, paged local tail loop, TMA V size, varlen dQ per-batch offset, full-block mask_mod clear, fused dK/dV mixed-headdim guard, TMA-selection shadowing, compile-cache key, sink-off-TMA, mask_mod varlen autograd, F401/F541/flake8/README — all confirmed against current code.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Integrates FA4 support for consumer Blackwell (
sm_120) and the dispatcher,kernel, correctness, and performance work needed to make the path usable on
RTX 5090-class GPUs.
Current branch:
sm120-integrateLatest performance/code keepers: D256 qpkv6/qpkv8/qpkv16 Q-in-regs paths,
packed-GQA valid-row fast path, and fused dK+dV guard cleanup.
Current State
This PR contains:
flash_attn.cute.flash_attn_func/flash_attn_varlen_funcAPIs.non-TMA path for shapes TMA cannot handle safely.
auto-downgrades for invalid packed-GQA cases.
head_dim <= 256; D192/D256 use the non-TMA64x64 path.
shared-memory alias path, with exact packed/fused/split keepers for repeatable
rows.
including D256 qpkv6/qpkv8/qpkv16 Q-in-regs long-sequence paths and the
B=2 qpkv6 S8192 noncausal short/mid path.
postprocess guard for mixed dK/dV head dims.
What This Adds On Top Of Dao-AILab#2553 + Dao-AILab#2349 + Dao-AILab#2389
pack_gqa=Trueworks where valid; invalid SM120 packed cases auto-downgrade.head_dim <= 256, validated against contiguous unpacked FA4.head_dim > head_dim_vmask_modautograd propagation, pack-GQA dQ offset, paged local loop bound, TMA V byte count, unsupported TMA learnable-sink gate, and fused dK/dV head-dim guard.Current Performance Vs FlashAttention 2
RTX 5090 / SM120 / bf16.
/tmp/sm120_longseq_qwen_gemma_after_qpkv16_qregs_20260529/tmp/sm120_current_long_miss_repeat_after_qpkv16_20260529/tmp/sm120_model_variants_after_qpkv5_notma_5x_20260529/tmp/sm120_sdpa_fa2_fa4_forward_qpkv4_patch_20260528AI/SM120_CURRENT_PERF_2026_05_26.md/tmp/sm120_bwd_d256_after_split3_10x_3b276ae_20260529Current long Qwen/Gemma rollup:
Performance Vs PyTorch SDPA
Matched forward run, RTX 5090 / SM120 / bf16 / batch=1. PyTorch SDPA is the
faster of
FLASH_ATTENTIONandCUDNN_ATTENTIONper cell. For GQA SDPA, K/Vare pre-expanded to Hq outside the timed region.
Artifact:
/tmp/sm120_sdpa_fa2_fa4_forward_qpkv4_patch_20260528Interpretation: FA4 is materially faster than FA2 on the matched 60-cell matrix.
Best PyTorch SDPA, especially cuDNN, remains stronger overall on that B=1 matrix.
Current Forward Keepers
/tmp/sm120_d256_qpkv6_qregs_default_auto_off_20260529showed auto beatingforced-off on all 8 long qpkv6 rows by +1.9% to +18.7%, and beating FA2
on all 8.
/tmp/sm120_qpkv6_b2_public_after_s8192_nc_qregs_20260529showed the newlydefaulted S8192 noncausal row at 1.026004x geomean vs FA2 with 5/5
wins. The exact default/off repeat
/tmp/sm120_qpkv6_b2_s8192_nc_default_qregs_validate_20260529showed+4.43% median versus forced-off with output diff 0.
/tmp/sm120_d256_qpkv8_causal_qregs_default_auto_off_20260529showedS32768/S65536/S131072 auto beating forced-off by
+9.36% / +5.96% / +6.72% and beating FA2 by
1.079x / 1.050x / 1.099x. S16384 repeat artifact:
/tmp/sm120_d256_qpkv8_s16384_causal_qregs_auto_off_95x_20260529./tmp/sm120_d256_qpkv16_causal_qregs_default_auto_off_20260529showedS16384/S32768/S65536/S131072 auto beating forced-off by
+3.92% / +5.33% / +5.36% / +9.08% and beating FA2 by
1.046x / 1.056x / 1.068x / 1.115x.
/tmp/sm120_d256_qregs128_default_validate_20260529and related repeatskeep exact B=1 Hkv=2 long noncausal rows on
128x64, 256 threads,Q_in_regs=True.validated tile gates; recent D128 qpkv8 S65536 256-thread recheck was only
+0.26% vs auto and was not shipped.
64x16local-window path. Recentlocal broad misses did not reproduce in focused repeats.
Important Fixes
score-mod argument order, unsupported
is_split_kv,vec_sizetypo,invalid head-dim routing, and varlen
cu_seqlensover-launch reads.dQ_single_wg,softmax_scale,atomic_add_fp32, dQ postprocess layout,8-warp causal masking compatibility, and v4 atomic layout.
memory and reloading Q/K inside the D256 mainloop.
head_dim <= 256.head_dim > head_dim_vaway from SM120 TMA to avoid GPU hangs.fused dK+dV postprocess.
Tuning Decisions To Preserve
Keep these negative results so future passes do not repeat them without fresh
paired evidence:
kv_stages=1and hook variants outside the accepted exactrows.
and S65536 256-thread tile switch.
check_inf=False.rP.qpkv6 fused-dKV, K/V load predicate elision,
acc_S_preguard, and nonpackedsplit4 as a default.
The full win/rejection ledger is in
win.md.Validation
Validated on RTX 5090 / SM120 / bf16:
tests/cute/test_flash_attn_sm120_local.py::test_sm120_qpkv16_d256_causal_qregs_matches_sdpa,test_sm120_qpkv8_d256_causal_qregs_matches_sdpa, andtest_sm120_qpkv6_d256_qregs_matches_sdpa: 4 passed.qpkv6 D256 Q-regs, B=2 Q-regs+hook, and S8192 noncausal default-Q-regs
subset: 5 passed.
head_dim > head_dim_vrouting suite: 11 cases pass.packed/nonpacked and fused-dKV paths during this PR.
Not In Scope
Older noisy dispatch variants are tracked as rejected tuning decisions, not as
feature-scope exclusions.
cc @coderabbitai