Hardware-Aware Attention: FlashAttention, PagedAttention, KV Cache Management

A GPU is not one kind of memory. On an A100:

SRAM is roughly 10× faster than HBM and thousands of times smaller. Compute throughput has grown far faster than memory bandwidth for a decade, so most kernels are now memory-bound: the arithmetic units idle while data moves.

Standard attention is a bandwidth disaster

Look at what a naive implementation writes to and reads from HBM:

S = Q @ K.T           # write S   (T, T)  to HBM
S = S / sqrt(d)       # read S, write S
S = S + mask          # read S, write S
P = softmax(S)        # read S, write P   (T, T)
P = dropout(P)        # read P, write P
O = P @ V             # read P, write O

The (T, T) matrix crosses the HBM boundary six times. At T = 4096 with 32 heads in bf16 that is ~1 GB per crossing, per layer.

The FLOPs are unavoidable. The memory traffic is not.

(Dao et al., 2022)

The insight: never materialise the (T, T) matrix in HBM at all. Tile the computation, keep tiles in SRAM, and fuse every step into one kernel.

Tiling

Split Q, K, V into blocks that fit in SRAM. For each Q block, loop over K/V blocks, accumulating the output incrementally:

for each block of Q (size B_r):
    load Q_i into SRAM
    initialize output accumulator O_i, running max m_i, running sum l_i
    for each block of K, V (size B_c):
        load K_j, V_j into SRAM
        S_ij = Q_i @ K_j.T                 # small, stays in SRAM
        update m_i, l_i, O_i               # online softmax
    write O_i / l_i to HBM                 # normalize weighted numerator

Only O(T·d) is written — the output — instead of O(T²).

Online softmax

The subtle part. Softmax needs a max and a sum over the whole row for numerical stability, but we only ever see one block at a time.

The fix is a running update. Having processed blocks up to j with running max m and sum l, a new block with max m' gives:

Here O is an unnormalized weighted numerator, not the final attention output: for the new tile, l'=sum(exp(scores-m')) and O'=sum(exp(scores-m') * values). m,l are scalars per query row; O has d_v components. Start with m=-inf,l=0,O=0, handle empty/all-masked tiles explicitly, and finally return O/l for rows with positive l. Rescaling preserves the same real-number function; floating-point reduction order can change the last bits.

This is exact

Worth emphasising, because it is unusual. FlashAttention is not an approximation. Unlike everything in module 10, it produces the same numbers as the naive implementation (up to floating-point associativity). There is no quality trade-off. It is a strictly better implementation of the same function.

Recomputation in the backward pass

Backprop normally needs the stored (T, T) attention matrix. FlashAttention stores only the softmax statistics (m, l) and recomputes scores on the fly during the backward pass.

That is more FLOPs and still faster, because the operation is memory-bound. This is the clearest demonstration of the module's thesis.

Results

FA-2 improved work partitioning: fewer non-matmul FLOPs, better parallelism across sequence length, better warp scheduling. Reached ~70% of peak on A100.

FA-3 targets Hopper specifically: asynchronous execution with the Tensor Memory Accelerator, overlapping GEMM and softmax, and FP8 support. ~75% of peak on H100.

Using it

You almost never write this yourself:

import torch.nn.functional as F

# dispatches to FlashAttention when shapes/dtype/hardware allow
out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

Eligibility depends on PyTorch/backend version, device, dtype, dimensions and mask support. Inspect backend diagnostics for the actual inputs. A dense mask may prevent a particular fused kernel; replacing it with is_causal=True is valid only if the two masks represent the same attention pattern. Kernel availability can motivate architecture choices, but never justifies changing a pretrained model's mask silently.

FlashAttention optimises training and prefill, where T is large and there is plenty of parallel work.

Decoding is different. Generating one token means T_q = 1 — a single query against the whole cache. FlashAttention parallelises over query blocks, and with one query there is nothing to parallelise. Most of the GPU sits idle.

FlashDecoding (2023) parallelises over the key/value dimension instead: split the cache into chunks, attend to each chunk independently, then combine using the same online-softmax rescaling.

Up to 8× faster decoding at long context. Also exact.

FlashDecoding++ adds a unified max value across partitions to avoid synchronisation, plus asynchronous softmax.

(2024–25)

A library rather than a single kernel — the recognition that serving needs many attention variants, and hand-writing each is unsustainable.

FlashInfer provides:

• Block-sparse and composable formats. KV cache stored as a block-sparse matrix, so paged layouts, sliding windows, and shared prefixes are all expressed in one framework.
• JIT compilation. Kernels generated per attention variant (GQA vs MLA, with or without sinks, custom masks) rather than written by hand.
• Load-balanced scheduling. Handles ragged batches — requests with wildly different context lengths — without wasting warps.
• CUDAGraph compatibility, cutting launch overhead for small decode kernels.

It is now the kernel backend for vLLM, SGLang, and others. The trend it represents: attention kernels have become a compiler problem.

Module 09 sized the cache. This section is about allocating it, which turns out to be where most memory was being wasted.

The fragmentation problem

Pre-vLLM systems allocated one contiguous buffer per request, sized to the maximum possible length:

Request A: max_len 2048, actually uses 300   -> 1748 slots wasted
Request B: max_len 2048, actually uses 1900  ->  148 slots wasted
Request C: max_len 2048, actually uses  50   -> 1998 slots wasted

Three sources of waste:

The vLLM paper measured real systems using 20–40% of KV cache memory for actual tokens. The rest was waste.

PagedAttention

(Kwon et al., SOSP 2023)

Apply the idea operating systems have used since the 1960s: virtual memory with paging.

Split the cache into fixed-size blocks (typically 16 tokens). A request gets a block table mapping its logical positions to arbitrary physical blocks. Blocks need not be contiguous.

What this buys:

• Near-zero waste. Blocks are allocated on demand as generation proceeds. Internal fragmentation is capped at one block (≤16 tokens) per request.
• No external fragmentation. All blocks are the same size, so any free block fits any request.
• Copy-on-write sharing. Requests with a common prefix — a shared system prompt, or n parallel samples from one prompt — share physical blocks. Only on divergence is a block copied.

Reported result: 2–4× higher throughput at the same latency, coming almost entirely from fitting more concurrent requests in the same memory.

Prefix caching

Shared blocks generalise into a major production optimisation. If many requests begin with the same long system prompt, its KV blocks are computed once and reused — the prefill for that prefix is skipped entirely.

For agentic workloads, where a large tool-definition preamble is resent on every turn, this is often the single biggest win available.

Cache eviction

When memory runs out, something must go:

The first two are exact; the last two trade quality for capacity. StreamingLLM exists specifically because of the attention-sink phenomenon from module 10 — dropping the first tokens is catastrophic even though they are oldest.

The scheduling counterpart to PagedAttention, and part of the same system.

Static batching groups requests and runs them to completion together. Because outputs have different lengths, short requests sit finished-but-blocked while the longest one runs. GPU utilisation collapses.

static:      [A....][B..........][C...]     all wait for B
             |------ blocked -----|

continuous:  A finishes -> D starts immediately in that slot

Continuous batching (a.k.a. in-flight batching) schedules at the level of individual decode steps. When a request emits [EOS], its slot is freed and a queued request takes it on the very next step.

Iteration-level scheduling does not require paged allocation. Orca (OSDI 2022) predates PagedAttention (2023). Paging makes variable-length cache allocation more efficient; scheduling decides which requests run. They complement one another but are distinct mechanisms.

Combined reported gains: up to 23× throughput over naive static batching with contiguous allocation.

More in module 14, where it sits alongside quantization and speculative decoding.

A framing that explains most serving-system design.

Because the two phases have opposite bottlenecks, running them on the same GPU means one always underuses the hardware. Hence chunked prefill (interleave prompt chunks with decode steps to keep both units busy) and disaggregated serving (separate GPU pools for each phase, shipping the cache between them).

Understanding that decode is memory-bandwidth-bound explains almost everything else in this course:

• why GQA/MLA help speed, not just capacity (module 09) — fewer bytes read
• why recomputation can be cheaper than extra memory traffic (§11.2)
• why quantization speeds up decoding (module 14) — fewer bytes again
• why speculative decoding works (module 14) — it converts memory-bound decode steps into compute-bound verification

For scores [0,1,2] split into [0,1] and [2], the first maximum is 1 and the first sum is exp(-1)+1. On the next tile the maximum becomes 2; scale the old sum and weighted numerator by exp(-1), add the new tile, then divide. For values [1,3,5], this yields the same weighted mean as ordinary softmax.

import torch
scores = torch.tensor([0., 1., 2.], dtype=torch.float64)
values = torch.tensor([[1.], [3.], [5.]], dtype=torch.float64)
m = torch.tensor(float('-inf'), dtype=torch.float64)
den = torch.tensor(0., dtype=torch.float64)
num = torch.zeros(1, dtype=torch.float64)
for start, end in [(0, 2), (2, 3)]:
    tile = scores[start:end]
    new_m = torch.maximum(m, tile.max())
    correction = (m - new_m).exp()
    p = (tile - new_m).exp()
    num = correction * num + p @ values[start:end]
    den = correction * den + p.sum()
    m = new_m
torch.testing.assert_close(num / den, scores.softmax(-1) @ values)
assert den > 0
print("Normalized online output:", (num / den).item())

The SRAM budget for a tile is per multiprocessor/workgroup, not the sum of all on-chip memory. A tensor fitting in aggregate GPU SRAM may not fit in the local workspace of the scheduled block.

For a weight matrix with N entries reused across M query tokens, dominant work is 2MN FLOPs and a single weight read costs sN bytes: weight-only intensity is 2M/s. At BF16 and M=1 this is 1 FLOP/byte; at M=128 it is 128. KV reads, activation writes and imperfect reuse lower full-operation intensity. Prefill often has larger M; batched decode also increases reuse. Neither phase is universally compute- or bandwidth-bound across all batch/context sizes.

Benchmark exercise: disclose device, versions, dtype, (B,H,T_q,T_k,d_head), mask, warmup count, repetitions and backend. Synchronize CUDA before/after timing or use CUDA events. Compare equal masks/functions and report numerical error alongside allocated/peak memory. CPU checks above verify mathematics only; they do not substantiate GPU speedups or backend-selection claims.