大模型之Flash attention V1

Flash attention：Fast and Memory Efficient Exact Attention with IO-Awareness

Fast with IO-Awareness：计算快。 通过降低对显存 HBM 的访问次数来加快 整体运算速度，这种方法 被称为 O-Awareness。具体是通过 分块计算 tiling 和 核函数融合 kernel fusion 来降低对显存的访问。

Memory Efficient： 节省显存。 通过 trick 2 使得 forward和backward阶段的存储压力 从 O(N方)降至 O(N)

Exact Attention：精准注意力。 做到了完全等同于 标准attention的实现方式。

问题归纳为：

（1）为什么能降低对显存HBM的访问次数：充分利用SRAM

（2）如何实现的trick2来节省显存：Tiling分块机制

Transformer的卡点

卡点在读写速度。

pi：硬件算力上限。每秒钟所能完成的浮点运算数，单位FLOP/s。1TFLOP/s指的是每秒万亿次，PFLOP/s指的是千万亿次每秒。
beta：硬件带宽上限。每秒所能完成的内存交换量。单位是Byte/s。
pi t：某个算法所需的总运算量
beta t：某个算法所需的总数据读取存储量

事实1：数据在读取的同时，可以同时计算；最终一个算法运行的总时间，取决于计算时间和数据读取时间中的最大值。即计算和读取哪个耗时长、哪个就是更限制。

Transformer中Attention矩阵计算分析的总结：

计算限制：大矩阵乘法（N和d）、通道数很大的卷积运算。

内存限制：逐点运算操作（激活函数、dropout、mask、softmax、BN和LN）。

flash attentin重点优化对象：在计算attention矩阵时受到了内存限制的softmax。

上图显示，当带宽上限beta高，而算力上限pi低时，计算强度就会因为内存限制而弱（没有发挥出高带宽的优势）。

背景：GPU的存储与计算

SRAM：Static Random Access Memory 静态随机存取存储器

HBM：High Bandwidth Memory 高带宽内存

on chip 片上内存：主要用于 Cache 缓存及少数特殊存储单元 texture，如SRAM，存储空间小、带宽大。

off chip 片下内存：主要用于全局存储，即显存，如HBM

硬件计算逻辑：会先从显存HBM中把数据加载到片上 SRAM，在片上由SM读取并进行计算，然后将计算结果再通过SRAM返回给显存。

所以：SRAM的存储为上限，尽量保证每次加载数据都把SRAM打满，节省数据读取时间。

背景：kernel融合

含义：在SRAM存储容许的情况下，能合并的计算我们尽量合并在一起，避免重复从显存读取数据。「把不超过SRAM内存的算子，无论是否可并行，一次性放进SRAM中，减少HBM和SRAM的IO次数」

flash attention 具体实现举例：将矩阵乘法、mask、softmax、dropout操作合并成一个kernel，做到了只读一次、只写会一次。

核心优化技术：Tiling 分块计算

分块的核心是 将大矩阵的乘法 转化为 小矩阵乘法的拼接（行列的分块 + 两组for循环）

分块的细节：

（1）分块后如何计算attention score： QK相乘之后的S矩阵，按照列分块。

（2）如何正确计算O：不断用当前最新的rowmax和rowsum去更新，直到遍历完最后一块「类比梯度下降的思想」

（3）如何优化IO、解决memory-bound问题：见下图中对伪代码的解释，用小矩阵就是抹去了对大矩阵S、P等的读写。

其他

具体运算量的计算，IO复杂度 详见 图解大模型计算加速系列：FlashAttention V1，从硬件到计算逻辑

效果