feat(kernels): Triton FlashAttention LSE export + varlen packing - #233
feat(kernels): Triton FlashAttention LSE export + varlen packing#233Billy1900 wants to merge 4 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.
|
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 (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds optional attention-domain LSE export to dense Triton FlashAttention, introduces packed variable-length attention, adds CUDA/Triton correctness tests and benchmarks, and documents the APIs, tensor contracts, backend scope, validation, and limitations. ChangesTriton attention enhancements
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant TestVarlenAttention
participant triton_flash_attention_varlen
participant TritonAttentionVarlenKernel
participant ref_attn
TestVarlenAttention->>triton_flash_attention_varlen: submit packed q, k, v, and cu_seqlens
triton_flash_attention_varlen->>TritonAttentionVarlenKernel: execute attention with return_lse
TritonAttentionVarlenKernel-->>triton_flash_attention_varlen: return packed output and LSE
TestVarlenAttention->>ref_attn: compute fp32 per-sequence reference
ref_attn-->>TestVarlenAttention: return reference output and LSE
TestVarlenAttention->>TestVarlenAttention: compare outputs, LSE, and gradients
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 functionality to the Triton FlashAttention fallback kernels to (1) optionally export attention-domain LSE and (2) support packed variable-length (cu_seqlens) attention, positioning Triton as the cross-platform semantic baseline for upcoming SM90/SM80/ROCm fused implementations.
Changes:
- Extend dense
triton_flash_attentionwithreturn_lse=Trueto return attention-domain log-sum-exp (non-differentiable). - Add
triton_flash_attention_varlenwith forward/backward kernels for packed[total_tokens, H, D]+cu_seqlensinputs, including causal decode-style masking. - Add dedicated tests and operator documentation for the new LSE export and varlen path.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
rl_engine/kernels/ops/triton/triton_attn.py |
Adds dense LSE export and introduces packed varlen FlashAttention (fwd+bwd) kernels/APIs. |
tests/test_triton_attention_varlen.py |
New CUDA+Triton-gated tests validating dense LSE + varlen fwd/bwd against an independent fp32 reference. |
docs/operators/attention-varlen.md |
New operator page documenting LSE export + packed varlen attention contract and limitations. |
docs/operators/README.md |
Adds the new operator doc page to the operators index. |
docs/.nav.yml |
Wires the new operator doc page into site navigation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/test_triton_attention_varlen.py (1)
54-54: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefix unused unpacked vars.
HandDfromH, Sq, D = q.shapeare unused (Ruff RUF059). Rename to_H/_Dor use_, Sq, _.🤖 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/test_triton_attention_varlen.py` at line 54, Update the shape unpacking near the attention test to mark the unused H and D dimensions as intentionally unused, using underscore-prefixed names or anonymous underscores while retaining Sq for subsequent logic.Source: Linters/SAST tools
rl_engine/kernels/ops/triton/triton_attn.py (1)
792-792: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider fp32 accumulation for
dqatomics.
dqis accumulated viatl.atomic_addin the input dtype (fp16), summing contributions from every K-block. fp16 atomic accumulation compounds rounding error; a fp32dqscratch buffer (cast toq.dtypeat the end) would be more robust. Current tests pass withinatol=0.05, so this is not blocking.🤖 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 `@rl_engine/kernels/ops/triton/triton_attn.py` at line 792, Update the dq accumulation flow around the tl.atomic_add call to use an fp32 scratch buffer for atomic accumulation, then cast the completed result back to q.dtype when writing the final dq output. Ensure all K-block contributions use the fp32 buffer and preserve the existing masking behavior.
🤖 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 `@rl_engine/kernels/ops/triton/triton_attn.py`:
- Around line 903-948: Fix varlen backward stride handling in the wrapper
invoking _bwd_preprocess_varlen and _bwd_kernel_varlen: do is contiguous, but
these kernels currently reuse out/q strides while q, k, and v may be
non-contiguous. Either make q, k, and v contiguous before launching both
kernels, or add and pass dedicated do strides through both launches, preserving
correct delta and gradient results for packed non-contiguous inputs.
---
Nitpick comments:
In `@rl_engine/kernels/ops/triton/triton_attn.py`:
- Line 792: Update the dq accumulation flow around the tl.atomic_add call to use
an fp32 scratch buffer for atomic accumulation, then cast the completed result
back to q.dtype when writing the final dq output. Ensure all K-block
contributions use the fp32 buffer and preserve the existing masking behavior.
In `@tests/test_triton_attention_varlen.py`:
- Line 54: Update the shape unpacking near the attention test to mark the unused
H and D dimensions as intentionally unused, using underscore-prefixed names or
anonymous underscores while retaining Sq for subsequent logic.
🪄 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: b4bd91fa-2f38-473f-ac41-d98a69fcfb0b
📒 Files selected for processing (5)
docs/.nav.ymldocs/operators/README.mddocs/operators/attention-varlen.mdrl_engine/kernels/ops/triton/triton_attn.pytests/test_triton_attention_varlen.py
_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.
|
Thank you for your contribution. Could you add some benchmark tests for the PR content? This will provide some convenience for the reviewers. Thanks! |
…varlen Compares triton_flash_attention (return_lse=True/False) against torch SDPA on dense tensors, and triton_flash_attention_varlen against a pad-mask-unpad SDPA baseline on packed cu_seqlens tensors -- the realistic naive alternative for RL rollout batches with skewed per-sequence lengths. Reports forward and forward+backward latency plus peak VRAM. Requested by a reviewer on RL-Align#233 for easier review of the perf tradeoffs. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Benchmark resultsAdded
Ran on H100 SXM5 ( Dense (causal):
Native SDPA (cuDNN/flash backend) is faster at small/medium shapes and roughly at parity at Packed varlen (causal, H=8, D=64):
This is where packing pays off — up to 20x forward speedup and several-x memory savings, growing with how skewed the batch's length distribution is. The zero-length-sequence case ( Reproduce with: |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
benchmarks/benchmark_triton_attention_varlen.py (1)
156-176: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBind loop-scoped variables as closure defaults (Ruff B023). The benchmark closures already bind
q/k/vvia default args to freeze them per iteration, but leavesm_scale(dense) andseqlens/cu_seqlens/max_seqlen(varlen) captured by reference. This is currently harmless because each closure is invoked within the same loop iteration through_time_ms, but it's inconsistent with the existing binding style, defends against future refactors that store/defer these closures, and clears the Ruff B023 warnings.
benchmarks/benchmark_triton_attention_varlen.py#L156-L176: addsm_scale=sm_scaletotriton_fwd,triton_fwd_lse, andtriton_fwd_bwdsignatures.benchmarks/benchmark_triton_attention_varlen.py#L235-L259: addseqlens=seqlenstopad_fwd/pad_fwd_bwdandcu_seqlens=cu_seqlens, max_seqlen=max_seqlentotriton_fwd/triton_fwd_bwd.🤖 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 `@benchmarks/benchmark_triton_attention_varlen.py` around lines 156 - 176, Add closure-default bindings for all loop-scoped variables: in benchmarks/benchmark_triton_attention_varlen.py lines 156-176, bind sm_scale=sm_scale in triton_fwd, triton_fwd_lse, and triton_fwd_bwd; in lines 235-259, bind seqlens=seqlens in pad_fwd and pad_fwd_bwd, and cu_seqlens=cu_seqlens plus max_seqlen=max_seqlen in triton_fwd and triton_fwd_bwd, preserving the existing q/k/v bindings.Source: Linters/SAST tools
🤖 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 `@benchmarks/benchmark_triton_attention_varlen.py`:
- Around line 156-176: Add closure-default bindings for all loop-scoped
variables: in benchmarks/benchmark_triton_attention_varlen.py lines 156-176,
bind sm_scale=sm_scale in triton_fwd, triton_fwd_lse, and triton_fwd_bwd; in
lines 235-259, bind seqlens=seqlens in pad_fwd and pad_fwd_bwd, and
cu_seqlens=cu_seqlens plus max_seqlen=max_seqlen in triton_fwd and
triton_fwd_bwd, preserving the existing q/k/v bindings.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 365bd06d-2e62-4c8c-a84b-f3b9722b3ef3
📒 Files selected for processing (1)
benchmarks/benchmark_triton_attention_varlen.py
|
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. |
Summary
First step of a stack toward "Fused FlashAttention with causal mask, varlen
packing, and exported attention LSE" (see discussion in the repo). Rather than
attempting SM90 WGMMA+TMA + SM80 + ROCm MFMA + Triton all in one PR, this
lands the Triton piece first: it's the only backend where a real kernel
already existed to extend (
rl_engine/kernels/ops/triton/triton_attn.py), andit becomes the cross-platform semantic baseline the SM90/SM80/ROCm kernels
get checked against later, rather than against each other.
What this adds
return_lse=Trueon the existing densetriton_flash_attention. No new kernel math —M(running max) andL(running sum-exp) were already computed online for the backward pass, just
never returned. LSE is marked non-differentiable (
ctx.mark_non_differentiable),matching
flash_attn's externalsoftmax_lsecontract. Default calls(
return_lse=False) are unchanged — no breaking change to the (currentlyuncalled) public API.
triton_flash_attention_varlen: a new packed variable-length path overcu_seqlens-style[total_tokens, H, D]tensors (same convention asflash_attn_varlen_funcand this repo'spackop, feat(kernels): add fused masking + variable-length pack-and-pad op #182), with forward +backward kernels. Causal masking uses the
Skv - Sqanchor fromNativeAttentionOp(docs/operators/attention.md) so prefill (Sq == Skv)and decode (
Sq < Skv) share one formula.Explicitly out of scope for this PR: SM90 WGMMA+TMA, SM80
mma.sync, andROCm MFMA kernels, and GQA support (the pre-existing dense Triton kernel
didn't support GQA either — not a regression). Not wired into
KernelRegistryyet.Correctness
Both paths are checked against an independent fp32 masked-softmax +
logsumexp reference (not against each other, so a shared bug in the
online-softmax loop can't hide). Run end-to-end on real H100 SXM5 hardware:
outmax-abs-diff ≈1.4e-3,lse≈1e-6vs. reference (fp16inputs, fp32 accumulation).
dq/dk/dv) checked across causal/non-causal,decode-style (
Sq < Skv, varying independently per sequence in the batch),sequence lengths deliberately not aligned to the kernel's 64-token block
size,
head_dim64 and 128, and a zero-length sequence in the batch (a realcase for fully-masked/empty RL responses — packed tensors have the next
sequence's data immediately after a short one, so this is a genuine
correctness boundary, not just an edge case to tolerate).
8/8 tests in
tests/test_triton_attention_varlen.pypass; the pre-existing26/26 in
tests/test_attention.pyare untouched and still pass.Docs
Added
docs/operators/attention-varlen.md(template-following operator page,wired into
docs/.nav.ymlanddocs/operators/README.md), including anexplicit note that this LSE is attention-domain, not the vocab-domain LSE the
fused_logp/linear_logpkernels already use internally — same name,different tensor, worth flagging to avoid confusion.
Test plan
pytest tests/test_triton_attention_varlen.py -v— 8 passed (H100 SXM5)pytest tests/test_attention.py -v— 26 passed, no regressionsruff check,black --check,isort --check-onlyall cleanSummary by CodeRabbit
New Features
Documentation
Tests
Benchmarks