Implementing Flash Attention’s Backward Pass in Raw CUDA: What I Learned

Most implementations of Flash Attention stop at the forward pass. The backward pass is harder to find, harder to understand, and much harder to get right. This is what I learned building it from scratch in raw CUDA C++.

Why I built this

I wanted to understand why Flash Attention exists — not just use it, but actually know what’s happening at the hardware level. So I built the whole thing from the ground up: naive softmax, optimized softmax, standard attention baseline, Flash Attention forward, and finally the backward pass.

The backward pass was the part I couldn’t find good raw CUDA implementations of. Most student repos use PyTorch autograd or CuPy. The original paper’s reference implementation is in Triton. I wanted raw CUDA C++ — the thing closest to the metal.

The memory problem Flash Attention solves

Standard scaled dot-product attention computes an N × N score matrix. At SEQ_LEN=4096, that matrix is 64MB per head. A 32-head model needs 4GB just for attention matrices — before weights, optimizer state, or gradients. That’s why long-context models were impractical before Flash Attention.

Flash Attention never writes that matrix to DRAM. It processes Q, K, V in tiles using the online softmax algorithm to accumulate the output without materializing the full matrix. Memory stays O(N·d) regardless of sequence length.

What the forward pass saves for backward

Instead of storing P (which would be O(N²)), the forward pass saves two small vectors:

• m[i]: the running max of scores for query row i
• l[i]: the softmax normalizer for query row i

With these two vectors we can recompute any attention weight on demand during backward: P[i][j] = exp(S[i][j] — m[i]) / l[i]

Deriving the gradients

Given the loss gradient dO, we need dQ, dK, and dV.

dV is the easiest: dV = P^T · dO

For dQ and dK we need the softmax backward. First compute a scalar correction per row: D[i] = dot(dO[i], O[i])

Then: dS[i][j] = P[i][j] * (dP[i][j] — D[i]) where dP[i][j] = dot(dO[i], V[j])

Finally: dQ[i] = sum_j dS[i][j] K[j] / sqrt(d) and dK[j] = sum_i dS[i][j] Q[i] / sqrt(d)

None of this requires storing the full N×N matrix at any point.

The bug that got me

My first version used dim3(D_MODEL, Bc) thread blocks — 64×32 = 2048 threads per block. CUDA’s hard limit is 1024. The kernel silently launched zero blocks. Every output was zero. The fix was switching to one block per output row with 64 threads, looping inside the kernel instead.

Results

All gradients verified against CPU reference:

Forward O: max error 9.69e-08 — PASSED dV: max error 3.58e-07 — PASSED dQ: max error 4.47e-08 — PASSED dK: max error 1.19e-07 — PASSED

Errors are float32 rounding noise. The math is exact.

Code

Full implementation: github.com/andreay99/cuda-flash-attention