feat(kernels): Triton FlashAttention LSE export + varlen packing#233
feat(kernels): Triton FlashAttention LSE export + varlen packing#233Billy1900 wants to merge 3 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.
📝 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
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
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