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.
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.
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.
| Time | Segment | Notes |
|---|---|---|
| 0–7 | §9.1 Three regimes | Board the three walls before any arithmetic. The overhead wall is the one they have never been told about. |
| 7–25 | §9.2 Attention's intensity | Protected — 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 ceiling | Short. Fusion deletes bytes that did not need to move; it cannot delete bytes that did. |
| 32–50 | §9.4 FlashAttention | Protected. Online softmax on the board, by hand, on two blocks. Then the IO bound. |
| 50–57 | §9.5 What it buys, honestly | The 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 recomputation | One idea: recompute is cheaper than store when you are bandwidth-bound. |
| 63–69 | §9.8 Decode is a different kernel | 1 FLOP/byte is unfixable. Then the 490-launch overhead and CUDA graphs. |
| 69–75 | §9.9 Wrap and the map | Point 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.
By the end of this class you should be able to:
S²d²/M) rather than O(S²), and say what M is on a specific GPU.S, and say at what context length it stops fitting on one card.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.
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.
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."
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.
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.
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.
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.
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×.
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.
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.
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).
Each update is exact: the factor exp(m − m_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.
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.
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.
There are two wins and they are wildly different in size. Take memory first, because it is the one that is categorical.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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 today | The quantity | Where it goes |
|---|---|---|
| One 129 ms prefill lands on a 22.6 ms decode step | interference between phases | Batching Oct 7, disaggregation Oct 14 |
| The batch is chosen per request, wasting every batch's tail | effective B over time | Batching and scheduling I, Oct 7 |
13.5 GB of weights re-read every step, at b = 2 | bytes per weight | Pruning and quantization, Nov 11 |
| The KV cache is the decode step's input and it is huge | bytes per cached token | KV-cache optimization, Oct 28 |
| A cache hit on the wrong replica is a cache miss | placement | Routing and load balancing, Oct 19 |
| Identical prefixes are prefilled twice | the 92% re-send of Lecture 5 | Prefix cache, Nov 9 |
| One pass over the weights yields one token | tokens per pass | Speculative decoding, Nov 18 |
| The engine cannot see the agent's structure | dependence between requests | Agent 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.
S = 512, and at S = 32,768?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?S²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.S·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.m − m_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.| Quantity | Value | Source |
|---|---|---|
| Unfused attention intensity | 64 FLOP/byte, independent of S | 512S² ÷ 8S²; §9.2 |
| Its ceiling on an H100 / A100 | 214 TFLOP/s (22%) / 130 TFLOP/s (42%) | 64 × bandwidth |
| Fused attention intensity | S/2 FLOP/byte | 512S² ÷ 1024S |
| Compute-bound crossover, fused | S = 590 (H100), 306 (A100) | S/2 = ridge |
| Score matrix, one head | 33.55 MB at 4K; 2.147 GB at 32K | S²·2 B |
| One layer, 32 heads, 32K, one sequence | 68.7 GB — does not fit 62.5 GB | × 32 |
| FlashAttention footprint at 32K | 33.5 MB — 64× smaller | 4·S·d·b |
| FlashAttention HBM traffic | O(S²d²/M), M ≈ 228 KB/SM (H100) | §3.2 of the paper |
| Prefill gain, 4K / 32K | 1.18× / 1.73× end-to-end | Amdahl on 8.8 of 64.0 TFLOP |
| 4K prefill, fused, at 50% peak | 129 ms — the figure all of Part II trades against | 64.0 TFLOP ÷ 494.5 |
| Backward: recompute vs store | 8.7 µs vs 20.0 µs (2.3×) | 4.29 GFLOP vs 67.1 MB |
| Decode attention intensity | 1 FLOP/byte = 2/b, unfixable by fusion | 4S·d ÷ 2S·d·b |
| Decode dispatch overhead | 490 launches × 5 µs = 2.45 ms on a 4.0 ms floor (38%) | Lecture 2 §2.9 |
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.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.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.m_old − m_new) applies exactly the missing factor, since exp(x − m_old)·exp(m_old − m_new) = exp(x − m_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.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.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×.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.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.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.N·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.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 → S ≤ 5,524. (b) One layer: 64S² ≤ 62.5e9 → S ≤ 31,250. (c) FlashAttention: 4·S·d·b·32 heads = 32,768·S bytes per layer → S ≤ 1.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.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?
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.