Attention kernel chains: tiling, the online softmax, merges and masks

Attention kernel questions are about where bytes live and which small state survives a tile. The two softmax chains use the same three scores, 1, 2 and 3, so the answer can be checked against a plain softmax: 0.090, 0.245, 0.665.

What FlashAttention removes from HBM

inputs   one head, N = 8,192 tokens, head dim 128, BF16
scores   N x N x 2 B = 128 MiB per head, written and read back by naive attention
naive    about 520 MiB of traffic per head, nearly all of it the score matrix
compulsory  read Q, K, V and write O once: 4 x 8,192 x 128 x 2 B = 8 MiB
FLOPs    4 x N^2 x d = 34.4 GFLOP, unchanged by tiling
sanity   34.4 GFLOP over 8 MiB is about 4,096 FLOP per byte: compute-bound

Tiling changes bytes moved, never FLOPs. Prefill attention becomes a compute problem once the score matrix stops going to memory.

Drilled in: FlashAttention HBM bytes for one head · FlashAttention generations

The online softmax on numbers

inputs   scores 1 and 2 arrive in the first block, 3 in the second
block 1  m = 2, l = e^(1-2) + e^(2-2) = 1.368
block 2  new max 3: alpha = e^(2-3) = 0.368
         l = 1.368 x 0.368 + e^(3-3) = 1.503
result   weight of the score 3 = 1 / 1.503 = 0.665
sanity   matches softmax(1, 2, 3); the accumulated output is rescaled by the
         same alpha whenever the max rises

The state per query row is the running max, the running sum and the unnormalized output. Dividing by the sum once at the end, rather than after each block, is the FlashAttention-2 change.

Drilled in: Derive the online softmax · Online softmax recurrence

The log-sum-exp merge

inputs   the same scores split across two workers: [1, 2] and [3]
lse      lse_1 = ln(e^1 + e^2) = 2.313; lse_2 = 3
merged   lse = ln(e^2.313 + e^3) = 3.408
weights  e^(2.313 - 3.408) = 0.335 and e^(3 - 3.408) = 0.665
sanity   0.335 is 0.090 + 0.245, the first worker's share of the softmax

Each worker writes a normalized partial output and its lse. The merge weights each partial by exp(lse_s - lse), so any split of the keys gives the exact result.

Drilled in: Split-KV split count and LSE merge · Split-KV flash decoding

How many splits for decode

inputs   Llama 3.1 8B decode at batch 1: 8 KV heads, 128K tokens of context
         H100 SXM with 132 SMs
unsplit  one thread block per KV head: 8 blocks, 6% of the SMs
split 16 16 chunks of 8K keys x 8 heads = 128 blocks, 0.97 of one wave
sanity   one wave keeps every SM busy without a second, mostly idle wave

Splitting pays when batch times KV heads is far below the SM count. At high batch there are already enough blocks, and the merge becomes pure overhead.

Drilled in: Split-KV split count and LSE merge

The causal mask for a decode step

inputs   T new queries over S cached keys; query row i is at position S - T + i
rule     key j is visible when j <= i + (S - T)
decode   T = 1, S = 4,096: the query sees keys 0 to 4,095
top-left a mask aligned to the top-left corner lets that query see key 0 only
sanity   the mask is lower-triangular and aligned to the bottom-right corner

PyTorch's is_causal=True aligns top-left when T < S. Using it for a decode or chunked-prefill step returns attention over the wrong keys without an error.

Drilled in: Causal mask when queries are shorter than keys

What the backward pass saves

inputs   32 heads, N = 8,192 tokens per layer
saved    Q, K, V, O and one FP32 log-sum-exp per query row:
         32 x 8,192 x 4 B = 1 MiB of LSE per layer
not saved  the probabilities P: 32 x 8,192^2 x 2 B = 4 GiB per layer in BF16
backward rebuilds each P tile as exp(S - LSE) on chip
sanity   one extra GEMM per tile buys back 4 GiB per layer

The softmax gradient needs one more vector, D = rowsum(dO x O), computed once before the main loop. Nothing else of size N x N is kept.

Drilled in: What FlashAttention-2's backward saves

GQA loads KV once per group

inputs   Llama 3.3 70B: 64 query heads, 8 KV heads, group g = 8; TP=8
         batch 64, 8K context, 320 KiB of KV per token in BF16
grouped  64 x 8,192 x 327,680 B / 8 = 20 GiB of KV read per GPU per step
per head without grouping, each query head re-reads its KV head: up to 160 GiB
intensity  1 FLOP per KV byte for one head, 8 when 8 heads share the tile
sanity   the cache is already 8x smaller than MHA; the kernel has to read
         it once to keep the saving

Stacking the group's query heads as rows of one tile also gives a Hopper matrix instruction 8 of its 64 rows to work on, where a single decode query would give it one.

Drilled in: GQA kernel that loads KV once per group

The block-table cost of paged attention

inputs   Llama 3.1 8B, 16-token blocks, BF16
block    16 x 8 KV heads x 128 x 2 (K, V) x 2 B = 64 KiB per layer
table    one int32 per block: 4 B per 64 KiB = 0.006%
sanity   the bytes are nothing; the cost is a dependent load before every
         block and tiles that span several pages

A 128-token tile over 16-token blocks touches 8 separate pages, so the kernel needs 8 table lookups before it can issue the loads that matter.

Drilled in: Paged attention block-table gather