Mamba一作再祭神作，H100利用率飙至75%！FlashAttention三代性能翻倍，比标准注意力快16倍

文章预览

来源 | 新智元 ID | AI-era FlashAttention又有后续了！ 去年7月， FlashAttention-2发布，相比第一代实现了2倍的速度提升 ，比PyTorch上的标准注意力操作快5～9倍，达到A100上理论最大FLOPS的50～73%，实际训练速度可达225 TFLOPS（模型FLOPs利用率为72%）。 然而，去年发布FlashAttenion-2尚未运用到硬件中的最新功能，在H100上仅实现了理论最大FLOPS 35%的利用率。 时隔一年，FlashAttention-3归来，将H100的FLOP利用率再次拉到75%，相比第二代又实现了1.5～2倍的速度提升，在H100上的速度达到740 TFLOPS。 论文地址： https:// tridao.me/publications/ flash3/flash3.pdf 值得一提的是，FlashAttention v1和v2的第一作者也是 Mamba的共同一作，普林斯顿大学助理教授Tri Dao ，他的名字也在这次FlashAttention-3的作者列表中。 Tri Dao师从于Christopher Ré和Stefano Ermon，去年6月在斯坦福大学获得计算机博士学位，毕业 ………………………………