FlashAttention and Optimized Kernels

I wrote exact block-wise attention in plain PyTorch, the algorithm that FlashAttention is built on. It gave the same answer as ordinary attention to within rounding. At 8,192 tokens it needed 16 MiB of extra memory where ordinary attention needed 3,148. And it was sixteen to eighteen times slower.

The version of the same algorithm that PyTorch ships as a fused kernel, running on the same inputs, needed 12 MiB and, depending on length and run, was between 5.5 and 13.5 times faster than ordinary attention.

Both have the same algorithm, the same arithmetic, the same exactness. One saved memory and lost time, the other saved both. The difference is what the chapter is about: block-wise attention is an algorithm, and making it fast is a property of how it is executed. If you read only the abstract of a paper, you can come away thinking the first is the contribution. It is a necessary part, and the second half is where the time goes.

There is one more thing to say at the start. I could not run FlashAttention. Its backend is not available on this machine, and the reason is not what I wrote in Chapter 8. Everything below about “optimised attention kernels” is measured on a close relative, and the chapter says so each time it matters.

Where this sits

Chapter 8 measured that the direct implementation of attention writes and reads a \(T \times T\) score tensor several times, and that the memory-efficient kernel avoids it. It estimated the traffic by hand. Chapter 11 provided the model for saying which of launch, bytes and arithmetic bounds an operation, and showed that fusing kernels reduces the byte term by the number of memory passes it removes. This chapter puts the specific algorithm under that model.

The Inference Lab piece added here is the reference implementation of block-wise attention, tiled_attention, together with the exact merge rule merge_partials, an analytic traffic model, and an availability probe that asks PyTorch why a backend was refused. The reference is slow on purpose. Its job is to be the oracle against which any fused kernel the Lab ever runs can be checked, in the same way that Chapter 4’s hand-written forward pass is the oracle for whole models.

For the Pi series, this is the kernel that makes long contexts affordable, which matters to every volume that stuffs retrieved text, memory or tool output into a prompt. It is also the first of several places where the book shows that a speed-up needs a kernel and a piece of hardware, and not only an idea.

The online softmax

The obstacle to computing attention a block at a time is the softmax. It needs the maximum and the sum over all the scores in a row before any of them can be normalised, and a block sees only some of the keys.

The remedy is to keep running statistics and correct them as new blocks arrive. For one query row and a set of keys \(J\), define a partial result by three quantities:

Chapter 3’s stable softmax subtracts the maximum before exponentiating. This is the same device made incremental: the running maximum plays the role of the global maximum, and the correction factor \(e^{m_{\text{old}} - m_{\text{new}}}\) repairs the sums computed under the old one. A direct consequence is that the algorithm does not overflow on large scores, as the tests confirm with queries scaled by 40.

The algorithm, for each tile of \(B_q\) queries:

m, l, acc = -inf, 0, 0
for each tile of B_k keys (skip tiles wholly in the future if causal):
    s   = q_tile @ k_tile^T / sqrt(d)           # [B_q, B_k], never written to main memory in a fused kernel
    m'  = max(m, rowmax(s))
    p   = exp(s - m')
    l   = l * exp(m - m') + rowsum(p)
    acc = acc * exp(m - m') + p @ v_tile
    m   = m'
output_tile = acc / l

The score tile is only \(B_q \times B_k\) and exists for one iteration.

It is exact

I compared the block-wise implementation with ordinary attention computed in float64, for 12 heads of dimension 64, at 512 and 2,048 tokens, for four block shapes, in both precisions.

The block shape makes no measurable difference. In fp16 the error is the rounding of the 16-bit output, the same level the library kernels showed in Chapter 8, and it is identical across block shapes. With queries scaled by 30 (scores in the hundreds), the fp32 error stayed at \(4.7 \times 10^{-5}\).

The tiles also let a causal kernel skip work. For blocks of 128 at 2,048 tokens, 53 per cent of the block pairs had to be computed; with blocks of 64, 52 per cent. Whole blocks that lie above the diagonal contribute nothing and are never visited. That is the saving that Chapter 8’s efficient kernel showed (causal time 0.58 to 0.71 of full), which a plain implementation can only get by writing the skip explicitly.

What it saves: traffic

The reason to tile is to keep the score tile in fast on-chip memory. The main-memory traffic can be counted per head and per sequence. The direct implementation moves, in elements,

The tile size is the lever. Doubling \(B_q\) halves the \(K\) and \(V\) traffic, and \(B_q\) is limited by how much a processor can hold on chip at once: the \(Q\) tile, the \(K\) and \(V\) tiles, the score tile and the accumulator. That is what “IO-aware” means in the title of the FlashAttention paper: the tile sizes are derived from the size of fast memory. I do not know the tile sizes the kernel I measured uses, and PyTorch does not report the capacity of this GPU’s on-chip memory, so the table is a model with \(B_q\) as a free parameter.

What it saves: the numbers

Three implementations at 1,024 to 8,192 tokens, fp16, 12 heads. Time is the median of ten runs (four for the unfused one), measured twice; peak memory is extra allocation over the inputs.

Both ways of avoiding the score tensor avoid its memory. The plain-PyTorch version is 13 to 23 times slower than the naive implementation it was supposed to improve. The fused kernel is 7.8 to 13.5 times faster than naive at 1,024 to 4,096 tokens, and 5.5 to 8.6 times at 8,192, where the two runs disagreed (7.8 and 13.0 ms for the fused kernel), so I give the range and not a value.

Why is the unfused version slow? Every block pair is still a sequence of separate operations: two matrix products, a scale, a mask, a maximum, two exponentials, several multiplications and sums. Each is a separate kernel launch, and each reads its inputs from and writes its results to main memory. The tile is small, but it still makes the round trip. At 8,192 tokens with blocks of 128 there are 2,080 block pairs; at about fourteen operations each that is some 29,000 launches, and at 40 microseconds per launch (the order of magnitude of Chapter 7’s floor, on a busier day) that alone is 1.2 seconds, which is where the measurement sits. This is an estimate from counts and a floor measured on another day, so I take it as the right order of magnitude and not as an accounting.

The point is not that PyTorch is slow. The point is that tiling at the level of the framework does not reduce memory traffic, because the framework materialises every intermediate. The traffic reduction only appears when the whole loop is one kernel whose intermediates live in registers and on-chip memory. That is fusion, the lever from Chapter 11, and it is why FlashAttention is distributed as a kernel and not as a recipe.

The traffic model against the measurement

The model predicts that for an assumed query tile of 128, naive attention moves 9.1, 10.3, 11.1 and 11.5 times the data of a fused kernel at 1,024 to 8,192 tokens. The measured ratio of times was 7.8, 13.5, 9.4 and 8.6 (first run), and 7.8, 13.5, 9.5 and 5.5 (second). For the range 5.5 to 13.5 against a prediction of 9 to 12, that is consistent. It is not a validation: I do not know the kernel’s tile size, one tile size was assumed, the naive implementation’s time also includes arithmetic and launches that the byte count ignores, and the runs themselves disagree by up to 67 per cent. What the comparison does support is the direction and the order of magnitude, which is the claim the model was built to make: the speed-up of a fused attention kernel over the direct one is about the ratio of the bytes they move.

Which kernel ran

FlashAttention is not available here. I asked PyTorch for its reasons for each backend rather than inferring them:

These are two different reasons. cuDNN attention is refused because of the GPU generation. FlashAttention is refused because this build of PyTorch does not contain it, which says nothing about whether the GPU could run such a kernel. Chapter 8 attributed both to the GPU’s compute capability. That was wrong for FlashAttention, and the Chapter 8 text and records have been corrected.

So there is no FlashAttention measurement in this book, and the claims I can make about it come from its authors. The FlashAttention-2 paper reports a speed-up of about two over the first version, reaching 50 to 73 per cent of the theoretical arithmetic peak on an A100. The abstract names no other GPU. Those figures are the authors’, on their hardware. I have not reproduced them, and nothing here tests whether they would hold on a different GPU.

What this chapter does measure is the memory-efficient kernel, a fused, block-wise, exact attention kernel of the same family. It reached 55 per cent of the best matrix-product rate in Chapter 8. It is a fair proxy for the approach and not for FlashAttention’s specific performance.

The decode shape: splitting the keys

Chapter 8 found that one query against a cache read the cache at only 50 GB/s at batch 1, a sixth of the sustainable bandwidth, and hypothesised that a single sequence offers too little parallelism: 12 heads, 34 multiprocessors. The remedy proposed in the literature as FlashDecoding is to split the key dimension into chunks, give each chunk to different processors, and combine the partial results with the merge rule above. I tested whether that helps.

The merge rule makes the splitting exact, so correctness is the easy part: the split implementation matched a float64 reference to between \(3 \times 10^{-5}\) and \(2 \times 10^{-4}\) for every split count from 1 to 64, and the tests check exactness for split counts from 1 to 96.

Speed was a different matter. The split implementation, in PyTorch operations, took between roughly 400 and 1,070 microseconds across every length and batch I tried, with no systematic dependence on the number of splits. Against the fused kernel:

The number of splits had no effect I could detect. One split was as fast as sixty-four. The split implementation was faster only for the longest cache (1.1 to 1.75 times), and slower, by up to eight times, for the short ones.

This result does not test the hypothesis, and it would be wrong to read it as refuting it. The implementation is a dozen separate operations with a launch cost of several hundred microseconds in total, which is larger than the memory time it was designed to reduce: 48 MiB of cache is 170 microseconds at the sustainable bandwidth. The experiment can only see the launch overhead. Whether a fused split-KV kernel would raise batch-1 bandwidth is untested here, and I do not have the ability to write and compile one on this machine. The honest status of Chapter 8’s hypothesis is “not established”.

That said, one fact is worth keeping. At 16,384 keys, the unfused split implementation with a single split beat the fused kernel at batch 1 (571 to 751 µs against 992), because it used the library’s batched matrix product, which reads the cache faster (67 to 88 GB/s) than the fused kernel does (50). That is a reminder that a “fused” kernel is not automatically the fastest way to do a particular job on a particular shape.

Other kernels

Attention is the best-known fused kernel, not the only one. The same argument, fewer memory passes and fewer launches, justifies fusing:

• the layer norm with the residual addition that precedes it;
• the bias and GELU after a matrix product;
• the splitting and reshaping of query, key and value;
• positional rotation of queries and keys (for models that use it);
• dequantisation with the matrix product, which Chapter 12 showed is the difference between a quantised model being slower and faster than fp16;
• the sampling step’s sort and cumulative sum.

I measured one instance in Chapter 11: a fused multiply-add took 1.5 to 1.6 times less time than two separate kernels, matching the ratio of memory passes, six to four. I did not measure the others, and they are listed as the design space and not as findings. Chapter 18 returns to launch cost, which fusion attacks from the other side.

What the numbers do not show

• No FlashAttention. Nothing here is a FlashAttention measurement, and nothing tests its published speed-ups.
• One fused kernel of unknown tile size. The traffic model’s agreement with measurement depends on an assumed query tile.
• Noisy fused timing. Two runs of the fused kernel at 8,192 tokens differed by 67 per cent.
• fp16, head dimension 64, 12 heads, batch 1. Other head dimensions, grouped key/value heads and masks can disable the fast paths.
• Forward pass only. The backward pass, where FlashAttention’s recomputation matters most, is not covered.
• The split-KV experiment could not test its hypothesis because it was limited by its own launch overhead.
• The launch-cost estimate for the unfused version is arithmetic from counts and a floor measured on another day.

Inference Lab increment

The Lab now has:

• tiled_attention(q, k, v, causal, block_q, block_k), the block-wise reference, which also reports how many block pairs it computed.
• merge_partials, the exact merge rule, tested for associativity and equivalence with a single softmax.
• splitkv_decode(q, k, v, splits), exact split-key decoding.
• naive_hbm_elements, tiled_hbm_elements and traffic_ratio, the traffic model, tested for its limiting behaviour.
• availability, which asks PyTorch to explain each backend’s status.

The rules this adds: a kernel claim names the kernel that ran and the reason the others did not; and an algorithm is not a speed-up until it is a kernel, so a result that shows only memory saved, or only arithmetic unchanged, has not shown that anything got faster.

What you should now be able to do

• Derive the online softmax and the merge rule, and say why the block order does not matter.
• Compute the main-memory traffic of direct and block-wise attention, and the ratio’s dependence on the query tile.
• Say what a framework-level tiling saves (memory) and what it cannot (time), and why.
• Read PyTorch’s refusal reasons for a backend, and distinguish a hardware limit from a build limit.
• Check a fused kernel against the block-wise reference with a stated tolerance.
• Explain why splitting the key dimension is exact and why that alone does not make it fast.
• Say which claims about FlashAttention in this chapter are measurements and which are reports.

Key terms

• Online softmax: computing a softmax over blocks with running maximum and sum.
• Partial result: the triple \((m, \ell, o)\) over a subset of keys.
• Tile: the block of queries or keys processed at once on chip.
• IO-aware: designed around the cost of moving data between main and on-chip memory.
• Fused kernel: one GPU program doing the work of several operations without writing intermediates to main memory.
• Split-KV decoding: dividing a cache across processors and merging partial results.
• Backend: one of the implementations behind a framework’s attention function.

Reader exercises

1. Prove that merge_partials is associative. What does that let a kernel do?
2. Run availability on your machine. Which backends run, and for each that does not, what is the stated reason?
3. If you have hardware where FlashAttention runs, repeat the traffic comparison. What query tile size does the measured speed-up imply, using \(6B_q/d\)?
4. Rewrite tiled_attention to launch fewer operations per block, for example by precomputing the scaled queries. How far does the unfused time fall?
5. Measure the unfused version with blocks of 512. Does it get closer to the naive time, and what does that say about launch cost?
6. Extend the split-KV experiment with larger caches (65,536 keys) or more sequences. At what size does the unfused version stop being limited by its own overhead?
7. Write down what a fused split-KV decode kernel would have to do, and predict its time from the bandwidth ceiling for 48 MiB of cache.

References

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — the algorithm and its IO analysis; not run in this chapter.
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning — reports roughly twice the speed of the first version, with figures given on an A100; not reproduced here.
• Online normalizer calculation for softmax — the running-normaliser recurrence for softmax, by Milakov and Gimelshein of NVIDIA.
• PyTorch: scaled_dot_product_attention — the backend selection rules.

Evidence in this chapter

Implementation status

• Manuscript: content/books/inference/13-chapter.md
• Runnable examples: examples/ch13/flash_attention.py (original), examples/ch13/tiled_attention.py (algorithm, traffic model and measurements)
• Tests: tests/test_ch13_flash.py, tests/test_ch13_tiled_attention.py (equivalence over block shapes, causal block skipping, large-score stability, merge associativity, exact split-KV, traffic-model limits; no GPU needed)
• Evidence: research/ch13/sources.md, research/ch13/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about FlashAttention’s own performance, other GPUs, training or grouped-query attention. Chapter 14 returns to the cache itself: the memory waste of one contiguous buffer per sequence, and the block-based design that removes it.