Triton FlashAttention

FlashAttention-2 forward and backward in Triton, tuned for the free Kaggle T4: 1.55-1.92× faster forward and 1.64-1.83× faster forward+backward than PyTorch's memory-efficient SDPA, using under 60% of its peak extra memory. 79 of 79 tests pass on the GPU.

Built with Python, Triton, PyTorch, CUDA, Kaggle.

attention(q, k, v, causal=...) runs fused Triton kernels on CUDA, the same kernels on the CPU under Triton's interpreter (which is how CI tests them), and a tiled PyTorch reference of the same algorithm anywhere else, with autograd support throughout. The forward kernel runs one program per query block and batch-head, streaming key and value blocks through shared memory with an online softmax (a running row max and sum in fp32, base-2 exponentials with scores pre-scaled by log2(e), and a rescale of the accumulator when the max moves); causal masking skips key blocks above the diagonal entirely, and the kernel saves one float per row for the backward pass. The backward is split the FlashAttention-2 way into a dK/dV kernel and a dQ kernel so nothing needs atomics, both rebuilding probabilities from the saved log-sum-exp instead of storing the N×N matrix, with rowsum(dO·O) computed by its own small kernel. Block sizes, warps and stages are autotuned per power-of-two length bucket for the T4's sm75 architecture, with configurations that can't fit its shared memory pruned first. Tests compare outputs and all three gradients against materialised float32 attention across fp16 and fp32, causal and non-causal, head dimensions 16-128, lengths that aren't block multiples, and non-contiguous inputs, and every autotune configuration is tested directly. On the T4, across sequence lengths 512-16,384, it runs 1.55-1.92× faster than SDPA's memory-efficient backend forward and 1.64-1.83× faster forward plus backward, with 130 MiB of peak extra memory against SDPA's 226, and a best forward rate of 22.6 TFLOP/s, 35% of the T4's nominal peak. Built by Archit with Claude as the AI coding assistant.

• Online softmax with exp2 and fp32 running statistics; causal block skipping
• Backward split into dK/dV and dQ kernels (no atomics), plus a fused rowsum(dO·O) kernel
• Autotuning per power-of-two length bucket, with shared-memory pruning for sm75
• fp16 and fp32, head dimensions 16-128, non-contiguous inputs; double backward refused rather than silently wrong
• A readable tiled PyTorch reference of the same algorithm, used as the fallback and as a test oracle
• CI runs the real kernels under Triton's CPU interpreter; a Kaggle job reruns the suite on the T4 before benchmarking and commits the logs

Every project hits walls. Here's what nearly broke Triton FlashAttention — and how it shipped anyway.

• Tensor-core instructions vanish on Turing in newer Triton: a sweep of 10 versions with PTX instruction counts showed 3.0-3.2 emit 256 mma instructions, 3.3 fails to compile and 3.4+ emit none (8,208 fp32 FMAs instead), so the project pins Triton 3.2.0 and warns on newer versions
• Turing has no TF32, so fp32 inputs need IEEE-precision dot products
• 64-bit program offsets to avoid int32 overflow at long sequence lengths
• Testing every autotune tile configuration, including unequal query and key block sizes, so a configuration only the tuner picks can't hide a bug
• Benchmark hygiene: median timers, accuracy checked at every length up to 16k, and an honest caveat on a power-throttled cuBLAS comparison

Every build leaves residue. The takeaways from Triton FlashAttention that carry forward.

• Compiler versions are part of the result: the same kernel ran at 1.3 instead of 22.2 TFLOP/s on a newer Triton
• Recomputing attention probabilities in the backward pass is cheaper than storing them, which is the whole point of FlashAttention
• A slow, readable reference implementation is the most useful test oracle a kernel can have
• Scope stated plainly: dense attention only (no dropout, bias, GQA, variable lengths or KV-cache decode), measured on one GPU in one run