ggaaooppeenngg

本文讨论了稀疏注意力的计算方法，其核心思想是通过选择矩阵的分块，将重要的矩阵挑选出来参与计算。这是因为注意力矩阵具有稀疏性，而如何选择这些重点矩阵是各类算法需要解决的主要问题。

Dynamo 发布以后，我大概速览了一些设计文档，并且提取了一些关键点，并对比一些其他方案的异同点。

参考 OpenAI Triton 主页

参考 Triton论文

参考 GPU MODE Lecture 14: Practitioners Guide to Triton

从Trinton主页引用的话

现代 GPU 的架构大致可以分为三个主要组件 ——DRAM、SRAM 和 ALU—— 在优化 CUDA 代码时必须考虑每个组件：

1. 来自 DRAM 的内存传输必须合并为更大的事务，以利用现代内存接口的大型总线宽度。
2. 数据在被再次使用之前，必须手动存储到SRAM中，并且要对数据进行管理，以便在检索数据时尽量减少共享内存存储体冲突的情况。
3. 计算必须在流式多处理器（SM）之间和内部仔细分区和调度，以促进指令 / 线程级并行性并利用专用 ALU（例如Tensor Core）。

这几句话可能比较抽象，下面给一下这几个组件的指标可能感受更直观，参考Which GPU(s) to Get for Deep Learning: My Experience and Advice for Using GPUs in Deep Learning。
其中指出：

• 全局内存访问（高达 80GB）：~380 个周期
• L2 缓存：~200 个周期
• L1 缓存或共享内存访问（每个流式多处理器高达 128 kb）：~34 个周期
• 融合乘法和加法，ab+c（FFMA）： 4 个周期
• Tensor Core 矩阵乘法：1 个周期

每个操作总是由 32 个线程组成的Warp执行，Warp中的线程必须相互等待。GPU 上的所有内存操作都针对warp进行了优化。

根据Simplifying CUDA kernels with Triton: A Pythonic Approach to GPU Programming的说法，GPU中的HBM（High Bandwidth Memory）等价于我们讲的Global Memory，SRAM对应的是L1和L2 Cache对应的是Shared Memory，这几个词在一些文档中可能会有不同的叫法，但是意思是一样的。

A100中的内存带宽约为 2TB/s，L1 缓存带宽：~100-200 TB/s 理论带宽，L2 缓存带宽：~4-7 TB/s 理论带宽。

再看OpenAI的三条说明的意思就是：

1. 因为DRAM很大，比较容易占满总线带宽，所以尽量合并传输的事务可以减少传输的时间，让高速公路跑满。
2. 如果数据要重复利用，反复参与计算，尽量让他们在SRAM当中能够缓存住，比如L1的读取只要34个cyle，能比从L2中快6到7倍。
3. 尽量跑满并行度，并且利用更高效的计算单元，比如Tensor Core。

这个是OpenAI给出的GPU架构的简图，我们需要明确不同内存，缓存，和执行单元的周期之间的关系就比较好理解GPU计算当中的性能瓶颈。

Triton的目标其实就是优化 HBM -> SRAM -> 寄存器 的带宽，这在Torch里面直接实现不了，通过一些融合算子是可以减少写回到HBM的。

Triton的文档给出的很多实现的代码，可能都不太奏效了，笔者自己测试下来并没有超过torch本身的实现，
可能torch本身也再不断改进吧，这些差别很快就超越了，但是在一些写自定义融合算子方面应该还是比较有优势的。