FA4 consumer Blackwell (sm_120) integration - #2
Open
thad0ctor wants to merge 242 commits into
Open
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>
… + 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.
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
CodeRabbit (PR #1) flagged that `seed=hash(page_table_pattern) & 0xFFFF` is non-reproducible across runs: Python's builtin `hash(str)` is randomized per process via PYTHONHASHSEED, so a tolerance-band failure on one run cannot be repro'd from the recorded seed. Replace both occurrences (test_page_table_patterns and test_d_gt64_page_table_patterns) with a fixed PATTERN_SEEDS mapping ("identity"->101, "permuted"->202, "shared"->303). The second test keeps its `d * 1000 + seed` derivation so d=96 and d=128 sweeps remain seed-disjoint. All 9 affected parametrized cases pass on RTX 5090 (sm_120). Closes-Upstream-Comment: #1
Add a subprocess-isolated benchmark harness for tuning the FA4 CuTe
backward kernel tile sizes on consumer Blackwell (sm_120) GPUs.
benchmarks/sm120_bwd_tuning/
measure_one_bwd.py - one cell + one tile config, prints JSON
bench_master_bwd.py - orchestrator: sweep + repro + analyze + validate
README.md - usage, env vars, output schema
For each (preset, seqlen, causal) cell, spawns a fresh Python process per
(tile_m, tile_n, num_stages) candidate so the CuTe JIT cache and CUDA
context start cold for every measurement.
The measurement script monkey-patches the SM120 branch of
flash_attn.cute.interface._flash_attn_bwd so the hard-coded tile
constants can be overridden per run without touching kernel sources.
No source changes to flash_attn/. Pure tooling; the harness is opt-in
via env vars (FA_BENCH_GPU_UUID, FA_BENCH_OUT_DIR, FA_BENCH_PYTHON,
FA_BENCH_CUTE_OVERRIDE) and falls back to inheriting CUDA_VISIBLE_DEVICES
and sys.executable when unset.
The 0.05 threshold rejected legitimate bf16 noise on backward gradients at longer sequences (causal=1, sl=2048 cells with dV diff=0.0625 -- ~8 bf16 ULPs at magnitude 1.0), causing the sweep to mis-flag the working baseline as a numerical_fail. 0.1 absolute is roughly 12 bf16 ULPs and still catches genuine corruption: when an MMA-unsupported tile (e.g. tile_n=32 on the SM80-base backward) compiles but produces garbage, the dK/dV diff is in the 0.4-0.75 range -- well past the 0.1 cutoff.
Phase 16c sweep + Phase 17C tightened paired validation (RTX 5090,
n_measure=30, interleaved trials) confirm ns=1 wins on d=64 backward.
ns=2 was inherited from the SM80-base default but the async pipeline
overhead exceeds the latency-hiding benefit at the small d=64 tile
size on consumer Blackwell.
Paired validation across 19 d=64 cells (5 model presets x 4 seqlens x
{causal, non-causal} from the qhead_per_kvhead = {1,4,7} matrix):
geomean ratio ns=2/ns=1 = 1.0558x
0 cells regress >2%
Code change is minimal: the SM120 backward branch now uses ns=1 for
all head_dim (the d>64 branch already used ns=1, so this is purely a
flip of the d<=64 case from 2 to 1). No effect on forward, non-sm_120
arches, or any other parameter.
Validation:
- E2E phase11_e2e/e2e.py: 34/34 pass
- tests/cute/test_paged_kv_sm120.py: 38/38 pass
- tests/cute/test_flash_attn_sm120_dgtdv.py: 11/11 pass
- tests/cute/test_flash_attn_bwd_sm120_postprocess.py: 10/10 pass
Phase 16a NCU profile attributed the 0.93x FA4/FA2 backward geomean to parallelism: FA4 SM120 backward ran 4 warps/SM, FA2 runs 8 warps/SM, both clamped to 1 block by ~82 KB SMEM. The 9pp compute-throughput gap (80.12% vs 88.57%) was the dominant explanation. This commit repartitions the SM120 backward kernel to 256-thread / 8-warp blocks at the SAME SMEM footprint. The bug surfaced by the prior 17A-config attempt (causal-path dQ/dK/dV errors of 5-20 vs the 0.004 noise floor) was traced to the R2P bitmask fast-path in flash_attn/cute/mask.py:r2p_bitmask_below + sm90_col_to_r2p_idx assuming the standard SM80/SM90 per-thread column pattern (col-pairs at stride 8). With AtomLayoutSdP = (4, 2, 1) the SM120 256-thread configuration has 2 N-warps and the per-thread cols interleave at stride 16 instead, so the R2P bitmask kept cells beyond the causal boundary. Fixed by adding an optional r2p_compatible field to AttentionMask (default True; preserves all SM80/SM90/SM100 behaviour) and gating the 4-warps-per-tile detection to sm_120 only via a new `arch = 120` marker on FlashAttentionBackwardSm120. The SM80 path is unaffected (no `arch` attribute -> getattr falls back to 80). Also fixes a postprocess invariant: the dq_accum / dk_accum / dv_accum byte buffers are written by the main kernel via a thread-major partition whose stride is num_threads; the postprocess reader must use the same num_threads or per-thread element->address mapping diverges. SM120 branch now mirrors num_threads_post_dQ/dKV = 256. SM80, SM90, SM100 paths untouched. NCU after on mistral-7b sl=4096 c=1 bwd: theoretical occupancy 8.33% -> 16.67% (matches FA2), compute throughput 80.12% -> 87.55% (1pp short of FA2's 88.57%). Phase 13 40-cell backward paired bench: geomean ~1.05x; the 6 flagged regressions are all in the documented small-seqlen / GQA bench-noise floor (10% CV) and flicker across re-runs. Validation: E2E 34/34, sm120 pytest 59/59, forward unchanged within bench noise.
…layout
Phase 17D-lite-v3: switch gmem_tiled_copy_dQaccum to a 128-bit copy atom
with val_layout=4 on SM120 only, so each thread owns 4 contiguous fp32 in
gdQaccum. With the post-17A-config 256-thread / 8-warp partition this is
the prerequisite for emitting red.global.add.v4.f32 in the dQ accumulator
write loop, cutting the atomic instruction count by 4x. The dK/dV GQA
atomic-add path (qhead_per_kvhead > 1) reuses the same V=4 copy and the
same v4 atomic helper.
The MMA m16n8k16 C-fragment for thread t holds 4 fp32 (c0,c1,c2,c3) in
2-contig col pairs at row offsets {r, r+8}. retile(acc_dQ) flattens
((2,2),1,N):(...) -> ((4,1),1,N):(...) without reordering registers — the
register-storage flat order is identity in both layouts. So writing the
4 per-atom registers to 4 contiguous gmem positions yields a consistent
flat encoding as long as the postprocess reads back in the same order;
flash_bwd_postprocess.py is updated to use num_s2r_copy_elems=4 for
SM120 to satisfy that invariant.
Architectural gating:
- SM80 path: untouched. The dQaccum gmem copy and atomic loop both
remain V=1 / scalar atomic.
- SM90/SM100: unaffected (separate files, not subclasses of
FlashAttentionBackwardSm80).
- SM120: V=4 / v4 atomic, gated via getattr(self, "arch", 80) == 120.
Adds utils.atomic_add_fp32_v4 as a v4 inline-asm wrapper (mirrors the
existing copy_utils.atomic_add_fp32x4 but lives next to atomic_add_fp32
so the backward kernel can import both from the same module).
Validation (RTX 5090 / sm_120, bf16):
- phase11_e2e/e2e.py: 34/34 pass
- tests/cute/test_flash_attn_bwd_sm120_postprocess.py: 8/8 pass
- tests/cute/test_flash_attn_sm120_dgtdv.py: 11/11 pass
- tests/cute/test_paged_kv_sm120.py: 40/40 pass (2 skipped)
- Smoke (mistral-7b/qwen2.5-7b sl=1024 c=1, llama2-7b sl=1024 c=0):
max diff vs SDPA <= 0.05 on all 3 cells.
Phase 16b audit found pack_gqa was forced to False in the backward
dispatcher with comment "not yet supported in bwd" - explaining the
largest backward FA4/FA2 regressions on GQA shapes (qwen2.5-7b 7-way:
0.79x, llama3-8b/mistral-7b/mixtral-8x7b 4/8-way: 0.83-0.94x; MHA
llama2-7b at 0.92-1.00x as control).
This commit ports pack_gqa from the forward path to the SM120 backward
kernel:
- flash_bwd.py:
* pack_gqa_layout remaps for mQ/mdO/mLSE/mdPsum (under arch=120
+ pack_gqa=True only, mirroring flash_fwd.py:701-706). mdQaccum
intentionally stays in the ORIGINAL (B, H_q, S*D) layout because
the postprocess kernel reads it back via an opaque per-thread MMA
register interpretation, and per-element address routing into the
correct head_q slot happens in the helper.
* post-pack mQ.shape[...] lookups fixed for num_head,
seqlen_q_static, m_block_max, d_head (the four fixes flagged by
the v2 POSTMORTEM).
* causal masking: pass qhead_per_kvhead_packgqa to AttentionMask so
the row index is divided by qh before computing the causal limit.
* m_block_min for causal scaled by qh under pack_gqa.
- pack_gqa.py:
* load_scalar_per_row: one-thread-per-row direct fp32 store from
gmem to sLSE/sdPsum (replaces the cp.async LSE load whose tile
layout doesn't fit the composite packed-mode-0 tensor).
* atomic_add_dQaccum: per-MMA-element atomic_add_fp32 that uses
partition_C to get each element's (mma_m, d), decomposes mma_m
to (h_in_kvgroup, m_actual), computes the absolute head_q
index, and inverts the postprocess's MMA layout to find the
canonical gmem position in head_q's slot where the postprocess
will emit at (m_actual, d). v4 atomic is NOT usable here because
the 4 vals in each MMA atom span 2 different m_block rows.
- interface.py:
* relax pack_gqa = False override to apply only when arch != 120
(line ~1562)
* thread pack_gqa through FlashAttnFunc.backward and
FlashAttnVarlenFunc.backward (was previously dropped between fwd
ctx and _flash_attn_bwd call)
dQ atomic reconciliation: dQ_mma uses per-element atomic_add_fp32 under
pack_gqa (vs the v4 atomic for non-pack); the postprocess and the
backward kernel both share the same per-thread MMA register layout
encoded by partition_C of tiled_mma_dq, so the inverse mapping is
exact for the SM120 config (AtomLayoutMdQ=4, m16n8k16, 256 threads, 4
warps in M x 2 in N).
SM80, SM90, SM100 paths bit-identical (verified by empty git diff on
flash_bwd_sm90.py and flash_bwd_sm100.py vs cee3b54).
Validation (smoke probe, all 4 brief shapes, max diff threshold 0.05):
- qwen2.5-7b non-causal (Hq=28 Hkv=4): dq=0.003 dk=0.005 dv=0.005 PASS
- llama3-8b causal (Hq=32 Hkv=8): dq=0.023 dk=0.016 dv=0.030 PASS
- mistral-7b causal (Hq=32 Hkv=8, D=128): dq=0.015 dk=0.024 dv=0.030 PASS
- llama2-7b MHA causal (Hq=Hkv=32, D=128): dq=0.015 dk=0.014 dv=0.020 PASS
Extra coverage (11 shapes spanning MHA/4-way/7-way/8-way GQA x causal
True/False x pack_gqa True/False): 11/11 PASS.
Forward unchanged: 4 brief shapes match SDPA reference within bf16
tolerance.
Known limitation (NOT exercised by brief shapes): non-divisible seqlens
under pack_gqa corrupt dK/dV because PackGQA.load_Q's row-OOB guard
skips the cute.copy entirely, leaving sQ smem stale. Fix is to use
cute.copy with a per-element predicate matching the non-pack path so
cp.async writes 0 for OOB rows. Brief shapes are all S in {64, 256,
1024} which are divisible by m_block_size=64.
Current winner after paired prior-commit sweep and patched rerun.\n\nSM120 FA4 forward 60-cell grid: 97.471 geomean TFLOPS, 60/60 ok, peak 182.0 TFLOPS.\nVs c43a7b4 fallback: 1.075x geomean, 39 wins, 16 ties, 5 losses.\nVs Phase 5c c8a1864: 1.018x geomean.\nVs prior sweep leader 55ab672: 1.016x geomean.\n\nKey recovered cells include hd64-small-gqa causal S=2048/4096/8192 at 81.4/133.0/166.3 TFLOPS.\n\nThis is an empty marker commit; implementation lives in 1b7db10.
|
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:
📝 WalkthroughWalkthroughThis PR adds comprehensive SM120 (consumer Blackwell) support to FlashAttention-4, including a new decode kernel, TMA-based forward implementation, backward kernel variants with packed-GQA handling, fused postprocessing, and host-side dispatch routing with compile-time caching. The implementation addresses SM120-specific constraints via vectorized atomics, dependent-grid control gating, and extensive regression test coverage. ChangesSM120 end-to-end implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
✨ Finishing Touches📝 Generate docstrings
🧪 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: 11
🧹 Nitpick comments (1)
flash_attn/cute/pack_gqa.py (1)
306-338: 💤 Low value
store_LSE_all_rows_validstill checks row bounds, inconsistent with the parametric version.The method
store_LSE_all_rows_validat line 336 still checksrow < seqlen * self.qhead_per_kvhead, whilestore_LSE(..., all_rows_valid=True)at lines 300-301 unconditionally stores. If both are intended to be "all rows valid" variants, they should have consistent behavior.If the row check in
store_LSE_all_rows_validis intentional (for partial blocks), consider renaming or documenting the distinction.🤖 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/pack_gqa.py` around lines 306 - 338, The function store_LSE_all_rows_valid is inconsistently still checking row bounds (row < seqlen * self.qhead_per_kvhead) even though it's the "all rows valid" variant; to fix, remove the row-bound check so behavior matches the parametric all_rows_valid path — i.e., in store_LSE_all_rows_valid (and specifically the if that reads "if taccOcO[0][1] == 0 and row < seqlen * self.qhead_per_kvhead"), either (A) delete the "row < seqlen * self.qhead_per_kvhead" clause so the write is gated only by the writer lane check (taccOcO[0][1] == 0) or (B) if truly intended to allow partial blocks, rename the function to reflect that and document the distinction; update store_LSE_all_rows_valid accordingly (refer to symbols: store_LSE_all_rows_valid, taccOcO, tLSErLSE, mLSE_copy, tPrLSEPtr).
🤖 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 `@benchmarks/sm120_bwd_tuning/bench_master_bwd.py`:
- Around line 406-415: The resume key built in the set assigned to existing only
includes (preset, sl, causal, mode, repeat) and must also include the tuned tile
to avoid treating entries with changed (tile_m, tile_n, num_stages) as cache
hits; update the set comprehension that constructs existing (the tuple produced
for r in read_jsonl(val_path)) to append the tuned identifier—either include
r.get("tuned") or the explicit fields r.get("tile_m"), r.get("tile_n"),
r.get("num_stages")—so the tuple becomes (r["preset"], r["sl"], r["causal"],
r.get("mode_label") or r.get("mode"), r.get("repeat", 0),
<tuned-or-tile-fields>).
- Around line 225-263: phase_repro currently appends to sweep_repro.jsonl on
every run causing duplicates; make it idempotent by either
truncating/overwriting repro_path at start or skipping entries already present:
read existing repro rows (if any) into a set keyed by (preset/name, sl, causal,
tile_m, tile_n, num_stages, repeat) and, before calling
run_measure/append_jsonl, skip runs whose key is present; update symbols:
phase_repro, repro_path, append_jsonl, run_measure, and by_cell (use cells_iter
to enumerate) so repeated --phase repro runs don't produce duplicate repeats
that bias analysis.
In `@flash_attn/cute/flash_bwd.py`:
- Around line 104-108: The reuse_qk_dov_smem flag currently enables the reuse
path based only on arch and padded head dims (in the assignment to
self.reuse_qk_dov_smem) but the reuse logic assumes single-stage kernels; update
the guard to also require single-stage operation by checking the stage counts
(e.g., self.num_stages_Q == 1 and self.num_stages_dO == 1 or the equivalent
smem_pipe_*/num_stages_* attributes used in this class) so that
reuse_qk_dov_smem is set to True only when arch==120, head_dim_padded==256,
head_dim_v_padded==256, and both stage counts are 1.
In `@flash_attn/cute/flash_fwd_combine.py`:
- Around line 56-63: The current predicate computes arch_int and sets
self.use_pdl = self.arch_int >= 90 which incorrectly enables PDL on SM120
(arch_int == 120); change the guard to explicitly exclude SM120 (e.g. set
self.use_pdl = (self.arch_int >= 90 and self.arch_int != 120) or check
arch.major != 12) so griddepcontrol_wait remains disabled on SM120; update the
code around BaseDSL._get_dsl().get_arch_enum(), the self.arch_int assignment,
and the self.use_pdl assignment accordingly and ensure any use of
griddepcontrol_wait respects the new self.use_pdl flag.
In `@flash_attn/cute/flash_fwd.py`:
- Around line 2020-2027: The paged-KV path can set unmasked_n_block_start beyond
valid range because n_block_min_causal_local_mask from
block_info.get_n_block_min_causal_local_mask(...) is not clamped; update the
block in the if const_expr(self.is_causal or self.is_local) branch to clamp
unmasked_n_block_start = min(n_block_min_causal_local_mask, n_block_max - 1)
(mirror the dense path) before entering the cutlass.range loop so the unmasked
loop never starts past or reprocesses blocks; reference symbols:
get_n_block_min_causal_local_mask, n_block_min_causal_local_mask,
unmasked_n_block_start, n_block_max, cutlass.range.
- Around line 1914-1919: In the paged-KV prologue where Q is loaded
(gmem_thr_copy_Q / gmem_tiled_copy_Q.get_slice(tidx)), handle the pack_gqa case
instead of always calling self.load_Q: if self.pack_gqa is true, call the
packed-GQA Q loader (use the class method that implements packed/composite
(qhead_per_kvhead, seqlen) layout — e.g. self.load_Q_packed or the existing
packed-Q loader method in this class) with the same arguments (gmem_thr_copy_Q,
gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1]); otherwise keep
calling self.load_Q. This ensures the paged-KV path uses the correct address
computation when pack_gqa is enabled.
In `@flash_attn/cute/interface.py`:
- Around line 1211-1249: The code currently allows fake/compile-only mode to
fall through when fp8_kv_decode is true but want_fp8_decode is false; change the
guard so fake mode is also blocked. Replace the condition in the
NotImplementedError check (the one using fp8_kv_decode, want_fp8_decode and
is_fake_mode()) with a check that does not exclude fake mode (e.g., if
fp8_kv_decode and not want_fp8_decode: raise ...), and make the identical change
at the second occurrence around FlashAttentionDecodeSm120 handling (the block at
~1274-1303) so both locations consistently prevent fake-mode fp8-KV decode from
falling through to the regular SM120 forward path.
In `@flash_attn/cute/pack_gqa.py`:
- Line 167: The parameter zero_oob_rows in load_Q is currently unused (dead
parameter); update load_Q to honor zero_oob_rows by using it when building the
out-of-bounds / row-valid predicate so rows marked OOB are zeroed instead of
treated as valid, or remove the parameter if zeroing is not desired. Concretely,
in the load_Q implementation adjust the logic that computes the row validity /
OOB mask (referencing load_Q and zero_oob_rows) so that when zero_oob_rows is
True you apply a mask that writes zeros for OOB rows (or short-circuits loads to
return zero vectors), otherwise preserve the existing behavior; ensure callers
in flash_fwd.py/flash_bwd.py retain compatibility.
In `@flash_attn/cute/seqlen_info.py`:
- Around line 111-126: In create(), several per-batch reads still use the raw
batch_idx and can OOB on over-launched SM80/SM120 tiles; clamp the index once
(e.g., compute clamped_idx = cutlass.min(batch_idx, tensor.shape[0] - 1)) and
use it for all per-batch accesses instead of the unclamped batch_idx or only
clamping mCuSeqlens[batch_idx+1]; update accesses to mSeqUsedQ[...],
mSeqUsedK[...], mCuTotalMBlocks[...], and mCuBlockIdxOffsets[...] (and any other
per-batch reads in create()) to use the clamped index (or batch_idx+1 clamped
where appropriate) to ensure all reads stay in-allocation.
In `@tests/cute/test_flash_attn_sm120_local.py`:
- Around line 17-35: The helper _ensure_worktree_cute_loaded currently deletes
every module named "flash_attn" or starting with "flash_attn.", which can create
a second live package and test-order flakiness; modify the cleanup to only
remove "flash_attn.cute" and modules whose names start with "flash_attn.cute."
(i.e., replace the sys.modules deletion loops to check for name ==
"flash_attn.cute" or name.startswith("flash_attn.cute.")), and when injecting
the shim pkg (pkg = types.ModuleType("flash_attn")) keep the rest of the package
intact or instead run this loader inside an isolated subprocess/pytest fixture
if isolation is required; adjust references in the function
(_ensure_worktree_cute_loaded, finder.MAPPING) accordingly.
In `@tests/cute/test_flash_attn.py`:
- Around line 865-877: The skip is too coarse: instead of skipping based on
local_enum, change the SM120-only skip to inspect the actual local window tuple
used in the test and only skip when the window has a negative-offset side (e.g.,
window == (None, -X) or window == (-X, None)). Replace the "if local_enum in (2,
3) and IS_SM120:" check with a condition that looks up the local window variable
(the one built from local_enum) and checks for a negative offset on either side,
keeping the existing IS_SM120 and pytest.skip usage so only true negative-offset
local windows are skipped on SM120.
---
Nitpick comments:
In `@flash_attn/cute/pack_gqa.py`:
- Around line 306-338: The function store_LSE_all_rows_valid is inconsistently
still checking row bounds (row < seqlen * self.qhead_per_kvhead) even though
it's the "all rows valid" variant; to fix, remove the row-bound check so
behavior matches the parametric all_rows_valid path — i.e., in
store_LSE_all_rows_valid (and specifically the if that reads "if taccOcO[0][1]
== 0 and row < seqlen * self.qhead_per_kvhead"), either (A) delete the "row <
seqlen * self.qhead_per_kvhead" clause so the write is gated only by the writer
lane check (taccOcO[0][1] == 0) or (B) if truly intended to allow partial
blocks, rename the function to reflect that and document the distinction; update
store_LSE_all_rows_valid accordingly (refer to symbols:
store_LSE_all_rows_valid, taccOcO, tLSErLSE, mLSE_copy, tPrLSEPtr).
🪄 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: 7ed35350-d2dc-42ab-b439-732cca820c4d
📒 Files selected for processing (27)
benchmarks/sm120_bwd_tuning/README.mdbenchmarks/sm120_bwd_tuning/bench_master_bwd.pybenchmarks/sm120_bwd_tuning/measure_one_bwd.pyflash_attn/cute/README.mdflash_attn/cute/block_sparse_utils.pyflash_attn/cute/flash_bwd.pyflash_attn/cute/flash_bwd_postprocess.pyflash_attn/cute/flash_bwd_sm120.pyflash_attn/cute/flash_fwd.pyflash_attn/cute/flash_fwd_combine.pyflash_attn/cute/flash_fwd_decode_sm120.pyflash_attn/cute/flash_fwd_sm120.pyflash_attn/cute/flash_fwd_sm120_tma.pyflash_attn/cute/interface.pyflash_attn/cute/mask.pyflash_attn/cute/pack_gqa.pyflash_attn/cute/paged_kv.pyflash_attn/cute/seqlen_info.pyflash_attn/cute/utils.pytests/cute/test_flash_attn.pytests/cute/test_flash_attn_bwd_sm120_pack_gqa.pytests/cute/test_flash_attn_bwd_sm120_postprocess.pytests/cute/test_flash_attn_sm120_dgtdv.pytests/cute/test_flash_attn_sm120_local.pytests/cute/test_fp8_decode_sm120.pytests/cute/test_mask_mod_varlen.pytests/cute/test_paged_kv_sm120.py
…+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.
Owner
Author
|
Addressed the CodeRabbit review in 5ff86ed (validated in-process vs SDPA on sm_120). Fixed (verified real on current HEAD):
Also fixed two correctness bugs found during review (replacing conservative guards):
Verified stale/already-fixed or false-positive (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, and the F401/F541/flake8/README items — all confirmed against current code. |
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
…te the sink to split 0) The conservative `learnable_sink -> num_splits=1` guard (scoped to non-sm120) is unnecessary: every SplitKV-capable forward already folds the sink into the LSE only in split 0, so it is counted exactly once across splits. - SM100 (flash_fwd_sm100.py:~2525): `not is_split_kv or split_idx == 0`, incl. the empty-split row_max==-inf case — so SM100 sink+SplitKV was already correct, and the guard was an unnecessary single-split regression I had introduced for it. - SM80-base / SM120 (compute_sink_val + softmax.finalize): split-0 gating added in the prior commit; verified in-process vs SDPA on sm_120 (5090). - SM90 has no SplitKV. Verification note: SM120 is run-verified; SM100/SM80 are verified by the split-0 gating in their kernels (no sm100/sm90 hardware available here to run).
…tream Dao-AILab#2592 insight) Upstream FA2 PR Dao-AILab#2592 routes sm_120 hd=128 forward through the sm_8x small-SMEM tile (128x32, 48 KB -> 2 blocks/SM) instead of the H100 128x64 (64 KB -> 1 block/SM), for +7-17% (biggest at small seqlen). That's the FA2 C++ kernel, not FA4 cute, so the code isn't portable — but the SMEM/occupancy principle is. Profiling our FA4 cute D128 forward (ncu, RTX 6000): the GQA pack-GQA path (FlashAttentionForwardSm120, num_stages=1) is SMEM-bound at 1 CTA/SM with the larger tiles at short seqlen; registers already allow 2 blocks. The long-seq GQA cells already use 128x32 (2 CTA/SM) — insight already captured there — and the MHA TMA path (kv_stages=2) is pinned near 99 KB and can't benefit. The only new, cross-batch-reliable wins are three S512 GQA cells, now switched to 2-CTA/SM tiles: - (128, 4, 512, 0): 128x64 -> 128x32 - (128, 8, 512, 0): (fallback 128x64) -> 128x32 [new entry] - (128, 4, 512, 1): 64x128 -> 64x64 Re-validated in-process (per-shape clock soak + round-wise ratios) at B=2 and B=16: +10.5/+4.9% (qpkv4 nc), +9.6/+4.8% (qpkv8 nc), +7.9/+10.7% (qpkv4 causal) over the prior tile; correctness rel <=3.2e-3 vs the old tile and vs SDPA. The causal cell also flips qpkv4 S512 from 0.98x (losing) to 1.09x vs the installed FA2. Longer seqlens either already use the small tile or regress with it (left unchanged).
The SM120 TMA forward's K/V pipeline depth (kv_stages) changes the compiled kernel (SMEM layout / pipeline), but it was absent from compile_key, so two kv_stages settings would share one cached binary. It's a constant (2) today so there's no live bug, but keying it now prevents a silent stale-binary reuse if kv_stages tuning is ever revisited. Verified MHA D128 forward still correct.
Keep this PR sm_120-scoped: 4 shared-kernel changes that also altered real SM80 (Ampere) are now gated so non-sm120 reproduces `main` byte-for-byte, while sm_120 keeps the fix. (Each was a pre-existing SM80 bug; those fixes are deferred to a separate SM80 PR.) The forward sm_120 specialization forces self.arch=sm_80 (to reuse the SM80 epilogue/MMA), so `arch==120` is always False there — forward items gate on a new `is_sm120` class marker (FlashAttentionForwardSm120); backward items gate on `arch==120` (FlashAttentionBackwardSm120 sets arch=120). Gated to sm_120 (non-sm120 == main, verified vs `git show main:`): - learnable_sink: restored `assert learnable_sink is None` for non-sm120 forward (so the softmax row_max_safe sink path is unreachable on real SM80). - pack_gqa forward num_head_kv + varlen q_offset (GQA R>1 indexing fix). - backward local-window m-block prune + window-aware mask (else: causal-only). - backward pack_gqa dQaccum per-element write (else: main's contiguous loop). (utils.atomic_add_fp32's red.global.add rewrite is numerically identical and left shared/un-gated; block-sparse SM80 is additive/asserted, untouched.) Cleanup: removed the multi-worktree dev shim from all 6 new sm_120 tests (plain top-level imports like tests/cute/test_flash_attn.py); neutralized internal-path references in test docstrings; aligned flash_fwd_decode_sm120.py copyright header and dropped 2 dead locals. Validated (RTX 6000, cutlass-dsl 4.5.1): sink+SplitKV 0 fails, negative-offset window bwd correct, pack_gqa fwd+bwd 3 passed, fp8 decode 1.69e-3, sm120_local 49 passed (shim-free import from repo root). sm_120 fixes all still active.
… (+5%) The dense D256 qpkv4 S512 backward runs non-packed with grid = ceil(S/64)*Hq*B = 8*num_head*batch CTAs. At small batch (num_head*batch <= 32, i.e. <=256 CTAs ~1.36 waves on the 188-SM RTX 6000) it underfills the SMs and the unsplit (split=1) default leaves >half the second wave idle. Splitting the M-loop x2 fills it to ~2.7 waves: +5.3% over split=1 (RTX 6000, interleaved config A/B), flipping this cell from below FA2 to ~parity-or-above. Gradients match SDPA (dQ/dK/dV ~3-5e-3). This shape is non-packed (pack_gqa=False), so the split-2 case is handled before the pack-only early-return in _sm120_bwd_pack_gqa_m_splits (whose result is the backward m-split regardless of pack_gqa). Tightly gated: D256, qpkv==4, S512, dense (causal uses the separate causal split policy), seqlen_q==seqlen_k, and num_head*batch<=32. Filled grids (B>=4 or Hq32 -> >32) keep the unsplit default (split regresses there); qpkv8 regresses ~25% even when underfilled -> excluded by qpkv==4. Added a batch_size param to the split helper for the underfill gate.
…2-20%) B=1 D256 backward grids (ceil(S/64)*Hq CTAs) underfill the 188 SMs for small-Hq shapes (<~1.4 waves), leaving the unsplit default idling SMs. Five exact cells win from an M-split (RTX6000 A/B vs the CURRENT dispatch, robust min across seeds): - causal qpkv8 Hq8/Hkv1 S512 -> split4 (+20%) - causal qpkv4 Hq8/Hkv2 S512 -> split3 (+18%) - dense qpkv4 Hq16/Hkv4 S1024 -> split2 (+20%) - dense qpkv4 Hq8/Hkv2 S2048 -> split2 (+18%) - dense qpkv6 Hq24/Hkv4 S1024 -> split3 (+12%) These run non-packed (B=1 doesn't auto-pack), so they're handled before the pack-only early-return in _sm120_bwd_pack_gqa_m_splits. Gated to batch_size==1 and the exact (qpkv, Hq/Hkv, S, causal): B>=2 is excluded because it either already auto-splits or the split is noise/regresses (re-validated vs current, not vs split1 — the split1 baseline overstated B>=2 gains). Gradients match SDPA (rel <=4.5e-3). Candidates that didn't robustly beat current (min<1.0) were dropped.
The forward decode-tile elif evaluated `sm120_seq_q <= 8` before the `cu_seqlens_q is None` guard. In the varlen path without an explicit max_seqlen_q, sm120_seq_q is None, so Python's left-to-right `and` raised `TypeError: '<=' not supported between NoneType and int` before the guard could short-circuit -- breaking the entire head_dim<=128 varlen forward (a regression from the decode-tile optimization in 512254c). Move the `sm120_seq_q <= 8` clause after the `... is None` guards: for varlen, `cu_seqlens_q is None` short-circuits to False before the None comparison; for the intended non-varlen decode case all guards are True so the comparison runs identically (no behavior change). Validated on RTX 5090: test_flash_attn_varlen.py head_dim=128 1728 passed/0 failed, head_dim=64 2160 passed/0 failed (all previously crashed); decode-tile path unregressed (seqlen_q<=8 non-varlen still selects FwdConfig(16,64)); varlen output matches per-sequence reference (max err 3.9e-3).
Two changes in the shared SM80-base forward leaked into non-sm120 archs vs main: 1. flash_fwd.py: the pack_gqa_layout folding (absent in main) was gated only on self.pack_gqa, so real SM80 + GQA (pack_gqa defaults True, no non-sm120 pack_gqa=False guard on the forward) ran a folded path that reads the wrong KV head. Gate it on `self.pack_gqa and is_sm120`. The dependent tile_sched_args (cute.size(mQ.shape[0/2]), is_split_kv/num_splits/seqlen_k) are byte-equivalent for the unfolded non-sm120 case (single scalar mode -> size==extent; split args inert at is_split_kv=False), so no further gating needed. 2. softmax.py: the finalize() row_max==-inf sink guard (absent in main) is reachable on SM90 (shares the base Softmax) via sink+block-sparse empty rows, changing SM90 output. Add an is_sm120 Constexpr param (default False = main's plain row_max[r]) and thread it from the three SM80-base callers; SM90/SM100/sm120-TMA callers are untouched so they keep main behavior. Non-sm120 (real SM80, SM90) now match main byte-for-byte; sm120 keeps both fixes. Validated on sm120 (RTX 5090): GQA forward + learnable_sink + SplitKV and the local-window GQA suite still pass. (No SM80/SM90 hardware available; non-sm120 equivalence is by static analysis of the const_expr branches.)
…rong) The paged-KV SplitKV mainloop (_paged_kv_mainloop) called compute_sink_val() and epilogue() WITHOUT split_idx, so it defaulted to 0: every split wrote its partial O/LSE into out_partial[0]/lse_partial[0]. Splits raced on slot 0 and the slots for splits >=1 were never written, so the combine merged uninitialized torch.empty garbage -> NaN (or, with luckier memory, silently wrong, max-diff 0.36 vs num_splits=1). Only paged was affected; the dense SplitKV mainloop already passes split_idx, which is why non-paged SplitKV+cache was fine. Secondary: the paged mainloop ran its first masked iteration unconditionally with no has_work guard, so a SplitKV split with no assigned KV blocks (cache_seqlens short relative to num_splits) processed a block belonging to a lower split and emitted a finite garbage partial the combine double-counts, instead of the clean O=0/LSE=-inf sentinel the combine drops. Fix: thread split_idx through _paged_kv_mainloop into compute_sink_val + epilogue; guard the first iteration with has_work = n_block_max > n_block_min (gated on is_split_kv). Non-split path is byte-identical (split_idx=0, has_work=True); this mainloop is sm120-only (SM80 rejects paged KV). Combine kernel untouched. Validated on RTX 5090: repro test_flash_attn_kvcache[...True...] passes; crafted paged SplitKV+cache 0.36->0.001 vs num_splits=1; empty-split 0.14->0.0; kvcache decode 528 passed/0 failed; test_paged_kv_sm120 + test_fp8_decode_sm120 76 passed; dense SplitKV num_splits=3 unregressed. (8 gqa kvcache failures are pre-existing -- the GQA+SplitKV+non-varlen TODO at interface.py:1196 -- verified on the clean tree.)
…arbage)
The "sliding-window large-right-window" failure was misdiagnosed: the real trigger
is varlen + pack_gqa + SplitKV (num_splits>1), independent of masking/window. When
is_sm120 and pack_gqa, the partial-O/LSE buffers are folded to the packed
(qhead_per_kvhead, seqlen_q) layout, but the SplitKV partial epilogue in flash_fwd.py
writes O/LSE via the unpacked local_tile path with an unpacked row predicate (never
calling pack_gqa.store_O/store_LSE), so packed rows scatter to the wrong partial
slots -> NaN/garbage (max-diff ~1.96 in the failing varlen GQA cell). The non-varlen
case was already disabled at interface.py:1197; the varlen path was left unguarded.
Fix: mirror that guard for varlen, gated to sm120 (arch//10==12) so SM100 -- which
uses its own pack_gqa+SplitKV kernel -- is unaffected. SplitKV still runs correctly
via the non-packed GQA path; only the pack_gqa optimization is skipped for this
unsupported combo (a packed-aware SplitKV epilogue would restore it -- future work).
Validated on RTX 6000/5090: the original failing test passes; varlen GQA pack_gqa
+num_splits=3 window sweep {(None,None),(None,64/357/1055/8000),(64/256,None),
(128,256),(512,512)} all max-diff <0.01 (were NaN/1.5+); dense pack_gqa×splits
unchanged; dv!=d varlen correct; test_flash_attn_sm120_local 49/49;
test_flash_attn_varlen_output local sample 8/8.
…sk_mod combo)
The block-sparse forward mainloop (mma_one_n_block_bs) reuses a single-stage smem
K/V buffer (smem_pipe_write=0) and reloads K/V per selected n-block, but did not
synchronize between blocks: the next block's load_V cp.async could overwrite sV[0]
before all warps finished the previous block's PV GEMM read of it (cp.async
commit/wait order cp.async ops among themselves, not the MMA's smem reads). This
WAR race produced nondeterministic wrong output (max-diff 0.10-0.23, run-to-run
varying) for tiles with >1 selected block AND enough heads/CTAs in flight -- e.g.
block-sparse + a within-tile mask_mod (mini_causal/causal/sliding_window/document)
at seqlen>=1024. The indexing (block list, n_block, mask_mod kv_idx) was verified
correct; block_diagonal (1 block/tile) and single-head never raced.
Fix: add cute.arch.barrier() before the K/V reload for every non-first block,
gated on `not is_first_n_block and is_sm120`. The first block has no in-tile
predecessor (its prologue already syncs after load_Q). sm120-gated; the latent
SM80-base race is deferred to the separate SM80 PR.
Validated on RTX 6000/5090: nondeterministic 0.10-0.23 -> deterministic 5.7e-3
(3 identical repeats); forward correct across mini_causal/block_diagonal/
sliding_window{128,256,512} x mha/gqa/mqa x seqlen{1023,1024,4096}; plain mask_mod
forward control unchanged; test_block_sparsity.py 4883 passed/0 failed. (Those
test IDs' residual failures are the pre-existing block-sparse-BACKWARD gate at
interface.py:2498, out of scope.)
…ut + IMA) A non-varlen block-sparse forward whose sparse Q block size exceeds the sm120 kernel tile_m (e.g. BlockMask built with tile_m=256 on the CC>=10 path, run by the 128-wide sm120 kernel) needs q_subtile_factor to map each kernel m_block to its owning sparse block (m_block // factor). normalize_block_sparse_config computes factor=2 correctly, and the SM80 and SM100 dispatch branches pass it -- but the arch//10==12 (FlashAttentionForwardSm120) branch omitted it, so the kernel defaulted to factor=1: kernel m_block=1 read sparse m-block index 1 when the metadata has only index 0, reading past mask_block_cnt/mask_block_idx -> wrong output (max-diff 0.77) and, for larger configs, cudaErrorIllegalAddress that poisoned the CUDA context (the 508-error suite cascade was this one crash). Fix: pass q_subtile_factor=q_subtile_factor to the FlashAttentionForwardSm120 constructor. It's already in the forward compile-cache key, sm120-only; other archs untouched. Validated on RTX 6000/5090: compute-sanitizer memcheck 0 errors on the formerly- IMA shapes; 20 mask x seqlen combos bit-exact (max_err 0.0); test_varlen_block_sparse non-varlen group 64 passed/0 failed (was 4 failed + IMA); full test_mask_mod_varlen.py 401 passed/500 skipped/0 failed/0 errors (cascade gone); test_block_sparsity.py 4883 passed (no regression).
The cute suite had ~1520 FAIL/ERROR on sm120 that are NOT bugs -- tests that exercise paths sm120 explicitly does not support, without arch-skipping, so the interface "not supported on SM 12.0" asserts surfaced as failures. Add surgical cc==12 guards (other archs untouched): - test_score_mod.py: skip the 3 score_mod BACKWARD tests on cc==12 (forward score_mod still runs+passes, 720 cases). - test_mask_mod.py: in _run_mask_test, skip the backward portion AFTER the forward is validated (forward mask_mod/block-sparse still exercised), skip the autograd block-sparse path (raises at forward-time building the bwd graph), and skip test_gqa_block_sparse_broadcast_pattern_recompilation (block-sparse bwd) and test_compact_block_sparse_indices (tile_n=128 vs sm120's default n_block_size=64). - test_flash_attn.py: relax the dQ allclose to bf16 tol on sm120 only (non-deterministic atomic-add dQ; the deterministic semaphore scheduler is SM90/SM100-only); dK/dV stay bit-exact on all archs. No over-skipping: forward score_mod/mask_mod/block-sparse paths still RUN and pass; vanilla backward still runs. RTX 5090 after: test_score_mod 1520 passed/1097 skipped/0 failed; test_mask_mod 145 passed/1405 skipped/0 failed; bwd_preallocated 8 passed; test_mask_mod_varlen 401/500/0 and test_block_sparsity 4883/0 unchanged.
…hmarks) - ruff format the 4 non-excluded sm120 files (flash_fwd_sm120_tma.py, flash_fwd_decode_sm120.py, mask.py, pack_gqa.py) -> satisfies the pre-commit ruff-format hook (whitespace/line-wrapping only, no logic change; py_compile OK). - Add the standard FA copyright header to block_sparse_utils.py. - Remove the vendor-attribution comment from flash_fwd_sm120_tma.py (attribution belongs in the PR description, not the file header). - Remove benchmarks/sm120_bwd_tuning/ (internal dev tuning harness with campaign-specific nomenclature; not appropriate for upstream).
…gression nvidia-cutlass-dsl 4.5.2 introduced a DSL codegen regression that breaks the sm120 fp8 (e4m3/e5m2) KV-cache decode kernel (nvgpu.cvt_fpext rejects a scalar f8E4M3FN operand -> compile failure); 4.5.1 compiles and runs correctly. Upstream CI on a different DSL version would otherwise surface this as a confusing compile error. Add `_fp8_decode_dsl_supported()` (+`_parse_dsl_version`): a reusable predicate guarding the half-open broken window [4.5.2, _DSL_FP8_DECODE_FIXED_VERSION=None). Unknown/unparseable versions are treated as supported (don't over-guard). When the DSL is fixed, set _DSL_FP8_DECODE_FIXED_VERSION to the first good release -- the only change needed. - interface.py: at the fp8-decode dispatch (gated by want_fp8_decode) raise a clear NotImplementedError on a broken DSL. No silent bf16 fallback -- the K/V cache is physically fp8, so a dtype switch would reinterpret bytes as garbage. - test_fp8_decode_sm120.py: pytest.skip on the broken version with an informative reason (in addition to the existing non-sm120 skip). Validated on 4.5.1: test_fp8_decode_sm120.py 25 passed / 0 skipped (guard does not over-fire); version-logic check confirms 4.5.0/4.5.1 -> ok, 4.5.2/4.5.3/4.6.0 -> guarded, unparseable -> ok; non-fp8 (bf16/fp16) paths unaffected (guard is strictly behind want_fp8_decode).
Neutralize device-serial and internal-campaign references in the sm120 tile-table / dispatch / M-split comments that don't resolve for upstream readers (23 comments): - device serials (RTX 6000 / RTX 5090 / 188-SM) -> "sm120" / "high-SM-count sm120 part", keeping the technical point; - internal report/phase/repro references (phase5c/REPORT.md, Phase 17C, agent_space paths) -> removed; - device-tied vs-FA2 absolute ratios -> dropped, keeping the relative speedups and the engineering rationale (occupancy/underfill/SMEM-cap reasoning, shape hints). Comments only: verified the code portion of every changed line (incl. the 7 dict-entry lines with trailing comments) is byte-identical; tuning logic, dict values, and conditions are untouched. py_compile + ruff clean.
…qa, +1.24x) GQA + SplitKV was forced onto the non-packed path because the SplitKV partial-O/LSE epilogue wrote the unpacked layout while pack_gqa folds O/LSE to the packed (qhead_per_kvhead, seqlen_q) layout -> scattered partial slots (the interface.py TODO for non-varlen, and the root cause of the varlen BUG 3). Correct but slower: GQA decode/long-context (where SplitKV matters most) lost the pack_gqa speedup. Add pack_gqa.store_O_partial / store_LSE_partial: scatter the fp32 MMA accumulator (and LSE) directly to each packed row's physical (h_idx, m_idx) partial slot via the composite stride. No combine-kernel change needed -- the combine reads partials in unpacked physical layout [split,batch,m,head,d], which is exactly the slot the packed scatter targets, so forward-write and combine-read are consistent end-to-end. Gated on is_sm120 + pack_gqa + is_split_kv; the unpacked path is unchanged for MHA / non-pack / non-sm120. Relax the interface guards: pack_gqa stays ENABLED for GQA+SplitKV on sm120 (non-varlen AND varlen); non-sm120 non-varlen stays disabled, SM100's own pack_gqa+SplitKV kernel untouched. This supersedes the varlen BUG 3 fallback (1a9d81d) with the real fix. Validated on RTX 5090: SplitKV(num_splits=3) == num_splits=1 baseline for GQA non-varlen AND varlen incl the BUG 3 window(None,1055) cell (rel <=6e-3), compute-sanitizer memcheck 0 errors, repo suites green (varlen GQA/MQA 2592, output GQA/MQA 9616, paged_kv 51). Perf A/B (5090, production auto-split path, pack-on vs pack-off): 1.02-1.56x, geomean 1.24x, scaling with GQA ratio; with num_splits held equal it is ~1.00x (no regression, confirms the win is real densification, not a split-count artifact). 6000 perf pending (cards were contended); the directional GQA-ratio-scaling win should hold.
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/cute/test_mask_mod.py (1)
2550-2558: ⚡ Quick winKeep SM120 forward coverage instead of skipping this test.
This is a forward-only test. You can avoid the SM120 skip by passing an explicit
tile_mn=(tile_m, tile_n)in both_flash_attn_fwdcalls sotile_nmatches the block-sparse tensors.Suggested change
- out_compact, _ = _flash_attn_fwd( + out_compact, _ = _flash_attn_fwd( q=tensors["q"], k=tensors["k"], v=tensors["v"], out=tensors["out"].clone(), lse=tensors["lse"].clone(), softmax_scale=1.0 / math.sqrt(headdim), + tile_mn=(tile_m, tile_n), causal=False, mask_mod=mask_mod_cute, block_sparse_tensors=block_sparse_compact, return_lse=True, ) ... - out_full, _ = _flash_attn_fwd( + out_full, _ = _flash_attn_fwd( q=tensors["q"], k=tensors["k"], v=tensors["v"], out=tensors["out"].clone(), lse=tensors["lse"].clone(), softmax_scale=1.0 / math.sqrt(headdim), + tile_mn=(tile_m, tile_n), causal=False, mask_mod=mask_mod_cute, block_sparse_tensors=block_sparse_full, return_lse=True, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/cute/test_mask_mod.py` around lines 2550 - 2558, The test currently skips SM 12.0 due to a mismatch between the default tile_n and the block-sparse tensor block_size; instead of skipping, update both calls to _flash_attn_fwd in this test to pass an explicit tile_mn=(tile_m, tile_n) so the forward uses the correct tile_n matching sparse_block_size[1]=128; locate the two _flash_attn_fwd invocations and add the tile_mn argument (using the existing tile_m and tile_n variables) so SM120 forward coverage is retained.flash_attn/cute/flash_fwd.py (1)
2067-2071: 💤 Low valueMinor style inconsistency with dense path.
The dense path (line 1061) uses Python
Truedirectly, while this paged path usescutlass.Boolean(True). PythonTruemay enable slightly better compile-time optimization since it's a constant literal. Not a correctness issue; functionally equivalent.🤖 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 2067 - 2071, The paged-path conditional for has_work uses cutlass.Boolean(True) while the dense path uses the Python literal True; change the paged-path branch to use the Python True literal for consistency and potential compile-time optimization by replacing cutlass.Boolean(True) with True in the has_work expression (referencing has_work, is_split_kv, n_block_max, n_block_min, and cutlass.Boolean in the diff).tests/cute/test_paged_kv_sm120.py (1)
344-401: 💤 Low valueMinor: Redundant import inside test function.
torch.nn.functional as Fis already imported at the module level (line 41). The import at line 381 is unnecessary and can be removed. Thesdpa_kernelandSDPBackendimports at line 382 are fine since they're only used in this test.Suggested cleanup
# Reconstruct logical K/V from page table for SDPA reference. - import torch.nn.functional as F from torch.nn.attention import sdpa_kernel, SDPBackend🤖 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 `@tests/cute/test_paged_kv_sm120.py` around lines 344 - 401, The test function test_d128_dv64_paged_varlen_correctness contains a redundant local import "import torch.nn.functional as F" — remove that local import line and keep using the module-level F import already provided; leave the subsequent "from torch.nn.attention import sdpa_kernel, SDPBackend" import in place since those are only used here. This change is limited to deleting the redundant import inside the function and verifying references to F, sdpa_kernel, and SDPBackend still resolve.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@flash_attn/cute/flash_fwd.py`:
- Around line 2067-2071: The paged-path conditional for has_work uses
cutlass.Boolean(True) while the dense path uses the Python literal True; change
the paged-path branch to use the Python True literal for consistency and
potential compile-time optimization by replacing cutlass.Boolean(True) with True
in the has_work expression (referencing has_work, is_split_kv, n_block_max,
n_block_min, and cutlass.Boolean in the diff).
In `@tests/cute/test_mask_mod.py`:
- Around line 2550-2558: The test currently skips SM 12.0 due to a mismatch
between the default tile_n and the block-sparse tensor block_size; instead of
skipping, update both calls to _flash_attn_fwd in this test to pass an explicit
tile_mn=(tile_m, tile_n) so the forward uses the correct tile_n matching
sparse_block_size[1]=128; locate the two _flash_attn_fwd invocations and add the
tile_mn argument (using the existing tile_m and tile_n variables) so SM120
forward coverage is retained.
In `@tests/cute/test_paged_kv_sm120.py`:
- Around line 344-401: The test function test_d128_dv64_paged_varlen_correctness
contains a redundant local import "import torch.nn.functional as F" — remove
that local import line and keep using the module-level F import already
provided; leave the subsequent "from torch.nn.attention import sdpa_kernel,
SDPBackend" import in place since those are only used here. This change is
limited to deleting the redundant import inside the function and verifying
references to F, sdpa_kernel, and SDPBackend still resolve.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 15984509-4462-4857-89ad-49e1a24faf46
📒 Files selected for processing (19)
flash_attn/cute/block_sparse_utils.pyflash_attn/cute/flash_bwd.pyflash_attn/cute/flash_fwd.pyflash_attn/cute/flash_fwd_decode_sm120.pyflash_attn/cute/flash_fwd_sm120.pyflash_attn/cute/flash_fwd_sm120_tma.pyflash_attn/cute/interface.pyflash_attn/cute/mask.pyflash_attn/cute/pack_gqa.pyflash_attn/cute/softmax.pytests/cute/test_flash_attn.pytests/cute/test_flash_attn_bwd_sm120_pack_gqa.pytests/cute/test_flash_attn_bwd_sm120_postprocess.pytests/cute/test_flash_attn_sm120_dgtdv.pytests/cute/test_flash_attn_sm120_local.pytests/cute/test_fp8_decode_sm120.pytests/cute/test_mask_mod.pytests/cute/test_paged_kv_sm120.pytests/cute/test_score_mod.py
💤 Files with no reviewable changes (2)
- tests/cute/test_flash_attn_bwd_sm120_postprocess.py
- tests/cute/test_flash_attn_sm120_local.py
🚧 Files skipped from review as they are similar to previous changes (9)
- flash_attn/cute/flash_fwd_sm120.py
- flash_attn/cute/mask.py
- flash_attn/cute/block_sparse_utils.py
- flash_attn/cute/flash_fwd_decode_sm120.py
- flash_attn/cute/softmax.py
- flash_attn/cute/flash_fwd_sm120_tma.py
- tests/cute/test_fp8_decode_sm120.py
- flash_attn/cute/flash_bwd.py
- flash_attn/cute/pack_gqa.py
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/cute/test_mask_mod.py (1)
2550-2558: ⚡ Quick winKeep SM120 forward coverage instead of skipping this test.
This is a forward-only test. You can avoid the SM120 skip by passing an explicit
tile_mn=(tile_m, tile_n)in both_flash_attn_fwdcalls sotile_nmatches the block-sparse tensors.Suggested change
- out_compact, _ = _flash_attn_fwd( + out_compact, _ = _flash_attn_fwd( q=tensors["q"], k=tensors["k"], v=tensors["v"], out=tensors["out"].clone(), lse=tensors["lse"].clone(), softmax_scale=1.0 / math.sqrt(headdim), + tile_mn=(tile_m, tile_n), causal=False, mask_mod=mask_mod_cute, block_sparse_tensors=block_sparse_compact, return_lse=True, ) ... - out_full, _ = _flash_attn_fwd( + out_full, _ = _flash_attn_fwd( q=tensors["q"], k=tensors["k"], v=tensors["v"], out=tensors["out"].clone(), lse=tensors["lse"].clone(), softmax_scale=1.0 / math.sqrt(headdim), + tile_mn=(tile_m, tile_n), causal=False, mask_mod=mask_mod_cute, block_sparse_tensors=block_sparse_full, return_lse=True, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/cute/test_mask_mod.py` around lines 2550 - 2558, The test currently skips SM 12.0 due to a mismatch between the default tile_n and the block-sparse tensor block_size; instead of skipping, update both calls to _flash_attn_fwd in this test to pass an explicit tile_mn=(tile_m, tile_n) so the forward uses the correct tile_n matching sparse_block_size[1]=128; locate the two _flash_attn_fwd invocations and add the tile_mn argument (using the existing tile_m and tile_n variables) so SM120 forward coverage is retained.flash_attn/cute/flash_fwd.py (1)
2067-2071: 💤 Low valueMinor style inconsistency with dense path.
The dense path (line 1061) uses Python
Truedirectly, while this paged path usescutlass.Boolean(True). PythonTruemay enable slightly better compile-time optimization since it's a constant literal. Not a correctness issue; functionally equivalent.🤖 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 2067 - 2071, The paged-path conditional for has_work uses cutlass.Boolean(True) while the dense path uses the Python literal True; change the paged-path branch to use the Python True literal for consistency and potential compile-time optimization by replacing cutlass.Boolean(True) with True in the has_work expression (referencing has_work, is_split_kv, n_block_max, n_block_min, and cutlass.Boolean in the diff).tests/cute/test_paged_kv_sm120.py (1)
344-401: 💤 Low valueMinor: Redundant import inside test function.
torch.nn.functional as Fis already imported at the module level (line 41). The import at line 381 is unnecessary and can be removed. Thesdpa_kernelandSDPBackendimports at line 382 are fine since they're only used in this test.Suggested cleanup
# Reconstruct logical K/V from page table for SDPA reference. - import torch.nn.functional as F from torch.nn.attention import sdpa_kernel, SDPBackend🤖 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 `@tests/cute/test_paged_kv_sm120.py` around lines 344 - 401, The test function test_d128_dv64_paged_varlen_correctness contains a redundant local import "import torch.nn.functional as F" — remove that local import line and keep using the module-level F import already provided; leave the subsequent "from torch.nn.attention import sdpa_kernel, SDPBackend" import in place since those are only used here. This change is limited to deleting the redundant import inside the function and verifying references to F, sdpa_kernel, and SDPBackend still resolve.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@flash_attn/cute/flash_fwd.py`:
- Around line 2067-2071: The paged-path conditional for has_work uses
cutlass.Boolean(True) while the dense path uses the Python literal True; change
the paged-path branch to use the Python True literal for consistency and
potential compile-time optimization by replacing cutlass.Boolean(True) with True
in the has_work expression (referencing has_work, is_split_kv, n_block_max,
n_block_min, and cutlass.Boolean in the diff).
In `@tests/cute/test_mask_mod.py`:
- Around line 2550-2558: The test currently skips SM 12.0 due to a mismatch
between the default tile_n and the block-sparse tensor block_size; instead of
skipping, update both calls to _flash_attn_fwd in this test to pass an explicit
tile_mn=(tile_m, tile_n) so the forward uses the correct tile_n matching
sparse_block_size[1]=128; locate the two _flash_attn_fwd invocations and add the
tile_mn argument (using the existing tile_m and tile_n variables) so SM120
forward coverage is retained.
In `@tests/cute/test_paged_kv_sm120.py`:
- Around line 344-401: The test function test_d128_dv64_paged_varlen_correctness
contains a redundant local import "import torch.nn.functional as F" — remove
that local import line and keep using the module-level F import already
provided; leave the subsequent "from torch.nn.attention import sdpa_kernel,
SDPBackend" import in place since those are only used here. This change is
limited to deleting the redundant import inside the function and verifying
references to F, sdpa_kernel, and SDPBackend still resolve.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 15984509-4462-4857-89ad-49e1a24faf46
📒 Files selected for processing (19)
flash_attn/cute/block_sparse_utils.pyflash_attn/cute/flash_bwd.pyflash_attn/cute/flash_fwd.pyflash_attn/cute/flash_fwd_decode_sm120.pyflash_attn/cute/flash_fwd_sm120.pyflash_attn/cute/flash_fwd_sm120_tma.pyflash_attn/cute/interface.pyflash_attn/cute/mask.pyflash_attn/cute/pack_gqa.pyflash_attn/cute/softmax.pytests/cute/test_flash_attn.pytests/cute/test_flash_attn_bwd_sm120_pack_gqa.pytests/cute/test_flash_attn_bwd_sm120_postprocess.pytests/cute/test_flash_attn_sm120_dgtdv.pytests/cute/test_flash_attn_sm120_local.pytests/cute/test_fp8_decode_sm120.pytests/cute/test_mask_mod.pytests/cute/test_paged_kv_sm120.pytests/cute/test_score_mod.py
💤 Files with no reviewable changes (2)
- tests/cute/test_flash_attn_bwd_sm120_postprocess.py
- tests/cute/test_flash_attn_sm120_local.py
🚧 Files skipped from review as they are similar to previous changes (9)
- flash_attn/cute/flash_fwd_sm120.py
- flash_attn/cute/mask.py
- flash_attn/cute/block_sparse_utils.py
- flash_attn/cute/flash_fwd_decode_sm120.py
- flash_attn/cute/softmax.py
- flash_attn/cute/flash_fwd_sm120_tma.py
- tests/cute/test_fp8_decode_sm120.py
- flash_attn/cute/flash_bwd.py
- flash_attn/cute/pack_gqa.py
Owner
Author
CodeRabbit review — dispositionAll 11 inline CodeRabbit comments (dated 2026-06-04) were re-validated against current HEAD (
Validation: |
thad0ctor
pushed a commit
that referenced
this pull request
Jun 9, 2026
* Bump aiter submodule commit Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com> * Bump aiter submodule to 3b2e6f48ce97e1d494e8b3f1af5c65f74e304b28 (#2) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com>
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.
FA4 consumer Blackwell (sm_120) integration
Adds FlashAttention-4 (CuTeDSL) support for consumer/workstation Blackwell — RTX 50-series / RTX PRO 6000 / DGX Spark (compute capability 12.x). Dispatch and tile selection are auto-tuned for the arch; no environment variables are required for normal use.
Background
Consumer/workstation Blackwell (
sm_120) ships 5th-gen tensor cores with new fp4/fp8 data types but — unlike datacenter Blackwell (sm_100) — keeps the Ampere-style warp-levelmma.syncprogramming model: no WGMMA, no tcgen05, no Tensor Memory (confirmed against the CUDA 13.x toolkit and ptxas SASS — the tensor-memory subsystem is absent in silicon, not just unexposed). So FA4'ssm_100path does not apply;sm_120reuses the SM80-base kernels (cp.async, plus a TMA + warp-specialized forward) and needs its own dispatch, tiling, and correctness work.This builds on the in-flight upstream FA4
sm_120kernels (Dao-AILab Dao-AILab#2553 / Dao-AILab#2349 / Dao-AILab#2389), which landed initial kernels but left public-API routing, correctness, and coverage gaps:sm_120paths via the public API, auto-tuned dispatchsm_120GQA shapessm_120correctlyhead_dim ≤ 256, validatedhead_dim > head_dim_vsm_120TMA hangsm_120byte-identical tomain); CodeRabbit follow-upsWhat's supported
score_mod/mask_mod; learnable sink. Head dims 64/96/128/192/256.seqlen_q == 1with a quantized K/V cache + bf16/fp16 query (pass fp8k/v+ per-(batch, kv_head)k_descale/v_descale). Auto-routes to a memory-efficient GEMV decode kernel; ~1.6–1.9× faster than bf16 at GQA ratio ≤ 4 while halving KV bandwidth, within ~2e-3 of an fp8-quantized reference.Performance vs FA2 (and SDPA)
Two
sm_120parts: RTX 5090 (consumer, 170 SM, 32 GB) and RTX PRO 6000 Blackwell (workstation, 188 SM, 96 GB). bf16. In-process round-wise interleaved A/B — each round times the impls back-to-back under the same clock state, so the ratios are clock-robust; medians reported.Legend.
D= head dim ·B= batch size ·r= GQA ratio = query-heads per KV-head (r1/MHA= one KV head per query head;r8= 8 query heads share 1 KV head;MQA= all heads share 1) · dense = full attention, causal = causal mask · prefill =seqlen_q == seqlen_k, decode =seqlen_q == 1(one new query against a length-SkKV cache) · ratio = FA4 / FA2 throughput, >1 = FA4 faster (so1.15= FA4 15% faster,0.40= FA4 60% slower). Config labels are(Hq/Hkv). Short seqlens have higher run-to-run variance — read the trend.Prefill forward — per seqlen (B=1)
The headline is D256; D128 is the other common case. (Full matrix incl. D64/D96/MHA and B=2 in the collapsibles below.)
D256 GQA (16/2), FA4/FA2:
D128 GQA (32/4), FA4/FA2:
Geomean rollup (over seqlen 512–32k, FA4/FA2)
Full per-seqlen prefill matrix — RTX 5090 (all configs, B=1 & B=2, fwd+bwd, dense+causal)
Full per-seqlen prefill matrix — RTX PRO 6000 (all configs, B=1 & B=2, fwd+bwd, dense+causal)
Prefill vs PyTorch SDPA (best torch backend)
FA4 matches or beats the best PyTorch SDPA on prefill (FA4/SDPA, >1 = FA4 faster); on the 6000 D256 it leads by 14–18%:
Decode (
seqlen_q == 1)Decode is memory-bandwidth-bound, so absolute latency (not a FLOP ratio) is the meaningful metric. fp8 K/V also halves cache memory (1 byte/elem vs 2). Latency at
Sk=16384(µs, lower = better):Across the full decode sweep (batch 16/64/128 ×
Sk4k–64k): fp8 decode wins at GQA r2/r4 with batch ≥ 64 — geomean ~1.4–2.0× vs FA2, up to ~2.2× — and scales up with context length, while halving KV memory. It loses at GQA r8 (~0.3–0.5×: the packed GQA-decode GEMV becomes compute-bound) and at tiny-batch + short-context. bf16 decode is at parity for MHA/r2 and trails at r4/r8.PyTorch SDPA is not a viable GQA-decode baseline — it has no GQA-decode path, so it materializes the expanded K/V (17–34 GB here): where it fits it is 23–45× slower than FA4-fp8 (e.g. 101 ms vs 3 ms), and otherwise OOMs. FA4 is the only efficient GQA-decode path of the three.
Achieved throughput (prefill forward, TFLOP/s)
Algorithmic TFLOP/s = standard attention FLOPs ÷ measured time; the FLOP count is ncu-verified exact (the hardware tensor-instruction count matches the analytical formula — dense/backward exact, causal +0.4%).
FA4 forward sustains ~180–191 TF/s on the 5090 (~86–91% of its ~209 TF/s bf16 dense peak) and ~245–261 TF/s on the 6000.
Honest sub-parity regimes (so they aren't a surprise to a reviewer)
Correctness
Validated against PyTorch SDPA across the full
tests/cutematrix, which is green onsm_120(with documentedskips for the unsupported configs below). Everysm_120change is arch-gated: SM80 / SM90 / SM100 / SM110 are byte-identical tomain(statically audited across the diff; their dedicated kernels are untouched).Real bugs found and properly fixed during development (not merely guarded):
learnable_sink+ SplitKV double-counted the sink across splits → wrong forward output. Fixed: the sink is folded into LSE only in split 0, so it's counted exactly once.window_size=(None,-X)/(-X,None)) produced wrong dK/dV (empty/inverted m-block prune). Fixed: the prune is skipped per-side for negative bounds (forward was already correct).sm_120(flash_fwd_combine.py):arch>=90wrongly enabledgriddepcontrolon the sm_80-compiled target. Fixed to excludesm_120.TypeError(head_dim ≤ 128withoutmax_seqlen_q): a decode-tile guard compared aNoneseqlen before its short-circuit, breaking the entire general varlen forward (a regression from a decode-tile optimization). Fixed (guard reorder).split_idx) → NaN/garbage withcache_seqlens+num_splits > 1. Fixed (threadsplit_idx+ empty-split guard).mask_modforward had a write-after-read shared-memory race (no inter-block barrier on the single-stage K/V buffer) → nondeterministic wrong output at multi-tile seqlens. Fixed (barrier).q_subtile_factoronsm_120→ out-of-bounds read (wrong output + illegal memory access) when the sparse block exceeded the kernel tile. Fixed.Independently double-checked: every fix re-verified cross-device (RTX 5090 + RTX PRO 6000),
compute-sanitizerclean on the formerly-IMA paths, and the arch-gating re-audited end-to-end.Known limitations / floors on sm_120 (documented in
flash_attn/cute/README.md)qhead_per_kvhead==1) atseqlen ≥ 8192is ~0.95× FA2 — an intrinsic register/occupancy wall (255 regs/thread → 1 CTA/SM). GQA (the common case) is at parity or faster.learnable_sinkruns single-split (no decode SplitKV speedup); negative-offset windows are forward-only;deterministic=Truebackward is unsupported. (Proper fixes for the last items are follow-ups.)m16n8k32thread-layout mismatch).Testing
New sm_120 test coverage:
test_fp8_decode_sm120.py,test_flash_attn_sm120_local.py,test_paged_kv_sm120.py,test_flash_attn_sm120_dgtdv.py,test_flash_attn_bwd_sm120_pack_gqa.py,test_flash_attn_bwd_sm120_postprocess.py. Thetest_flash_attn.pysuite is green on sm_120 modulo documentedskips for the unsupported configs above.Supersedes #1 — rebased and cleaned for upstream: experimental tuning env-flags collapsed to their shipping defaults, internal campaign notes/artifacts dropped, history squashed into logical commits, all changes strictly
sm_120-gated, and the correctness + hardening work above completed and double-checked.Summary by CodeRabbit
New Features
Documentation
Bug Fixes
Tests