feat(kernels): SM90 TMA + mma.sync FlashAttention, forward + backward (causal, varlen, LSE) - #237
Open
Billy1900 wants to merge 6 commits into
Open
feat(kernels): SM90 TMA + mma.sync FlashAttention, forward + backward (causal, varlen, LSE)#237Billy1900 wants to merge 6 commits into
Billy1900 wants to merge 6 commits into
Conversation
Extends the Triton FlashAttention fallback with an attention-domain LSE output (M + log(L), already computed online for the backward pass but previously discarded) and a packed variable-length path operating on cu_seqlens-style [total_tokens, H, D] tensors instead of a padded dense batch. This is the cross-platform semantic baseline the planned SM90 WGMMA+TMA, SM80 mma.sync, and ROCm MFMA fused-attention kernels will be checked against. Both paths are validated against an independent per-sequence fp32 masked-softmax + logsumexp reference (not against the dense kernel itself), covering causal/non-causal, decode-style (Sq < Skv), sequence lengths not aligned to the kernel's block size, and empty sequences in a packed batch -- run end-to-end on H100 hardware.
_bwd_preprocess_varlen indexed DO with Out's strides, and _bwd_kernel_varlen indexed DO with Q's strides. do.contiguous() only guarantees DO's own contiguous layout -- it does not make DO's strides equal to Out's or Q's, which may be non-contiguous (e.g. a transposed view). The mismatch produced wrong gradients (or out-of-bounds reads) whenever Q/Out were non-contiguous. Fixes by plumbing do.stride(*) through both kernels, matching how the pre-existing dense _bwd_preprocess/_bwd_kernel already keep DO's strides separate from Out's. Added a regression test building q/k/v as transposed (non-contiguous) views; confirmed it fails against the prior code (85% of dq elements wrong, max abs diff 2.3) and passes with this fix.
…rlen, LSE) Adds the CUDA SM90 (Hopper) forward-only FlashAttention kernel from the roadmap: causal masking, packed cu_seqlens variable-length input, and attention-domain LSE export, matching the semantics of triton_flash_attention_varlen (the cross-platform baseline). Compute path is TMA global->shared loads + PTX mma.sync.aligned.m16n8k16 -- the same combination every other "SM90" kernel in this repo already uses, not literal wgmma.mma_async. Forward only this milestone (bf16, head_dim=128, contiguous, no GQA); falls back to triton_flash_attention_varlen for anything outside that support surface. Backward, SM80, ROCm, and KernelRegistry wiring are explicit follow-ups tracked in docs/operators/attention-varlen.md. Validated against both an independent fp32 masked-softmax+LSE reference and triton_flash_attention_varlen across uneven/block-aligned seqlens, non-causal, decode-style, zero-length sequences, long multi-tile sequences, single-head and 8-entry batches, and the input-validation error paths (bad head_dim, dtype, GQA mismatch, cu_seqlens batch mismatch, non-contiguous input). Build with KERNEL_ALIGN_FORCE_SM90=1 pip install -e . on a Hopper GPU.
Billy1900
requested review from
EthanZero2Hero,
Flink-ddd,
KJLdefeated,
bitborne and
inaniloquentee
as code owners
July 20, 2026 03:28
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds attention-domain LSE output to Triton dense and varlen attention, introduces forward and backward SM90 packed varlen FlashAttention, wires Python fallback dispatch, and adds tests, build integration, and documentation. ChangesVariable-length FlashAttention
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant flash_attention_sm90_varlen
participant FlashAttentionVarlenSM90Op
participant SM90_CUDA_extension
participant triton_flash_attention_varlen
Caller->>flash_attention_sm90_varlen: packed Q/K/V and cu_seqlens
flash_attention_sm90_varlen->>FlashAttentionVarlenSM90Op: forward request
FlashAttentionVarlenSM90Op->>SM90_CUDA_extension: supported bf16 D=128 request
SM90_CUDA_extension-->>FlashAttentionVarlenSM90Op: out and lse
FlashAttentionVarlenSM90Op->>triton_flash_attention_varlen: unsupported-input fallback
triton_flash_attention_varlen-->>FlashAttentionVarlenSM90Op: out and lse
FlashAttentionVarlenSM90Op-->>Caller: requested output
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Adds a new FlashAttention forward implementation for CUDA SM90 (Hopper) using TMA→shared loads and PTX mma.sync, including causal masking, packed cu_seqlens varlen inputs, and attention-domain LSE export, with Triton as the semantic baseline (and fallback).
Changes:
- Extend Triton FlashAttention to optionally return attention-domain LSE (
return_lse=True) for dense and add a packed varlen FlashAttention path (forward + backward). - Add an SM90 forward-only packed varlen FlashAttention kernel (TMA +
mma.sync) with a Python wrapper that falls back to Triton outside the supported surface. - Add docs + tests covering LSE correctness/non-diff behavior, varlen correctness (including decode-style and zero-length sequences), and SM90-vs-reference comparisons.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
rl_engine/kernels/ops/triton/triton_attn.py |
Adds return_lse support for dense Triton attention and introduces packed varlen Triton attention (fwd/bwd). |
rl_engine/kernels/ops/cuda/attention/flash_attn_sm90.py |
Provides a Python wrapper for the SM90 fused varlen attention op with Triton fallback. |
csrc/cuda/attention/flash_attention_varlen_sm90.cu |
Implements the SM90 TMA + mma.sync forward-only packed varlen attention kernel and C++ entrypoint. |
csrc/ops.cpp |
Exposes the SM90 fused attention op via PyBind. |
setup.py |
Adds the SM90 attention CUDA source to the SM90 extension sources list. |
tests/test_triton_attention_varlen.py |
New tests for Triton dense LSE export and Triton varlen attention correctness (fwd/bwd). |
tests/test_flash_attention_sm90.py |
New tests for SM90 fused varlen attention correctness vs fp32 reference and Triton baseline. |
docs/operators/attention-varlen.md |
New operator doc describing LSE export + varlen packing and backend status. |
docs/operators/README.md / docs/.nav.yml |
Wires the new operator doc into the docs index/navigation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
1000
to
+1014
| if sm_scale is None: | ||
| sm_scale = 1.0 / (q.shape[-1] ** 0.5) | ||
|
|
||
| return _TritonAttention.apply(q, k, v, causal, sm_scale) | ||
| out, lse = _TritonAttentionVarlen.apply( | ||
| q, | ||
| k, | ||
| v, | ||
| cu_seqlens_q, | ||
| cu_seqlens_k, | ||
| max_seqlen_q, | ||
| max_seqlen_k, | ||
| causal, | ||
| sm_scale, | ||
| return_lse, | ||
| ) |
Comment on lines
+284
to
+292
| #pragma unroll | ||
| for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) | ||
| #pragma unroll | ||
| for (int kd = 0; kd < DIM / MMA_K; kd += 2) { | ||
| uint32_t addr = K_smem_thread + buf_off; | ||
| addr += nkv * MMA_N * DIM * sizeof(nv_bfloat16); | ||
| addr += kd * MMA_K * sizeof(nv_bfloat16); | ||
| ldmatrix_x4(K_rmem[nkv][kd], addr); | ||
| } |
Comment on lines
+386
to
+394
| #pragma unroll | ||
| for (int nkv = 0; nkv < BLOCK_KV / MMA_K; nkv++) | ||
| #pragma unroll | ||
| for (int d = 0; d < DIM / MMA_N; d += 2) { | ||
| uint32_t addr = V_smem_thread + buf_off; | ||
| addr += nkv * MMA_K * DIM * sizeof(nv_bfloat16); | ||
| addr += d * MMA_N * sizeof(nv_bfloat16); | ||
| ldmatrix_x4_trans(V_rmem[nkv][d], addr); | ||
| } |
| int64_t max_seqlen_q, | ||
| int64_t max_seqlen_k, | ||
| bool causal, | ||
| double sm_scale) { |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
csrc/cuda/attention/flash_attention_varlen_sm90.cu (1)
509-518: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd error checks for
cudaFuncSetAttributeand the kernel launch.The return value of
cudaFuncSetAttributeis ignored, and there is no launch-error check after the kernel. If the dynamic smem request exceeds the device limit or the launch config is otherwise invalid (e.g., running on non-Hopper hardware where the symbol was compiled in), the failure surfaces later at an unrelated synchronization point, making it hard to diagnose.♻️ Suggested error handling
if (smem > 48 * 1024) { - cudaFuncSetAttribute(flash_attention_varlen_sm90_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + C10_CUDA_CHECK(cudaFuncSetAttribute( + flash_attention_varlen_sm90_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); } dim3 grid(cdiv(static_cast<int>(max_seqlen_q), BLOCK_Q), batch * H, 1); flash_attention_varlen_sm90_kernel<<<grid, TB_SIZE, smem, at::cuda::getCurrentCUDAStream()>>>( q_tmap, k_tmap, v_tmap, cu_seqlens_q.data_ptr<int>(), cu_seqlens_k.data_ptr<int>(), reinterpret_cast<nv_bfloat16 *>(out.data_ptr<at::BFloat16>()), lse.data_ptr<float>(), H, static_cast<float>(sm_scale), causal); + C10_CUDA_KERNEL_LAUNCH_CHECK();🤖 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 `@csrc/cuda/attention/flash_attention_varlen_sm90.cu` around lines 509 - 518, In the launch path around flash_attention_varlen_sm90_kernel, check and propagate the return status from cudaFuncSetAttribute when dynamic shared memory is configured, then check the kernel launch error immediately after the launch using the project’s established CUDA error-handling pattern. Ensure both invalid shared-memory configuration and invalid launch configurations fail at this call site with useful diagnostics.
🤖 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 `@csrc/cuda/attention/flash_attention_varlen_sm90.cu`:
- Around line 416-436: The epilogue’s Lse and Out address calculations overflow
32-bit integers for large packed batches. In
csrc/cuda/attention/flash_attention_varlen_sm90.cu lines 416-436, promote the
row/head offset to int64_t for both Lse indexing and Out pointer arithmetic;
apply the same promotion to the degenerate num_kv_iter == 0 Lse/Out writes at
lines 217-231. Preserve the existing bounds checks and stored values.
In `@docs/operators/attention-varlen.md`:
- Around line 82-86: Update the documentation around FlashAttentionVarlenSM90Op
to clarify that the Python wrapper checks contiguity via _supported and routes
non-contiguous or transposed inputs to triton_flash_attention_varlen; reserve
the RuntimeError statement for direct backend invocation or unsupported paths
that bypass the wrapper. Keep the existing supported-input and fallback behavior
unchanged.
---
Nitpick comments:
In `@csrc/cuda/attention/flash_attention_varlen_sm90.cu`:
- Around line 509-518: In the launch path around
flash_attention_varlen_sm90_kernel, check and propagate the return status from
cudaFuncSetAttribute when dynamic shared memory is configured, then check the
kernel launch error immediately after the launch using the project’s established
CUDA error-handling pattern. Ensure both invalid shared-memory configuration and
invalid launch configurations fail at this call site with useful diagnostics.
🪄 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: 0f101ce7-cf85-4c65-9d3b-3ea9d95c8eed
📒 Files selected for processing (10)
csrc/cuda/attention/flash_attention_varlen_sm90.cucsrc/ops.cppdocs/.nav.ymldocs/operators/README.mddocs/operators/attention-varlen.mdrl_engine/kernels/ops/cuda/attention/flash_attn_sm90.pyrl_engine/kernels/ops/triton/triton_attn.pysetup.pytests/test_flash_attention_sm90.pytests/test_triton_attention_varlen.py
Comment on lines
+416
to
+436
| if (row0 < seqlen_q) | ||
| Lse[(q_start + row0) * H + h] = lse0; | ||
| if (row0 + 8 < seqlen_q) | ||
| Lse[(q_start + row0 + 8) * H + h] = lse1; | ||
| } | ||
|
|
||
| #pragma unroll | ||
| for (int d = 0; d < DIM / MMA_N; d++) { | ||
| const int col = d * MMA_N + (lane_id % 4) * 2; | ||
| float *regs = O_rmem[mi][d]; | ||
| regs[0] *= inv0; | ||
| regs[1] *= inv0; | ||
| regs[2] *= inv1; | ||
| regs[3] *= inv1; | ||
|
|
||
| if (row0 < seqlen_q) | ||
| *reinterpret_cast<nv_bfloat162 *>(Out + ((q_start + row0) * H + h) * DIM + col) = | ||
| __float22bfloat162_rn({regs[0], regs[1]}); | ||
| if (row0 + 8 < seqlen_q) | ||
| *reinterpret_cast<nv_bfloat162 *>(Out + ((q_start + row0 + 8) * H + h) * DIM + col) = | ||
| __float22bfloat162_rn({regs[2], regs[3]}); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
32-bit integer overflow in Out/Lse global address computation. Both write sites compute the element offset as ((q_start + row0) * H + h) * DIM + col entirely in int. For large packed batches this product exceeds INT_MAX (e.g., total_q≈524288, H=32, DIM=128 → offset > 2^31), wrapping to a negative index and causing an out-of-bounds/corrupting store. Promote the offset to 64-bit before pointer arithmetic at both sites.
csrc/cuda/attention/flash_attention_varlen_sm90.cu#L416-L436: cast the row/head index toint64_twhen indexingLse[...]andOut[...]in the epilogue, e.g.Out + ((int64_t)(q_start + row0) * H + h) * DIM + col.csrc/cuda/attention/flash_attention_varlen_sm90.cu#L217-L231: apply the sameint64_tpromotion to the degenerate (num_kv_iter == 0)Lse/Outwrites.
📍 Affects 1 file
csrc/cuda/attention/flash_attention_varlen_sm90.cu#L416-L436(this comment)csrc/cuda/attention/flash_attention_varlen_sm90.cu#L217-L231
🤖 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 `@csrc/cuda/attention/flash_attention_varlen_sm90.cu` around lines 416 - 436,
The epilogue’s Lse and Out address calculations overflow 32-bit integers for
large packed batches. In csrc/cuda/attention/flash_attention_varlen_sm90.cu
lines 416-436, promote the row/head offset to int64_t for both Lse indexing and
Out pointer arithmetic; apply the same promotion to the degenerate num_kv_iter
== 0 Lse/Out writes at lines 217-231. Preserve the existing bounds checks and
stored values.
…lly differentiable Adds the CUDA backward kernel for the SM90 TMA + mma.sync varlen FlashAttention forward from the previous commit: a Delta preprocess kernel plus the main dQ/dK/dV kernel (grid over KV-tiles, looping Q-tiles, recompute-based using the forward's saved lse directly -- algebraically identical to Triton's separate-M/L form, so no forward-kernel changes were needed). dQ is atomic-accumulated directly in bf16 (native SM90 atom.add.noftz.bf16, not an emulated CAS loop -- verified against cuda_bf16.hpp); dK/dV are written once per KV-tile CTA with no atomics. Rewrites the Python wrapper as a real torch.autograd.Function -- previously flash_attention_sm90_varlen called the raw pybind op directly with no autograd wrapping at all, so gradients silently did not flow to q/k/v. Dispatch now requires both the forward and backward native symbols before routing to the hardware path; missing either falls back to triton_flash_attention_varlen entirely (never mixes an SM90 forward with a Triton backward). Two indexing bugs were caught by re-derivation before ever building (mixing up per-warp vs whole-tile loop bounds for the dK/dV transpose-read and epilogue store; a shared-memory transpose scratch indexed by a global instead of local column, correct only for the first KV-tile). A third finding during testing turned out not to be a bug: seqlen_q > seqlen_k under causal=True produces NaN, but triton_flash_attention_varlen produces NaN on the identical out-of-contract input too, confirmed directly. Tests extended to 24 cases: the existing forward shape matrix now also exercises .backward() and checks dQ/dK/dV against an fp32 autograd reference (Triton's own backward can't run on bf16 at all in this Triton version -- tl.atomic_add doesn't support bf16, a compiler limitation, not numerical), plus dedicated single-KV-tile, multi-batch dQ atomic-accumulation stress, return_lse=False, partial requires_grad, non-contiguous do, and backward-specific validation-guard cases.
Billy1900
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/cuda/attention/flash_attention_varlen_sm90.cu`:
- Around line 912-919: The backward attention offsets currently multiply in
32-bit integers and can overflow for large packed batches. In
csrc/cuda/attention/flash_attention_varlen_sm90.cu at lines 912-919, cast
q_start + row_a/row_b to int64_t before computing DQ atomic-add addresses; at
lines 456-457, compute preprocess-kernel base with an int64_t row/head term; and
at lines 936-946, cast k_start + row and k_start + row8 to int64_t before DK/DV
writeout address multiplication.
🪄 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: 8ce5e181-7acc-498f-971d-ade3847f921e
📒 Files selected for processing (5)
csrc/cuda/attention/flash_attention_varlen_sm90.cucsrc/ops.cppdocs/operators/attention-varlen.mdrl_engine/kernels/ops/cuda/attention/flash_attn_sm90.pytests/test_flash_attention_sm90.py
Comment on lines
+912
to
+919
| atomicAdd(&DQ[((q_start + row_a) * H + h) * DIM + col], __float2bfloat16_rn(regs[0])); | ||
| atomicAdd(&DQ[((q_start + row_a) * H + h) * DIM + col + 1], | ||
| __float2bfloat16_rn(regs[1])); | ||
| } | ||
| if (row_b < seqlen_q) { | ||
| atomicAdd(&DQ[((q_start + row_b) * H + h) * DIM + col], __float2bfloat16_rn(regs[2])); | ||
| atomicAdd(&DQ[((q_start + row_b) * H + h) * DIM + col + 1], | ||
| __float2bfloat16_rn(regs[3])); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
32-bit overflow in backward global element offsets (* DIM products). Every offset of the form (row * H + h) * DIM (+ col) is computed in int; for large packed batches total_q * H * DIM (or total_k * H * DIM) exceeds INT_MAX (e.g. total_q≈524288, H=32, DIM=128 → 2^31), wrapping to a negative index and causing out-of-bounds / corrupting reads and stores. Promote the row/head term to int64_t before the * DIM multiply at each site. (Same failure class previously flagged on the forward Out/Lse epilogue.) Note: the Lse/Delta reads at L752-L755 use * H only and stay within int range.
csrc/cuda/attention/flash_attention_varlen_sm90.cu#L912-L919: cast(q_start + row_*)toint64_tin theDQatomic-add addresses.csrc/cuda/attention/flash_attention_varlen_sm90.cu#L456-L457: computebaseas((int64_t)row * H + h) * DIMin the preprocess kernel.csrc/cuda/attention/flash_attention_varlen_sm90.cu#L936-L946: cast(k_start + row)/(k_start + row8)toint64_tin theDK/DVwriteout addresses.
📍 Affects 1 file
csrc/cuda/attention/flash_attention_varlen_sm90.cu#L912-L919(this comment)csrc/cuda/attention/flash_attention_varlen_sm90.cu#L456-L457csrc/cuda/attention/flash_attention_varlen_sm90.cu#L936-L946
🤖 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 `@csrc/cuda/attention/flash_attention_varlen_sm90.cu` around lines 912 - 919,
The backward attention offsets currently multiply in 32-bit integers and can
overflow for large packed batches. In
csrc/cuda/attention/flash_attention_varlen_sm90.cu at lines 912-919, cast
q_start + row_a/row_b to int64_t before computing DQ atomic-add addresses; at
lines 456-457, compute preprocess-kernel base with an int64_t row/head term; and
at lines 936-946, cast k_start + row and k_start + row8 to int64_t before DK/DV
writeout address multiplication.
Collaborator
|
please resolve the CI error first, Thank you. |
… comment CI linting job failed: black wanted to reformat flash_attn_sm90.py and test_flash_attention_sm90.py (multi-arg calls collapsed onto one line), and flake8 flagged a 104-char line in setup.py (E501). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Author
fix it. waiting for further review. |
3 tasks
|
This pull request has been inactive for 30 days. To keep things moving, please update it or let us know if it's still in progress. Otherwise it will be automatically closed in 30 days. |
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
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
cu_seqlensvariable-length input, attention-domain LSE export, anddQ/dK/dV, matching the semantics oftriton_flash_attention_varlen(the cross-platform baseline from feat(kernels): Triton FlashAttention LSE export + varlen packing #233).mma.sync.aligned.m16n8k16-- the same combination every other "SM90" kernel in this repo already uses, not literalwgmma.mma_async.lsedirectly rather than separate M/L state.dQis atomic-accumulated directly in bf16 (native SM90atom.add.noftz.bf16, not an emulated CAS loop -- verified againstcuda_bf16.hpp);dK/dVare written once per KV-tile CTA, no atomics needed there.torch.autograd.Function-- previouslyflash_attention_sm90_varlencalled the raw op directly with no autograd wrapping at all, so gradients silently did not flow toq/k/v. Dispatch requires both the forward and backward native symbols before routing to the hardware path; missing either falls back totriton_flash_attention_varlenentirely (never mixes an SM90 forward with a Triton backward).head_dim=128, contiguous, no GQA. Falls back totriton_flash_attention_varlenfor anything outside that support surface rather than erroring.mma.syncfallback, ROCm MFMA, andKernelRegistrywiring remain explicit, separately tracked follow-ups (see Known Limitations indocs/operators/attention-varlen.md).Depends on #233 (Triton varlen/LSE baseline) -- this PR is stacked on top of that branch, so its diff includes those commits until #233 merges.
Test plan
KERNEL_ALIGN_FORCE_SM90=1 MAX_JOBS=2 pip install -e .builds clean on an H100 (SM90) --MAX_JOBS=2needed after a full-parallelism build hitcicc died due to signal 9on constrained build resources.python -m pytest tests/test_flash_attention_sm90.py -v-- 24/24 passing: the forward shape matrix (uneven/block-aligned/long seqlens, non-causal, decode-style, zero-length sequences, single-head, 8-entry mixed batch) now also exercises.backward()and checksdQ/dK/dV; plus dedicated single-KV-tile, multi-batchdQatomic-accumulation stress,return_lse=False, partialrequires_grad, non-contiguousdo, and forward+backward input-validation error paths (bad head_dim/dtype/GQA/cu_seqlens mismatch/non-contiguous).python -m pytest tests/test_triton_attention_varlen.py -v-- 9/9 passing, no regression to the baseline.triton_flash_attention_varlen; gradients checked against the fp32 reference only (Triton's own backward can't run on bf16 tensors in this Triton version --tl.atomic_adddoesn't support bf16, a compiler limitation confirmed directly, not a numerical one).Summary by CodeRabbit
return_lse) for both dense and varlen paths; LSE is non-differentiable while output gradients still work.