CS2680 Modern AI Systems: Agents and System Optimizations
Lecture 9 — Efficient LLM computing: GPU kernels

Wednesday opened the machine and left one number on the board: a ridge point of 295 FLOP/byte, against which almost everything a transformer does is memory-bound. Today we stop describing that fact and start beating it, on the one kernel where the win is largest and best documented. The claim is specific and it is arithmetic, not opinion: attention as normally written sits at about 64 FLOP/byte and is therefore capped at 22% of an H100's peak, and the only thing standing between it and the compute-bound regime is a matrix nobody needs to keep. Fuse the attention pipeline so the S×S score matrix never reaches HBM and intensity becomes S/2 — which crosses the ridge at 590 tokens, meaning every sequence you care about. By the end you should be able to compute a kernel's arithmetic intensity and place it in one of three regimes, derive FlashAttention's intensity and IO complexity from the tiling alone, explain why the backward pass recomputes the scores instead of storing them, and say precisely why none of this rescues decode.

Date: Monday, October 5, 2026 · 11:15am – 12:30pm · SEC LL2.221 · Instructor-led. Assignment 2 (build an agent) was due last night, Sun Oct 4, and Assignment 3 (optimize the agent) goes out today, due Sun Oct 25.

Required FlashAttention (Dao et al., 2022) — the main text, and the paper to read as accounting rather than as an algorithm. Read §2's background on the memory hierarchy, then §3.1 (tiling and the online-softmax rescaling) until you can write the update rule for the running maximum and the running denominator from memory, and §3.2's IO-complexity theorem until you can say what M is and why the bound has M in the denominator. Read the backward-pass discussion for the recomputation trade. Skim the block-sparse extension and the experiments; treat their multipliers as measurements on their shapes. Hold one question: the paper's title says IO-awareness, and its central result is a bound on HBM accesses — where in §3 does the FLOP count rise, and why is that the right trade?

Required Roofline (Williams, Waterman, and Patterson, 2009) — six pages, and Lecture 2 §2.12 already borrowed its result. Read it properly now, because today uses it as a decision procedure rather than a picture. Pay attention to what the original paper does that the popular version of it does not: it draws ceilings below the roof for optimizations you have not applied yet, which turns the diagram into a to-do list. Hold one question: what would the ceilings be for a transformer prefill on an H100?

Optional FlashAttention-2 — the same algorithm, rewritten for how a GPU actually distributes work; read the work-partitioning section and count what fraction of the gain comes from doing fewer non-matmul operations. FlashAttention-3 — Hopper-specific asynchrony: TMA, warp specialization, and overlapping softmax with matmul. Read for the shape of the argument, which is that on modern hardware the scarce resource is not bandwidth or arithmetic but overlap. Making deep learning go brrr — not a paper, and the clearest statement anywhere of §9.1's three regimes; read it first if the roofline has never quite clicked.

Where this sits

Wednesday (GPU programming, Sep 30) built the machine: SIMT execution, the memory hierarchy, occupancy and latency hiding, coalescing, tensor cores, and Triton as the language in which a kernel is a readable object. It ended on the observation that bytes are the budget. Lecture 2 §2.12 supplied the ridge point that turns that observation into a number, and Lecture 2 §2.3 supplied the specific scandal this lecture repays: the elementwise operations in a transformer block are 0.19% of its FLOPs and move fifteen times more bytes than the layer's weights.

Today is the payoff, and it is the last lecture in Part II that looks inside a single forward pass. From Oct 7 onward every lecture is about which requests run together and where their state lives; the kernel stops being the object of study and becomes a fixed cost that scheduling has to work around. That makes today the right place to establish the habit the rest of the semester depends on: before optimizing anything, compute its intensity and find out which of three walls you are actually standing against.

Standing assumptions, from Lecture 2 §2.12–§2.15. One H100 SXM: 3,350 GB/s, 989 TFLOP/s dense BF16, ridge point 295 FLOP/byte, 132 SMs with 228 KB of shared memory each. A100: 2,039 GB/s, 312 TFLOP/s, ridge 153, 108 SMs at 164 KB. Reference 7B: L = 32, d_model = 4096, 32 heads, d_head = 128, N = 6.74B, 13.5 GB of bf16 weights. Prefill at 50% of dense peak = 494.5 TFLOP/s. Kernel launch ≈5 µs; a forward-only pass over 32 layers issues about 490 launches. Where a number here disagrees with an older note, Lecture 2 is the one that is right.

Instructor notes — Timing plan

75-minute class (Mon/Wed 11:15am–12:30pm, SEC LL2.221). Instructor-led, two required readings, one of which is six pages. The lecture has exactly one centerpiece and everything else is scaffolding for it.

TimeSegmentNotes
0–7§9.1 Three regimesBoard the three walls before any arithmetic. The overhead wall is the one they have never been told about.
7–25§9.2 Attention's intensityProtected — never cut. Derive 64 FLOP/byte and then S/2 live. The 590-token crossover is the lecture's thesis.
25–32§9.3 Fusion and its ceilingShort. Fusion deletes bytes that did not need to move; it cannot delete bytes that did.
32–50§9.4 FlashAttentionProtected. Online softmax on the board, by hand, on two blocks. Then the IO bound.
50–57§9.5 What it buys, honestlyThe 68.7 GB that does not fit is the number they will remember. Say the 1.18× out loud too.
57–63§9.6 Backward and recomputationOne idea: recompute is cheaper than store when you are bandwidth-bound.
63–69§9.8 Decode is a different kernel1 FLOP/byte is unfixable. Then the 490-launch overhead and CUDA graphs.
69–75§9.9 Wrap and the mapPoint at Oct 7. End on "the kernel is now a fixed cost."

Reading-only, not scheduled: §9.7 in full (FA-2 and FA-3), and the split-K derivation in §9.8 — state the result, assign the arithmetic.

If running long: compress §9.6 to its one sentence and drop §9.7 entirely; the reading guide carries them. Never cut §9.2 or §9.4 — the intensity derivation and the online softmax are the class. If you must choose between them, keep §9.2: a student who can compute intensity can read FlashAttention alone, but a student who has memorized FlashAttention cannot analyze the next kernel.

Learning objectives

By the end of this class you should be able to:

  1. Classify a kernel as compute-bound, bandwidth-bound, or overhead-bound, and name the measurement that distinguishes the third from the second.
  2. Compute the arithmetic intensity of unfused attention and of fused attention from the tensor shapes alone, and find the sequence length at which fused attention becomes compute-bound.
  3. State what fusion can and cannot remove, and give an example of each in a transformer block.
  4. Write down the online-softmax rescaling rule and explain why it makes a single pass over the keys sufficient.
  5. Derive why FlashAttention's HBM traffic is O(S²d²/M) rather than O(S²), and say what M is on a specific GPU.
  6. Explain why the backward pass recomputes the score matrix rather than storing it, and price both choices.
  7. Compute the peak activation memory of unfused attention for a given S, and say at what context length it stops fitting on one card.
  8. Explain why fusion does not rescue batch-1 decode, and compute the launch overhead of a decode step against its memory floor.

9.1 Three regimes, not two

The roofline gives two regimes and most kernels live in a third that the diagram cannot show. Name all three now, because the first question to ask about any kernel is which one it is in, and the three have completely different fixes.

Compute-bound. Intensity is above the ridge point; the arithmetic units are the constraint. The fix is to do less arithmetic, or to do it in a cheaper precision, or to use a unit that does more of it per cycle — which is what tensor cores are. A big prefill GEMM lives here: Lecture 2 §2.14 put the 4K prefill at ≈4,100 FLOP/byte, fourteen times past the ridge.

Bandwidth-bound. Intensity is below the ridge; the memory system is the constraint and the arithmetic units idle. The fix is to move fewer bytes — fuse, tile, quantize, or restructure the algorithm so that data reaching fast memory is used more before it leaves. Batch-1 decode lives here at exactly 1 FLOP/byte, using 0.34% of the machine's arithmetic.

Overhead-bound. The kernel is so short that fixed per-launch costs dominate the work. Neither axis of the roofline describes this, because the roofline assumes the kernel is running. The fix is to launch fewer kernels — fusion again, but for a completely different reason — or to stop paying dispatch cost per launch at all, which is what CUDA graphs do. Lecture 2 §2.9 found this regime below about 75 tokens per step, which is precisely decode.

attainable throughput = min(peak FLOP/s, intensity × bandwidth) — and 0 if you never launch

The three are not mutually exclusive and the same optimization often addresses two of them, which is why "fusion" is a confusing word: fusing a normalization into a GEMM removes bytes and removes a launch, and in decode the second saving is the larger one. Keep the two effects separate in your accounting or you will attribute a win to the wrong cause and then fail to reproduce it.

One methodological note that is worth more than any single number here. Intensity is a ratio, so it is invariant to problem size — a huge kernel and a tiny kernel can have identical intensity and identical efficiency. That is what makes it the right diagnostic: it tells you what to change rather than how big the problem is.

Instructor notes

Minutes: 7. Board: Three walls in a row: ARITHMETIC / BANDWIDTH / DISPATCH, with the fix under each — "less math", "fewer bytes", "fewer launches". Leave it up all class. Ask the room: "Which wall is a batch-1 decode step standing against?" Most will say bandwidth. The honest answer is both bandwidth and dispatch, and §9.8 shows dispatch is 38% of the step. Expect confusion: Students conflate intensity with FLOP count. Say: "Intensity is a ratio. Doubling the problem does not change it."

9.2 Attention looks like arithmetic and behaves like memory

This is the section the lecture exists for. Take one attention head in prefill, sequence length S, d = d_head = 128, bf16 so b = 2 bytes. Count FLOPs and bytes independently and take the ratio.

The arithmetic is two matmuls. Scores QKᵀ is (S×d)·(d×S), costing 2S²d FLOPs. The weighted sum PV is (S×S)·(S×d), another 2S²d. The softmax itself is O(S²) elementwise work, negligible in FLOPs and — as we are about to see — not negligible at all in bytes.

attention FLOPs per head = 2`S`²`d` + 2`S`²`d` = 4`S`²`d` = 512·`S`² (at `d` = 128)

Now the bytes, written the way a naive implementation writes them — as a sequence of separate kernels each of which must hand its result to the next through HBM, because that is the only memory they share. Q, K, V come in once: 3·S·d·b = 768S bytes. Then the score matrix is written (S²·b), read by the softmax (S²·b), the probabilities are written (S²·b), and read by the PV matmul (S²·b). That is 4S²·b = 8S² bytes of traffic for a matrix that exists only to be consumed immediately.

Unfused attention: intensity, and the ceiling it implies

FLOPs: 512·S² Bytes: 8S² + 768S

For any S large enough that the quadratic term dominates — and 768S is under 5% of the total by S = 2,000 — the S² cancels:

intensity → 512S² ÷ 8S² = 64 FLOP/byte, independent of S

Against the H100's ridge of 295, attention is bandwidth-bound by a factor of 4.6, so its ceiling is 64 × 3,350 GB/s = 214 TFLOP/s — 22% of the machine's 989. On an A100: 64 × 2,039 = 130 TFLOP/s, 42% of 312, because the older card's ridge is lower.

Attention is not slow because attention is expensive. It is slow because a matrix nobody wants is being written to HBM and read back, twice.

Note the shape of that result, because it is the reason the fix is so clean. Intensity came out independent of S: making the sequence longer makes attention quadratically more expensive without making it one bit more efficient. And the 64 is not a property of attention as mathematics — it is a property of attention as a pipeline of separate kernels. The FLOPs are irreducible; the 8S² is an implementation choice.

So make the other choice. Suppose the score matrix never leaves the chip: it is produced in fast memory, consumed there, and discarded. Then the only HBM traffic is the inputs and the output.

Fused attention: the same FLOPs, a different denominator

Bytes: Q, K, V in and O out = 4·S·d·b = 1024·S

intensity = 512S² ÷ 1024S = S/2 FLOP/byte

Now it grows with S. Set it equal to the H100's ridge to find where attention becomes compute-bound:

S/2 = 295 → S = 590 tokens

On an A100, S/2 = 153 → S = 306.

Every sequence length this course cares about is past that crossover. At S = 4,096 fused attention sits at 2,048 FLOP/byte, seven times past the ridge; at S = 32,768 it is 16,384. The same computation, with the same FLOP count, moves from 22% of the machine to compute-bound — and it does so because of a decision about where an intermediate lives. This is the single most valuable worked example in Part II, and the reason is not the speedup. It is that the analysis found the optimization: nothing about attention had to be understood mathematically to see that 8S² was the problem.

Instructor notes

Minutes: 18. Protected — never cut. Board: Two columns, FLOPs and BYTES. Fill FLOPs first (4S²d, one line). Then fill bytes term by term, saying "written, read, written, read" out loud while writing 4 × S²b. Then divide and let the S² cancel in front of them. Write 64 and circle it. Then erase only the bytes column, write 1024S, divide again, and write S/2. Finally 590. Ask the room: before the second derivation — "Which of those two columns can we change?" Push until someone says the score matrix does not need to be stored. Do not supply it. Expect confusion: "Isn't the softmax the expensive part?" It is 0.02% of the FLOPs and roughly half the bytes. That inversion is the lecture. Common wrong answer: "So attention is O(S²) and that's the problem." The complexity is unchanged by everything today. What changes is the constant on the memory term, and that is worth 4.6×.

9.3 What fusion can remove, and what it cannot

Generalize the move, because it is not specific to attention. Two kernels that communicate through HBM can sometimes be rewritten as one kernel that communicates through registers or shared memory. Whenever that is possible, the intermediate's bytes disappear from the traffic entirely.

Lecture 2 §2.3 measured the opportunity for the elementwise operations in a transformer block: they are 0.19% of the block's FLOPs and they move about fifteen times more bytes than the layer's weights, because each one reads and writes a full activation tensor at an intensity of roughly 1 FLOP/byte. A normalization, a residual add, and an activation function are, from a bandwidth standpoint, three full passes over the activations to do almost no arithmetic. Fusing them into their neighbouring GEMMs removes essentially all of that traffic.

But be precise about the limit, because "fuse everything" is not a strategy.

Fusion removes bytes that did not need to move. An intermediate consumed exactly once by the next operation is pure overhead in HBM, and fusing it is free.

Fusion cannot remove bytes that must move. The weights must be read: 13.5 GB per forward pass over the reference 7B, and no fusion changes that. The KV cache must be read during decode. The inputs and outputs of the whole computation must cross the boundary. When those dominate, fusion buys nothing — which is why fusing kernels inside a big prefill GEMM is not where the win is, and why §9.8's decode step is immune.

Fusion is limited by what fits. The fused kernel's working set must live in registers and shared memory, and there is not much: 228 KB per SM on an H100. That is the whole reason FlashAttention is a paper rather than a compiler pass — the score matrix does not fit, so it cannot merely be kept in fast memory. It has to be produced, consumed, and thrown away in pieces, and the softmax stands in the way of doing that naively. §9.4 is how.

Instructor notes

Minutes: 7. Board: Three lines: "fuse: bytes that need not move" / "cannot fuse: bytes that must" / "limited by: 228 KB". Then Lecture 2's 0.19%-and-15× pair. Ask the room: "Name one byte in a decode step that fusion cannot remove." The weights. Then: "So what is left to optimize there?" Nothing in the kernel — which is why Oct 7 is about batching. Expect confusion: Students believe torch.compile makes this a solved problem. It handles the elementwise cases well and the attention case not at all, which is exactly why the attention kernel is hand-written in every serious stack.

9.4 FlashAttention: tiling, and the softmax problem it has to solve

The goal is set: compute attention without materializing S×S in HBM. Tiling is the obvious approach — split the keys and values into blocks, and for each block of queries walk the key blocks, accumulating the output. The obstacle is the softmax, and it is a real one: the softmax denominator is a sum over all keys, so a naive tiling cannot normalize anything until it has seen the last block, which means keeping every partial score — the matrix we were trying to avoid.

The resolution is online softmax: maintain a running maximum and a running denominator, and rescale the accumulated output whenever the maximum moves. Let block j have scores x_j, and carry three running quantities — m (max seen so far), (sum of exponentials, relative to m), and O (the accumulated weighted output).

`m_new` = max(`m`, max `x_j`) `ℓ_new` = `ℓ`·exp(`m` − `m_new`) + Σ exp(`x_j` − `m_new`) `O_new` = `O`·exp(`m` − `m_new`) + exp(`x_j` − `m_new`)·`V_j`

Each update is exact: the factor exp(mm_new) retroactively corrects everything accumulated under the old maximum. After the last block, divide O by once. The subtraction of the running maximum is not an optimization — it is what keeps exp() from overflowing, and it is why the algorithm is numerically better than the naive one rather than merely faster.

Online softmax by hand, two blocks

Scores for one query, in two blocks: x₁ = [1, 3], x₂ = [5, 2]. True softmax denominator is e¹+e³+e⁵+e² = 2.718 + 20.09 + 148.41 + 7.389 = 178.6.

Block 1: m = 3, = e^(1−3) + e^(3−3) = 0.135 + 1 = 1.135. Block 2: max x₂ = 5, so m_new = 5. ℓ_new = 1.135 · e^(3−5) + (e^(5−5) + e^(2−5)) = 1.135·0.135 + (1 + 0.0498) = 0.1536 + 1.0498 = 1.2034. Check: ℓ_new · e^m_new = 1.2034 · 148.41 = 178.6. Exact.

One pass over the keys, three scalars carried, no S×S matrix anywhere.

With the softmax handled, the IO accounting follows from the tiling. Let M be the size of fast memory available to the kernel — shared memory plus registers, so on the order of 228 KB per SM on an H100. Blocks are sized so that a K/V block and a Q block co-reside, which makes the block dimension O(M/d). Each of the O(S·d/M) query-block passes reads the whole K/V stream once, giving O(S·d) · O(S·d/M) = O(S²d²/M) HBM accesses, against O(S²) for the unfused version.

Read what the M in the denominator means: making fast memory larger makes this algorithm asymptotically better, which is not true of the unfused version, and is the sense in which the algorithm is IO-aware rather than merely fused. And note the direction of the trade — the tiling recomputes exponentials it could have stored, so the FLOP count goes up slightly while the byte count goes down by orders of magnitude. That is only a good trade below the ridge point, which is exactly where §9.2 found attention sitting. The optimization is licensed by the arithmetic.

Instructor notes

Minutes: 18. Protected. Board: Do the two-block example by hand with those exact numbers, and finish by multiplying by e^m to show 178.6. The moment the check lands is the moment the algorithm becomes obvious. Ask the room: "Why subtract the maximum at all?" Overflow. Many will not have met the trick. Expect confusion: Students think tiling alone is the contribution. Tiling is obvious; tiling through a softmax is the contribution. Say that explicitly. If short on time: Do the hand example and state the IO bound without deriving the block count.

9.5 What it buys, stated honestly in both currencies

There are two wins and they are wildly different in size. Take memory first, because it is the one that is categorical.

The matrix that does not fit

Unfused attention materializes S×S per head, in bf16, for one sequence.

S = 4,096: 4,096² · 2 B = 33.55 MB per head → × 32 heads = 1.07 GB for one layer. S = 32,768: 32,768² · 2 B = 2.147 GB per head → × 32 heads = 68.7 GB for one layer, for one sequence — against a 62.5 GB KV budget on an 80 GB card. It does not fit, at batch 1, with the weights already loaded.

FlashAttention's footprint is O(S): Q, K, V, O for a head at 32,768 tokens is 4 · 32,768 · 128 · 2 B = 33.5 MB, a factor of 64× smaller, and it falls further as S grows because the ratio is S/64.

That is the real result. Long-context attention is not slow without FlashAttention; it is impossible. Every 32K-context model you can run exists because this matrix stopped being materialized. Sequence-length limits that looked like modelling decisions were activation-memory limits.

Now the speed win, which is smaller than the folklore suggests and worth being honest about, because Amdahl applies to kernels too. Attention is only part of a prefill.

Prefill time at 4K and at 32K, reference 7B, one H100

Weight GEMMs at 50% of peak = 494.5 TFLOP/s. Unfused attention is capped at 214 TFLOP/s (§9.2); give fused attention the same 50% of peak as the GEMMs.

S = 4,096. Weights 2·6.74e9·4,096 = 55.2 TFLOP → 111.6 ms. Attention 4·32·32·128·4,096² = 8.8 TFLOP. Unfused: 8.8 ÷ 214 = 41.1 ms → total 152.7 ms Fused: 8.8 ÷ 494.5 = 17.8 ms → total 129.4 ms — a 1.18× end-to-end gain.

S = 32,768. Weights 2·6.74e9·32,768 = 441.7 TFLOP → 893 ms. Attention 4·32·32·128·32,768² = 563 TFLOP. Unfused: 563 ÷ 214 = 2,631 ms → total 3.52 s Fused: 563 ÷ 494.5 = 1,139 ms → total 2.03 s — a 1.73× gain.

So the kernel-level gain is 4.6× and the end-to-end gain is 1.18× at 4K and 1.73× at 32K, rising with S because attention's share rises with S. Both numbers are true; quoting the first as though it were the second is the most common way this result is misreported. The 129.4 ms figure is worth holding onto for another reason: it is where Lecture 2 §2.16's and Lecture 10 §10.5's 129 ms prefill for a 4K sequence comes from, so every break-even in Part II that trades against a prefill is implicitly assuming a fused attention kernel.

Instructor notes

Minutes: 7. Board: 68.7 GB against 62.5, big. Then the 1.18× and 1.73× pair beside the 4.6×. Ask the room: "The kernel got 4.6× faster and the prefill got 1.18× faster. Where did it go?" Amdahl — attention is 14% of the FLOPs at 4K. Expect confusion: The internet says "2–4× faster attention". That is a kernel measurement on particular shapes, not an end-to-end one. Both can be right.

9.6 The backward pass, and why recomputing beats storing

Training needs the score matrix again in the backward pass, to compute the gradients of Q, K, and V. The obvious approach is to save it in the forward pass — which reintroduces exactly the O(S²) memory the forward pass just eliminated. FlashAttention instead saves only the running statistics (m and , which are O(S)) and recomputes the score blocks during the backward pass.

Price it. Recomputing the scores costs one extra QKᵀ, so about 2S²d extra FLOPs — call it a 25% increase over the backward pass's own 4× forward cost for that operation. Storing them instead costs S²·b bytes written in the forward pass and read in the backward. At S = 4,096 for one head that is 33.55 MB each way, 67.1 MB of traffic, which at 3,350 GB/s is 20.0 µs; the recomputation is 2·4,096²·128 = 4.29 GFLOP, which at 494.5 TFLOP/s is 8.7 µs. Recomputing is 2.3× cheaper in time and infinitely cheaper in memory.

The general principle is worth extracting because it recurs: when you are bandwidth-bound, recomputation is a bandwidth optimization. Gradient checkpointing in Lecture 2 §2.7 is the same trade at the level of whole layers, and the fetch-versus-recompute decisions of Lecture 10 §10.5 and of Nov 9 are the same trade at the level of a KV cache. The question is always the ratio of the recompute's FLOPs to the transfer's bytes, measured against the machine's ridge.

Instructor notes

Minutes: 6. Board: "store: 67.1 MB = 20.0 µs" over "recompute: 4.29 GFLOP = 8.7 µs". One comparison. Ask the room: "When would storing win?" When the link is fast relative to the arithmetic — i.e. above the ridge, or when the tensor is small. That framing is what makes Nov 9's break-even obvious when they meet it. Expect confusion: Recomputation sounds wasteful. It is, in FLOPs, and FLOPs are the resource you have spare.

9.7 FlashAttention-2 and -3: what was left on the table

Reading-only; not scheduled in class.

FlashAttention-2 keeps the algorithm and rewrites the implementation around how a GPU actually schedules work. Three changes matter. The parallelism moves to the sequence dimension as well as batch and heads, so a single long sequence can occupy the whole card rather than a few SMs — which matters precisely in the long-context regime §9.5 says is the point. Work is partitioned between warps so that the shared-memory traffic between them nearly disappears. And the non-matmul operations — the rescalings — are reduced, which matters far more than their FLOP count suggests, because on an H100 a non-tensor-core FLOP is worth roughly an order of magnitude less throughput than a tensor-core FLOP. The lesson generalizes: once a kernel is compute-bound, the composition of its arithmetic starts to matter, not just the amount.

FlashAttention-3 targets Hopper specifically, and its argument is that the scarce resource has moved again. With tensor cores this fast, the goal is to never let them wait: the Tensor Memory Accelerator moves tiles asynchronously so loads overlap with math; warp specialization dedicates some warps to data movement and others to computation, turning the kernel into a software pipeline; and the softmax of one block is overlapped with the matmul of the next. It also exploits fp8, where the interesting problem becomes accuracy rather than scheduling.

The trajectory is the useful thing to notice. FlashAttention-1 was an algorithmic insight about IO. FlashAttention-2 was a work-partitioning fix. FlashAttention-3 is a latency-hiding fix. Each generation solved the bottleneck the previous one exposed — which is the same pattern as the lectures on either side of today, and the reason that "optimize the kernel" is never finished, only handed forward.

9.8 Decode attention is a different kernel, and fusion cannot save it

Everything above is about prefill, where there are S tokens to amortize a pass over the weights. In decode there is one token per sequence, and every conclusion changes.

The attention a decode step performs is one query vector against the whole KV cache. FLOPs: for one sequence at context length S, the two matmuls become 2·S·d + 2·S·d = 4S·d. Bytes: the cache itself, 2·S·d·b. The ratio is 4S·d ÷ 2S·d·b = 2/b = exactly 1 FLOP/byte at bf16 — Lecture 2 §2.15's result, and now visibly a property of the shape rather than of any implementation.

There is no intermediate to fuse away. The bytes being moved are the KV cache, and they are the input to the computation, not a temporary. Fusion removes temporaries; here there are none. This is the sharpest illustration of §9.3's limit, and it is why decode is attacked by every lecture from Oct 7 onward but never by a better kernel: you cannot fuse your way out of reading your own state. Batching does not rescue it either, because no two sequences share a cache, so the cache traffic grows exactly as fast as the arithmetic — which is why Lecture 2 §2.15's fully-batched step still reaches only ~24 FLOP/byte at GQA-8, and 6 at MHA, against a ridge of 295.

What a kernel can still do in decode is two things, and both are about parallelism rather than bytes. Split-K (FlashDecoding). At small batch there is not enough work to fill 132 SMs: one sequence's decode attention is a handful of thin matmuls. Splitting the KV cache along the sequence dimension lets many SMs each reduce over a slice, followed by a combine step that merges the partial softmax statistics — the same m/ rescaling as §9.4, used for a completely different purpose. This does not improve intensity at all; it improves occupancy, converting an underfilled machine into a full one. Paged gathers. Lecture 10 §10.3 handed attention a block table, and the kernel now has to gather through it at close to contiguous speed; FlashInfer exists because that turned out to be a real engineering problem rather than an indexing detail.

And then the third wall. A forward-only eager pass over 32 layers issues about 490 kernel launches at ≈5 µs each — 2.45 ms of dispatch — against a batch-1 decode floor of 4.0 ms.

The dispatch wall in decode

Naive eager step: 4.0 ms of memory time + 2.45 ms of dispatch ≈ 6.45 ms, so 38% of the step is the CPU enqueuing work. Fuse or graph-capture the launch count down by 3×, to ~163: dispatch falls to 0.82 ms and the step to 4.82 ms, a 25% cut — within 20% of the floor.

Nothing goes below 4.0 ms: the 13.5 GB of weights is read once per token regardless. Dispatch is the meter you can drive to zero; bytes are the meter you cannot.

CUDA graphs are the standard answer — record the launch sequence once, replay it as a single submission — and they are the reason every production engine graph-captures its decode step. Note what this does to the three regimes of §9.1: decode starts overhead-bound and bandwidth-bound, and removing the overhead does not make it fast, it makes it honestly bandwidth-bound at last.

Instructor notes

Minutes: 6. Board: "1 FLOP/byte — the bytes ARE the input" then the 4.0 / 2.45 / 6.45 stack. Ask the room: "What would you fuse in a decode attention kernel?" Let them search for a temporary and fail to find one. That failure is the lesson. Expect confusion: "FlashAttention speeds up decode." It is used in decode (as split-K/paged variants) for occupancy and memory layout, not for intensity.

9.9 What a kernel cannot fix

Today bought a factor of 4.6 on one kernel, a categorical memory win that made long context possible, and a 25% cut in a decode step's dispatch. Put the rest of the semester against what is left, because every remaining lecture attacks a quantity no kernel can touch.

Still broken after todayThe quantityWhere it goes
One 129 ms prefill lands on a 22.6 ms decode stepinterference between phasesBatching Oct 7, disaggregation Oct 14
The batch is chosen per request, wasting every batch's taileffective B over timeBatching and scheduling I, Oct 7
13.5 GB of weights re-read every step, at b = 2bytes per weightPruning and quantization, Nov 11
The KV cache is the decode step's input and it is hugebytes per cached tokenKV-cache optimization, Oct 28
A cache hit on the wrong replica is a cache missplacementRouting and load balancing, Oct 19
Identical prefixes are prefilled twicethe 92% re-send of Lecture 5Prefix cache, Nov 9
One pass over the weights yields one tokentokens per passSpeculative decoding, Nov 18
The engine cannot see the agent's structuredependence between requestsAgent serving, Nov 18, Nov 23

The habit to carry forward is §9.1's, not FlashAttention's. Compute the intensity, find the wall, then choose the tool. Everyone in the room can now look at a kernel and say which of three things is wrong with it, which is a more durable skill than knowing what one famous kernel does — and it is the skill Nov 18 will ask whether a language model has.

Discussion seeds

  1. The ceilings you were not shown. The roofline paper draws intermediate ceilings for optimizations not yet applied. Draw the roofline for a reference-7B prefill on an H100 and add ceilings for: no fusion, no tensor cores, and no tiling. Which ceiling is the binding one at S = 512, and at S = 32,768?
  2. Where the analysis fails. §9.2's intensity argument assumes bytes and FLOPs are the only currencies. Name a kernel that is fast by both measures and still slow, and say what the third currency is.
  3. M in the denominator. FlashAttention's IO bound improves as fast memory grows. Hopper gave 228 KB per SM against Ampere's 164. What does that predict for the speedup, and does the prediction hold in the FA-3 numbers?
  4. The trade you would not make. Recomputation beat storage by 2.3× in §9.6 on an H100. Construct a machine on which storing wins, and say whether such a machine is plausible.
  5. Decode's floor, one more time. Given §9.8, is there any kernel-level change that improves a batch-1 decode step by more than 40%? Defend your answer with intensity, not intuition.
  6. What you would ask an agent to write. Nov 18 asks whether a model can generate kernels. Which of today's four ideas — tiling, online softmax, recomputation, split-K — would you expect a model to find, and which would you expect it to miss? What does your answer imply about how to design the benchmark?

Key takeaways

  • There are three walls, not two: arithmetic, bandwidth, and dispatch. The roofline shows the first two and cannot show the third, which is the one decode hits first.
  • Unfused attention does 4S²d FLOPs and moves ~8S² bytes, so its intensity is 64 FLOP/byte independent of S — capping it at 214 TFLOP/s, 22% of an H100. The S² cancels: longer sequences are quadratically more expensive and not one bit more efficient.
  • Keep the score matrix off HBM and the bytes become 4S·d·b, so intensity is S/2 and crosses the 295 ridge at S = 590 tokens. The same FLOPs, a different denominator, and every sequence in this course is past the crossover.
  • Tiling through a softmax requires online softmax: carry a running max and denominator, and rescale the accumulator by exp(mm_new) when the max moves. Exact, one pass, three scalars. HBM traffic drops from O(S²) to O(S²d²/M) — an algorithm that gets better as fast memory grows.
  • The memory win is the categorical one: unfused attention needs 68.7 GB for one 32K sequence's layer against a 62.5 GB budget, and FlashAttention needs 33.5 MB — 64× less. Long context was an activation-memory problem masquerading as a modelling limit.
  • The speed win is real but obeys Amdahl: 4.6× on the kernel, 1.18× on a 4K prefill, 1.73× at 32K. Quoting the kernel number as the end-to-end number is the standard misreport.
  • Recomputing the scores in the backward pass beats storing them 2.3× in time and entirely in memory. When you are bandwidth-bound, recomputation is a bandwidth optimization — the same trade as gradient checkpointing and as Nov 9's fetch-versus-recompute.
  • Fusion cannot save decode. At 1 FLOP/byte the bytes are the KV cache, which is the input, not a temporary. What a kernel can still do is occupancy (split-K) and dispatch: 490 launches at 5 µs is 2.45 ms against a 4.0 ms floor, so graph capture cuts a step 25% and then stops.

Numbers worth memorizing

QuantityValueSource
Unfused attention intensity64 FLOP/byte, independent of S512S² ÷ 8S²; §9.2
Its ceiling on an H100 / A100214 TFLOP/s (22%) / 130 TFLOP/s (42%)64 × bandwidth
Fused attention intensityS/2 FLOP/byte512S² ÷ 1024S
Compute-bound crossover, fusedS = 590 (H100), 306 (A100)S/2 = ridge
Score matrix, one head33.55 MB at 4K; 2.147 GB at 32KS²·2 B
One layer, 32 heads, 32K, one sequence68.7 GB — does not fit 62.5 GB× 32
FlashAttention footprint at 32K33.5 MB — 64× smallerS·d·b
FlashAttention HBM trafficO(S²d²/M), M ≈ 228 KB/SM (H100)§3.2 of the paper
Prefill gain, 4K / 32K1.18× / 1.73× end-to-endAmdahl on 8.8 of 64.0 TFLOP
4K prefill, fused, at 50% peak129 ms — the figure all of Part II trades against64.0 TFLOP ÷ 494.5
Backward: recompute vs store8.7 µs vs 20.0 µs (2.3×)4.29 GFLOP vs 67.1 MB
Decode attention intensity1 FLOP/byte = 2/b, unfixable by fusion4S·d ÷ 2S·d·b
Decode dispatch overhead490 launches × 5 µs = 2.45 ms on a 4.0 ms floor (38%)Lecture 2 §2.9

Self-check

  1. Unfused attention's intensity came out independent of S. Why, and what does that imply for long context?Both FLOPs (4S²d) and the dominant byte term (~8S²) are quadratic in S, so the ratio cancels to 4d/(4b) = 64 at d = 128, bf16. The implication is bleak and precise: doubling the context quadruples attention's cost while leaving its efficiency at 22% of the machine. You cannot grow out of the problem, which is why it had to be solved rather than out-scaled.
  2. Fused attention's intensity is S/2. Where does the 2 come from, and why is the answer independent of d?Bytes are 4·S·d·b (three inputs and one output) and FLOPs are 4S²d, so intensity = 4S²d ÷ 4S·d·b = S/b = S/2 at bf16. The d cancels because both the arithmetic and the traffic are linear in the head dimension — which means the crossover at 590 tokens holds for any head width, and halving b to fp8 halves the crossover to 295 tokens.
  3. Why can't you simply keep the S×S score matrix in shared memory instead of tiling?Because it does not fit. At S = 4,096 one head's score matrix is 33.55 MB against 228 KB of shared memory per SM — a factor of 147. Fusion by co-residency works only when the intermediate fits; attention needs the harder thing, which is to produce and consume the intermediate in pieces. The softmax's global denominator is what makes that hard, and online softmax is what makes it possible.
  4. What exactly does the online-softmax rescaling correct, and why is the result exact rather than approximate?When a later block contains a larger score, every exponential accumulated under the old maximum was computed with the wrong offset. Multiplying the running denominator and the running output by exp(m_oldm_new) applies exactly the missing factor, since exp(xm_old)·exp(m_oldm_new) = exp(xm_new). It is an algebraic identity, so no accuracy is lost — and subtracting a maximum at all is what prevents overflow, so the tiled version is numerically better than the naive one.
  5. FlashAttention's IO bound is O(S²d²/M). What is M, and what does its presence in the denominator tell you?M is the size of the fast memory the kernel tiles into — shared memory plus registers, of order 228 KB per SM on an H100. Its presence in the denominator means the algorithm's asymptotic traffic improves as on-chip memory grows, which is not true of the unfused version (stuck at O(S²) regardless). That is the precise content of "IO-aware": the algorithm is parameterized by the memory hierarchy rather than merely respecting it.
  6. The kernel got 4.6× faster and the 4K prefill got 1.18× faster. Reconcile those.Amdahl. At S = 4,096 attention is 8.8 TFLOP of a 64.0 TFLOP prefill — 14% of the work — so even making it free would cap the gain at about 1.16×; the measured 1.18× reflects attention going from 41.1 ms to 17.8 ms against an unchanged 111.6 ms of weight GEMMs. The kernel ratio and the end-to-end ratio are different quantities, and attention's share grows with S, which is why 32K gives 1.73×.
  7. Why does fusion do nothing for a batch-1 decode step, and what does help?Because there is no temporary to eliminate: the bytes being moved are the weights and the KV cache, which are inputs rather than intermediates, so intensity is pinned at 1 FLOP/byte = 2/b. What helps is not intensity but occupancy and overhead — split-K/FlashDecoding to fill 132 SMs from one thin sequence, and graph capture to remove the 2.45 ms of dispatch from a 6.45 ms step. Both leave the 4.0 ms memory floor untouched, which is why Oct 7 onward attacks bytes and batching instead.

Exercises

  1. Move the crossover. Recompute §9.2's fused-attention crossover for fp8 (b = 1) and for fp32 (b = 4), on both an H100 and an A100. Then state which of the four numbers changes the engineering conclusion and why. Solution sketch: Fused intensity is S/b, so the crossover is S = b · ridge. H100 (ridge 295): fp8 295, bf16 590, fp32 1,180. A100 (ridge 153): fp8 153, bf16 306, fp32 612. None of them changes the conclusion, and that is the point: every crossover is in the hundreds of tokens while every sequence of interest is in the thousands, so fused attention is compute-bound in all eight combinations. The number that would matter is the fp8 H100 figure if you served very short sequences — a 128-token classification workload at fp8 sits at intensity 128, below the 295 ridge, and is bandwidth-bound even when fused. Short-sequence serving is a different regime, which is worth knowing before you reuse a long-context kernel there.
  2. Price the unfused ceiling properly. §9.2 assumed 4 S²·b bytes of score traffic. Suppose a moderately good implementation fuses the softmax with the second matmul, so the probabilities are never written or read. Recompute the byte count, the intensity, the H100 ceiling, and the 4K prefill time. What fraction of FlashAttention's win does this cheap fix capture? Solution sketch: Traffic falls from 4S²·b to 2S²·b = 4S² bytes (write scores, read scores). Intensity = 512S² ÷ 4S² = 128 FLOP/byte, still below 295, so still bandwidth-bound; ceiling 128 × 3,350 = 429 TFLOP/s, 43% of peak. Attention time at 4K: 8.8 ÷ 429 = 20.5 ms, so total prefill 111.6 + 20.5 = 132.1 ms against 152.7 unfused and 129.4 fully fused. The cheap fix captures (152.7 − 132.1)/(152.7 − 129.4) = 88% of the available end-to-end gain. The moral is uncomfortable and worth sitting with: most of FlashAttention's time win at moderate S is available from ordinary fusion, and what the full algorithm uniquely delivers is the O(S) memory — which is why §9.5 leads with 68.7 GB rather than with a speedup.
  3. When does attention dominate? Find the sequence length at which attention's FLOPs equal the weight GEMMs' FLOPs for the reference 7B, and compare it with Lecture 2 §2.5's crossover of 24,704. Then compute, at that length, the end-to-end prefill gain from fusing attention. Solution sketch: Weight GEMMs cost 2N·S = 13.48e9·S; attention costs 4·32·32·128·S² = 524,288·S². Equal when S = 13.48e9 ÷ 524,288 = 25,713 — within 4% of Lecture 2's 24,704, the difference being that §2.5 compared per-layer parameters against the per-layer attention term while this compares the totals including embeddings. At S = 25,713: weights 346.6 TFLOP → 701 ms; attention 346.6 TFLOP → unfused 1,620 ms, fused 701 ms; totals 2.32 s vs 1.40 s = 1.66×. So even at the crossover the end-to-end gain is under 1.7×, because fusing a term that is half the work and making it 2.3× faster is bounded by 1.75× — a useful sanity check on any claim above 2× for a prefill.
  4. The dispatch budget. A decode step for the reference 7B issues ~490 launches at 5 µs against a 4.0 ms floor. Compute the step time and dispatch fraction for: (a) eager, (b) fully graph-captured (treat dispatch as one 10 µs submission), (c) eager after quantizing weights to int4. Then say which of the three is the most misleading benchmark to publish. Solution sketch: (a) 4.0 + 2.45 = 6.45 ms, dispatch 38%. (b) 4.0 + 0.01 = 4.01 ms, dispatch 0.2% — a 1.61× gain from nothing but submission mechanics. (c) int4 quarters the weight bytes so the floor falls to 1.0 ms, but dispatch is unchanged: 1.0 + 2.45 = 3.45 ms, dispatch 71%. The most misleading benchmark is (c): quantization was supposed to give 4× and delivered 6.45/3.45 = 1.87×, and a reader would blame the quantization when the fault is that the harness never graph-captured. Combining (b) and (c) gives 1.01 ms, the honest 6.4×. The lesson is that overhead-bound regimes make every other optimization look broken, which is why §9.1 insists on identifying the wall before choosing the tool.
  5. Size a long-context deployment by activation memory. Ignore the KV cache entirely and ask only what unfused attention's score matrices cost. For the reference 7B in bf16, compute the largest S that fits in 62.5 GB at batch 1 if (a) all 32 layers' score matrices are live at once, (b) only one layer's are live, (c) FlashAttention is used. State which assumption real frameworks make and what that implies. Solution sketch: Per layer, all heads: 32 · S² · 2 B = 64S² bytes. (a) All 32 layers: 2,048·S² ≤ 62.5e9 → S5,524. (b) One layer: 64S² ≤ 62.5e9 → S31,250. (c) FlashAttention: 4·S·d·b·32 heads = 32,768·S bytes per layer → S1.9e6, i.e. not the binding constraint by three orders of magnitude. Real frameworks are case (b) — autograd frees a layer's scores once the next layer has consumed them in the forward pass, but must retain them for the backward pass during training, which pushes training back toward case (a). So unfused attention caps inference context at ~31K and training context at ~5.5K on one card, and both ceilings are activation-memory ceilings that look like model limits. That is the concrete sense in which FlashAttention did not speed up long context so much as create it.

Reading guide

Required — FlashAttention. Read §2 for the memory-hierarchy framing, then §3.1 slowly: the tiling is easy and the online-softmax rescaling is the whole paper, so do not move on until you can reproduce §9.4's two-block hand calculation without looking. Then §3.2's IO-complexity theorem, where the thing to extract is not the proof but the form of the bound — O(S²d²/M) against O(S²) — and what it says about a kernel whose asymptotics depend on the machine. Read the backward-pass subsection for the recomputation argument and check it against §9.6's 2.3×. Skim the block-sparse extension entirely, and read the experiments with §9.5 in hand: find which of their reported multipliers are kernel-level and which are end-to-end, because the paper is careful about the distinction and most citations of it are not. Hold this question: the paper's contribution is a reduction in HBM accesses at the cost of extra FLOPs — write down the ratio of the two for the reference 7B at S = 4,096 and say why an H100 is happy to make that trade.

Required — Roofline. Six pages, and read them properly rather than trusting the version of the model that has reached folklore. Two things to take. First, the ceilings: the original draws horizontal and diagonal lines below the roof for optimizations not yet applied — no ILP, no SIMD, no memory-affinity — which turns the diagram from a diagnosis into an ordered work list, and discussion seed 1 asks you to draw the transformer version. Second, the insistence that intensity is measured at the DRAM boundary, which is exactly what makes §9.2's question ("does the score matrix reach HBM?") the decisive one. Skip the specific multicore machines. Hold this question: the paper predates GPUs as compute devices entirely — what about it survived, and what would you add for a machine whose ridge point is 295?

Optional — FlashAttention-2. Read the work-partitioning section and try to attribute its gain between the three changes (sequence-dimension parallelism, inter-warp traffic, fewer non-matmul ops). Question: which of the three would have been unnecessary if the first paper had targeted long sequences from the start?

Optional — FlashAttention-3. Read for the argument that the scarce resource is overlap. Question: warp specialization makes the kernel a software pipeline — what does that imply about how much of a modern kernel's performance is scheduling rather than arithmetic, and where else in this course have you seen that claim?

Optional — Making deep learning go brrr. The clearest statement of §9.1's three regimes anywhere, written for exactly this audience. Read it first if the roofline has never clicked. Question: it argues most practitioners misattribute overhead-bound behavior to bandwidth — how would you tell the two apart with one measurement?

Looking ahead

Wednesday (batching and scheduling I, Oct 7) leaves the inside of the forward pass for good. From there the kernel is a fixed cost and the questions are about which requests run together: Orca's iteration-level scheduling re-chooses the batch every step rather than every request, which is only practical because Lecture 10's allocator made adding and removing a sequence a table edit. It will also inherit today's unfinished business — the 129 ms prefill computed in §9.5 is exactly the thing that lands on a 22.6 ms decode step and spikes it 6.7×, and Wednesday (disaggregation, Oct 14), which is a student-led discussion, is the two answers to that.

Assignment 3 (make the agent faster) goes out today, due Sun Oct 25, and it is still outside the API — so nothing today is a lever you can pull on it yet. Assignment 4 (serve your own agent) goes out Nov 11, and there today's material becomes operational: you will choose an attention backend, and now you know what the choice is about.

One thing to carry out of the room. The 4.6× was not found by understanding attention better than anyone else. It was found by writing down FLOPs in one column and bytes in the other and noticing that the second column contained a matrix nobody wanted. Do that for every kernel, every cache, and every request stream for the rest of the semester — it is the same move Lecture 10 made on the allocator, and it is the move Nov 9 will make on the whole prompt.