How Flash Attention Runs on a GPU

This is the fourth post in a series on LLM internals. Part 1 covered attention, Part 2 covered generation, Part 3 covered the Flash Attention algorithm. This post takes Part 3’s numpy simulation and puts it on a GPU where the SRAM/HBM distinction is real.