Attention from the kernel's side: FlashAttention, online softmax and fused dequant

Attention is three lines of math:

S = Q·Kᵀ / √d        scores, one row per query
P = softmax(S)       row by row
O = P·V              weighted sum of values

The math is easy. The memory isn’t. Here’s what the kernels actually do about it, from the textbook version to the fused 4-bit path my inference box runs.

The N×N problem

Written naively, S and P are N×N matrices for every head in every layer. During a 32K-token prefill that’s 32,768² scores. At FP16 that’s 2 GiB per head per layer, written to memory and read straight back. No GPU has the memory, and even if it did, the traffic would dominate everything else.

The fix is to never build the matrix. You compute it in tiles that live in on-chip memory and throw each one away once it’s used. The obstacle is the softmax: it needs the maximum and the sum of a whole row before it can normalise any element of it.

The online softmax

Softmax subtracts the row maximum before exponentiating. That isn’t optional: in FP16, exp(x) overflows past x ≈ 11.1, because 65,504 is the largest FP16 value. The online softmax keeps a running maximum m, a running sum l and a running, unnormalised output o, and fixes them up as each new tile of scores arrives:

m_new = max(m, max_j s_j)
l     = exp(m − m_new) · l  +  Σ_j exp(s_j − m_new)
o     = exp(m − m_new) · o  +  Σ_j exp(s_j − m_new) · v_j
m     = m_new
                                         … after the last tile:  O = o / l

When a larger maximum shows up, everything accumulated so far is rescaled by exp(m − m_new). The result is exactly the softmax, not an approximation. It just never needs the whole row at once.

FlashAttention: what never touches memory

FlashAttention (Dao et al., 2022) is the online softmax turned into a kernel:

1. Load a tile of queries into shared memory.
2. Stream tiles of K and V past it.
3. For each pair, compute the scores in registers and update m, l and o.
4. Write O once, at the end.

S and P exist only in registers and shared memory, tile by tile. The kernel reads Q, K and V and writes O, and nothing N×N ever goes to VRAM. The arithmetic doesn’t change; the memory traffic does, and on a GPU that’s the whole game.

Prefill and decode are different problems

Prefill processes the whole prompt at once. There are many queries, so every K and V tile is reused across many rows. The work is big matrix multiplies on tensor cores, and it’s compute-bound.

Decode has one query per sequence. Each cached key is used for one dot product with that query, and each cached value for one multiply-add. At FP16 that’s about 1 FLOP per byte read, against the hundred-plus FLOPs per byte a modern GPU can sustain. Decode attention isn’t a matmul; it’s a stream over the KV cache, and it runs at memory speed.

Two things change the picture:

• Grouped-query attention. When g query heads share one KV head, the kernel loads each K/V tile once and uses it g times. That’s g× the work per byte read.
• Flash-decoding. With a single query there’s nothing to parallelise across rows, so the kernel splits the KV sequence across thread blocks instead. Each block computes a partial (m, l, o) over its slice, and a final step merges them with the same rescale rule shown above.

Why a quantised cache needs a fused kernel

At a full 262K context, my model’s cache is 17.2 GB at FP16 and 4.4 GB as 4-bit TBQ4 (the KV cache note has the derivation). Decode attention reads all of it on every step, so the format sets the bytes.

There are two ways to use a quantised cache:

• A separate dequantise pass. Read the 4-bit cache, write an FP16 copy, then let attention read the copy. For the full cache that’s 4.4 + 17.2 + 17.2 ≈ 39 GB of traffic per step. That’s more than just keeping the cache in FP16. You save memory and lose speed.
• Dequantise inside the tile load. Read the 4-bit bytes straight into the kernel and turn them into values in registers. Traffic is the 4.4 GB, and that’s all.

The upstream TBQ4 kernel in Indras-Mirror’s fork, which my engine builds on, does the second, with one more trick. The cache is stored in a Hadamard-rotated domain, and dot products survive the rotation. So the query is rotated once, keys and values are read directly as centroid × norm, and the output is rotated back once. The entire dequantise is two table lookups per byte:

// per byte = 2 KV values
const uint8_t byte = __ldg(&blk->qs[b]);
const half lo = __float2half(d_tbq4_centroids[byte & 0xF] * norm);
const half hi = __float2half(d_tbq4_centroids[byte >> 4] * norm);
tile[...] = __halves2half2(lo, hi);

At 262K, that makes attention about 21% of the bytes each decode step reads, next to 17 GB of weights. With an FP16 cache it would be half. The decode note turns that into a tokens-per-second ceiling.

Rules I use

1. Profile prefill and decode separately. They have different bottlenecks, and an average hides both.
2. At long context, decode attention is a bandwidth problem. Bytes per cached value matter more than any FLOP optimisation.
3. A quantised cache only pays off if dequant is fused. A separate dequant pass can move more data than the FP16 cache you were trying to shrink.
4. Check the GQA ratio. It sets how much reuse the attention kernel gets out of each byte it loads.

Sources: Milakov and Gimelshein, Online normalizer calculation for softmax (2018). Dao et al., FlashAttention (2022) and FlashAttention-2 (2023). Dao, Haziza, Massa and Sizov, Flash-Decoding for long-context inference (2023). The TBQ4 kernel snippet is from the upstream fork’s README.