Hacker News 中文摘要

RSS订阅

为5090编写CUDA C++中的光速Flash Attention -- Writing Speed-of-Light Flash Attention for 5090 in CUDA C++

原文链接 | HN讨论 | 2025-08-24 03:59:27

文章摘要

本文介绍了作者如何学习在CUDA C++中实现Flash Attention for 5090的过程，旨在填补现有博客中关于快速注意力机制实现的空白。作者建议读者熟悉CUDA C++和NVIDIA GPU的Tensor核心使用，并提供了相关学习资源。文章还提供了完整的实现代码链接。

文章总结

在这篇文章中，作者详细介绍了如何在CUDA C++中实现Flash Attention算法，并逐步优化其性能。以下是文章的主要内容：

1. 背景与目标

  • 作者的目标是学习如何在CUDA C++中实现Flash Attention算法，因为Triton等工具不支持某些特性（如MXFP8/NVFP4 MMA）。
  • 作者认为这是学习矩阵乘法内核后的自然下一步，并且目前关于如何实现快速注意力机制的博客较少，因此决定填补这一空白。

2. 读者要求

  • 读者需要熟悉CUDA C++和如何在NVIDIA GPU上使用Tensor Core。
  • 作者推荐了一些学习资源，如GPU-MODE系列讲座和关于矩阵乘法的优秀博客。

3. 实现与优化

  • 作者提供了完整的实现代码，并展示了在不同优化版本下的性能对比。
  • 通过逐步优化，作者从基础版本（v1）开始，逐步引入了共享内存交换（v2）、双缓冲流水线（v3）、使用ldmatrix.x4加载K和V（v4）以及更好的流水线设计（v5），最终将性能提升到了接近理论极限的94.39%。

4. Flash Attention算法

  • 作者首先介绍了注意力机制的基本实现，并引用了Flash Attention 2论文中的算法。
  • 每个线程块负责处理Q的一个分块，并在KV序列长度上进行迭代。作者用Python伪代码展示了算法的高层设计。

5. 版本1：基础实现

  • 作者首先实现了基础版本，使用cp.async从全局内存加载数据到共享内存，并使用ldmatrix将数据从共享内存加载到寄存器，最后调用mma.m16n8k16进行矩阵乘法。
  • 为了减少错误，作者没有引入复杂的优化技巧，如共享内存交换和流水线。

6. 在线Softmax

  • 作者详细解释了在线Softmax的实现，包括如何计算行最大值、如何重新缩放输出以及如何将tile_S打包为tile_P。
  • 在线Softmax的关键思想是通过逐步更新注意力状态来避免重复计算。

7. 版本2：共享内存交换

  • 在版本2中，作者引入了共享内存交换技术，解决了共享内存加载时的bank冲突问题，显著提升了性能。

8. 版本3：双缓冲流水线

  • 在版本3中，作者引入了双缓冲流水线技术，通过重叠全局内存操作和计算操作，进一步提升了性能。

9. 版本4：使用ldmatrix.x4加载K和V

  • 在版本4中，作者使用ldmatrix.x4来加载K和V，减少了指令数量，进一步提升了性能。

10. 版本5：更好的流水线设计

  • 在版本5中，作者优化了流水线设计，减少了共享内存的使用，并增加了分块大小，最终将性能提升到了94.39%。

11. 未来工作

  • 作者提出了几个未来的优化方向，包括实现反向传播、量化/低比特注意力机制、使用TMA（cp.async.bulk）以及实现PagedAttention等。

12. 总结

  • 作者希望通过这篇文章帮助更多人理解和实现高效的CUDA内核，并鼓励大家继续探索和优化。

这篇文章详细记录了作者在实现和优化Flash Attention算法过程中的思考和经验，适合对CUDA编程和深度学习优化感兴趣的读者。

评论总结

评论内容总结：

  1. 对NVIDIA生态系统的贡献与Triton的使用

    • 主要观点：评论者质疑为何不直接为Triton贡献代码，而非为NVIDIA等大公司免费开发产品。
    • 关键引用：
      • "Why not contribute to Triton, they accept PRs?"
      • "Like so what if you do free product ecosystem development for NVIDIA and giant corporations by contributing to Triton?"

  2. 硬件兼容性与性能问题

    • 主要观点：评论者提到5090显卡在Flash Attention上的兼容性问题，并询问5080的情况。
    • 关键引用：
      • "I had a 5090 some months ago but couldnt get flash attention to work."
      • "Does it now work natively? What about 5080?"

  3. 技术内容的深度与复杂性

    • 主要观点：评论者认为文章内容非常深入，需要多次阅读才能理解。
    • 关键引用：
      • "Damn awesome. This going to take me 3 reads and a week to digest."

  4. 显卡性能与性价比的讨论

    • 主要观点：评论者指出5090的理论性能与服务器级显卡相比差距较大，且NVIDIA对游戏显卡的ML训练性能进行了限制，认为游戏显卡不再适合作为“廉价FLOPs”替代方案。
    • 关键引用：
      • "That's not even 10% of the server Blackwell."
      • "It doesn't seem worth it to use NVIDIA gaming cards as a 'cheaper FLOPs' alternative anymore."

  5. 工作站ML使用中的功耗问题

    • 主要观点：评论者对5090的功耗表示担忧，认为其TDP高于4090且功耗限制不如4090灵活。
    • 关键引用：
      • "My issue with upgrading to the 5090 for workstation ML use is that it both has higher TDP than the 4090."
      • "It can only be limited to 70% power (not 50% like the 4090)."

总结：评论主要围绕NVIDIA显卡在ML应用中的性能、兼容性、功耗及性价比展开，既有对技术细节的深入讨论，也有对硬件选择的实际考量。