Forcing Flash Attention onto a TPU and Learning the Hard Way
Summary
A technical exploration of forcing a GPU oriented flash attention algorithm onto a TPU, with a detailed discussion of the TPU programming model, memory architecture, and performance benchmarks. The author finds that on a single head and modest dimensions, the native TPU behavior plus XLA fusion outperforms handwritten flash attention, and shows how using vmap to expose independence can unlock significant speedups. The piece concludes that hardware and compiler design shapes the best approach more than raw code cleverness.