1. 从显存瓶颈说起Attention 到底卡在哪做 Transformer 推理或训练的朋友大概率都遇到过这样的场景模型结构没变、参数量也没变但序列长度一拉长显存就爆了。尤其是自注意力模块序列长度从 512 涨到 2048显存占用不是线性增长而是平方级膨胀。这个现象背后的原因就是标准 Attention 实现里那个巨大的中间矩阵。标准自注意力的计算过程可以拆成三步先把 Query 和 Key 做矩阵乘法得到注意力分数矩阵再对分数矩阵做 Softmax 归一化最后用归一化后的分数矩阵去加权 Value。问题就出在第一步和第二步之间——那个注意力分数矩阵的大小是序列长度的平方。序列长度 4096 时单头单批次的分数矩阵就有 1600 万个元素如果用 FP32 存储光这一个矩阵就占 64MB。多头、多批次、多层叠加之后显存消耗会迅速失控。更麻烦的是这个中间矩阵在标准实现里需要被完整写入显存然后再从显存读回来做 SoftmaxSoftmax 的结果又要写回去最后再读出来和 Value 相乘。一次 Attention 计算中间矩阵至少经历三次显存读写。GPU 的算力其实早就过剩了真正拖后腿的是显存带宽。这就是所谓的“内存墙”问题——计算单元等数据而不是数据等计算单元。FlashAttention 的核心思路就是冲着这个痛点去的既然中间矩阵这么大、读写这么频繁那能不能干脆不存它答案是能。通过把 Attention 计算拆成小块在片上高速缓存里完成 Softmax 和加权求和中间结果永远不落到显存里。这就是“不存中间矩阵”的字面含义也是它能把 Attention 显存占用从平方级降到线性级的关键。这篇文章适合两类人看一类是想搞明白 FlashAttention 到底怎么做到不存中间矩阵的算法爱好者另一类是准备用 Triton 自己动手实现一版工业级 Attention 的工程师。我会从原理拆解讲到 Triton 实战把每一步的“为什么”都说清楚代码可以直接抄作业。2. 标准 Attention 的显存账本与计算账本2.1 一次 Attention 到底读写了多少数据先把标准 Attention 的数据流捋一遍。假设输入 Q、K、V 的形状都是 [N, d]其中 N 是序列长度d 是每个头的维度。标准实现的计算流程是S Q K^T得到 [N, N] 的分数矩阵P softmax(S)逐行归一化还是 [N, N]O P V得到 [N, d] 的输出从显存读写角度看S 需要被写入一次、读出一次做 SoftmaxP 需要被写入一次、读出一次做加权求和。如果考虑 Softmax 内部还要读 S、写 P那 S 和 P 各自至少经历两次完整的显存往返。以 N4096、d64、FP16 存储为例S 和 P 各占 4096×4096×2 字节 32MB单头单批次下中间矩阵的显存往返量就是 32MB × 4 128MB。这还没算 Q、K、V 本身的读取。对比一下计算量QK^T 是 2×N×N×d 次浮点运算PV 也是 2×N×N×d总共 4×N×N×d。N4096、d64 时大约是 4.3 GFLOP。现代 GPU 的 FP16 算力动辄上百 TFLOPS算这些乘法理论上只要几十微秒。但显存带宽通常只有几百 GB/s 到一两 TB/s搬运 128MB 数据需要上百微秒。计算等数据算力利用率被显存带宽死死压住。这里有个容易被忽略的点Softmax 本身的计算量很小但它强制要求把整行分数都拿到才能算归一化。这个“整行依赖”是标准实现必须物化中间矩阵的根本原因。2.2 为什么不能简单地把矩阵分块算有人可能会想既然矩阵太大那我把它切成小块分别算 Softmax 不就行了问题在于 Softmax 的分母是整行的指数和你只拿到一行的一部分算不出完整的分母。如果强行分块每块算自己的局部 Softmax结果和全局 Softmax 不一致Attention 的输出就错了。这就是 FlashAttention 要解决的核心矛盾Softmax 需要整行信息但整行信息太大存不下。它的解法是“在线 Softmax”——一边遍历行内的块一边增量更新分母和当前输出最后再统一校正。这样只需要维护一个运行中的最大值和指数和不需要把整行都存下来。2.3 显存占用对比平方级 vs 线性级把标准 Attention 和 FlashAttention 的显存占用列个表差距一目了然。假设头维度 d64FP16 存储只算中间矩阵部分序列长度 N标准 Attention 中间矩阵FlashAttention 中间缓存5120.5 MB约 0.06 MB10242 MB约 0.06 MB20488 MB约 0.06 MB409632 MB约 0.06 MB8192128 MB约 0.06 MB标准 Attention 的中间矩阵随 N 平方增长FlashAttention 的片上缓存只和块大小有关与 N 无关。这就是“不存中间矩阵”带来的直接收益。实际训练中这意味着同样显存可以支持更长的序列或者同样序列可以用更大的批次。3. FlashAttention 的核心原理拆解3.1 在线 Softmax 的数学推导在线 Softmax 是 FlashAttention 的数学基石值得把推导过程写清楚。标准 Softmax 对一行分数 x 的计算是softmax(x_i) exp(x_i - max(x)) / sum_j exp(x_j - max(x))减去最大值是为了数值稳定避免 exp 溢出。现在假设这一行被分成若干块我们逐块处理。维护两个运行变量当前见过的最大值 m 和对应的指数和 l。处理新块时新块的最大值是 m_new更新后的全局最大值是 m_updated max(m, m_new)。原来的指数和 l 是基于旧最大值 m 算的需要按比例缩放l_updated l × exp(m - m_updated) sum(exp(x_new - m_updated))输出 O 也要同步缩放O_updated O × exp(m - m_updated) exp(x_new - m_updated) V_new这样每处理一个块就用新的最大值校正之前累积的结果。全部块处理完后O 除以 l 就得到最终输出。整个过程只需要维护 m、l、O 三个运行变量不需要存储任何中间分数矩阵。这个推导的关键在于exp(a - c) exp(a - b) × exp(b - c)。缩放因子 exp(m - m_updated) 就是利用这个恒等式把基于旧最大值的累积结果换算到新最大值下。3.2 分块策略与片上缓存管理有了在线 Softmax接下来要把 Q、K、V 分块。典型的分块方式是外层循环遍历 K 和 V 的块内层循环遍历 Q 的块。对于每个 Q 块初始化 m、l、O然后遍历所有 K、V 块用在线 Softmax 增量更新。为什么外层是 K、V内层是 Q因为每个 Q 块需要和所有 K、V 块交互把 Q 块放在内层可以让它常驻片上缓存减少重复加载。K、V 块在外层被反复读取但它们可以从显存流式加载不需要同时驻留。片上缓存的容量决定了块大小。以常见的 GPU 为例共享内存通常有 100KB 到 200KB 量级。Q 块、K 块、V 块、输出块、以及 m 和 l 的运行变量都要放在共享内存或寄存器里。假设 d64FP16 存储一个块大小 B64 的 Q 块占 64×64×2 8KBK 块和 V 块各 8KB输出块 8KB加上运行变量总共不到 40KB完全放得下。块大小可以适当调大以提高计算密度但要注意不要超出共享内存上限。3.3 重计算换显存反向传播怎么处理前向不存中间矩阵反向传播怎么办标准实现里反向需要用到前向的分数矩阵 P 来计算梯度。FlashAttention 的做法是反向时重新计算一遍前向用重计算换显存。具体来说反向时重新加载 Q、K、V 块重新走一遍在线 Softmax 的前向过程顺便把需要的中间量算出来用于梯度计算。这个策略的代价是反向多了一次前向计算计算量增加约三分之一。但省下的显存是平方级的对于长序列场景非常划算。而且重计算的过程也是分块的中间结果同样不落显存显存占用依然是线性级。重计算不是 FlashAttention 首创但它在 Attention 场景下收益特别大因为中间矩阵的平方级增长是显存的主要矛盾。用计算换显存在算力过剩、带宽受限的当下是明智的取舍。4. Triton 实战手写一版分块 Attention4.1 Triton 为什么适合写这类 KernelTriton 是 OpenAI 开源的一门类 Python 的 GPU 编程语言它的抽象层级介于 CUDA 和 PyTorch 之间。用 Triton 写 Attention Kernel 有几个好处一是它自动处理共享内存分配和线程调度你只需要关注块级别的数据流二是它的编程模型天然适合分块计算tl.load、tl.dot、tl.store 这些原语都是块级别的三是它和 PyTorch 生态无缝集成写好的 Kernel 可以直接在 Python 里调用。相比手写 CUDATriton 省去了大量样板代码比如线程索引计算、共享内存 bank conflict 处理、warp 级别同步等。对于 FlashAttention 这种分块逻辑复杂的 Kernel用 Triton 实现能把代码量压缩到 CUDA 的三分之一左右而且可读性好很多。4.2 前向 Kernel 的完整实现下面是一版简化但可运行的 FlashAttention 前向 Kernel。为了便于理解我假设头维度 d 不超过 128块大小 BLOCK_M 和 BLOCK_N 取 64。import triton import triton.language as tl triton.jit def flash_attn_fwd( Q, K, V, O, stride_qm, stride_qd, stride_kn, stride_kd, stride_vn, stride_vd, stride_om, stride_od, N, scale, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, D: tl.constexpr, ): pid_m tl.program_id(0) offs_m pid_m * BLOCK_M tl.arange(0, BLOCK_M) offs_d tl.arange(0, D) q_ptrs Q offs_m[:, None] * stride_qm offs_d[None, :] * stride_qd q tl.load(q_ptrs, maskoffs_m[:, None] N, other0.0) m_i tl.full([BLOCK_M], float(-inf), dtypetl.float32) l_i tl.zeros([BLOCK_M], dtypetl.float32) acc tl.zeros([BLOCK_M, D], dtypetl.float32) for start_n in range(0, N, BLOCK_N): offs_n start_n tl.arange(0, BLOCK_N) k_ptrs K offs_n[:, None] * stride_kn offs_d[None, :] * stride_kd k tl.load(k_ptrs, maskoffs_n[:, None] N, other0.0) qk tl.dot(q, tl.trans(k)) * scale qk tl.where(offs_n[None, :] N, qk, float(-inf)) m_new tl.maximum(m_i, tl.max(qk, axis1)) alpha tl.exp(m_i - m_new) p tl.exp(qk - m_new[:, None]) l_i l_i * alpha tl.sum(p, axis1) acc acc * alpha[:, None] v_ptrs V offs_n[:, None] * stride_vn offs_d[None, :] * stride_vd v tl.load(v_ptrs, maskoffs_n[:, None] N, other0.0) acc tl.dot(p.to(v.dtype), v) m_i m_new acc acc / l_i[:, None] o_ptrs O offs_m[:, None] * stride_om offs_d[None, :] * stride_od tl.store(o_ptrs, acc.to(O.dtype.element_ty), maskoffs_m[:, None] N)这段代码的核心逻辑和前面推导的在线 Softmax 完全对应。m_i 和 l_i 是运行中的最大值和指数和acc 是累积输出。每处理一个 K、V 块先算 qk 分数更新 m_i 和 l_i缩放 acc再加上当前块的贡献。循环结束后 acc 除以 l_i 得到最终输出。有几个细节值得注意。qk 计算后乘了 scale这是 1/sqrt(d) 的缩放因子防止点积结果过大导致 Softmax 饱和。mask 处理是为了应对序列长度不是块大小整数倍的情况超出部分置为负无穷Softmax 后权重为零。acc 用 FP32 累积最后再转回输出类型这是为了保证数值精度。4.3 块大小怎么选参数计算与实测对比块大小 BLOCK_M 和 BLOCK_N 的选择直接影响性能。块太小计算密度低GPU 利用率上不去块太大共享内存放不下会触发寄存器溢出或共享内存溢出。我实测过几组配置在 d64、N2048 的场景下BLOCK_MBLOCK_N共享内存占用相对吞吐3232约 16KB0.726464约 40KB1.0012864约 72KB1.08128128约 120KB0.9525664约 140KB0.88可以看到 64×64 是一个比较稳的基线128×64 略有提升但共享内存占用翻倍。128×128 开始出现下降因为共享内存接近上限occupancy 降低。256×64 更差寄存器压力太大。实际选型时建议从 64×64 起步根据具体 GPU 的共享内存容量和寄存器数量微调。块大小的选择没有万能答案和 GPU 架构、头维度、数据类型都有关。我的经验是先用 64×64 跑通然后用 Triton 的 autotune 功能在几组候选配置里自动搜索通常能找到比手工调优更好的组合。4.4 反向传播的 Triton 实现要点反向 Kernel 比前向复杂不少核心是要重算前向的中间量然后根据链式法则求 Q、K、V 的梯度。关键步骤包括重新加载 Q、K、V 块重算分数矩阵和 Softmax 输出计算输出梯度对分数矩阵的梯度再分别对 Q、K、V 求梯度。反向的块循环方向和前向相反通常外层遍历 Q 块内层遍历 K、V 块。每个 Q 块需要累积对 K、V 的梯度贡献所以 K、V 的梯度要用原子加或者分块归约的方式写回。Triton 里可以用 tl.atomic_add 实现但原子操作有性能开销更好的做法是每个 Q 块独立计算对 K、V 的贡献然后用一个单独的归约 Kernel 合并。重计算带来的额外开销在反向里体现为多了一次前向的分数计算。实测下来反向耗时大约是前向的 2.5 到 3 倍比标准实现的反向慢一些但显存占用大幅降低长序列场景下整体吞吐反而更高。5. 踩坑记录与性能调优实录5.1 数值精度问题FP16 累积的陷阱最开始我用 FP16 累积 acc结果发现长序列下输出误差明显偏大。原因是 acc 在循环里不断累加FP16 的尾数只有 10 位累加几百次之后精度损失严重。改成 FP32 累积后误差降到可接受范围。同样l_i 和 m_i 也必须用 FP32否则指数和的精度不够。还有一个隐蔽的坑是 exp 的输入范围。如果 qk 的值很大exp 之后可能溢出。虽然在线 Softmax 减了最大值但在更新 m_i 之前qk 本身可能已经很大。我的做法是在算 qk 之后先做一次 clamp把值限制在合理范围内再参与后续计算。这个 clamp 的阈值可以根据 scale 和 d 估算一般取 80 左右比较安全。5.2 共享内存溢出与寄存器压力Triton 在编译时会根据块大小和数据类型自动分配共享内存和寄存器。如果块太大编译会报共享内存溢出或者虽然编译通过但 occupancy 极低。我遇到过 BLOCK_M128、BLOCK_N128、D128 时编译失败的情况报错信息是共享内存超出限制。解决办法是减小块大小或者把 D 维度也分块。寄存器压力是另一个隐形杀手。Triton 的 tl.dot 会占用大量寄存器如果同时有多个大块在寄存器里会触发寄存器溢出到本地内存性能断崖式下跌。用 Triton 的编译日志可以看到寄存器使用量如果接近上限就要考虑减小块大小或者拆分计算。5.3 常见问题速查表问题现象可能原因排查方向解决方法输出误差大FP16 累积精度不足检查 acc、l_i 类型改用 FP32 累积编译报共享内存溢出块大小过大查看编译日志减小 BLOCK_M/BLOCK_N性能远低于预期occupancy 低检查寄存器使用量减小块大小或拆分 D长序列结果异常exp 溢出检查 qk 范围加 clamp 或调大 scale反向梯度不对重计算逻辑有误对比标准实现逐块核对在线 Softmax边界序列出错mask 处理遗漏检查非整数倍长度补全 mask 逻辑5.4 实测性能对比与调优心得在 N2048、d64、批次 8、头数 8 的配置下我对比了几种实现。标准 PyTorch Attention 的显存占用约 1.2GBFlashAttention Triton 实现约 0.3GB显存降到四分之一。吞吐方面标准实现约 120 样本/秒Triton 实现约 280 样本/秒提升超过一倍。序列越长差距越大N4096 时 Triton 实现的吞吐优势接近三倍。调优过程中最大的心得是不要迷信大块。很多人直觉上觉得块越大计算密度越高但实际上共享内存和寄存器的限制会让大块适得其反。我最后稳定的配置是 BLOCK_M64、BLOCK_N64配合 num_warps4、num_stages2。num_stages 控制软件流水线的级数设成 2 可以在加载下一块的同时计算当前块进一步掩盖显存延迟。还有一个容易被忽略的点是 Q、K、V 的内存布局。如果它们不是连续存储的tl.load 的效率会大打折扣。我建议在调用 Kernel 之前先用 contiguous() 把张量整理成连续布局虽然多了一次拷贝但 Kernel 内部的加载效率提升明显整体还是划算的。6. 这套方案还能怎么扩展FlashAttention 的思路不止能用在标准自注意力上。因果掩码场景下只需要在分数矩阵上加上因果 mask让未来位置的分数为负无穷在线 Softmax 的逻辑完全不用改。多头注意力可以把头维度合并到块维度里一次处理多个头提高计算密度。变长序列场景下可以用块级别的序列长度信息来控制 mask避免 padding 带来的浪费。Triton 这边也有不少可以深挖的方向。比如用 tl.dot 的输入精度控制来做混合精度计算用 tl.inline_asm 嵌入特定指令优化关键路径或者用 Triton 的 autotune 做更细粒度的参数搜索。我最近在试的是把 KV 缓存也做成分块的推理时只加载当前需要的块进一步降低显存占用。这套东西我自己在几个项目里用过从最初的磕磕绊绊到后来能稳定跑长序列踩的坑基本都写在上面了。如果你也在做类似的事情建议先把前向 Kernel 跑通用一个小规模输入和标准实现对比数值确认无误后再上反向和调优。块大小从 64×64 起步精度用 FP32 累积这两条能帮你避开大部分初级问题。