Research
FlashAttention (FA1–FA4) in PyTorch - educational implementations focused on algorithmic differences [P]
This r/MachineLearning post presents educational PyTorch implementations of FlashAttention versions 1 through 4, designed to highlight the key algorithmic differences across each iteration rather than
This r/MachineLearning post presents educational PyTorch implementations of FlashAttention versions 1 through 4, designed to highlight the key algorithmic differences across each iteration rather than maximize raw performance. The implementations trace how each version improved upon the last — from FA1's foundational tiled IO-aware attention and online softmax, to FA2's parallelism and reduced non-matmul operations, to FA3's pipelined compute for Hopper GPUs, and FA4's deeply asynchronous warp-specialized pipeline targeting Blackwell hardware. The project serves as a learning resource for practitioners and researchers wanting to understand why each FlashAttention version was designed the way it was, grounded in the evolving constraints of GPU memory hierarchies and hardware architectures.
Related
- Parax: Parametric Modeling in JAX + Equinox [P]
- AI Systems Performance Engineering by Chris Fregly - is it worth it? [D]
- [[p-building-a-llm-from-scratch-with-mary-shelleys-frankenstei|[P] Building a LLM from scratch with Mary Shelley's 'Frankenstein' (on Kaggle)]]
- [[p-citracer-a-small-cli-tool-to-trace-where-a-concept-comes-f|[P] citracer: a small CLI tool to trace where a concept comes from in a citation graph]]
Source: r/MachineLearning | 2026-04-11