FlashAttention-3发布，带来了哪些新的优化点？

整体而言，它是针对Hopper架构的新特性进行了相应的优化。这样的改进迟早会出现，当我看到Hopper架构中的GPC、WGMMA、TMA以及异步事务屏障技术时，就预感到会有类似更新。TurboAttention也在开展相关优化工作，主要聚焦于适配H20的特性。接下来，我们可以先梳理一下Hopper架构在这些方面具体实现了哪些优化内容，这有助于更清晰地理解FlashAttentionV3和相关技术的发展方向与改进之处。先简单看看FlashAttentionV3的核心优化，再深入理解Hopper架构的改进，后续将结合代码进行详细讲解。FlashAttentionV3相较于前一代：GPU 利用率显著提升：新技术能发挥 H100 GPU 高达 75% 的性能，远超此前的 35%。这一进步大幅加快了大语言模型（LLM）的训练与推理速度，相比旧版提升 1.5 至 2 倍，效率优势明显。更佳性能与更低精度：FlashAttention-3 支持使用 FP8 低精度数值，同时确保计算准确性。这不仅提升了运行速度，还可能降低内存消耗，帮助执行大规模人工智能任务的用户节省成本、提高运算效率。

FlashAttentionV3在Hopper架构中主要应用了WGMMA、TMA和Fp8技术。WGMMA是一种基于GPC技术的MMA指令，与以往架构中的MMA指令相比，其最大特点是支持warp-group功能。传统MMA指令仅限于单warp操作，而WGMMA借助Hopper的线程块集群（Thread Blocks Cluster）技术以及图形处理集群（Graphics Processing Clusters，GPC）和分布式共享内存实现突破。这项创新使得WGMMA能够在更大范围的线程组内协同工作，从而提升计算效率和灵活性，为复杂任务提供更强大的支持。这种设计充分利用了现代GPU架构的优势，进一步优化了矩阵运算性能。TMA是Hopper架构引入的一项技术，旨在提升从全局内存加载数据到共享内存的速度。与上一代的async_cp不同，TMA拥有独立的硬件支持，因此参与调度的线程数量可以很少，同时能够处理更大规模的数据加载。例如，可以通过TMA轻松加载1KB的数据量，显著提高效率和性能。这项技术为大规模数据传输提供了更优的解决方案。FP8大家应该比较熟悉，相较于INT8，其主要优势在于数值范围更广。作为FP8的最后一块高地，注意力机制在FlashAttention V3支持FP8后，必将推动业界对FP8的广泛应用。训练和推理均可直接使用FP8，从而减少对量化的依赖，这或许会对量化技术造成一定影响。接下来是最核心的技术：跨线程组重叠技术与乒乓调度相结合。这一方法有效提升了性能表现，通过巧妙的并行处理和调度优化，实现了更高效的运算能力。这项技术的背景在于：在现代加速器中，非矩阵运算的速度远低于矩阵运算。例如，像用于软最大值计算的指数函数等特殊函数，其吞吐量甚至不及浮点乘加运算；这些特殊函数由独立于浮点乘加或矩阵乘加单元的多功能单元来评估。以 H100 GPU SXM5 为例，其 FP16 矩阵乘法的运算能力高达 989 TFLOPS，但特殊函数的运算能力仅为 3.9 TFLOPS，相差约 256 倍。对于 head-size 为 128 的情况，matmul 的 FLOPS 数量是指数运算的 512 倍，这意味着指数运算所花费的时间会比 matmul 多出约 50%。而在 FP8 模式下，问题更加严重，因为 matmul 的 FLOPS 性能是指数运算的两倍，而指数运算的性能却没有提升。因此，在理想情况下，我们希望 matmul 和 softmax 能够并行执行。也就是说，当张量核心忙于处理 matmul 运算时，多功能单元应当同时进行指数运算，从而最大化资源利用率并提高整体效率。在此背景下，FlashAttention V3采用了Pingpong Scheduling技术，如下所示：Warp级专属持久化乒乓调度。该方法通过优化调度策略，提升计算资源利用率，实现更高效的并行处理，显著改善性能表现。

GEMM运算能够与Softmax相互叠加，这是基于Hopper架构下的一项新技术——Warp-Specialized Persistent Cooperative。这项技术在Cutlass中已有详细介绍，若想深入理解，可参考文献3和4。它与Warp-Specialized kernel（同样为Hopper架构下Cutlass 3.0引入的技术）类似，在协作设计中依旧保留了warp-group之间屏障同步的概念。而Warp-Specialized Persistent Cooperative kernel的特殊之处在于以下几个方面：首先，它通过持久化的方式优化线程调度，减少线程启动开销；其次，利用warp粒度的专门化分工提高计算资源利用率；通过更高效的跨warp协作机制，进一步提升数据复用率和整体性能表现，从而在大规模矩阵运算场景下展现出显著优势。这些特点使其成为针对特定任务优化的理想选择。专为持久乒乓操作优化的Warp特化内核设计特点表现如下：更多细节可参见Cutlass 3.0文档（参考文献），后续笔者也将深入解析这些技术。