feat(kernels): Triton FlashAttention LSE export + varlen packing#233
Open
Billy1900 wants to merge 2 commits into
Open
feat(kernels): Triton FlashAttention LSE export + varlen packing#233Billy1900 wants to merge 2 commits into
Billy1900 wants to merge 2 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.
Billy1900
requested review from
EthanZero2Hero,
KJLdefeated and
inaniloquentee
as code owners
July 19, 2026 09:53
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds optional attention-domain LSE export to dense Triton FlashAttention, introduces packed variable-length attention, adds CUDA/Triton correctness tests, 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.
Comment on lines
+891
to
+892
| do = do.contiguous() | ||
| dq = torch.zeros_like(q) |
Comment on lines
+826
to
+831
| batch = cu_seqlens_q.numel() - 1 | ||
| assert cu_seqlens_k.numel() - 1 == batch | ||
|
|
||
| cu_seqlens_q = cu_seqlens_q.to(device=q.device, dtype=torch.int32) | ||
| cu_seqlens_k = cu_seqlens_k.to(device=q.device, dtype=torch.int32) | ||
|
|
![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.
Open
5 tasks
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
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