I recently updated my FlashAttention-PyTorch repo so it now includes educational implementations of FA1, FA2, FA3, and FA4 in plain PyTorch. The main goal is to make the progression across versions easier to understand from code. This is not meant to be an optimized kernel repo, and it is not a hardware-faithful recreation of the official implementations. The point is to expose the algorithmic ideas and design changes without immediately going deep into CUDA/Hopper/Blackwell-specific details. Roughly, the repo now shows: FA1: tiled online softmax baseline FA2: split-Q / query-tile ownership, deferred normalization FA3: explicit staged pipeline with ping-pong tile buffers, plus a simplified educational FP8 forward path FA4: explicit scheduler with main / softmax / correction phases, and conditional/selective rescaling So the same exact attention math is preserved, but the orchestration changes version by version. I wrote it for people who want to understand: “What actually changed from FA1 → FA2 → FA3 → FA4?”" without having to start from highly optimized CUDA kernels. Repo: https://github.com/shreyansh26/FlashAttention-PyTorch Would be interested in feedback on whether the code makes the version-to-version differences intuitive. submitted by /u/shreyansh26
Originally posted by u/shreyansh26 on r/ArtificialInteligence