Attention and FlashAttention

Attention is where every idea in this series meets at once. The online-softmax recurrence from Part 4, the tiling from Part 5, the tensor-core layouts and warp specialization from Part 6 — FlashAttention is those four things assembled into one kernel. And then, at inference time, it turns into something else entirely.

The operation itself is three lines:

Count the arithmetic first. Per head, forward: 2N²d for QK^T plus 2N²d for PV = 4N²d FLOPs. The compulsory memory is Q, K, V, O — four N×d tensors. So the arithmetic intensity is

Seven times above the H100 ridge point. Attention should be comfortably compute-bound.

It is not, in the naive implementation, because S and P are materialised. Those are N×N matrices, written to HBM and read back — twice, in most implementations, once for the softmax and once for the second matmul. And N² grows a great deal faster than Nd.

The obstacle is that softmax needs the whole row before it can normalize anything, and the whole row is exactly what we are refusing to materialise. Part 4 already dissolved that: the online-softmax recurrence lets you consume a row in pieces, keeping only a running maximum m and a running sum ℓ.

FlashAttention adds one thing: the output accumulator has to be rescaled too. Here is the forward pass, verbatim in structure from FlashAttention-2, Algorithm 1:

That bolded line is the entire difference from Part 4. Whenever a new block reveals a larger maximum, the accumulated output is scaled by the same factor as the accumulated sum, so the two stay consistent. Both factors are e to a non-positive power, so neither can overflow.

Two FA2-specific refinements are visible in the last two lines. Dividing by ℓ is deferred out of the inner loop, because a division is a non-matmul FLOP and those are 16× more expensive than matmul FLOPs on an A100. And only the single scalar L = m + \log ℓ is stored for the backward pass, not both m and ℓ.

# The inner loop, in the shape Triton's tutorial 06 uses.
for start_n in range(0, hi, BLOCK_N):
    k = tl.load(K_block_ptr)
    qk = tl.dot(q, k) * qk_scale                     # S tile — stays in registers
    if CAUSAL:
        qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, -1.0e6)

    m_new = tl.maximum(m_i, tl.max(qk, 1))            # running max
    qk    = qk - m_new[:, None]
    p     = tl.math.exp2(qk)                         # exp2, not exp — see Part 4
    alpha = tl.math.exp2(m_i - m_new)                # THE rescale factor

    l_i = l_i * alpha + tl.sum(p, 1)
    acc = acc * alpha[:, None]                       # rescale the accumulator
    v   = tl.load(V_block_ptr)
    acc = tl.dot(p.to(v.dtype), v, acc)              # second GEMM, straight from registers
    m_i = m_new
    K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
    V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

acc = acc / l_i[:, None]                             # normalize once

# What it must equal. Note this allocates the N×N matrix.
S = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(d))
if causal:
    S = S.masked_fill(mask, float('-inf'))
P = torch.softmax(S, dim=-1)
O = P @ v

Notice what does not appear: any store of p to HBM. The tile stays on chip and feeds the second GEMM; the compiler and kernel handle whatever register-layout conversion the chosen dtype and MMA instruction require. That qualification matters: FlashAttention-3's FP8 wgmma path, for example, explicitly permutes values because its accumulator and A-operand layouts differ.

The 2022 version proved the idea and reached 25–40% of an A100's peak. Getting from there to 75% of an H100's peak took two more papers, and each one is a lesson from an earlier part of this series.

7.3.1 FA2: parallelism and warp partitioning

• Fewer non-matmul FLOPs. On an A100, 312 TFLOP/s of matmul against 19.5 TFLOP/s of FP32 non-matmul — every non-matmul FLOP costs 16×. Hence deferring the division and storing only L.
• Parallelize over sequence length, not just over batch×heads. With long sequences and small batches there simply were not enough thread blocks to fill the GPU — the wave-quantization problem from Part 1.
• Split-Q instead of split-K. FA1 split K and V across the four warps of a block, so each warp computed a slice of QK^T and then had to write intermediates to shared memory, synchronize, and sum. FA2 splits Q across warps and keeps K and V shared. In the paper's words: "There is no need for communication between warps."

Result: about 2× over FA1, reaching 50–73% of theoretical peak on an A100 (FA1 managed 25–40%).

7.3.2 FA3: asynchrony, and the exponential problem

On Hopper the bottleneck moves. From Part 4: the SFU produces 3.9 TFLOP/s of exp against 989.5 TFLOP/s of BF16 matmul — a 254× gap, which the paper rounds to 256×. In an attention forward pass with head dimension 128 there are roughly 512× more matmul FLOPs than exponentials, so the exponentials take about half as long as the matmuls themselves. They are on a different execution unit, so the fix is to run them concurrently.

• Producer/consumer warp specialization with TMA and wgmma, exactly the structure in Part 6 §6.7. The producer warpgroup drops to 40 registers and issues TMA loads into an s-stage circular buffer; consumer warpgroups take 232 registers and run the two GEMMs.
• Ping-pong scheduling. bar.sync instructions force warpgroup 1's GEMMs to be scheduled ahead of warpgroup 2's, so WG1's softmax runs during WG2's GEMMs and then the roles swap. Measured: 570 → 620–640 TFLOP/s for FP16, head dim 128, sequence 8192. A second overlap runs inside each consumer warpgroup: a 2-stage register pipeline issues block j+1's QK^T while block j's softmax is still in flight — that is the "GEMM–softmax pipelining" row of the ablation below.
• FP8, with two accuracy devices: per-block quantization (one scale per B_r×d block, fused into the preceding rotary embedding for free), and incoherent processing — multiply Q and K by a random orthogonal matrix M before quantizing. Since MM^T = I, (QM)(KM)^T = QK^T, so the result is unchanged, but outliers get spread across dimensions.

Peak reported by FlashAttention-3: 740 TFLOP/s FP16 (75% of theoretical) and close to 1.2 PFLOP/s in FP8. Speedup over FA2: 1.5–2.0× forward, 1.5–1.75× backward.

With a causal mask, query i only attends to keys j ≤ i. In a tiled kernel that means roughly half the (i,j) tile pairs are entirely masked out, and you can skip computing them altogether — not mask them, skip them. Only the tiles on the diagonal need an actual elementwise mask applied.

FA2 measures this at 1.7–1.8× versus non-causal. It also creates a load-imbalance problem: block i does work proportional to i, so a naive assignment leaves half the SMs finishing early. Production kernels reverse or interleave the block order so that long and short rows are paired.

At inference, after the prompt has been processed, each step has one query row per sequence. The arithmetic changes character completely:

Decode attention is not a matmul problem. It is a KV-cache read, at an arithmetic intensity of 1 against a ridge point of 295. And FlashAttention's parallelism — batch × heads × query blocks — loses its third factor entirely.

The fix is the same one as split-K in Part 5: find parallelism along the reduction dimension. Flash-Decoding splits K and V along the sequence, runs an independent FlashAttention against each chunk, stores one extra log-sum-exp scalar per chunk, and then combines them — using the very same (m, ℓ) merge operator from Part 4, applied one level up.

Flash-Decoding reports near-constant time up to ~32k tokens, and up to 8× end-to-end on long-context decoding — with the attention step itself up to 50× faster than FlashAttention. Note the middle column at 131k: FlashAttention-2 is slower than eager PyTorch there, because it has traded parallelism for memory efficiency in exactly the regime where parallelism is what you need.

Everything a decode step reads is the KV cache, so its size is the whole story:

For Llama-3-70B (80 layers, 8 KV heads, head dim 128, BF16) that is 320 KiB per token. At 8k context and batch 32 the cache alone is 85.9 GB — more than an 80 GB H100 — and every decode step must re-read all of it on top of the 140 GB of weights.

7.6.1 Grouped-query attention, and what batching cannot amortise

With g query heads sharing one KV head, each KV element is loaded once and consumed g times, so the intensity of decode attention is

All still one to two orders of magnitude below the ridge point. And here is the asymmetry that shapes serving systems:

7.6.2 PagedAttention

The second KV problem is not bandwidth but fragmentation. A request's cache grows unpredictably, so a serving system that pre-allocates a contiguous maximum-length buffer wastes most of it. The vLLM paper (Kwon et al., 2023) measured the fraction of allocated KV memory actually holding token state: 20.4–38.2% for existing max-length allocators, 96.3% for PagedAttention.

The mechanism is virtual memory, applied to the cache: split it into fixed-size blocks, keep a per-request block table mapping logical to physical blocks, and let the blocks be scattered. Because every block has the same size there is no external fragmentation at all, and the waste that remains is internal — the unfilled slots of one final block per request, under 4% in the paper's measurements. Blocks can then be shared — between beam-search branches, parallel samples, or requests with a common prefix. The reported end-to-end gain is 2–4× throughput at the same latency.

The kernel pays for it with an extra indirection: every K/V access goes through the block table. vLLM's cache layout is shaped around keeping the resulting loads 16-byte aligned:

That trailing x exists purely so a thread group's 16-byte vector load lands contiguously — the rule from Part 2 §2.3, reappearing as a tensor layout decision.

• Training / prefill: use FlashAttention-4 on supported Hopper and Blackwell systems, with FA3 or FA2 as compatibility fallbacks. Alternatively use F.scaled_dot_product_attention, which selects among PyTorch's documented backends, including FlashAttention-2 where eligible; it does not promise to dispatch to external FA3 or FA4 packages.
• Decode: a split-KV kernel — Flash-Decoding, vLLM's paged kernel, or SGLang's. The right choice depends on your batching and cache-sharing strategy more than on the kernel.
• Writing your own: reasonable in Triton if you have an unusual mask, a custom score modification, or a fused pre/post-op. Tutorial 06 is a working starting point. On standard shapes, compare against the current FlashAttention implementation before claiming a win.

Next: the last big lever on both of those — using fewer bits.