I'm writing a short series of tutorials on FlashAttention: from theory to efficient CUDA kernels.
Part 1 is the theoretical foundation. It walks through a modern algebraic formalism showing that FlashAttention is an associative operation, which lets you treat it as a regular reduction on the GPU and apply all the same scheduling optimizations. Some recent MLSys and CVPR (ELSA) papers lean on this framing, and I find it much more powerful than the original.
This framework is particularly useful in ML compilers, where you should implement general optimizations applicable to many operations rather than writing specialized kernels. This article shows that attention actually belongs to a large family of "secretly-associative" operations, walks through a handful of examples, and links a few concepts from abstract algebra that let you identify whether an operation is secretly associative.
Overview:
- Safe softmax, Welford's variance, and FlashAttention are the same secretly-associative operation
- The twisted monoid (transport of structure), why the max-rescale coupling doesn't break associativity
- The qk_scale = log2(e)/√D like in FA-2 derived from scratch
- Numerical analysis: overflow bounds, error limits, and why tiling never amplifies error
- Bird's 3rd Homomorphism Theorem as a test for whether any loop is secretly associativ
I'm writing a short series of tutorials on FlashAttention: from theory to efficient CUDA kernels.
Part 1 is the theoretical foundation. It walks through a modern algebraic formalism showing that FlashAttention is an associative operation, which lets you treat it as a regular reduction on the GPU and apply all the same scheduling optimizations. Some recent MLSys and CVPR (ELSA) papers lean on this framing, and I find it much more powerful than the original.
This framework is particularly useful in ML compilers, where you should implement general optimizations applicable to many operations rather than writing specialized kernels. This article shows that attention actually belongs to a large family of "secretly-associative" operations, walks through a handful of examples, and links a few concepts from abstract algebra that let you identify whether an operation is secretly associative.
Overview:
- Safe softmax, Welford's variance, and FlashAttention are the same secretly-associative operation
- The twisted monoid (transport of structure), why the max-rescale coupling doesn't break associativity
- The qk_scale = log2(e)/√D like in FA-2 derived from scratch
- Numerical analysis: overflow bounds, error limits, and why tiling never amplifies error
- Bird's 3rd Homomorphism Theorem as a test for whether any loop is secretly associativ