最近,有个关于英伟达B200 GPU算力浪费的新闻火了。普林斯顿大学团队发现,这个GPU浪费了高达60%的计算资源。为啥会这样?主要是因为软硬件配合不太好。这让很多依赖这个GPU的开发者们心里很不是滋味,难道他们的钱都打水漂了吗?为了解决这个问题, FlashAttention-4 就出来了。它是为英伟达Blackwell架构GPU设计的一种注意力算法,成功把GPU的利用率从20%到30%提升到了71%。这个算法是由TriDao领头搞出来的,还拉上了Meta、TogetherAI和英伟达自己一起干。Blackwell B200虽然算力很强,理论上能帮我们加快很多工作,但实际上它在某些关键单元上还是有点问题。特别是那个负责指数运算的MUFU单元,它和上一代Hopper架构差不多。这就导致大模型计算的时候虽然矩阵乘法挺快了,但指数运算和共享内存读写反而拖后腿了。针对这些问题,FlashAttention-4提出了三大优化策略。第一个就是用软件模拟和多项式近似来提升指数运算的速度,同时还引进了条件性softmax rescaling策略来减少不必要的计算步骤。第二个就是重新设计计算流水线,让核心环节完全重叠起来干活,不让算力闲着没事干。第三个就是给未来留条后路,为了适应未来硬件升级准备好优化空间。 FlashAttention-4 跟以前的 FlashAttention-3 比起来进步可不小。它用的是CuTe-DSL框架写的代码,编译速度快多了。以前编译一次得55秒或者45秒,现在只要4秒或者2.5秒。实测一下发现,在B200上它的前向传播能跑到1613TFLOPS/s,真正达到了71%的理论峰值利用率。跟主流框架比起来表现还是很不错的,在处理长序列或者因果掩码这些大模型任务的时候特别强。这个消息不仅给开发者们带来了希望也说明咱们科技界在解决问题上还是挺有创新能力的。以后我也期待看到更多这样的解决方案出来把高性能计算的道路铺平吧。