在RTX GPU上探索FlashAttention-3/4优化

我很好奇FA-3/4的优化是否能在RTX GPU上发挥作用。vLLM/SGLang的注意力机制在消费级显卡上回退到FA-2（FA-3和FA-4仅限数据中心），因此我想知道是否还有性能提升空间，于是从头重建了注意力内核。该内核达到了与FA-2相同的性能（在RTX5090上，batch=1, heads=8, seq_len=4096, head_dim=64时为206微秒），但遗憾的是，FA-3/4的优化在消费级显卡上要么不适用，要么没有帮助。看来FA-2就是天花板了。总结一下：更快的张量核心指令（WGMMA）是FA-3背后的主要杠杆，但RTX GPU并不支持。TMA（张量内存加速器）在sm_120（RTX 5XXX）上可用。理论上它有帮助（LSU减少），但数据传输并非瓶颈，因此最终数值几乎没有变化。Warp特化也可用，但它主要是一种调度优化，即有助于消除流水线气泡并更好地利用张量核心，但缺少异步张量核心指令。结果是负面的：213微秒对比206微秒。在FA-4中，他们还使用FMA指令模拟了exp函数，因为B200上的张量核心速度极快，整个流水线变成了SFU（特殊函数单元）受限。RTX 5090是张量核心受限，因此这项优化也无意义。事实上，即使使用更快的exp2f替代expf进行softmax这种常规优化，数值也纹丝不动。我还尝试了一些可能在消费级芯片上起效的其他优化，例如更深的流水线和寄存器乒乓。但没有成功。由于整个流水线是张量流水线受限，我相信FA-2就是天花板，任何有意义的杠杆都需要牺牲一些精度来利用更快、更低精度的张量核心。请注意，这是对常规注意力机制的探索，它主导了预填充和计算受限阶段。针对大型KV缓存的解码则是另一个内存受限的故事，其中split-KV/Flash-Decoding比上述任何优化都更重要。完整文章：https://riftstack.ai/research/learning-flashattention-the-hard-way-part-2 Github：github.com/cloudrift-ai/emmy