From 3d5364b5c15d252ae39575c82cc681c135255290 Mon Sep 17 00:00:00 2001 From: Billy1900 Date: Sun, 19 Jul 2026 05:52:24 -0400 Subject: [PATCH 1/3] feat(kernels): Triton FlashAttention LSE export + varlen packing 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. --- docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/attention-varlen.md | 138 ++++++ rl_engine/kernels/ops/triton/triton_attn.py | 524 +++++++++++++++++++- tests/test_triton_attention_varlen.py | 194 ++++++++ 5 files changed, 850 insertions(+), 8 deletions(-) create mode 100644 docs/operators/attention-varlen.md create mode 100644 tests/test_triton_attention_varlen.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 8794c731..3b77709b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -16,6 +16,7 @@ nav: - operators/README.md - operators/activation.md - operators/attention.md + - operators/attention-varlen.md - operators/fused-logp.md - operators/linear-logp.md - operators/linear-logp-tp-test.md diff --git a/docs/operators/README.md b/docs/operators/README.md index bf38f368..e06cccf1 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -20,6 +20,7 @@ Every operator page should include: - [SiLU / SwiGLU Activation](activation.md) - [Standard Attention](attention.md) +- [FlashAttention: LSE Export + Variable-Length Packing (Triton)](attention-varlen.md) - [Fused LogP](fused-logp.md) - [Fused Linear LogP](linear-logp.md) - [Fused Linear LogP TP Test Runbook](linear-logp-tp-test.md) diff --git a/docs/operators/attention-varlen.md b/docs/operators/attention-varlen.md new file mode 100644 index 00000000..b8d4cff0 --- /dev/null +++ b/docs/operators/attention-varlen.md @@ -0,0 +1,138 @@ +# FlashAttention: LSE Export + Variable-Length Packing (Triton) + +## Summary + +Extends the existing pure-Triton FlashAttention fallback +(`rl_engine/kernels/ops/triton/triton_attn.py`) with two capabilities needed for +long-context RL rollout/training workloads: + +- **Attention-domain LSE export** — the per-query-row log-sum-exp of the scaled, + masked `QKᵀ` logits, in the same fixed online-softmax reduction order the + kernel already used internally for the backward pass. Useful for backward + recomputation, diagnostics, and rollout/training attention alignment checks. + **This is not the vocab-domain LSE produced by the `fused_logp` / + `linear_logp` kernels** — same name, different tensor domain (per key-length + reduction vs. per-vocab reduction). Do not conflate the two. +- **Variable-length (packed) attention** — operates directly on `cu_seqlens`-packed + `[total_tokens, H, D]` tensors instead of a padded `[B, H, S, D]` batch, so RL + batches with wildly different response lengths (or empty/fully-masked + responses) don't pay for padding in either compute or memory. + +This module is the **cross-platform semantic baseline**: the planned SM90 +WGMMA+TMA, SM80 `mma.sync`, and ROCm MFMA fused-attention kernels are checked +against it (and, transitively, against `NativeAttentionOp`) rather than against +each other. See `docs/operators/attention.md` for the pure-PyTorch WS1 +ground-truth reference this whole family is validated against. + +## Entry Point + +```python +from rl_engine.kernels.ops.triton.triton_attn import ( + triton_flash_attention, + triton_flash_attention_varlen, +) + +# Dense — unchanged default behavior, opt into LSE: +out = triton_flash_attention(q, k, v, causal=True) # [B, H, S, D] +out, lse = triton_flash_attention(q, k, v, causal=True, return_lse=True) # lse: [B, H, S] fp32 + +# Packed variable-length: +out = triton_flash_attention_varlen( + q, k, v, # [total_q, H, D], [total_k, H, D], [total_k, H, D] + cu_seqlens_q, cu_seqlens_k, # int32 [batch + 1], cu_seqlens[0] == 0 + max_seqlen_q, max_seqlen_k, # host ints, size the launch grid + causal=True, +) +out, lse = triton_flash_attention_varlen(..., return_lse=True) # lse: [total_q, H] fp32 +``` + +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| Triton (CUDA) | `triton_flash_attention` / `triton_flash_attention_varlen` | `_fwd_kernel[_varlen]`, `_bwd_kernel[_varlen]` | Cross-platform semantic baseline. | +| CUDA SM90 (WGMMA+TMA) | — | — | Planned; validates against this Triton path. | +| CUDA SM80 (`mma.sync`) | — | — | Planned; validates against this Triton path. | +| ROCm (MFMA) | — | — | Planned; validates against this Triton path. | + +Not yet wired into `KernelRegistry` — this is the standalone kernel-development +stage (mirrors how `prefix_shared_attention` and the SM90 logp kernels were +built before registry integration). Registry dispatch is future work once a +production op_type contract for causal/varlen attention is settled. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `q` (dense) | `[B, H, Sq, D]` | fp16/bf16/fp32 | `D ∈ {16,32,64,128,256}`. | +| `k`, `v` (dense) | `[B, H, Skv, D]` | same as `q` | **No GQA**: `k`/`v` head count must equal `q`'s (matches the pre-existing dense kernel's limitation, not new). | +| `q` (varlen) | `[total_q, H, D]` | fp16/bf16/fp32 | Packed, no padding. | +| `k`, `v` (varlen) | `[total_k, H, D]` | same as `q` | Same no-GQA constraint. | +| `cu_seqlens_q`, `cu_seqlens_k` | `[batch + 1]` | int32 | Cumulative offsets, `cu_seqlens[0] == 0`. Same convention as `flash_attn_varlen_func` and this repo's `pack` op (#182). | +| `max_seqlen_q`, `max_seqlen_k` | — | Python `int` | Host-side; sizes the launch grid. | +| `causal` | — | bool | Per-sequence anchor `Skv - Sq` (see Accuracy) — identical formula for prefill (`Sq == Skv`) and decode (`Sq < Skv`). | +| `return_lse` | — | bool | If `True`, also returns the attention-domain LSE, float32, **non-differentiable** (`ctx.mark_non_differentiable`) — diagnostics/backward-recompute only, matching `flash_attn`'s `softmax_lse` external contract. | +| output | dense: `[B, H, Sq, D]`; varlen: `[total_q, H, D]` | input dtype | — | + +## Dispatch Behavior + +Direct function calls only (see Backends). No registry op_type yet. + +## Accuracy + +Both paths are checked against an **independent** fp32 masked-softmax + LSE +closed form (not against each other, so a shared online-softmax bug can't +hide): + +```python +scores = einsum("hqd,hkd->hqk", q.float(), k.float()) * scale +if causal: + mask = torch.triu(torch.ones(Sq, Skv, bool), diagonal=Skv - Sq + 1) + scores = scores.masked_fill(mask, -inf) +probs = softmax(scores, dim=-1) +out = einsum("hqk,hkd->hqd", probs, v.float()) +lse = torch.logsumexp(scores, dim=-1) +``` + +Measured on an H100 SXM5 (fp16 inputs, fp32 accumulation in-kernel): + +- Dense: `out` max-abs-diff ≈ `1.4e-3`, `lse` max-abs-diff ≈ `1e-6`. +- Varlen: `out`/`dq`/`dk`/`dv` max-abs-diff in the `5e-4`–`1.1e-2` range across + uneven, non-block-aligned sequence lengths; `lse` ≈ `1e-6`. + +Varlen boundary handling: since packed sequences sit back-to-back in memory (no +padding), reading past a sequence's own `cu_seqlens` bound is a genuine +correctness bug (you'd read the *next* sequence's tokens), not just wasted +compute. All loads/stores in the varlen kernels are therefore explicitly masked +against `seqlen_q`/`seqlen_k` (not `boundary_check` on dynamically-shaped block +pointers) — this is deliberate and was chosen specifically to make the +padding-vs-next-sequence distinction unambiguous. + +## Performance Notes + +No standalone benchmark script yet (unlike `fused-logp.md` / `linear-logp.md`). +Follow-up: add `benchmarks/benchmark_attention_varlen.py` measuring +packed-vs-padded latency/VRAM at representative RL group sizes, alongside the +existing `benchmarks/benchmark_attention.py`. + +## Tests + +```bash +python -m pytest tests/test_triton_attention_varlen.py -v +``` + +Covers: dense LSE vs. independent reference, LSE non-differentiability with +backward still populating `q`/`k`/`v` grads, default-call backward +compatibility (bare tensor when `return_lse=False`), varlen forward+backward +across uneven non-block-aligned seqlens, non-causal, decode-style (`Sq < Skv` +varying per sequence), `head_dim=128`, and a zero-length sequence in the batch +(a real occurrence for fully-masked/empty RL responses). + +## Known Limitations + +- No GQA support on either path (`Hk` must equal `Hq`) — same constraint the + pre-existing dense kernel already had; not introduced by this change. +- Not wired into `KernelRegistry`; no automatic backend selection yet. +- SM90 WGMMA+TMA, SM80 `mma.sync`, and ROCm MFMA kernels are not implemented — + this page covers the Triton baseline only. +- No standalone performance benchmark yet (see Performance Notes). diff --git a/rl_engine/kernels/ops/triton/triton_attn.py b/rl_engine/kernels/ops/triton/triton_attn.py index c1bd4809..398c8e34 100644 --- a/rl_engine/kernels/ops/triton/triton_attn.py +++ b/rl_engine/kernels/ops/triton/triton_attn.py @@ -333,7 +333,7 @@ def _bwd_kernel( class _TritonAttention(torch.autograd.Function): @staticmethod - def forward(ctx, q, k, v, causal, sm_scale): + def forward(ctx, q, k, v, causal, sm_scale, return_lse): # [batch, num_heads, seq_len, head_dim] # Triton tutorial standard layout requires specific contiguity Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -390,10 +390,20 @@ def forward(ctx, q, k, v, causal, sm_scale): ) ctx.save_for_backward(q, k, v, out, L, M) - return out + + if not return_lse: + return out, None + + # Attention-domain LSE: log-sum-exp of the (scaled, masked) QK^T logits per + # query row, in the same fixed reduction order the fwd kernel already used to + # accumulate M (running max) and L (running sum-exp) online. This is distinct + # from the vocab-domain LSE produced by the logp/linear_logp kernels. + lse = M + torch.log(L) + ctx.mark_non_differentiable(lse) + return out, lse @staticmethod - def backward(ctx, do): + def backward(ctx, do, _dlse): q, k, v, out, L, M = ctx.saved_tensors do = do.contiguous() @@ -468,7 +478,7 @@ def backward(ctx, do): num_stages=1, ) - return dq, dk, dv, None, None + return dq, dk, dv, None, None, None def triton_flash_attention( @@ -477,6 +487,7 @@ def triton_flash_attention( v: torch.Tensor, causal: bool = True, sm_scale: float | None = None, + return_lse: bool = False, ): """ Universal backup Triton FlashAttention (support Forward / Backward) @@ -484,12 +495,509 @@ def triton_flash_attention( Args: q, k, v: Tensors of shape [batch, num_heads, seq_len, head_dim]. It is recommended to switch to contiguous memory (.contiguous()). - causal: Whether to turn on causal masking. - sm_scale: Softmax scaling factor, default to 1.0 / sqrt(head_dim). + causal: Whether to turn on causal masking. + sm_scale: Softmax scaling factor, default to 1.0 / sqrt(head_dim). + return_lse: If True, also return the per-query-row attention-domain LSE + (log-sum-exp of the scaled, masked QK^T logits), shape + [batch, num_heads, seq_len], float32. Not a vocab-domain logprob LSE. + The returned LSE is non-differentiable (diagnostics / backward-recompute + use only), matching the external contract of `flash_attn`'s `softmax_lse`. + Returns: + `out` of shape [batch, num_heads, seq_len, head_dim] if `return_lse` is False, + else `(out, lse)`. + """ + if sm_scale is None: + sm_scale = 1.0 / (q.shape[-1] ** 0.5) + + out, lse = _TritonAttention.apply(q, k, v, causal, sm_scale, return_lse) + return (out, lse) if return_lse else out + + +# --------------------------------------------------------------------------- +# Variable-length (packed) attention. +# +# Q/K/V are packed along the token dimension: [total_tokens, H, D], with no +# per-sequence padding. `cu_seqlens_{q,k}` are int32 [batch + 1] cumulative +# sequence-length offsets (cu_seqlens[0] == 0), the same convention as +# `flash_attn_varlen_func` and this repo's `pack` op (#182). Each program +# handles one (query-block, batch, head) triple; `causal_offset = seqlen_k - +# seqlen_q` reproduces the `Skv - Sq` causal anchor from the WS1 +# `NativeAttentionOp` reference (docs/operators/attention.md) so prefill +# (Sq == Skv) and decode (Sq < Skv) share one formula. Boundary handling is +# via explicit masks (not `boundary_check`) because a block that runs past a +# sequence's length would otherwise read into the *next* packed sequence. +# --------------------------------------------------------------------------- + + +@triton.jit +def _fwd_kernel_varlen( + Q, + K, + V, + sm_scale, + cu_seqlens_q, + cu_seqlens_k, + L, + M, + Out, + stride_qm, + stride_qh, + stride_qk, + stride_kn, + stride_kh, + stride_kk, + stride_vn, + stride_vh, + stride_vk, + stride_om, + stride_oh, + stride_ok, + H, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, +): + start_m = tl.program_id(0) + pid_bh = tl.program_id(1) + b = pid_bh // H + h = pid_bh % H + + q_start = tl.load(cu_seqlens_q + b) + q_end = tl.load(cu_seqlens_q + b + 1) + seqlen_q = q_end - q_start + if start_m * BLOCK_M >= seqlen_q: + return + + k_start = tl.load(cu_seqlens_k + b) + k_end = tl.load(cu_seqlens_k + b + 1) + seqlen_k = k_end - k_start + causal_offset = seqlen_k - seqlen_q + + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + valid_m = offs_m < seqlen_q + + q_ptrs = ( + Q + (q_start + offs_m[:, None]) * stride_qm + h * stride_qh + offs_d[None, :] * stride_qk + ) + q = tl.load(q_ptrs, mask=valid_m[:, None], other=0.0) + + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + if IS_CAUSAL: + hi = tl.minimum(seqlen_k, (start_m + 1) * BLOCK_M + causal_offset) + else: + hi = seqlen_k + + for start_n in range(0, hi, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + valid_n = offs_n < seqlen_k + + k_ptrs = ( + K + + (k_start + offs_n[:, None]) * stride_kn + + h * stride_kh + + offs_d[None, :] * stride_kk + ) + k = tl.load(k_ptrs, mask=valid_n[:, None], other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + + mask = valid_n[None, :] + if IS_CAUSAL: + mask = mask & (offs_m[:, None] >= (offs_n[None, :] - causal_offset)) + qk = tl.where(mask, qk, float("-inf")) + + m_ij = tl.max(qk, 1) + m_i_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_i_new) + beta = tl.exp(qk - m_i_new[:, None]) + l_i_new = alpha * l_i + tl.sum(beta, 1) + + p_scale = beta / l_i_new[:, None] + acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + + v_ptrs = ( + V + + (k_start + offs_n[:, None]) * stride_vn + + h * stride_vh + + offs_d[None, :] * stride_vk + ) + v = tl.load(v_ptrs, mask=valid_n[:, None], other=0.0) + p = p_scale.to(v.dtype) + acc += tl.dot(p, v) + + l_i = l_i_new + m_i = m_i_new + + acc = acc.to(Out.dtype.element_ty) + o_ptrs = ( + Out + (q_start + offs_m[:, None]) * stride_om + h * stride_oh + offs_d[None, :] * stride_ok + ) + tl.store(o_ptrs, acc, mask=valid_m[:, None]) + + l_ptrs = L + (q_start + offs_m) * H + h + m_ptrs = M + (q_start + offs_m) * H + h + tl.store(l_ptrs, l_i, mask=valid_m) + tl.store(m_ptrs, m_i, mask=valid_m) + + +@triton.jit +def _bwd_preprocess_varlen( + Out, + DO, + Delta, + stride_om, + stride_oh, + stride_ok, + total_q, + H, + BLOCK_M: tl.constexpr, + D_HEAD: tl.constexpr, +): + pid_m = tl.program_id(0) + h = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, D_HEAD) + valid_m = offs_m < total_q + + o_ptrs = Out + offs_m[:, None] * stride_om + h * stride_oh + offs_d[None, :] * stride_ok + do_ptrs = DO + offs_m[:, None] * stride_om + h * stride_oh + offs_d[None, :] * stride_ok + + o = tl.load(o_ptrs, mask=valid_m[:, None], other=0.0) + do = tl.load(do_ptrs, mask=valid_m[:, None], other=0.0).to(o.dtype) + + delta = tl.sum(o * do, axis=1) + tl.store(Delta + offs_m * H + h, delta, mask=valid_m) + + +@triton.jit +def _bwd_kernel_varlen( + Q, + K, + V, + sm_scale, + DO, + DQ, + DK, + DV, + L, + M, + Delta, + cu_seqlens_q, + cu_seqlens_k, + stride_qm, + stride_qh, + stride_qk, + stride_kn, + stride_kh, + stride_kk, + stride_vn, + stride_vh, + stride_vk, + H, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, +): + start_n = tl.program_id(0) + pid_bh = tl.program_id(1) + b = pid_bh // H + h = pid_bh % H + + k_start = tl.load(cu_seqlens_k + b) + k_end = tl.load(cu_seqlens_k + b + 1) + seqlen_k = k_end - k_start + if start_n * BLOCK_N >= seqlen_k: + return + + q_start = tl.load(cu_seqlens_q + b) + q_end = tl.load(cu_seqlens_q + b + 1) + seqlen_q = q_end - q_start + causal_offset = seqlen_k - seqlen_q + + offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + valid_n = offs_n < seqlen_k + + k_ptrs = ( + K + (k_start + offs_n[:, None]) * stride_kn + h * stride_kh + offs_d[None, :] * stride_kk + ) + v_ptrs = ( + V + (k_start + offs_n[:, None]) * stride_vn + h * stride_vh + offs_d[None, :] * stride_vk + ) + k = tl.load(k_ptrs, mask=valid_n[:, None], other=0.0) + v = tl.load(v_ptrs, mask=valid_n[:, None], other=0.0) + + dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + + # Causal: query rows before `offs_n - causal_offset` never attend to this K/V + # block, so the M-loop can start there instead of at row 0. + lo = tl.maximum(0, start_n * BLOCK_N - causal_offset) if IS_CAUSAL else 0 + lo = (lo // BLOCK_M) * BLOCK_M + + for start_m in range(lo, seqlen_q, BLOCK_M): + offs_m = start_m + tl.arange(0, BLOCK_M) + valid_m = offs_m < seqlen_q + + q_ptrs = ( + Q + + (q_start + offs_m[:, None]) * stride_qm + + h * stride_qh + + offs_d[None, :] * stride_qk + ) + do_ptrs = ( + DO + + (q_start + offs_m[:, None]) * stride_qm + + h * stride_qh + + offs_d[None, :] * stride_qk + ) + q = tl.load(q_ptrs, mask=valid_m[:, None], other=0.0) + do = tl.load(do_ptrs, mask=valid_m[:, None], other=0.0) + + m_i = tl.load(M + (q_start + offs_m) * H + h, mask=valid_m, other=0.0) + l_i = tl.load(L + (q_start + offs_m) * H + h, mask=valid_m, other=1.0) + delta = tl.load(Delta + (q_start + offs_m) * H + h, mask=valid_m, other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + + mask = valid_n[None, :] & valid_m[:, None] + if IS_CAUSAL: + mask = mask & (offs_m[:, None] >= (offs_n[None, :] - causal_offset)) + qk = tl.where(mask, qk, float("-inf")) + + p = tl.exp(qk - m_i[:, None]) / l_i[:, None] + p = tl.where(mask, p, 0.0) + + dv += tl.dot(tl.trans(p.to(do.dtype)), do) + + dp = tl.dot(do, tl.trans(v)) + ds = p * (dp - delta[:, None]) * sm_scale + ds = tl.where(mask, ds, 0.0) + + dq_val = tl.dot(ds.to(q.dtype), k) + dq_ptrs = ( + DQ + + (q_start + offs_m[:, None]) * stride_qm + + h * stride_qh + + offs_d[None, :] * stride_qk + ) + tl.atomic_add(dq_ptrs, dq_val.to(q.dtype), mask=valid_m[:, None]) + + dk += tl.dot(tl.trans(ds.to(q.dtype)), q) + + dk_ptrs = ( + DK + (k_start + offs_n[:, None]) * stride_kn + h * stride_kh + offs_d[None, :] * stride_kk + ) + dv_ptrs = ( + DV + (k_start + offs_n[:, None]) * stride_vn + h * stride_vh + offs_d[None, :] * stride_vk + ) + tl.store(dk_ptrs, dk.to(k.dtype), mask=valid_n[:, None]) + tl.store(dv_ptrs, dv.to(v.dtype), mask=valid_n[:, None]) + + +class _TritonAttentionVarlen(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + sm_scale, + return_lse, + ): + total_q, H, head_dim = q.shape + assert k.shape[1] == H and v.shape[1] == H, "GQA is not supported by the varlen path yet" + assert head_dim in {16, 32, 64, 128, 256} + + 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) + + out = torch.empty_like(q) + # Packed [total_q, H] layout (not [batch, H, seq_len]): total_q varies per + # batch, so there is no fixed per-sequence stride to lay these out densely. + M = torch.empty((total_q, H), device=q.device, dtype=torch.float32) + L = torch.empty((total_q, H), device=q.device, dtype=torch.float32) + + BLOCK_M = 64 + BLOCK_N = 64 if head_dim > 64 else 128 + + grid = (triton.cdiv(max_seqlen_q, BLOCK_M), batch * H) + _fwd_kernel_varlen[grid]( + q, + k, + v, + sm_scale, + cu_seqlens_q, + cu_seqlens_k, + L, + M, + out, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + H, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=head_dim, + IS_CAUSAL=causal, + num_warps=4, + num_stages=2, + ) + + ctx.sm_scale = sm_scale + ctx.causal = causal + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.total_q = total_q + ctx.save_for_backward(q, k, v, out, L, M, cu_seqlens_q, cu_seqlens_k) + + if not return_lse: + return out, None + + lse = M + torch.log(L) + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + def backward(ctx, do, _dlse): + q, k, v, out, L, M, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + + do = do.contiguous() + dq = torch.zeros_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + + total_q, H, head_dim = q.shape + delta = torch.empty_like(L) + + BLOCK_M = 64 + BLOCK_N = 64 + + grid_prep = (triton.cdiv(total_q, BLOCK_M), H) + _bwd_preprocess_varlen[grid_prep]( + out, + do, + delta, + out.stride(0), + out.stride(1), + out.stride(2), + total_q, + H, + BLOCK_M=BLOCK_M, + D_HEAD=head_dim, + ) + + batch = cu_seqlens_q.numel() - 1 + grid_bwd = (triton.cdiv(ctx.max_seqlen_k, BLOCK_N), batch * H) + _bwd_kernel_varlen[grid_bwd]( + q, + k, + v, + ctx.sm_scale, + do, + dq, + dk, + dv, + L, + M, + delta, + cu_seqlens_q, + cu_seqlens_k, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + H, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=head_dim, + IS_CAUSAL=ctx.causal, + num_warps=4, + num_stages=1, + ) + + return dq, dk, dv, None, None, None, None, None, None, None + + +def triton_flash_attention_varlen( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool = True, + sm_scale: float | None = None, + return_lse: bool = False, +): + """ + Packed variable-length FlashAttention (Triton), the cross-platform semantic + baseline for RL rollout/training batches where sequences are concatenated + rather than padded. + + Args: + q: [total_q, H, D] packed queries. + k, v: [total_k, H, D] packed keys/values. GQA (Hk != Hq) is not supported + by this path yet (matches the existing dense Triton kernel's limitation). + cu_seqlens_q, cu_seqlens_k: int32 [batch + 1] cumulative sequence-length + offsets, cu_seqlens[0] == 0 (the `flash_attn_varlen_func` convention; + also what `pack`, #182, produces). + max_seqlen_q, max_seqlen_k: max per-sequence length in the batch (host + ints), used to size the launch grid. + causal: causal masking, anchored per-sequence via `Skv - Sq` (same + convention as `NativeAttentionOp` in docs/operators/attention.md, so + Sq == Skv is prefill and Sq < Skv is decode-with-cache). + sm_scale: defaults to `1/sqrt(D)`. + return_lse: if True, also return the packed `[total_q, H]` float32 + attention-domain LSE (non-differentiable; see `triton_flash_attention`). Returns: - Attention Output tensor of shape [batch, num_heads, seq_len, head_dim] + `out` of shape [total_q, H, D] if `return_lse` is False, else `(out, lse)`. """ 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, + ) + return (out, lse) if return_lse else out diff --git a/tests/test_triton_attention_varlen.py b/tests/test_triton_attention_varlen.py new file mode 100644 index 00000000..787cff9c --- /dev/null +++ b/tests/test_triton_attention_varlen.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the Triton FlashAttention LSE export and packed varlen path. + +Extends `rl_engine.kernels.ops.triton.triton_attn` (the "cross-platform semantic +baseline" the SM90/SM80/ROCm fused-attention kernels will validate against) with: + + * `return_lse=True` on the existing dense kernel: exposes the attention-domain + LSE (log-sum-exp of the scaled, masked QK^T logits) already accumulated + online as `M`/`L` for the backward pass, but previously discarded. + * `triton_flash_attention_varlen`: packed `[total_tokens, H, D]` + `cu_seqlens` + attention, for RL rollout/training batches that are concatenated rather than + padded (same convention as `flash_attn_varlen_func` and this repo's `pack` + op, #182). + +Both are checked against an independent per-sequence fp32 reference (`_ref_attn`, +a masked-softmax + logsumexp closed form matching the causal-offset convention +`Skv - Sq` documented for `NativeAttentionOp` in docs/operators/attention.md), +not against the dense Triton kernel itself, so a shared bug in the online-softmax +loop cannot hide from these tests. + +The exported LSE is attention-domain (per query row, over the key dimension), +not the vocab-domain LSE produced by the logp/linear_logp kernels. +""" + +import math + +import pytest +import torch + +from rl_engine.kernels.ops.triton.triton_attn import ( + triton_flash_attention, + triton_flash_attention_varlen, +) + +try: + import triton # noqa: F401 + + _HAS_TRITON = True +except ImportError: # pragma: no cover + _HAS_TRITON = False + +requires_triton_cuda = pytest.mark.skipif( + not (_HAS_TRITON and torch.cuda.is_available()), + reason="Triton attention requires a CUDA device and Triton.", +) + +_ATOL_OUT = 0.05 # fp16 accumulation tolerance, dense and varlen paths alike +_ATOL_LSE = 1e-2 + + +def _ref_attn(q, k, v, causal, sm_scale): + """[H, Sq, D] / [H, Skv, D] fp32 masked-softmax reference with LSE.""" + H, Sq, D = q.shape + Skv = k.shape[1] + scores = torch.einsum("hqd,hkd->hqk", q, k) * sm_scale + if causal: + mask = torch.triu( + torch.ones(Sq, Skv, dtype=torch.bool, device=q.device), diagonal=Skv - Sq + 1 + ) + scores = scores.masked_fill(mask, float("-inf")) + probs = torch.softmax(scores, dim=-1) + out = torch.einsum("hqk,hkd->hqd", probs, v) + lse = torch.logsumexp(scores, dim=-1) + return out, lse + + +def _run_varlen_case(seqlens_q, seqlens_k, heads, head_dim, causal, dtype=torch.float16): + device = "cuda" + batch = len(seqlens_q) + total_q = sum(seqlens_q) + total_k = sum(seqlens_k) + cu_q = torch.tensor( + [0, *torch.tensor(seqlens_q).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + cu_k = torch.tensor( + [0, *torch.tensor(seqlens_k).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + + q = torch.randn(total_q, heads, head_dim, device=device, dtype=dtype, requires_grad=True) + k = torch.randn(total_k, heads, head_dim, device=device, dtype=dtype, requires_grad=True) + v = torch.randn(total_k, heads, head_dim, device=device, dtype=dtype, requires_grad=True) + + sm_scale = 1.0 / math.sqrt(head_dim) + out, lse = triton_flash_attention_varlen( + q, + k, + v, + cu_q, + cu_k, + max(seqlens_q), + max(seqlens_k), + causal=causal, + sm_scale=sm_scale, + return_lse=True, + ) + do = torch.randn_like(out) + out.backward(do) + dq, dk, dv = q.grad.clone(), k.grad.clone(), v.grad.clone() + + q_ref = q.detach().clone().float().requires_grad_() + k_ref = k.detach().clone().float().requires_grad_() + v_ref = v.detach().clone().float().requires_grad_() + out_ref = torch.empty(total_q, heads, head_dim, device=device, dtype=torch.float32) + lse_ref = torch.empty(total_q, heads, device=device, dtype=torch.float32) + + qs = ks = 0 + for b in range(batch): + sq, sk = seqlens_q[b], seqlens_k[b] + o, lval = _ref_attn( + q_ref[qs : qs + sq].transpose(0, 1), + k_ref[ks : ks + sk].transpose(0, 1), + v_ref[ks : ks + sk].transpose(0, 1), + causal, + sm_scale, + ) + out_ref[qs : qs + sq] = o.transpose(0, 1) + lse_ref[qs : qs + sq] = lval.transpose(0, 1) + qs += sq + ks += sk + out_ref.backward(do.float()) + + torch.testing.assert_close(out.float(), out_ref, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(lse, lse_ref, atol=_ATOL_LSE, rtol=0.0) + torch.testing.assert_close(dq.float(), q_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dk.float(), k_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dv.float(), v_ref.grad, atol=_ATOL_OUT, rtol=0.0) + + +@requires_triton_cuda +class TestDenseLSEExport: + def test_lse_matches_reference(self): + device = "cuda" + batch, heads, seq, head_dim = 2, 4, 256, 64 + sm_scale = 1.0 / math.sqrt(head_dim) + q = torch.randn(batch, heads, seq, head_dim, device=device, dtype=torch.float16) + k = torch.randn(batch, heads, seq, head_dim, device=device, dtype=torch.float16) + v = torch.randn(batch, heads, seq, head_dim, device=device, dtype=torch.float16) + + out, lse = triton_flash_attention(q, k, v, causal=True, sm_scale=sm_scale, return_lse=True) + assert lse.shape == (batch, heads, seq) + assert lse.dtype == torch.float32 + + ref_out = torch.empty_like(q, dtype=torch.float32) + ref_lse = torch.empty(batch, heads, seq, device=device) + for b in range(batch): + o, lval = _ref_attn(q[b].float(), k[b].float(), v[b].float(), True, sm_scale) + ref_out[b], ref_lse[b] = o, lval + + torch.testing.assert_close(out.float(), ref_out, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(lse, ref_lse, atol=_ATOL_LSE, rtol=0.0) + + def test_lse_is_non_differentiable_and_backward_still_works(self): + device = "cuda" + q = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16, requires_grad=True) + k = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16, requires_grad=True) + v = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16, requires_grad=True) + + out, lse = triton_flash_attention(q, k, v, causal=True, return_lse=True) + assert not lse.requires_grad + + out.float().sum().backward() + assert q.grad is not None and k.grad is not None and v.grad is not None + + def test_default_call_is_backward_compatible(self): + device = "cuda" + q = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16) + k = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16) + v = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16) + + out = triton_flash_attention(q, k, v, causal=True) + assert isinstance(out, torch.Tensor) + + +@requires_triton_cuda +class TestVarlenAttention: + def test_prefill_uneven_seqlens_not_block_aligned(self): + # BLOCK_M/BLOCK_N default to 64/128; deliberately not multiples of either. + _run_varlen_case([37, 128, 200, 5], [37, 128, 200, 5], heads=4, head_dim=64, causal=True) + + def test_non_causal(self): + _run_varlen_case([37, 130, 61], [37, 130, 61], heads=2, head_dim=64, causal=False) + + def test_decode_style_sq_less_than_skv(self): + # Small new-query chunk (e.g. rollout decode step) against a longer KV cache, + # varying independently per sequence in the batch. + _run_varlen_case([1, 3, 1], [50, 91, 17], heads=4, head_dim=64, causal=True) + + def test_head_dim_128(self): + _run_varlen_case([100, 260], [100, 260], heads=2, head_dim=128, causal=True) + + def test_zero_length_sequence_in_batch(self): + # A fully-masked / empty response is a real occurrence in packed RL batches. + _run_varlen_case([0, 128, 0, 37], [50, 128, 0, 37], heads=2, head_dim=64, causal=True) From 8276596fff27bf945c6fa8bbb613b13870ea33f4 Mon Sep 17 00:00:00 2001 From: Billy1900 Date: Sun, 19 Jul 2026 06:04:46 -0400 Subject: [PATCH 2/3] fix(kernels): use DO's own strides in varlen backward, not Q's/Out's _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. --- rl_engine/kernels/ops/triton/triton_attn.py | 20 +++++-- tests/test_triton_attention_varlen.py | 63 +++++++++++++++++++++ 2 files changed, 79 insertions(+), 4 deletions(-) diff --git a/rl_engine/kernels/ops/triton/triton_attn.py b/rl_engine/kernels/ops/triton/triton_attn.py index 398c8e34..69430afa 100644 --- a/rl_engine/kernels/ops/triton/triton_attn.py +++ b/rl_engine/kernels/ops/triton/triton_attn.py @@ -654,6 +654,9 @@ def _bwd_preprocess_varlen( stride_om, stride_oh, stride_ok, + stride_dom, + stride_doh, + stride_dok, total_q, H, BLOCK_M: tl.constexpr, @@ -667,7 +670,7 @@ def _bwd_preprocess_varlen( valid_m = offs_m < total_q o_ptrs = Out + offs_m[:, None] * stride_om + h * stride_oh + offs_d[None, :] * stride_ok - do_ptrs = DO + offs_m[:, None] * stride_om + h * stride_oh + offs_d[None, :] * stride_ok + do_ptrs = DO + offs_m[:, None] * stride_dom + h * stride_doh + offs_d[None, :] * stride_dok o = tl.load(o_ptrs, mask=valid_m[:, None], other=0.0) do = tl.load(do_ptrs, mask=valid_m[:, None], other=0.0).to(o.dtype) @@ -700,6 +703,9 @@ def _bwd_kernel_varlen( stride_vn, stride_vh, stride_vk, + stride_dom, + stride_doh, + stride_dok, H, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, @@ -755,9 +761,9 @@ def _bwd_kernel_varlen( ) do_ptrs = ( DO - + (q_start + offs_m[:, None]) * stride_qm - + h * stride_qh - + offs_d[None, :] * stride_qk + + (q_start + offs_m[:, None]) * stride_dom + + h * stride_doh + + offs_d[None, :] * stride_dok ) q = tl.load(q_ptrs, mask=valid_m[:, None], other=0.0) do = tl.load(do_ptrs, mask=valid_m[:, None], other=0.0) @@ -907,6 +913,9 @@ def backward(ctx, do, _dlse): out.stride(0), out.stride(1), out.stride(2), + do.stride(0), + do.stride(1), + do.stride(2), total_q, H, BLOCK_M=BLOCK_M, @@ -938,6 +947,9 @@ def backward(ctx, do, _dlse): v.stride(0), v.stride(1), v.stride(2), + do.stride(0), + do.stride(1), + do.stride(2), H, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, diff --git a/tests/test_triton_attention_varlen.py b/tests/test_triton_attention_varlen.py index 787cff9c..45c7d165 100644 --- a/tests/test_triton_attention_varlen.py +++ b/tests/test_triton_attention_varlen.py @@ -192,3 +192,66 @@ def test_head_dim_128(self): def test_zero_length_sequence_in_batch(self): # A fully-masked / empty response is a real occurrence in packed RL batches. _run_varlen_case([0, 128, 0, 37], [50, 128, 0, 37], heads=2, head_dim=64, causal=True) + + def test_non_contiguous_q_k_v_backward(self): + # Regression test (review finding): `do` was forced `.contiguous()` in + # backward, but the kernels indexed DO using Q's/Out's strides instead of + # DO's own -- silently correct only when Q/Out/DO all happen to share the + # same contiguous layout, as every other case in this file does by + # construction (fresh `torch.randn(...)`). + # + # Building q/k/v as `[H, total, D].transpose(0, 1)` gives a non-contiguous + # but non-overlapping-and-dense view. `torch.empty_like` (used for `out` + # and `dq`) preserves that exact stride pattern rather than falling back to + # a contiguous layout, so `out`'s strides differ from `do`'s (`do` is a + # plain contiguous tensor here) -- exactly the mismatch the review flagged. + device = "cuda" + heads, head_dim = 2, 64 + seqlens = [37, 61] + total = sum(seqlens) + cu = torch.tensor( + [0, *torch.tensor(seqlens).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + sm_scale = 1.0 / math.sqrt(head_dim) + + def make(seed): + gen = torch.Generator(device=device).manual_seed(seed) + base = torch.randn( + heads, total, head_dim, device=device, dtype=torch.float16, generator=gen + ) + return base.transpose(0, 1).detach().requires_grad_() + + q, k, v = make(1), make(2), make(3) + assert not (q.is_contiguous() or k.is_contiguous() or v.is_contiguous()) + + out = triton_flash_attention_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale + ) + assert not out.is_contiguous() # confirms empty_like preserved q's layout + + do = torch.randn(total, heads, head_dim, device=device, dtype=out.dtype) # plain contiguous + assert do.stride() != out.stride() + out.backward(do) + dq, dk, dv = q.grad.clone(), k.grad.clone(), v.grad.clone() + + q_ref = q.detach().clone().float().requires_grad_() + k_ref = k.detach().clone().float().requires_grad_() + v_ref = v.detach().clone().float().requires_grad_() + out_ref = torch.empty(total, heads, head_dim, device=device, dtype=torch.float32) + qs = 0 + for sq in seqlens: + o, _ = _ref_attn( + q_ref[qs : qs + sq].transpose(0, 1), + k_ref[qs : qs + sq].transpose(0, 1), + v_ref[qs : qs + sq].transpose(0, 1), + True, + sm_scale, + ) + out_ref[qs : qs + sq] = o.transpose(0, 1) + qs += sq + out_ref.backward(do.float()) + + torch.testing.assert_close(out.float(), out_ref, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dq.float(), q_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dk.float(), k_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dv.float(), v_ref.grad, atol=_ATOL_OUT, rtol=0.0) From 466be55251a6f1aec63f887ae222ebb578607daa Mon Sep 17 00:00:00 2001 From: Billy1900 Date: Wed, 22 Jul 2026 00:40:04 -0400 Subject: [PATCH 3/3] bench(kernels): add benchmark for Triton FlashAttention LSE export + 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 #233 for easier review of the perf tradeoffs. Co-Authored-By: Claude Sonnet 5 --- .../benchmark_triton_attention_varlen.py | 348 ++++++++++++++++++ 1 file changed, 348 insertions(+) create mode 100644 benchmarks/benchmark_triton_attention_varlen.py diff --git a/benchmarks/benchmark_triton_attention_varlen.py b/benchmarks/benchmark_triton_attention_varlen.py new file mode 100644 index 00000000..851cd884 --- /dev/null +++ b/benchmarks/benchmark_triton_attention_varlen.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark the Triton FlashAttention LSE export + packed varlen path. + +Covers the two additions from `rl_engine.kernels.ops.triton.triton_attn` +(docs/operators/attention-varlen.md): + + * `triton_flash_attention(..., return_lse=True)` vs. `torch.nn.functional. + scaled_dot_product_attention` on dense `[B, H, S, D]` tensors -- reports + the LSE-export overhead (return_lse=True vs. False) alongside the native + baseline. + * `triton_flash_attention_varlen` on packed `[total_tokens, H, D]` tensors + vs. a pad-mask-unpad SDPA baseline -- the realistic alternative for RL + rollout/training batches with skewed per-sequence response lengths. This + is where packing's compute/memory savings actually show up, since the + padded baseline pays for the longest sequence in the batch on every row. + +Reports forward and forward+backward latency plus peak extra VRAM for the +forward pass. + +Usage: + python benchmarks/benchmark_triton_attention_varlen.py + python benchmarks/benchmark_triton_attention_varlen.py --dense-only + python benchmarks/benchmark_triton_attention_varlen.py --varlen-only +""" + +import argparse + +import torch +from tabulate import tabulate + +from rl_engine.kernels.ops.triton.triton_attn import ( + triton_flash_attention, + triton_flash_attention_varlen, +) +from rl_engine.platforms.device import device_ctx +from rl_engine.utils.logger import logger + +# (batch, heads, seq_len, head_dim) +DENSE_CONFIGS = [ + (2, 8, 512, 64), + (2, 8, 2048, 64), + (1, 16, 4096, 128), + (1, 16, 8192, 128), +] + +# Per-sequence lengths (query == key/value, self-attention prefill), one +# entry per batch config. Deliberately skewed / not block-aligned, mimicking +# RL rollout batches where responses stop at wildly different lengths (and +# occasionally immediately, e.g. an empty completion). +VARLEN_CONFIGS = [ + ("uniform_512", [512] * 8), + ("skewed_short_long", [32, 64, 96, 128, 1536, 2048, 3072, 4000]), + ("many_short_few_long", [16, 24, 40, 48, 64, 80, 96, 4096]), + ("with_empty_seq", [0, 128, 256, 512, 1024, 2048, 3000, 3500]), +] + +HEADS = 8 +HEAD_DIM = 64 + + +def _time_ms(fn, warmup, iters): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def _peak_vram_gb(fn, warmup=3, iters=5): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + torch.cuda.empty_cache() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - baseline) / (1024**3) + + +def _sdpa_dense(q, k, v, causal): + # q, k, v: [B, H, S, D] + return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=causal) + + +def _pad_mask_unpad_sdpa_varlen(q, k, v, seqlens, causal): + """Pad-mask-unpad SDPA baseline for packed varlen input. + + q, k, v: [total_tokens, H, D] packed along dim 0, split by `seqlens`. + This is the naive alternative to `triton_flash_attention_varlen`: pad + every sequence to the batch max, run dense masked SDPA, then gather the + valid rows back out. Included so the benchmark reflects what packing is + actually saving, not just kernel-vs-kernel throughput. + """ + batch = len(seqlens) + max_len = max(seqlens) if seqlens else 0 + H, D = q.shape[1], q.shape[2] + device, dtype = q.device, q.dtype + + q_pad = torch.zeros(batch, H, max_len, D, device=device, dtype=dtype) + k_pad = torch.zeros(batch, H, max_len, D, device=device, dtype=dtype) + v_pad = torch.zeros(batch, H, max_len, D, device=device, dtype=dtype) + attn_mask = torch.zeros(batch, 1, max_len, max_len, device=device, dtype=torch.bool) + + offset = 0 + for b, s in enumerate(seqlens): + if s == 0: + continue + q_pad[b, :, :s, :] = q[offset : offset + s].transpose(0, 1) + k_pad[b, :, :s, :] = k[offset : offset + s].transpose(0, 1) + v_pad[b, :, :s, :] = v[offset : offset + s].transpose(0, 1) + valid = torch.ones(s, s, device=device, dtype=torch.bool) + if causal: + valid = torch.tril(valid) + attn_mask[b, 0, :s, :s] = valid + offset += s + + out_pad = torch.nn.functional.scaled_dot_product_attention( + q_pad, k_pad, v_pad, attn_mask=attn_mask + ) + + out = torch.zeros_like(q) + offset = 0 + for b, s in enumerate(seqlens): + if s == 0: + continue + out[offset : offset + s] = out_pad[b, :, :s, :].transpose(0, 1) + offset += s + return out + + +def run_dense_benchmark(args): + device = device_ctx.device + dtype = torch.float16 + rows = [] + + for batch, heads, seq, head_dim in args.dense_configs: + q = torch.randn(batch, heads, seq, head_dim, device=device, dtype=dtype) + k = torch.randn(batch, heads, seq, head_dim, device=device, dtype=dtype) + v = torch.randn(batch, heads, seq, head_dim, device=device, dtype=dtype) + sm_scale = 1.0 / (head_dim**0.5) + + def sdpa_fwd(q=q, k=k, v=v): + with torch.no_grad(): + _sdpa_dense(q, k, v, causal=True) + + def triton_fwd(q=q, k=k, v=v): + with torch.no_grad(): + triton_flash_attention(q, k, v, causal=True, sm_scale=sm_scale) + + def triton_fwd_lse(q=q, k=k, v=v): + with torch.no_grad(): + triton_flash_attention( + q, k, v, causal=True, sm_scale=sm_scale, return_lse=True + ) + + qg = q.clone().requires_grad_(True) + kg = k.clone().requires_grad_(True) + vg = v.clone().requires_grad_(True) + + def sdpa_fwd_bwd(q=qg, k=kg, v=vg): + out = _sdpa_dense(q, k, v, causal=True) + torch.autograd.grad(out, (q, k, v), torch.ones_like(out)) + + def triton_fwd_bwd(q=qg, k=kg, v=vg): + out = triton_flash_attention(q, k, v, causal=True, sm_scale=sm_scale) + torch.autograd.grad(out, (q, k, v), torch.ones_like(out)) + + s_fwd = _time_ms(sdpa_fwd, args.warmup, args.iters) + t_fwd = _time_ms(triton_fwd, args.warmup, args.iters) + t_fwd_lse = _time_ms(triton_fwd_lse, args.warmup, args.iters) + s_fb = _time_ms(sdpa_fwd_bwd, args.warmup, args.iters) + t_fb = _time_ms(triton_fwd_bwd, args.warmup, args.iters) + s_vram = _peak_vram_gb(sdpa_fwd) + t_vram = _peak_vram_gb(triton_fwd) + + rows.append( + [ + f"{batch}x{heads}x{seq}x{head_dim}", + f"{s_fwd:.3f}", + f"{t_fwd:.3f}", + f"{t_fwd_lse:.3f}", + f"{s_fwd/t_fwd:.2f}x", + f"{s_fb:.3f}", + f"{t_fb:.3f}", + f"{s_fb/t_fb:.2f}x", + f"{s_vram*1024:.0f}", + f"{t_vram*1024:.0f}", + ] + ) + + headers = [ + "shape (B x H x S x D)", + "sdpa fwd ms", + "triton fwd ms", + "triton fwd+lse ms", + "fwd speedup", + "sdpa f+b ms", + "triton f+b ms", + "f+b speedup", + "sdpa fwd MB", + "triton fwd MB", + ] + print("\nDense attention: triton_flash_attention vs. torch SDPA (causal)") + print(tabulate(rows, headers=headers, tablefmt="github")) + + +def run_varlen_benchmark(args): + device = device_ctx.device + dtype = torch.float16 + heads, head_dim = args.heads, args.head_dim + sm_scale = 1.0 / (head_dim**0.5) + rows = [] + + for name, seqlens in args.varlen_configs: + total = sum(seqlens) + max_seqlen = max(seqlens) + cu_seqlens = torch.tensor( + [0, *torch.tensor(seqlens).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + + q = torch.randn(total, heads, head_dim, device=device, dtype=dtype) + k = torch.randn(total, heads, head_dim, device=device, dtype=dtype) + v = torch.randn(total, heads, head_dim, device=device, dtype=dtype) + + def pad_fwd(q=q, k=k, v=v): + with torch.no_grad(): + _pad_mask_unpad_sdpa_varlen(q, k, v, seqlens, causal=True) + + def triton_fwd(q=q, k=k, v=v): + with torch.no_grad(): + triton_flash_attention_varlen( + q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, + causal=True, sm_scale=sm_scale, + ) + + qg = q.clone().requires_grad_(True) + kg = k.clone().requires_grad_(True) + vg = v.clone().requires_grad_(True) + + def pad_fwd_bwd(q=qg, k=kg, v=vg): + out = _pad_mask_unpad_sdpa_varlen(q, k, v, seqlens, causal=True) + torch.autograd.grad(out, (q, k, v), torch.ones_like(out)) + + def triton_fwd_bwd(q=qg, k=kg, v=vg): + out = triton_flash_attention_varlen( + q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, + causal=True, sm_scale=sm_scale, + ) + torch.autograd.grad(out, (q, k, v), torch.ones_like(out)) + + p_fwd = _time_ms(pad_fwd, args.warmup, args.iters) + t_fwd = _time_ms(triton_fwd, args.warmup, args.iters) + p_fb = _time_ms(pad_fwd_bwd, args.warmup, args.iters) + t_fb = _time_ms(triton_fwd_bwd, args.warmup, args.iters) + p_vram = _peak_vram_gb(pad_fwd) + t_vram = _peak_vram_gb(triton_fwd) + + padded_tokens = len(seqlens) * max_seqlen + waste_pct = 100.0 * (1.0 - total / padded_tokens) + + rows.append( + [ + name, + f"{len(seqlens)}", + f"{total} ({waste_pct:.0f}% pad waste)", + f"{p_fwd:.3f}", + f"{t_fwd:.3f}", + f"{p_fwd/t_fwd:.2f}x", + f"{p_fb:.3f}", + f"{t_fb:.3f}", + f"{p_fb/t_fb:.2f}x", + f"{p_vram*1024:.0f}", + f"{t_vram*1024:.0f}", + ] + ) + + headers = [ + "config", + "batch", + "total tokens (vs. padded)", + "pad+sdpa fwd ms", + "triton varlen fwd ms", + "fwd speedup", + "pad+sdpa f+b ms", + "triton varlen f+b ms", + "f+b speedup", + "pad+sdpa fwd MB", + "triton varlen fwd MB", + ] + print( + f"\nPacked varlen attention: triton_flash_attention_varlen vs. " + f"pad-mask-unpad SDPA (causal, H={heads}, D={head_dim})" + ) + print(tabulate(rows, headers=headers, tablefmt="github")) + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--heads", type=int, default=HEADS) + parser.add_argument("--head-dim", type=int, default=HEAD_DIM) + parser.add_argument("--dense-only", action="store_true") + parser.add_argument("--varlen-only", action="store_true") + parser.add_argument( + "--dense-configs", + type=str, + default=None, + help="Semicolon-separated 'batch,heads,seq_len,head_dim' tuples.", + ) + args = parser.parse_args() + if args.dense_configs: + args.dense_configs = [ + tuple(int(x) for x in tup.split(",")) for tup in args.dense_configs.split(";") + ] + else: + args.dense_configs = DENSE_CONFIGS + args.varlen_configs = VARLEN_CONFIGS + return args + + +def main(): + args = parse_args() + if device_ctx.device_type != "cuda": + raise RuntimeError( + "Triton attention benchmark requires a CUDA device (Triton kernels are CUDA-only)." + ) + + logger.info(f"Triton attention benchmark on {device_ctx.device}") + + if not args.varlen_only: + run_dense_benchmark(args) + if not args.dense_only: + run_varlen_benchmark(args) + + +if __name__ == "__main__": + main()