Part II — Systems

Kernels — fusion and FlashAttention

Lesson 06 left you with a verdict about the chip: most of what a language model does — every norm, every softmax, every elementwise add, and attention during decode — is memory-bound. The tensor cores sit idle while the kernel waits on HBM. But that verdict was about a single op in isolation. Run a real model in naive PyTorch and it is far worse: each op is its own kernel, its own launch, its own round trip to HBM. This lesson closes that gap two ways — by fusing chains of memory-bound ops into one trip, and by the headline case of fusion, FlashAttention, which refuses to ever write the T×T score matrix to HBM at all.

1 · The cost of a kernel: launch + a round trip to HBM

A kernel is one GPU program launched from the host: the CPU tells the GPU "run this function over this grid of threads." Two costs come with every launch. First, launch overhead — on the order of a few microseconds (≈3–10 μs) to queue and dispatch the grid. Second, and usually dominant, the HBM round trip: a typical elementwise kernel reads its input tensor from HBM, does a handful of FLOPs per element, and writes the result back to HBM. The math is trivial; the bytes are not.

Put numbers on it. Take a residual-stream activation of shape (B, T, d) = (8, 4096, 4096) in bf16: that is 8 · 4096 · 4096 · 2 ≈ 268 MB. On an H100 at ≈3.3 TB/s, just reading it costs 268 MB / 3.3 TB/s ≈ 81 μs; a read-plus-write op costs ≈162 μs. The arithmetic — say a bias add and a dropout, two FLOPs per element — is 2 · 134M ≈ 0.27 GFLOP, which the tensor cores would finish in well under a microsecond. The op is ≈99% waiting on memory. This is the §06 roofline restated as a wall-clock: the kernel's time is set by bytes / bandwidth, and the FLOPs are free.

Now chain three such ops — the classic dropout(bias(x)) followed by a LayerNorm. Run eagerly, that is three launches and three round trips: read x, write; read, write; read, write — roughly 6× the tensor's bytes crossing HBM, ≈3 launches. Fusion rewrites the chain as one kernel: load each tile of x from HBM into fast on-chip SRAM/registers once, apply bias, dropout, and the norm while the data is hot, write the final result once. One launch, one read, one write — ≈2× the bytes instead of 6×. For a memory-bound chain that is a ≈3× speedup, and it comes purely from moving fewer bytes, not from doing less math.

The principle generalizes: for memory-bound work, value is measured in FLOPs per byte loaded. Fusion raises that ratio by doing more work per trip to HBM. torch.compile and Triton exist largely to find and emit these fused kernels automatically. The one place where the payoff is so large it gets its own algorithm — not just a fused epilogue — is attention.

2 · Naive attention: the T×T matrix that eats HBM

Recall the attention math from lesson 2. For each head you compute scores, mask, softmax, then weight the values:

A naive implementation does this literally: it computes S, writes the full (B, h, T, T) tensor to HBM, reads it back to apply the mask and softmax, writes P, reads it back to multiply by V. The killer is the shape of S: it is quadratic in the sequence length T, and it lives in HBM.

Make it concrete. Take B = 8, h = 32 heads, bf16 (2 bytes). The score matrix alone is B · h · T² · 2 bytes:

Two separate problems, both fatal. Memory: the score matrix grows as O(T²) and OOMs at modest context — long-context training is simply impossible this way. Traffic: even when it fits, every byte of that O(T²) tensor is written to HBM and read back two or three times, so attention is bandwidth-bound — the matmuls themselves are cheap relative to the scores being shuttled in and out. The output we actually want, O, is only (B, h, T, d_head) — linear in T. We are paying O(T²) memory to produce an O(T) result. The entire trick is to never write the big intermediate down.

3 · The online-softmax recurrence

The obstacle is softmax. Softmax over a row needs the whole row at once: you subtract the row max (for numerical stability — e^large overflows), exponentiate, and divide by the row sum. That seems to demand having all T scores in hand, which is exactly the O(T²) matrix we want to avoid. Online softmax dissolves that demand: it computes the same softmax incrementally as scores arrive in tiles, carrying just two running scalars per query row.

Process the key/value sequence in blocks. Keep, for each query, a running maximum m (the largest score seen so far), a running denominator ℓ (the sum of exponentials, shifted by m), and a running output accumulator o. When a new block of scores s arrives:

The single subtle step is the rescale. When a later block contains a larger score, the running max jumps, and every exponential computed against the old max is now off by a factor e^{(m_old − m_new)}. So before folding in the new block you multiply the old ℓ and o by that correction, bringing the running totals onto the new reference. At the end, the true softmax-weighted output is just o / ℓ. Critically, m, ℓ, and o are tiny — a couple of scalars and one d_head-vector per query — so you can stream arbitrarily many K, V tiles past a query while holding constant memory. This is the same running-reduction pattern Triton teaches for blocked reductions (cross-link below); FlashAttention is online softmax plus a value accumulator.

for each query row q:
    m = -inf;  l = 0;  o = 0                 # running max, denom, output
    for each K,V tile (Kj, Vj):              # stream tiles; never store all of S
        s   = q @ Kj.T / sqrt(d_head)        # this tile's scores  (block-wide)
        m_n = max(m, max(s))                 # update running max
        c   = exp(m - m_n)                   # correction for the jump
        p   = exp(s - m_n)                   # rescaled exponentials, this tile
        l   = l * c + sum(p)                 # rescale old denom, add new
        o   = o * c + p @ Vj                 # rescale old output, add new
        m   = m_n
    out = o / l                              # exact softmax(QK^T)V, O(T) memory

4 · FlashAttention: tile, accumulate in SRAM, never write S

FlashAttention (Dao et al., 2022; FlashAttention-2/3 since) is the online-softmax recurrence laid out as one fused CUDA/Triton kernel that is aware of the memory hierarchy. The dataflow:

That last move is the key trade in the systems half made local to one kernel: it is activation checkpointing for attention (lesson 5). Recomputing S in the backward costs extra FLOPs, but FLOPs are exactly what was free for this memory-bound op — so you spend the cheap resource to save the scarce one (HBM bytes and bandwidth). Because the score matrix never touches HBM, attention memory is O(T) and the HBM traffic collapses from the O(T²) of shuttling S to the O(T·d) of streaming K, V once.

The headline is not just "less memory" — it is that FlashAttention moves attention across the roofline. Naive attention is memory-bound; the fused, tiled version keeps the data on-chip and finally puts the tensor cores to work, so it is both more memory-frugal and 2–4× faster in wall-clock at long context. This is the difference between 2k-token toy models and the 128k-context models that ship.

5 · How you actually write these: the Triton tile model

You will not write FlashAttention in raw CUDA by hand in this course — you will write it in Triton, the Python-embedded kernel language whose whole abstraction is the tile. In Triton you write the program for one block: a @triton.jit function receives pointers and a program_id, computes which tile of the output it owns, loads input tiles with tl.load into what becomes SRAM/registers, does block-level matmuls and reductions (tl.dot, tl.max, tl.sum) on that hot data, and writes its tile back with tl.store. The compiler handles the threads-within-a-block; you reason at the granularity of tiles and HBM↔SRAM movement — exactly the level at which fusion and FlashAttention are designed.

@triton.jit
def flash_fwd(Q, K, V, O, L, sm_scale, T, d, BLOCK_M, BLOCK_N):
    pid = tl.program_id(0)                    # this block owns one Q tile
    q   = tl.load(Q + q_tile_offsets)         # (BLOCK_M, d) into SRAM
    m   = tl.full([BLOCK_M], -float('inf'))   # running max  (per query row)
    l   = tl.zeros([BLOCK_M])                 # running denom
    acc = tl.zeros([BLOCK_M, d])              # running output
    for j in range(0, T, BLOCK_N):            # stream K,V tiles
        k = tl.load(K + j_offsets)            # (BLOCK_N, d)
        v = tl.load(V + j_offsets)
        s = tl.dot(q, tl.trans(k)) * sm_scale # (BLOCK_M, BLOCK_N) in SRAM
        m_new = tl.maximum(m, tl.max(s, 1))
        c     = tl.exp(m - m_new)             # rescale factor
        p     = tl.exp(s - m_new[:, None])
        l     = l * c + tl.sum(p, 1)
        acc   = acc * c[:, None] + tl.dot(p, v)
        m     = m_new
    tl.store(O + o_offsets, acc / l[:, None]) # write O once; S never hit HBM
    tl.store(L + l_offsets, m + tl.log(l))    # save logsumexp for the backward

This is exactly CS336 Assignment 2: you implement fused kernels and a FlashAttention forward + backward in Triton, then benchmark them against the naive PyTorch version and watch the memory and time curves separate. For the full line-by-line derivation see the deeper treatments — triton_kernels · 12 FlashAttention for the kernel itself, triton_kernels · 06 reductions for the online-softmax/running-reduction pattern, and gpu_kernels · 12 the vLLM kernel stack for how fused attention and paged KV are wired into a production serving engine.

6 · Drive it: naive vs Flash memory

The widget below plots the one number that decides everything: HBM bytes for the attention score path, as you slide context length T (and batch B, heads h). Naive grows as O(T²) — it materializes S. FlashAttention grows as O(T) — it only ever writes O plus a couple of scalars per row. The dashed line is one 80 GB H100. Watch where each curve crosses it: naive OOMs at around ten thousand tokens (~12k), while Flash is still a rounding error at 100k. Toggle to relative HBM traffic to see the same gap in bandwidth, which is the speed story.

The shape of the two curves is the whole lesson. Below modest T they look comparable; then the naive parabola lifts off and slams into the 80 GB ceiling while the Flash line stays nearly on the floor. Long context is not a clever inference trick — it is a direct consequence of a kernel that refuses to write down an O(T²) intermediate.

Failure modes & checklist

Checkpoint

Where this points next

You have now made a single GPU efficient: fused epilogues keep memory-bound chains to one HBM round trip, and FlashAttention turns attention from an O(T²) memory hog into a compute-bound, long-context-capable kernel. The chip is no longer the thing wasting your budget. But efficiency on one device does not change the ledger from lesson 5: a 7B model's training state — bf16 params + fp32 master + Adam's two moments + grads — is ≈112 GB, and that overflows an 80 GB H100 before a single activation is stored, no matter how fast every kernel runs. Fusion shrank the activation term and made attention fit; it cannot shrink the optimizer state below one device. The model still does not fit. The only move left is to stop fitting it on one GPU and spread the work — parameters, gradients, optimizer state, and the batch — across many. That is lesson 8, Data parallelism — DDP → ZeRO → FSDP.

Interview prompts

• Why does kernel fusion help memory-bound ops but barely help a big GEMM? (§1 — for a memory-bound chain wall-clock ≈ bytes/bandwidth, and fusion cuts the round trips (≈6×→2×); a GEMM is already compute-bound, so its time is set by FLOPs and fusing the epilogue saves little traffic relative to the matmul.)
• What exactly does naive attention write to HBM, and how does it scale? (§2 — the full (B·h·T·T) score matrix S, written then reread to mask/softmax/multiply; it is O(T²) in both memory and traffic, so it OOMs at modest T even though the output O is only O(T).)
• Derive the online-softmax update and explain the rescale. (§3 — track running max m, denom ℓ, output o; on a new tile set m_new=max(m,max(s)), correction c=e^(m−m_new), then ℓ←ℓ·c+Σe^(s−m_new), o←o·c+Σe^(s−m_new)·v; the rescale corrects every earlier exponential when the max jumps so the final o/ℓ is the exact softmax.)
• How does FlashAttention get away with not storing S for the backward? (§4 — it saves only O and the per-row logsumexp (m, ℓ), O(T); in the backward it recomputes S tile-by-tile from Q, K. It spends cheap FLOPs (attention is memory-bound) to save scarce HBM bytes — activation checkpointing localized to one kernel.)
• Is FlashAttention an approximation? Is it faster, and why? (§4 — no, it computes exact softmax attention to float precision; it is 2–4× faster because it moves attention from memory-bound to compute-bound by keeping the scores in SRAM, so HBM traffic drops from O(T²) to O(T·d) and the tensor cores stop starving.)
• You implement FlashAttention and memory still scales as T². What went wrong? (§Failure modes — you stored S (or its softmax) for the backward instead of recomputing it, reintroducing the O(T²) HBM tensor; save only O and (m, ℓ) and regenerate S from Q, K.)
• Fusion and FlashAttention made one GPU efficient — why isn't that enough? (Where this points next — efficiency doesn't change lesson 5's ledger; a 7B's ≈112 GB training state overflows one 80 GB GPU regardless of kernel speed, which forces sharding across many GPUs in lesson 8.)