Conversation
| batch_idx=batch_size, | ||
| seqlen_q_static=mQ.shape[0], | ||
| seqlen_k_static=mK.shape[0], | ||
| mCuSeqlensQ=mCuSeqlensQ, |
There was a problem hiding this comment.
Can you rebase on main? Varlen was already merged in #2333, as you know.
| # Unpack sparse block lists for this (batch, head, m_block) tile. | ||
| # mask_blocks = partially-masked KV blocks (need mask applied). | ||
| # full_blocks = fully-unmasked KV blocks (no masking needed). | ||
| _bs_mask_cnt, _bs_mask_idx, _bs_full_cnt, _bs_full_idx = blocksparse_tensors |
There was a problem hiding this comment.
Please keep variable names consistent with the other kernels. Is there a reason not to use the same design pattern for handling blocksparsity as in the others, e.g.
?| # Indices stored in decreasing order: highest n_block first. | ||
| # First mask block: is_first=True, mask_seqlen=True, mask_mod=self.mask_mod. | ||
| # Remaining mask blocks: is_first=False, mask_seqlen=False, mask_mod=self.mask_mod. | ||
| # CuTe DSL forbids closures in dynamic CF, so block processing is inlined at |
There was a problem hiding this comment.
Unnecessary comment on CuTe DSL
| # === Process full blocks (no masking) === | ||
| # When no mask blocks preceded: first full block is the first block overall | ||
| # (is_first=True) and gets mask_seqlen=True for seqlen_k boundary check. | ||
| # When mask blocks preceded: all full blocks have is_first=False. |
There was a problem hiding this comment.
First full block must always get seqlen masking, since it is possible that all mask blocks are further left than some full block.
6935828 to
a904872
Compare
|
Thanks for the review! Here's a summary of what changed in the v2 rewrite: 1. Rebased on main ✅ 2. Design pattern matching SM90 ✅ The block-sparse mainloop now follows the same
SM80 can't use the exact TMA producer/consumer split since it has no warpgroup specialization, but the 3. CuTe DSL comment removed ✅ (was in the old version) 4. First full block seqlen masking ✅ In Also fixed two bugs found during validation on SM121a:
Validated: |
377f1bd to
4dd1d4d
Compare
… + 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.
FlashAttentionForwardSm80 (which SM120 inherits) accepted blocksparse_tensors but the mainloop was dense-only; block-sparse forward existed on SM90/SM100 only. Add it to the non-warp-specialized cp.async path. - block_sparse_utils.run_block_sparse_mainloop_sm80: visits the active mask/full blocks (mask blocks first with mask_mod + seqlen masking, then full blocks with seqlen masking only). The first full block always gets seqlen masking even after mask blocks, since a full block may sit at the seqlen_kv boundary regardless of mask-block positions. Mirrors the masking contract of consume_block_sparse_loads (SM90/SM100). - FlashAttentionForwardSm80.mma_one_n_block_bs: per-block load+compute. Unlike the dense compute_one_n_block it does not prefetch the next block, since sparse blocks are not contiguous (no producer warp on this pipeline to run the warp-specialized produce/consume helpers). - Guard the dense contiguous prologue prefetch behind use_block_sparsity; the block-sparse path drains async copies and makes Q available instead. - Wire blocksparse_tensors from __call__ through kernel() (was dropped).
4dd1d4d to
78957f5
Compare
|
Thanks for the review @reubenconducts, and apologies for the long gap. Rebased onto current Rebased on main (varlen #2333, etc.) — done; the PR is now a single forward-only addition on top of current Design pattern / consistency with
Naming now matches the other kernels ( First full block must always get seqlen masking — fixed, and it matches Unnecessary CuTe DSL comment — gone in the reworked version. Validation — on DGX Spark (GB10, |
|
Status update: our SM120 hardware is currently allocated to other work, so I am unable to re-validate this for the time being. The branch currently applies cleanly against main, and the gap it fills is still present (block sparsity is asserted off for SM 12.0 on current main). Glad to see it adopted by anyone with SM120-class hardware; for design questions I am also reachable at github.com/bledden. |
This branch has not been deployed
3 participants
Summary
flash_fwd.py), which is shared by SM80 and SM120mask_blocks(partially masked, applymask_modper element) andfull_blocks(fully unmasked, skip masking entirely)FlashAttentionForwardSm120) inherits the SM80 base class and gets block sparsity for free; removes theassert not use_block_sparsityguard ininterface.pymma_one_n_blockcallback pattern as SM90/SM100Design
mma_one_n_block callback pattern (mirrors SM90)
A new
mma_one_n_block_bsmethod handles one KV block: load K, load V, GEMM QK, optional score_mod, mask, online softmax, GEMM PV. This mirrors SM90'smma_one_n_block.A new
run_block_sparse_mainloop_sm80utility inblock_sparse_utils.pytakes the callback and iterates mask blocks (applyingmask_mod) then full blocks (seqlen masking only). This mirrors the non-intra-wg-overlap path ofconsume_block_sparse_loads.SM80 can't use the exact TMA producer/consumer split (no warpgroup specialization), but the
mma_one_n_blockcallback abstraction is the same.First full block seqlen masking
The first full block always receives
mask_seqlen=Trueinrun_block_sparse_mainloop_sm80, since full blocks may sit at a higher n position than any mask block (and thus need seqlen-boundary masking).SM120 arch fix
FlashAttentionForwardSm120.__init__now forcesself.arch = Arch.sm_80. Without this,FlashAttentionForwardBase.__init__setsself.archfrom the real GPU arch (sm_121a on DGX Spark), which causes the SM80 epilogue to incorrectly attempt TMA-O — crashing sincetma_atom_O=Nonein this kernel variant.Dense mainloop unchanged
The original dense mainloop is guarded by
if const_expr(blocksparse_tensors is None):and is completely unmodified.Validation
Validated on SM121a (DGX Spark GB10):
test_block_sparsity.py: 4621 passed, 40 skipped (causal, sliding window, block diagonal; bf16/fp16; D=64/128/256; various seqlens)Contributed by Second Nature Computing (https://joinsecondnature.com)