一项关于将 FlashAttention-3 和 FlashAttention-4 优化应用于消费级 RTX GPU 的探索显示,这些面向数据中心的改进并未带来性能提升。作者从头重建了注意力内核,以测试 FA-3/4 技术是否能在消费级硬件上超越现有的 FA-2 回退方案。

  • FA-3 核心的更快张量核心指令 (WGMMA) 在 RTX GPU 上不可用。
  • Tensor Memory Accelerator (TMA) 存在于 sm_120 中,但由于传输并非瓶颈,因此没有帮助。
  • Warp 专业化导致性能下降,延迟从 206us 增加到 213us。
  • FA-4 的模拟指数优化无关紧要,因为 RTX 5090 仍然受限于张量核心而非 SFU。

该研究得出结论,FlashAttention-2 代表了消费级芯片上常规注意力的性能上限,因为进一步的提升将需要牺牲精度以换取较低的精度。

此分析侧重于预填充和计算密集型场景;针对大型 KV 缓存的解码仍然是一个独立的内存瓶颈挑战,其中 split-KV 技术更为相关。