Why does the TPU flash attention backward pass run two kernels instead of one?

Because the two halves of the gradient want opposite loop orders. The gradient of a key block is a sum over every query block, so queries have to be the innermost, fastest-varying grid axis. The gradient of a query block is a sum over every key block, so keys have to be innermost. One grid cannot have both, and a Pallas grid carries its running state in scratch along its innermost axis, so the file launches two grids: one shaped batch, heads, keys, queries that produces the key and value gradients, and one shaped batch, heads, queries, keys that produces the query gradient. The price is that the attention scores are recomputed twice, once in each launch.

What does block_k_major mean versus block_k in Pallas flash attention?

The major block is what gets copied from the chip’s main memory into fast memory in one transfer, and it is what the grid steps over. The minor block is what a single matrix multiply inside the kernel body consumes, and the body loops over the major block in minor-sized steps with the loop unrolled. The dataclass enforces the only relationship that makes sense between them: the minor must be no larger than the major and must divide it exactly, so the unrolled trip count is a whole number. Setting them equal gives one large multiply per grid step; setting the minor to a quarter of the major gives four smaller multiplies and four rounds of the running-maximum correction.

Does a causal Pallas kernel skip the blocks above the diagonal?

It skips the arithmetic, not the grid step and not the copy. The grid stays rectangular, every cell of it is invoked, and a four-line predicate decides whether the body runs by asking whether the last row index of the query block exceeds the first column index of the key block. When it does not, the body is suppressed by a device-side conditional. The block’s copy still happens, which is why the index map for that input does not return the requested block: it returns block zero, so that the transfer the machine was going to make anyway carries something the next grid step can use.

Why must a Pallas TPU block size be a multiple of 128?

Because a TPU vector register is 128 lanes wide, and the kernel stores per-row quantities across a full register. The running maximum and the running sum are one number per query row held as 128 identical lanes, and they have to be tiled sideways to the width of the score block before they can be subtracted from it, which only works if the score block is a whole number of registers wide. The rule is not enforced by the block-size dataclass or by the wrapper, both of which only check divisibility against each other and against the sequence length. It is enforced inside the kernel body, so a block size of 64 passes every argument check and then raises a NotImplementedError while Pallas is tracing your kernel.

What does the TPU flash attention kernel save for the backward pass?

Eight things: the queries, keys, values, the attention bias, the segment ids, the output, and two residual arrays holding the per-row running sum and running maximum that the forward softmax produced. It deliberately does not save the attention weights or the scores, which are the large objects. The backward kernels recompute the scores from the queries and keys and rebuild the softmax by subtracting the saved maximum and dividing by the saved sum, which is two lines standing in for a whole forward pass. One more quantity, a row-wise dot product of the output with the incoming gradient, is computed once outside both kernels and passed in, so neither kernel has to re-read the output array.

Flash Attention on a TPU, Line by Line

The flash attention algorithm fits on a napkin. The Pallas kernel that runs it on a TPU is 1715 lines: four kernel bodies, three separate launches, eleven block sizes, and a four-line diagonal test consulted from ten places. This walks the whole file at a pinned commit of jax-ml/jax.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Flash Attention on a TPU, Line by Line — built to be understood, not skimmed.

Questions

Why does the TPU flash attention backward pass run two kernels instead of one?: Because the two halves of the gradient want opposite loop orders. The gradient of a key block is a sum over every query block, so queries have to be the innermost, fastest-varying grid axis. The gradient of a query block is a sum over every key block, so keys have to be innermost. One grid cannot have both, and a Pallas grid carries its running state in scratch along its innermost axis, so the file launches two grids: one shaped batch, heads, keys, queries that produces the key and value gradients, and one shaped batch, heads, queries, keys that produces the query gradient. The price is that the attention scores are recomputed twice, once in each launch.
What does block_k_major mean versus block_k in Pallas flash attention?: The major block is what gets copied from the chip’s main memory into fast memory in one transfer, and it is what the grid steps over. The minor block is what a single matrix multiply inside the kernel body consumes, and the body loops over the major block in minor-sized steps with the loop unrolled. The dataclass enforces the only relationship that makes sense between them: the minor must be no larger than the major and must divide it exactly, so the unrolled trip count is a whole number. Setting them equal gives one large multiply per grid step; setting the minor to a quarter of the major gives four smaller multiplies and four rounds of the running-maximum correction.
Does a causal Pallas kernel skip the blocks above the diagonal?: It skips the arithmetic, not the grid step and not the copy. The grid stays rectangular, every cell of it is invoked, and a four-line predicate decides whether the body runs by asking whether the last row index of the query block exceeds the first column index of the key block. When it does not, the body is suppressed by a device-side conditional. The block’s copy still happens, which is why the index map for that input does not return the requested block: it returns block zero, so that the transfer the machine was going to make anyway carries something the next grid step can use.
Why must a Pallas TPU block size be a multiple of 128?: Because a TPU vector register is 128 lanes wide, and the kernel stores per-row quantities across a full register. The running maximum and the running sum are one number per query row held as 128 identical lanes, and they have to be tiled sideways to the width of the score block before they can be subtracted from it, which only works if the score block is a whole number of registers wide. The rule is not enforced by the block-size dataclass or by the wrapper, both of which only check divisibility against each other and against the sequence length. It is enforced inside the kernel body, so a block size of 64 passes every argument check and then raises a NotImplementedError while Pallas is tracing your kernel.
What does the TPU flash attention kernel save for the backward pass?: Eight things: the queries, keys, values, the attention bias, the segment ids, the output, and two residual arrays holding the per-row running sum and running maximum that the forward softmax produced. It deliberately does not save the attention weights or the scores, which are the large objects. The backward kernels recompute the scores from the queries and keys and rebuild the softmax by subtracting the saved maximum and dividing by the saved sum, which is two lines standing in for a whole forward pass. One more quantity, a row-wise dot product of the output with the incoming gradient, is computed once outside both kernels and passed in, so neither kernel has to re-read the output array.

Flash Attention on a TPU, Line by Line

Questions

Related explainers