Module 5 · Kernels & memory

I/O-Aware Kernels

Attention's real cost is not the arithmetic — it is shuttling data between slow, roomy HBM and fast, tiny SRAM. FlashAttention rewrites the kernel so the giant k×k score matrix is never written to memory at all.

The memory bottleneck: SRAM vs HBM

step-through

Why attention is I/O-bound — the cores can compute far faster than memory can feed them.

SRAM / on-chip compute HBM (main memory) slow HBM traffic

The GPU holds model weights and the KV cache in HBM — gigabytes, but comparatively slow. The cores actually work out of SRAM: only tens of megabytes, yet with an order of magnitude more bandwidth. Any kernel that keeps bouncing data between the two is limited by the slow link, not by raw math. Attention, done naively, is exactly that kind of kernel.

Naive attention: materializing the k×k score matrix

step-through

Each stage writes a full matrix to HBM and reads it back — and the score matrix is O(k²).

HBM resident score matrix S = QKᵀ on-chip compute expensive HBM round-trip
The key idea The score matrix S = QKᵀ is only an intermediate. There is no reason to send it all the way out to HBM and read it back. FlashAttention fuses the whole computation into a single kernel and keeps that intermediate on-chip — using an online softmax so it never needs the full row of scores at once.

FlashAttention: tiling + online softmax

step-through · centerpiece

One Q-tile in SRAM consumes K/V tiles in a stream, updating a running max m and sum ℓ — the k×k matrix is never written out.

HBM (Q, K, V blocks) SRAM tile compute partial scores Sij running stats m, ℓ output accumulator O

Because the softmax is computed incrementally, FlashAttention gets the exact same numbers as the naive version — it just never spills the score matrix to HBM. The K and V tiles are re-read once per Q-tile, but that is cheap next to round-tripping an O(k²) matrix. The win grows with sequence length, which is why it especially accelerates the long-context prefill stage.

FlashInfer: I/O-aware attention for serving

auto-play

The same fused kernel, but reading K/V from a paged cache across a variable-length batch — the bridge to Modules 4 & 8.

paged KV blocks (HBM) FlashInfer kernel (SRAM) output per sequence

FlashInfer is a set of inference-optimized kernels that implement FlashAttention while gathering keys and values from paged KV layouts — the PagedAttention and RadixAttention schemes from Module 4 — and handling ragged, variable-length batches plus multi-head and grouped-query attention. Same I/O-aware core; serving-grade plumbing.