Triton使用和Softmax实现
参考 Triton论文
参考 GPU MODE Lecture 14: Practitioners Guide to Triton
从Trinton主页引用的话
现代 GPU 的架构大致可以分为三个主要组件 ——DRAM、SRAM 和 ALU—— 在优化 CUDA 代码时必须考虑每个组件：
- 来自 DRAM 的内存传输必须合并为更大的事务，以利用现代内存接口的大型总线宽度。
- 数据在被再次使用之前，必须手动存储到SRAM中，并且要对数据进行管理，以便在检索数据时尽量减少共享内存存储体冲突的情况。
- 计算必须在流式多处理器（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的三条说明的意思就是：
- 因为DRAM很大，比较容易占满总线带宽，所以尽量合并传输的事务可以减少传输的时间，让高速公路跑满。
- 如果数据要重复利用，反复参与计算，尽量让他们在SRAM当中能够缓存住，比如L1的读取只要34个cyle，能比从L2中快6到7倍。
- 尽量跑满并行度，并且利用更高效的计算单元，比如Tensor Core。
这个是OpenAI给出的GPU架构的简图，我们需要明确不同内存，缓存，和执行单元的周期之间的关系就比较好理解GPU计算当中的性能瓶颈。
Triton的目标其实就是优化 HBM -> SRAM -> 寄存器 的带宽，这在Torch里面直接实现不了，通过一些融合算子是可以减少写回到HBM的。
Triton的文档给出的很多实现的代码，可能都不太奏效了，笔者自己测试下来并没有超过torch本身的实现，
可能torch本身也再不断改进吧，这些差别很快就超越了，但是在一些写自定义融合算子方面应该还是比较有优势的。