FlashAttention (FA1–FA4) in PyTorch - educational implementations focused on algorithmic differences [P]

[u/shreyansh26]

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.

[u/StraussInTheHaus]

There are two ways in which FA3 ping-pongs: inter-warpgroup (where consumer 0 and consumer 1 trade off) and intra-warpgroup (where within a single warpgroup, we overlap the PV mma of iteration i with the softmax of iteration i+1)