Study note
Writing Speed-of-Light Flash Attention for 5090 in CUDA C++
Properties
- Type
- Blogs
- Status
- 待读
- Domain
- AI / ML
- Category
- GPU、CUDA 与内核优化
- Source
- gau-nernst.github.io
- Vault note
library/articles/ai_ml/Writing-Speed-of-Light-Flash-Attention-for-5090-in-CUDA-C++-6327ae77202d3522.md
Summary
gau-nernst 以 FlashAttention-2 的分块算法和 online softmax 为基础，在 CUDA C++ 中完整实现 attention kernel，并按 v1-v5 依次加入 shared-memory swizzle、两级流水线、面向 K/V 的 `ldmatrix.x4` 和更紧密的流水线；正文同时拆解 global/shared/register 数据搬运、BF16 Tensor Core MMA、warp 划分和数值稳定的在线归一化。
Highlights
文章把“算法正确”到“逼近硬件极限”的优化过程、代码和 benchmark 放在同一条线上。作者报告最终版本在其 RTX 5090、CUDA 12.9 和指定 shape 下达到 197.74 TFLOPS，约为该卡 BF16 理论峰值的 94.39%，很适合在学过 tiled matmul、Tensor Core、shared-memory bank conflict 和基础 profiling 后精读。
Notes
从基础实现逐步优化 RTX 5090 FlashAttention CUDA kernel。
Comments
Loading comments...