很多人在长文本训练里都撞过这堵墙序列长度从 2048 提到 4096显存直接爆掉。看日志先怀疑 Q、K、V但把中间变量一列真正吃显存的是两个看起来不起眼的矩阵——S 和 P。这一讲我们就把 FlashAttention 掰开揉碎看看它怎么做到不存这两个中间矩阵再用 Triton 手写一个能上生产的前向 kernel顺便说说我在实际跑 kernel 时踩过的坑。FlashAttention 这几年几乎成了 GPU 上跑 Transformer 的默认选项很多人把它简单理解成“省显存”。其实它的核心贡献远不止显存更重要的是它把注意力机制里那些反复写回显存、再从显存读出来的中间结果全部压缩在一个 kernel 内部流转省下了大量显存读写。如果你正在做长文本训练、推理加速或者想自定义 attention 变体这篇文章会把原理、公式、Triton 实现和常见问题一次讲清楚看完可以直接动手复现。1. 标准 Attention 的显存都去哪儿了中间矩阵的真实开销1.1 标准实现的核心计算与存储标准 Attention 的计算流程写出来非常简洁S QK^T / sqrt(d)P softmax(S)O PV三行公式但 PyTorch 在背后做的事情远不止三行。每一步都会在显存里分配一个新的张量S 算完被写进显存P 算完又被写进显存最后 O PV 还要再写一次。反向传播时为了算梯度PyTorch 的自动求导机制还保留了 S 和 P或者说保留了未 mask 的 P这些中间变量全部要占用显存。真正的问题在于S 和 P 的大小是 B × H × N × N其中 B 是 batchH 是头数N 是序列长度。Q、K、V、O 的大小是 B × H × N × DD 一般是 64 或 128。N 一变大S 和 P 会迅速膨胀它们的规模远超 QKV 本身。我把这个账算给同学听的时候很多人一开始都不信觉得 Q、K、V 怎么着也有三份S 就一份能占多少其实你把数量级算一遍就明白了。1.2 用数字算一笔账假设一个不太夸张的配置B8H32N2048D128权重用 fp16每元素 2 字节。Q、K、V、O 四个张量加起来元素个数是 8 × 32 × 2048 × 128 × 4算下来每个张量 6700 万左右四个总共约 2.68 亿元素每个 2 字节约 536MB。这个量级在现在的显卡上还算轻松。再看 S 矩阵8 × 32 × 2048 × 2048约 10.7 亿元素fp16 下约 2.1GB。P 矩阵和 S 一样大又是约 2.1GB。两个加起来就是 4.2GB比 QKV 全套还多出近 8 倍。如果序列长度拉到 4096S 单张就是 8.6GBP 再来 8.6GB单是这两个矩阵就把一张主流加速卡的显存吃干净了。这就是为什么你在长文本任务里一扩序列就 OOM根本不是 Q、K、V 的问题是中间矩阵在作祟。1.3 为什么这个 N² 如此致命关键在 S 和 P 的维度是 N×N序列每翻一倍中间矩阵占的显存翻四倍。D 只有 128QKV 本身是线性增长根本扛不住 N² 的爆炸。这也是 FlashAttention 想要解决的核心问题能不能把 N×N 的中间矩阵彻底从显存里拿掉拿掉了之后显存占用就从 O(N²) 降回 O(N)长序列场景才能继续往下做。值得注意的是FlashAttention 并没有减少计算量。S 该算还是要算P 该算也还是要算只是它不再把 S 和 P 完整地写进显存。这种“计算不变、存储和 IO 大减”的思路让它在长序列上收益极大但在短序列上优势反而不明显这一点后面实战部分会再提到。2. 不存中间矩阵的数学前提online softmax 与分块累加2.1 softmax 的朴素实现为什么必须扫两遍先看 softmax 的朴素实现。我们要对 S 的每一行做归一化m max(S[i, :])l sum(exp(S[i, :] - m))P[i, :] exp(S[i, :] - m) / l这里 max 和 sum 都需要看完整行才能确定。所以标准实现必须先把 S 整行算完再读一遍做 exp 和累加。这迫使 S 必须完整存在显存里。但 FlashAttention 利用了一个数学上的小技巧softmax 的归一化过程可以拆成多个局部累加只要每次合并时把之前块的最大值和累加和做一次修正就行。这个技巧通常叫 online softmax是 FlashAttention 能成立的数学根基。2.2 online softmax 的合并公式假设我们已经处理了前 j-1 个 block得到当前行的最大值 m_old 和累加分母 l_old同时输出还没除分母的中间结果 O_old。现在来了第 j 个 block它局部的最大值是 m_block局部 exp 和是 l_block。合并方式如下m_new max(m_old, m_block)alpha exp(m_old - m_new)l_new l_old × alpha l_blockO_new O_old × alpha O_block_newO_old 在之前已经用 m_old 做了局部归一化分子 exp 减的是 m_old现在全局最大值从 m_old 变成了 m_new那么旧块的所有 exp 项都要乘上 exp(m_old - m_new) 这个缩放因子把分母对齐。l_old 同样要乘 alpha。新块本来就用 m_new 做归一化所以 l_block 和 O_block_new 不需要再乘。这个公式的本质是最大值和分母都是可以分块合并的不需要完整的一行数据才能算。每个 block 只要把自己的局部统计量传给下一个 block累计过程中的输出 O 始终是“未完成分母归一化”的中间值直到整行扫完最后再统一除以 l 即可。2.3 FlashAttention 前向分块流程有了 online softmaxFlashAttention 的前向就变成一个双层循环外层循环遍历 Q 的行块每个 block 是 BLOCK_M 行。内层循环遍历 K、V 的列块每个 block 是 BLOCK_N 列。对于每个 Q_i 块把 K、V 按列分成多块依次计算局部 S_ij更新该 Q_i 块的 m、l 和输出累加。具体到内层循环每一步先算局部 S_ij Q_i × K_j^T × scale再和当前行的最大值比较得到 m_new。然后算局部 exp 矩阵 P_ij exp(S_ij - m_new)这一步作用在二维 block 上。接着更新全局统计量l_i l_i × exp(m_old - m_new) rowsum(P_ij)O_i O_i × exp(m_old - m_new) P_ij × V_j。关键是S_ij 和 P_ij 只在这个内层循环里存活。算完 P_ij马上乘 V_j累加进 O_i。这个二维 block 不会写回显存它的生命周期只在缓存和寄存器里。最终每个 Q_i 块处理完所有 K、V 块后用 l_i 对 O_i 做一次除法得到最终输出。由于整个过程把 Q、K、V、O 读完、中间矩阵算完、输出写回全部压在一个 kernel 里完成S 和 P 根本没有机会落回显存。这就是“可以不存中间矩阵”的直接原因。2.4 反向为什么也能不存重计算换显存前向不存中间矩阵还不够因为反向传播时标准实现还需要知道 S 和 P 才能算出 dS、dQ、dK、dV。FlashAttention 的做法是反向时不读显存里的 S/P而是重新算一遍前向用重计算的方式恢复 P。反向需要保存的东西被压缩到极致——只需要保存每行的 m 和 l 两个统计向量大小是 B × H × N相比 B × H × N × N 的中间矩阵几乎可以忽略。重计算的额外计算量在硬件上通常可以接受。因为注意力机制的计算密度很高重算一次前向所增加的 FLOPs换来的是 N² 显存的释放这笔交易在长序列场景下非常划算。这也解释了为什么不少工程同学会发现FlashAttention 不仅省显存训练速度还更快。它的 IO 复杂度从标准实现的 O(N²) 降到了 O(N² × D² / M)其中 M 是缓存大小。当序列变长时IO 瓶颈比计算瓶颈更严峻FlashAttention 正好打在七寸上。3. Triton 实战手写 FlashAttention Forward Kernel3.1 为什么选 Triton 而不是手写底层 GPU 指令FlashAttention 论文里的参考实现是用底层 GPU 指令手写的性能极限但代码量很大普通人很难直接上手修改。工业界更常用的方案是 Triton一种面向 GPU 的类 Python 编程语言写起来像伪代码但生成的性能可以非常接近手写水平。Triton 最大的优势在于它会自动处理 block 调度、内存布局、缓存复用这些脏活。你只要用tl.dot、tl.load、tl.store这样的高层算子描述计算逻辑它就会帮你映射到底层指令。对于 attention 这种形状规则但访存模式复杂的算子开发效率高得不是一点半点。更重要的是Triton 很容易改。你想加一个自定义的 score mask、想支持 window attention、想调 block 大小做自动调优只需要改几行代码。相比之下手写底层实现每次改动都要重新处理很多底层细节维护成本很高。3.2 环境准备与快速验证如果你的环境还没有 Triton装一下很快pip install triton装完可以验证一下版本和基本可用性import triton import triton.language as tl print(triton.__version__)Triton 对 Python 版本和 GPU 后端有要求但主流加速卡基本都能跑。需要注意 Triton 的版本差异部分 API 在不同大版本里名字会变。比如tl.trans在 3.x 里可用早期版本里可能没有。遇到 API 报错先查当前版本的文档别拿旧示例直接套。3.3 教学版 Kernel单头、单 batch、直接可读下面这个代码是教学简化版默认输入是一块 [N, D] 的 Q、K、V只处理一个 batch、一个 head。生产环境要处理多头和 batch我会在 3.5 里说明扩展思路。import torch import triton import triton.language as tl triton.jit def flash_attn_fwd_single_head( Q, K, V, O, M, L, scale, N_CTX, stride_q, stride_k, stride_v, stride_o, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, CAUSAL: tl.constexpr, ): pid tl.program_id(0) offs_m pid * BLOCK_M tl.arange(0, BLOCK_M) offs_n tl.arange(0, BLOCK_N) offs_d tl.arange(0, HEAD_DIM) # 加载 Q 的行块 q tl.load( Q offs_m[:, None] * stride_q offs_d[None, :], maskoffs_m[:, None] N_CTX, other0.0, ) m_i tl.full((BLOCK_M,), float(-inf), tl.float32) l_i tl.zeros((BLOCK_M,), tl.float32) acc tl.zeros((BLOCK_M, HEAD_DIM), tl.float32) for start_n in range(0, N_CTX, BLOCK_N): offs_n_cur start_n offs_n # 加载 K 的列块注意这里是 [BLOCK_N, HEAD_DIM] k tl.load( K offs_n_cur[None, :] * stride_k offs_d[None, :], maskoffs_n_cur[None, :] N_CTX, other0.0, ) # 局部 QK^T用 fp32 做累加避免 fp16 精度问题 qk scale * tl.dot(q, tl.trans(k), out_dtypetl.float32) if CAUSAL: qk tl.where(offs_m[:, None] offs_n_cur[None, :], qk, float(-inf)) m_ij tl.maximum(m_i, tl.max(qk, axis1)) p tl.exp(qk - m_ij[:, None]) l_ij tl.sum(p, axis1) alpha tl.exp(m_i - m_ij) l_i l_i * alpha l_ij acc acc * alpha[:, None] v tl.load( V offs_n_cur[:, None] * stride_v offs_d[None, :], maskoffs_n_cur[:, None] N_CTX, other0.0, ) # p 是 fp32V 可能是 fp16/bf16转成同类型再做矩阵乘 acc tl.dot(p.to(v.dtype), v) m_i m_ij acc acc / l_i[:, None] tl.store( O offs_m[:, None] * stride_o offs_d[None, :], acc.to(O.dtype.element_ty), maskoffs_m[:, None] N_CTX, ) tl.store(M offs_m, m_i, maskoffs_m N_CTX) tl.store(L offs_m, l_i, maskoffs_m N_CTX)这个 kernel 的逻辑和伪代码几乎一样不需要额外状态管理。代码里的 mask 是为了处理序列长度不是 BLOCK_M 或 BLOCK_N 整数倍的情况。other0.0会把越界位置填零对 softmax 的结果没有实质影响。调用方式def flash_attn_forward(q, k, v, causalTrue): assert q.dim() 2, 教学版只处理单 batch、单 head 输入 N, D q.shape out torch.empty_like(q) m torch.empty(N, dtypetorch.float32, deviceq.device) l torch.empty(N, dtypetorch.float32, deviceq.device) BLOCK_M 64 BLOCK_N 64 grid (triton.cdiv(N, BLOCK_M),) flash_attn_fwd_single_head[grid]( q, k, v, out, m, l, 1.0 / (D ** 0.5), N, q.stride(0), k.stride(0), v.stride(0), out.stride(0), BLOCK_MBLOCK_M, BLOCK_NBLOCK_N, HEAD_DIMD, CAUSALcausal, ) return out3.4 用实测脚本对比显存和速度要验证这个 kernel 是否真的有效不要用眼睛看跑脚本对比。下面是最基础的验证流程import torch.nn.functional as F def torch_attention(q, k, v, causalTrue): s q k.transpose(-2, -1) * (q.shape[-1] ** -0.5) if causal: mask torch.triu(torch.ones_like(s), diagonal1).bool() s s.masked_fill(mask, float(-inf)) p torch.softmax(s, dim-1) return p v N, D 2048, 128 q torch.randn(N, D, devicecuda, dtypetorch.float16) k torch.randn(N, D, devicecuda, dtypetorch.float16) v torch.randn(N, D, devicecuda, dtypetorch.float16) # Triton 版 torch.cuda.reset_peak_memory_stats() out_tri flash_attn_forward(q, k, v) tri_ms triton.testing.do_bench(lambda: flash_attn_forward(q, k, v)) tri_peak torch.cuda.max_memory_allocated() # 标准版 torch.cuda.reset_peak_memory_stats() out_torch torch_attention(q, k, v) torch_ms triton.testing.do_bench(lambda: torch_attention(q, k, v)) torch_peak torch.cuda.max_memory_allocated() print(fTriton : {tri_ms:.3f} ms, peak {tri_peak / 1e6:.1f} MB) print(fPyTorch : {torch_ms:.3f} ms, peak {torch_peak / 1e6:.1f} MB) print(fmax diff: {(out_tri.float() - out_torch.float()).abs().max().item():.5e})跑这个脚本你会看到 Triton 版的峰值显存低很多尤其把 N 拉到 4096 时差距更夸张。速度方面短序列可能差距没那么大但序列越长Triton 版的优势越明显。这正是 IO 优化在长序列场景下的实际体现。3.5 从教学版到工业版Batch、多头与自动调优教学版只处理一个 [N, D] 张量生产环境里你的输入通常是 [B, H, N, D]。扩展方式有两种第一种是外层循环把每个 batch 和 head 作为一个独立任务喂给同一个 kernel简单粗暴但会引入额外循环开销。第二种是把 grid 改成二维一个维度编号 batch/head另一个维度编号行块。kernel 内部用 program_id 拿到对应的 batch/head 索引再通过 stride 计算正确的内存偏移。推荐第二种因为 Triton 的 block 调度会做得更好。具体写的时候你需要在 kernel 参数里额外传入 batch_head 总数然后pid_bh tl.program_id(0) pid_m tl.program_id(1)Q、K、V、O 的指针偏移都要加上pid_bh * head_seq_len的部分。这部分改动看起来小但写错过一次就容易越界或者踩到别的位置的数据调试的时候要注意。另外工业级 kernel 一般会加自动调优。Triton 的triton.autotune可以直接对 block 大小、num_warps、num_stages 做搜索triton.autotune( configs[ triton.Config({BLOCK_M: 64, BLOCK_N: 64}, num_warps4, num_stages2), triton.Config({BLOCK_M: 128, BLOCK_N: 64}, num_warps4, num_stages3), triton.Config({BLOCK_M: 128, BLOCK_N: 128}, num_warps8, num_stages3), ], key[N_CTX], )你只需要在一开始挂上装饰器Triton 会在第一次调用时自动尝试多组配置后续使用它选出的最优配置。不同 GPU 上最优 block 大小差异很大自动调优是保证“换机器不降速”的稳妥做法。4. 常见问题与排查技巧精度、性能与正确性4.1 精度问题fp16 累加和 softmax 的坑写 FlashAttention 最容易踩的坑是矩阵乘累加精度。QK^T 的结果 S 里数值范围随着序列长度增大可能很宽。如果直接用 fp16 做累加很容易出现数值饱和。解决方法是把tl.dot的out_dtype设为tl.float32让输出 S 和 P 计算都在 fp32 下进行。另一个常见问题是 causal mask 里的-inf。在 fp16 下float(-inf)会和正常值相互作用产生 NaN但只要在 softmax 之前做 mask再对整行做 max通常不会出问题。怕的话可以换成很小的负数比如-1e4在 fp16 下也足够把对应项压到接近 0但不会引发 NaN。验证精度的方式是把 Triton 输出和 PyTorch 输出对比看 max abs error。一般来说fp16 输入下 max error 应该在 1e-2 以内如果误差大优先检查是不是累加精度不够或者 mask 加错了位置。4.2 性能不达标先看这五个方向我见过不少同学照着网上代码写完一测速度反而比 PyTorch 慢第一反应是“FlashAttention 是不是被吹过头了”。其实大部分性能问题都出在参数和边界处理上。先看 BLOCK_N 是否太大。如果序列长度固定 2048BLOCK_N 选 256尾部可能有一半是无效计算。Triton 的 mask 会保护内存边界但计算量不减少空跑一圈等于浪费算力。把 BLOCK_M 和 BLOCK_N 调成 64 或 128 往往更稳。再看 num_warps。BLOCK_M 比较小时num_warps 取 4 通常比 8 好因为每个 block 能并行的工作有限。BLOCK_M 拉到 128 或 256 时num_warps8 可以更好利用线程资源。这个参数不是越大越好。还有一个容易被忽略的地方causal 模式下的循环边界。标准 FlashAttention 在 causal 下可以跳过完全在右上三角的 K/V 块让内层循环从合适的 start_n 开始而不是从 0 扫到最后。这个优化对小 BLOCK_N 尤其明显能直接把计算量砍一半。最后确认一下你的 kernel 是不是真的“融合”了。如果代码里在循环内用了 Python 的torch操作或者把中间量return到 Python 端那它就不是一个纯 Triton kernel性能根本起不来。整个 inner loop 必须全在triton.jit函数内部完成。4.3 正确性验证别只测一组数据我拿到一个新 kernel不会只测一个 N 就收工。至少要测这些场景序列长度对齐 BLOCK_M 的、不对齐的、N 很小的、N 很大的、causal 和非 causal 的、fp16 和 bf16 的。每个场景都跑一遍和 PyTorch 的对比。不对齐的场景是最容易出 bug 的。mask 如果写错尾巴上的数据可能被填 0 或者被重复计算。测试方法很简单构造 N 1000 这种不是 64 整数倍的输入对比输出差异。如果 N 对齐时没问题、不对齐时出错基本可以断定 mask 处理有问题。还有一个检查点是 M 和 L 这两个统计向量。前向结束保存的 m 和 l 要传给反向重算使用。如果你发现 M 里出现-inf或者 L 里出现 NaN说明 mask 可能覆盖了整行或者某一行所有 QK 值都被 mask 成了-inf。正常时每一行至少有一个位置有效m 会是有限值。4.4 关于反向传播的一点经验本文只写了前向 kernel但生产环境训练一定需要反向。自己写反向 kernel 的工作量比前向大不少因为要重算 P、再算四个梯度。如果你只是想把 FlashAttention 用起来可以考虑先用 Triton 接 PyTorch 的自动求导通过自定义 autograd Function 包装前向和反向。具体做法是在 forward 里调用 Triton kernel并保存 M 和 Lbackward 里写一个反向 Triton kernel或者调用重算逻辑的 kernel返回 dQ、dK、dV。由于前向已经验证过正确性反向可以分步验证每个梯度再和 PyTorch 的 gradcheck 对比。gradcheck 对 fp16 可能不稳定建议先在 fp32 下验证逻辑再切回 fp16 看数值。我在实际调这个 kernel 时的体会是原理看起来很简单无非是分块 online softmax但自己动手写一遍你才会真正理解为什么每个 block 的 m/l 要这样合并为什么输出在最后才除分母。纸上得来终觉浅调一次精度、跑一次性能数据比读十遍论文都管用。如果你的线上模型还在用标准 attention我建议先别急着上各种稀疏注意力或者状态空间模型。把“不存中间矩阵”这件事做好把标准的 softmax attention 融合成一个 kernel往往是性价比最高的一步。后续要扩展也可以在这个 kernel 基础上继续叠加 causal、window mask、flash decode 这些方案越往后越顺。