CUDA基础入门系列文章创作历程是怎样的？

C++

新年打算开始撰写CUDA基础入门系列文章，旨在帮助初学者入门CUDA，内容从易到难。实际上，早在24年3月就有这个想法了，只是我的笔记和资料一直未达到理想状态，而且我对CUDA的理解与应用水平也有待提升。所以，我决定对之前整理的CUDA - Learn - Notes笔记进行扩展。经过近一年的业余时间努力，我觉得现在已经可以开始撰写CUDA基础系列笔记了。资料已准备好，25年将继续更新笔记并保持学习。开篇第一篇，简单介绍一下CUDA - Learn - Notes的内容。在2024年的时候，我把自己那些CUDA笔记（这里要说明一下，是我学得不好，笔记比较垃圾，而不是CUDA本身不好）整理成文章发布到了网络上，同时也将其开源到了CUDA - Learn - Notes仓库。这里面涉及的kernel有warp reduce、block all reduce、dot - product、softmax、layer - norm、rms - norm、element - wise、sgemv以及sgemm等。其中涵盖的CUDA优化技巧主要包含合并访存、向量化、减少bank conflicts、warp shuffle、warp sgemv以及sgemm双缓冲等。上一篇文章的链接如下：CUDA - Learn - Notes的第一版得到了一些赞同，也收获了一些stars。不过很明显，这一版笔记存在诸多问题，像没有给出数值验证的案例；仅支持FP32数值类型，不支持FP16、BF16、INT8、FP8等类型；纯CUDA/C++代码让Python/PyTorch用户难以对比性能和数值结果；缺乏Tensor Cores应用案例；没有TF32相关内容；没有HGEMM/FlashAttention的Tensor Cores实现；也没有从易到难的难度分级等等。CUDA - Learn - Notes不能只是个cu文件，而要成为有学习价值的笔记。过去这一年的业余时间里，我一直在对它进行维护与更新。增添了许多内容，这些内容包含CUDA入门的常见主题，难易程度不一，从简单到困难都有所涉及。并且，充分顾及Python/PyTorch的使用习惯，给每个主题的kernel实现都添加了PyTorch绑定。于是，现在它的全名改成了：Modern CUDA Learn Notes with PyTorch for Beginners。目前，这个笔记已经收获了2k个星星。

CUDA - Learn - Notes依据不同主题对kernel加以实现与整理，共实现了近200个CUDA kernels，其中涉及CUDA Cores、Tensor Cores的运用，还涵盖了FP32/TF32/FP16/BF16/FP8/INT8等多数常见的数据类型。这些kernel有的极为简单（像elementwise这种），有的则极具挑战性（例如手动编写、性能基本与官方FA2的MMA版FlashAttention相当的）。每个主题的工作流程是：自定义CUDA kernel实现→PyTorch Python绑定→运行测试。其难度等级划分为5个级别，分别为：简单 ? 中等 ?? 困难 ??? 困难+ ???? 困难++ ?????简单与中等部分包含了元素级操作、矩阵转置（mat_trans）、warp/block reduce、非极大值抑制（nms）、ReLU、GELU、Swish、层归一化（layer - norm）、RMS归一化（rms - norm）、在线Softmax、点积（dot - prod）、嵌入（Embedding）以及FP32、FP16、BF16和FP8的基本用法。而困难、困难 + 和困难 ++ 部分则深入钻研高级主题，重点聚焦于sgemv、sgemm、hgemv、hgemm和flash - attention等操作（op）。这些部分还提供了许多运用纯Tensor Cores MMA PTX实现的内核（kernels）。HGEMM的最优实现，性能可达cuBLAS的98% - 100%。手搓FlashAttention - 2 MMA，在MMA Acc F32时，其性能能达到官方FA - 2的95% - 99%左右。

在这个过程里，尝试了众多SRAM和registers的优化方法，最终有了一个针对FA - 2的改进方案，即FFPA。它有着O(1)的SRAM复杂度，能将headdim扩展到1024，同时保持80%以上的TFLOPS利用率，速度比SDPA EA快2到3倍。由于FFPA在工程思路上与FA2和SDPA EA都有所不同，并且对FA有进一步拓展的意义，所以从CUDA - Learn - Notee里单独拿出来，作为一个单独的repo进行维护。考虑到headdim>256的使用场景较为少见，FFPA目前仅提供一些实验性的kernel和benchmark以供参考。有兴趣的同学，可以跳转至GitHub链接：（本文不详细介绍FFPA，之后会单独写文章介绍）我除了撰写CUDA笔记与示例，还整理出100多篇自己钟爱的技术博客。从中我获益良多。每次看完，都不禁感叹写得太棒了，自愧不如。于是，我把它们整理到CUDA - Learn - Notes里推荐给大家。这些博客我都按主题分类了，大家可按需查看，例如：

Python

本小节来谈一谈CUDA - Learn - Notes的工作流程。CUDA - Learn - Notes把kernels依据主题加以划分，针对每个主题的kernel实现，均做了PyTorch的Python绑定，能够直接运用Python脚本对性能和数值进行验证。就像下面这个block all reduce的例子：在这个实例里，涵盖了大多数常见的数据类型。现有的CUDA笔记基本都不怎么提及FP16和Tensor Cores，与之不同的是，CUDA - Learn - Notes在FP16和Tensor Cores方面投入了很多精力，给出了HGEMM、FlashAttention - MMA等诸多案例。而且，对于BF16/FP8等数据类型，同样提供了应用实例，像本示例中的block all reduce，它支持FP32/FP16/BF16/FP8/INT8等数据类型。在经过PyTorch绑定之后，验证数值结果和性能就变得非常容易，直接运行Python脚本就可以了，无需去配置C++编译环境。日志输出示例如下：（block all reduce的示例）主打上手就能运行，运行完再读代码，自己修改一下，就会自然学会编写block all reduce了。HGEMM和FlashAttention也是这样。HGEMM是CUDA优化中必然会涉及的部分，不亲手实践一下就感觉少了点什么。CUDA - Learn - Notes里也有许多HGEMM+Tensor Cores MMA的实现范例。其最优实现的性能，能达到cuBLAS性能的98% - 100%。它所实现的众多特性，大体如下：

当前，CUDA - Learn - Notes里所提供的HGEMM Kernels具备了一些主要特性，像对K进行循环、分块（Tile Block）、分块MMAs（更多线程）、分块Warp（更多数值）、打包LDST、共享内存填充（SMEM Padding）、共享内存混洗（SMEM Swizzle）、块混洗（Block Swizzle）、线程束混洗（Warp Swizzle）、多阶段（Mulit - Stages）、寄存器双缓冲（Registers Double Buffers）、NN/TN布局、集体存储（线程束混洗与寄存器复用）等。而其中已实现的HGEMM CUDA Kernel包括：我还顺便试了下SGEMM TF32，结果性能和cuBLAS比差远了，真惨。这儿给大家贴一个HGEMM的性能数据，要是想查看更多benchmark，可以到CUDA - Learn - Notes上去看。

Modern CUDA - Learn - Notes肯定不能缺少FlashAttention呀。在CUDA - Learn - Notes里，我通过Tensor Cores MMA PTX亲手打造了FlashAttention - 1（Split - KV）和FlashAttention - 2（Split - Q）。在这个过程中，我尝试了许多节省SRAM以及优化寄存器的方法。在MMA Acc F32这种情况下，最优的实现能够达到FlashAttention - 2官方实现大约95% - 98%的性能。我还实现了MMA Acc F16版本，在4090之类的机器上，能够达到FlashAttention - 2原生性能的1.5倍，不过，这个版本可能会发生溢出，这样就直接导致效果不行了。（另外，FFPA中实现了混合MMA Acc精度，即Q@K^T MMA Acc F32+P@V MMA Acc F16，这样既能确保精度，又能有1.2 - 1.3倍的性能提升。）

FlashAttention - 2 MMA所支持的特性有：多阶段（Multi - Stages）、分块矩阵乘法加速（Tile MMA）、分块扭曲（Tile Warp）、共享键值暂存器（Shared KV SMEM）、完全共享查询键值暂存器（Fully Shared QKV SMEM）、预取查询（Prefetch Q s2r）、预取键值（Prefetch K/V g2s）、查询键值细粒度分块（QKV Fine - grAIned Tiling）、集体存储（Collective Store）等。

当下，就小规模注意力机制（其中B≤4，H≤48，SeqLen≤8192，D≤64）而言，在部分设备上其运行速度能超过FA2/SDPA。比如在NVIDIA RTX 3080笔记本电脑上，采用Split Q + Fully Shared QKV SMEM方法，当D = 64时，能达到55 TFLOPS，差不多比FA2快1.5倍。在NVIDIA L20上，运用ffpa - attn - mma方法，当D = 512时，可达到104 TFLOPS，这几乎比SDPA（EFFICIENT ATTENTION）快1.8倍。不过，针对大规模注意力机制，性能方面仍存在差距。请持续留意后续更新情况。在FlashAttention - 2 MMA的实现进程里，进行了诸多节省SRAM以及优化寄存器的尝试。其中一个核心思考的问题为：为何FlashAttention的头维（headdim）会受SRAM限制，然而通用的HGEMM/GEMM的K维度却不会（只要显存足够，K可以无限大）？经过思考与尝试，最终得出了FFPA：又一个针对大头维（headdim）具有O(1)??GPU SRAM复杂度的更快的Flash预填充注意力（Faster Flash Prefill Attention）。FFPA把FA - 2的SRAM复杂度直接变为O(1)，并把头维（headdim）扩展到了1024（甚至可以更大）。除FFPA之外，还有许多其他尝试，大致如下：

CUDA - Learn - Notes对kernel按不同主题加以实现与整理，总共实现了近200个CUDA kernels，其中涉及CUDA Cores、Tensor Cores的运用，还涵盖了FP32/TF32/FP16/BF16/FP8/INT8等大部分常见数据类型。这些kernel有的极为简单（像elementwise这种），有的则极具挑战性（例如手动编写性能基本与官方FA2的MMA版FlashAttention相当的程序）。每个主题的工作流程是：先自定义CUDA kernel实现，接着进行PyTorch Python绑定，最后运行测试。其难度等级分为5个级别，分别为：简单 ??、中等 ????、困难 ??????、困难+ ????????、困难++ ????????。简单与中等部分包含了诸多内容，像元素级操作、矩阵转置（mat_trans）、warp/block reduce、非极大值抑制（nms）、ReLU、GELU、Swish、层归一化（layer - norm）、RMS归一化（rms - norm）、在线Softmax、点积（dot - prod）、嵌入（Embedding），还有FP32、FP16、BF16和FP8的基本用法。而困难、困难 + 和困难 ++ 部分则深入钻研高级主题，重点聚焦在sgemv、sgemm、hgemv、hgemm以及flash - attention等运算操作（op）上。这些部分也提供了许多用纯Tensor Cores MMA PTX实现的内核（kernels）。题内话，要是为了面试的话，简单??、中等????、困难??????这三个难度级别大体上就足够用了。在困难??????这个级别里，包含了HGEMM八股的各种实现以及Tensor Cores MMA的应用。而困难+????????是FlashAttention - 2 MMA的各个版本，困难++??????????是FFPA的MMA kernels，毕竟其性能比SDPA EA好太多了，给5个星也不为过。由于Kernels数量接近200，所以每个难度级别只截取部分来展示，若要看完整列表，请跳转到CUDA - Learn - Notes仓库。

我除了撰写CUDA笔记与示例，还整理出100多篇自己钟爱的技术博客。从这些博客里我获益匪浅。每次阅读后都觉得写得太棒了，自愧不如。所以我把它们整理到CUDA - Learn - Notes中推荐给大家。这些博客我都按主题分类了，大家可按需查看，例如：

本文全面介绍了CUDA - Learn - Notes v3.0升级后的新特性。CUDA - Learn - Notes v3.0包含200多个CUDA/Tensor Cores Kernels、PyTorch Python绑定、100多篇LLM/CUDA技术博客推荐，还有HGEMM - MMA和FlashAttention - 2 MMA的高性能实现，以及FFPA对FlashAttention - 2（1K headdim~）的扩展。相关的repo链接如下：兄弟们，欢迎来点赞（星标）呀~