在展开正式多模态前会再插入几篇现代工程优化技术的相关内容作为支撑。一直以来针对 Attention 的计算量问题主要有两条路线近似注意力如稀疏注意力、低秩近似等总结来说就是压缩用质量换速度。更好的硬件利用充分利用 GPU 的计算单元让显存的数据搬运不再成为瓶颈这也是现代 LLM 的主流。22 年的论文FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness 走的就是第二条路。它不仅没有牺牲精度反而在注意力计算上实现了 2-4 倍的实际加速并将注意力中间结果的显存占用从降到了。 其核心观点是注意力计算慢不是计算量太大而是显存搬运太频繁。这一技术涉及到 GPU 相关的硬件知识因此先行补充GPU 上的注意力#1.1 计算指标和 GPU 核心构建#我们知道支撑 LLM 庞大算力的基础设施是 GPU因此要理解 FlashAttention首先需要理解 GPU 是怎么工作的。先展开两个核心指标FLOPsFloating Point Operations指浮点运算次数代表“计算量”。也就是 GPU 的计算单元需要做多少次乘法和加法衡量的是算法的算力消耗。IOInput/Output指数据在不同内存之间的搬运次数与字节量代表“访存量”。也就是数据从慢速显存搬到快速缓存计算完再搬回去的开销。很显然前者是我们的算法账而后者则是硬件相关的开销在训练和推理中二者都缺一不可。我们已经十分熟悉在代码上进行算法优化的逻辑而 FlashAttention 是首个把 IO 优化思想引入注意力机制的工作。由此我们继续展开现代 GPU 有两种主要的存储空间HBMHigh Bandwidth Memory高带宽显存容量大A100 有 40/80GB但带宽有限我们在市面上听说的多少 GB 的显卡指的都是 HBM。SRAM片上共享内存容量极小A100 每个 SM 只有 192KB但带宽极高速度极快它的层次和 CPU 中的 cache 相似但更加受我们操作者控制。为了获得最高计算效率GPU 的计算单元往往只直接读写 SRAM不会直接访问 HBM。所以要从 HBM 读取数据到 SRAM再从 SRAM 做计算最后把结果写回 HBM这一过程涉及大量 IO 操作。abbb133b-ef69-4e58-83ce-9fc133efaa3f.png1.2 注意力中的 IO 操作#我们对标准注意力计算的过程已经十分熟悉了这里我们扩展开来看看其计算在 GPU 上的具体 IO 过程a1537138-5570-488c-abf1-9a5d3c93f83a.png注意这里反复进行了 6 次 HBM ↔ SRAM 的数据搬运。举个例子对于序列长度、维度多头总和、批量大小的情况我们进行一个估计:总搬运量搬运趟数每趟搬运的元素量字节数字节假设 FP16最终大约的数据量就是现在扩展上下文到这个数字就会暴涨而一个 GPU 的 HBM 带宽大约是 2 TB/s也就是说仅仅数据搬运IO就要 40 多毫秒。注意这只是一次注意力计算。相比之下A100 的算力高达 312 TFLOPS完成这些矩阵乘法的实际计算FLOPs可能只需要不到 1 毫秒。计算单元极快。于是我们发现了问题标准注意力的瓶颈不在计算FLOPs而在访存IO。是数据搬运太慢导致 GPU 大量时间在“空转等数据”。由此FlashAttention 开始了优化FlashAttention 的核心思路#FlashAttention 的想法很直接与其把整个 Q、K、V 搬来搬去不如每次只加载一小块到 SRAM 上进行完所有计算再写回 HBM。这样原本需要在 HBM 和 SRAM 之间来回搬运多次的数据现在就只需要一次完整的读取和一次写入。但这里有一个硬性问题Softmax 是一个全局操作要计算某个位置的 softmax需要知道所有位置的分数。而如果要使用刚刚说的分块计算显然每个块只能看到自己的局部信息怎么算全局 softmax 得到注意力权重答案是 18 年 NVIDIA 研究者的论文 Online normalizer calculation for softmax提出的 Online Softmax在线 Softmax。2.1 Online Softmax#我们知道对于一个包含个元素的向量Softmax 函数将其转换为概率分布的标准公式为但在实际工程中我们往往并不会使用这个公式而是使用 Safe Softmax其中这是因为指数函数增长极快。如果比较大比如 1000会直接超出计算机浮点数的表示范围变成 inf无穷大导致最终计算结果是 NaN。因此Safe Softmax 在分子分母同除以一个常数通常取向量中的最大值。因为指数相减等于相除不会改变最终的相对比例但把所有指数都拉到了的范围避免了上溢出。但这仍然没有解决我们现在的问题因为无论是取最大值还是计算分母我们还是需要遍历所有元素。而 Online Softmax 的思想是这样的当数据无法一次性全部读入内存或者需要分块计算时可以使用流式更新的方式计算最大值和分母即。假设我们正在逐个读取向量的元素到第个元素时更新最大值然后更新分母最终结果其核心技巧是的公式中的这一项的作用是当遇到更大的新最大值时把之前累加的分母“按比例缩小”使其基准与新的最大值对齐。我们来看一个简单实例假设向量为初始化第一步我们读入 2更新最大值和分母此时继续第二步读入 1再次更新现在继续第三步读入 3注意这里最大值发生变化因此之前累加的分母会重新缩放最终得到与标准 Safe Softmax 的结果完全一致。 整个过程中我们只需要维护两个标量即可而不需要等全部数据读完后再计算。但很显然这种 Online Softmax 是逐个更新并不符合 FlashAttention 分块计算逻辑因此我们还要再进行适配2.2 把 Online Softmax 嵌入注意力#在注意力中我们不仅要算 softmax还要用 softmax 的结果乘以 V得到最终输出由此FlashAttention 的最终做法是固定一个Q块不断流过所有KV块累积输出并修正。d3ad6edf-f62c-423b-ac1a-113d2ebee69b.png再展开一下这个过程对于某个 Q 块读取当前的 K 块和 V 块到 SRAM 并进行以下步骤对当前的 Q 块与 K 块算分数矩阵。对分数矩阵做局部 softmax得到局部权重。用局部权重乘以 V 块得到局部输出。基于当前已处理的块维护全局和并用它们修正之前累积的输出。为了方便理解我们只看单个 Query Token。对于这个 Token 来说它需要和所有 Key Token 计算注意力分数然后做一次完整的 Softmax。因此它会维护三个状态目前见过的最大注意力分数。Softmax 的归一化分母。当前已经累积得到的输出向量。当 KV 被分成多个块依次加载时这三个量会随着块的处理不断更新假设已经处理了前个 KV 块我们得到如下累积量前个块中的最大分数。前个块的指数修正和。前个块的部分输出。现在开始处理第个 KV 块当前 Query 与这一块 Key 计算注意力分数求出当前块中的最大分数更新截至目前的全局最大值按照 Online Softmax 的思想修正之前累计的分母这里第一项对应历史 KV 块的贡献第二项对应当前 KV 块的贡献。当新的最大值出现时历史部分会自动重新缩放到新的基准下。4c25839e-068b-4cb9-9486-f5a0e1f50d95 (1).png修正之前的输出并加上当前块的贡献0282d607-638c-46cd-8403-0f0dc3fd8d46.png这样处理完当前块后就得到了新的随后继续读取下一个 KV 块重复这一过程直到所有 KV 块都被遍历完成。需要强调的是上述推导是针对单个 Query Token 进行说明的。实际 FlashAttention 中会同时处理一整个 Q 块因此这里的、和实际上并不是标量而是对应每行 Query Token 分别维护的一组向量状态。总结来说FlashAttention 可以看成代码里的嵌套循环其中外层循环遍历 KV 块内层循环遍历 Q 块。整个过程中只有 QKV 块从HBM 加载最终的输出 O 写入 HBM中间的所有结果都只存在于 SRAM 中不写回 HBM。这就是 FlashAttention 在注意力中引入 IO 优化思想得到的突破。FlashAttention 的局限和后续发展#FlashAttention 虽然解决了 IO 瓶颈但它仍然没有完全发挥 GPU 的计算能力论文报告FlashAttention 在 A100 上只达到了 25-40% 的理论峰值 FLOPs。其原因大体分析如下非矩阵乘法的开销softmax 的指数运算、逐元素的乘法和加法等操作不如矩阵乘法高效。线程块和 warp 的调度FlashAttention 在并行化上还有优化空间部分线程块可能处于空闲状态。共享内存的读写虽然减少了 HBM 读写但 SRAM 内部的通信也存在开销。但这也同时说明了该方向仍然存在极大的优化空间。在此之后23 年的论文FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning提出了 FlashAttention-2对 FlashAttention 的工作分配进行了系统性的优化其核心内容是通过应用更多矩阵乘法优化线程并行线程内分工的方式以更好的并行得到加速效果。FlashAttention-2 中的各种“榨干 GPU” 的优化其实更符合现代工业标准。但这里涉及到较多的操作系统和硬件内容就不再详细展开了。FlashAttention-2 将 GPU 利用率从 FA1 的 25-40% 提升到了 50-73%相比 FA1 实现了大约 2 倍的速度提升。还没结束24 年的论文FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision 将目光投向了新一代的 Hopper 架构 GPUH100。