惨痛教训：为何GPU上的Flash Attention优化策略在TPU上彻底失效？

本文记录了作者尝试将基于GPU Triton编写的Flash Attention算法移植到Google TPU (JAX)上的失败与顿悟过程。起初，直接移植的性能不仅没有提升，反而比标准的XLA融合内核慢了数十倍。通过深入剖析TPU架构和JAX编译机制，作者发现使用`fori_loop`掩盖了并行性，而改用`vmap`显式告知编译器数据独立性后，性能惊人地提升了45倍。文章揭示了TPU凭借巨大的片上内存（VMEM）和强大的XLA编译器融合能力，在大多数情况下并不需要像GPU那样激进地进行手动分块优化，这是对“软硬件协同设计”的一堂生动教育课。