★ FEATURED ARTICLE Attention 机制几乎是所有现代大模型推理和训练的绝对核心但很多人第一次听到 FlashAttention 的时候都会有一个疑问算 Attention 不是必须先把 (QK^T) 这个大矩阵算出来、存下来再 softmax、再乘 (V) 吗不存中间矩阵那到底是怎么算出来的我最早接触这个方向的时候也卡在这个点上很久直到自己用 Triton 手写了一遍 FlashAttention 的 kernel才真正把“不存中间矩阵”这件事想明白。这篇内容就把我从原理到工程落地的完整思路拆开讲包括分块计算、在线 softmax、Triton 里的具体实现细节以及实际跑起来会踩的那些坑。适合已经了解基本 Attention 公式、想深入理解 FlashAttention 内部机制并且希望自己动手用 Triton 复现一遍的读者。读完你应该能自己写出一个能跑通的 FlashAttention kernel并且知道每一步为什么这么设计。1. 从“必须存中间矩阵”说起Attention 的显存困境1.1 标准 Attention 到底存了什么先把最朴素的 Attention 计算流程摆出来。给定查询 (Q)、键 (K)、值 (V)形状分别是 ([N, d])这里先忽略 batch 和 head 维度方便讨论标准计算是三步计算得分矩阵 (S QK^T)形状 ([N, N])对 (S) 做 softmax得到注意力权重 (P \text{softmax}(S))形状还是 ([N, N])计算输出 (O PV)形状 ([N, d])。问题就出在第一步和第二步。(S) 和 (P) 都是 ([N, N]) 的矩阵。当序列长度 (N) 是 4096、8192 甚至更长的时候这个矩阵的大小是 (N^2) 级别的。以 (N8192)、单精度浮点4 字节为例一个 ([8192, 8192]) 的矩阵就是 (8192 \times 8192 \times 4 \approx 268) MB。这还只是一个 head、一个 batch 的量。如果 batch 是 32、head 是 32那显存直接爆炸。更关键的是这个中间矩阵在标准实现里是要被写回显存的。GPU 的显存带宽是有限的写一个 (N^2) 的矩阵再读回来做 softmax再写回去再读回来乘 (V)这一来一回的 IO 开销往往比计算本身还大。这就是所谓的memory-bound问题算力没跑满卡在数据搬运上了。1.2 为什么大家默认要存中间矩阵这里要澄清一个认知误区。很多人以为“算 Attention 必须存中间矩阵”是数学上的要求其实不是。数学上 (O \text{softmax}(QK^T)V) 这个式子完全可以不显式构造 (S) 和 (P) 就得到结果。之所以标准实现要存是因为它把计算拆成了三个独立的 kernel或者三个独立的步骤每个步骤之间需要通过显存传递数据。第一个 kernel 算 (S QK^T)结果写显存第二个 kernel 读 (S)做 softmax结果 (P) 写显存第三个 kernel 读 (P) 和 (V)算 (O)。这种拆分的逻辑很自然因为每个步骤都是独立的、可并行的。但它带来的代价就是中间结果的显存占用和反复的显存读写。当 (N) 不大的时候这个代价可以接受当 (N) 上千甚至上万的时候这个代价就成了瓶颈。所以“不存中间矩阵”的本质不是数学上做不到而是工程上要避免把 (N^2) 的中间结果落到显存里。FlashAttention 的核心思想就是把这个计算融合成一个 kernel在片上高速缓存SRAM里完成分块计算中间结果只在 SRAM 里短暂存在不写回显存。1.3 显存墙与计算墙的双重夹击要理解 FlashAttention 的价值得先理解现代 GPU 的存储层次。GPU 的存储大致分几层存储层级典型容量访问速度说明寄存器每线程几十到几百个极快线程私有共享内存/SRAM每 SM 几十到几百 KB很快块内共享L2 缓存几 MB 到几十 MB较快全芯片共享显存/HBM几十 GB慢全局可见FlashAttention 的关键洞察是SRAM 的带宽比显存高一个数量级但容量小。所以策略就是把 (Q)、(K)、(V) 分块每次只把一小块加载到 SRAM在 SRAM 里算完这一块的贡献累加到输出上然后丢弃中间结果。这样显存里永远只有 (Q)、(K)、(V)、(O) 这几个 ([N, d]) 级别的矩阵没有 ([N, N]) 的中间矩阵。但这里有个数学上的难点softmax 是全局的它需要知道整行的最大值和指数和才能算。如果分块计算你怎么在只看一块 (K) 的情况下算出正确的 softmax这就引出了 FlashAttention 最精妙的部分——在线 softmaxonline softmax。2. 核心原理拆解不存中间矩阵是怎么做到的2.1 在线 softmax 的数学推导标准 softmax 对一行向量 (x) 的计算是[ \text{softmax}(x_i) \frac{e^{x_i - m}}{\sum_j e^{x_j - m}} ]其中 (m \max_j x_j)减去最大值是为了数值稳定。问题在于要算这个 (m) 和分母的求和你得先看完整行。在线 softmax 的思路是边看边更新。假设我们分块处理先看了第一块得到一个局部的最大值 (m_1) 和局部的指数和 (l_1)。然后看第二块得到 (m_2) 和 (l_2)。这时候新的全局最大值是 (m_{new} \max(m_1, m_2))新的指数和需要把之前的 (l_1) 按新的最大值重新缩放[ l_{new} l_1 \cdot e^{m_1 - m_{new}} l_2 \cdot e^{m_2 - m_{new}} ]这个缩放因子 (e^{m_1 - m_{new}}) 就是关键。因为 (m_{new} \geq m_1)所以这个因子小于等于 1不会溢出。同理输出的累加也要按这个因子缩放。把这个逻辑套到 Attention 上对于输出 (O) 的每一行我们分块遍历 (K) 和 (V)。每处理一块 (K_j, V_j)就算出这一块的得分 (S_j Q_i K_j^T)更新局部的最大值和指数和然后把这一块的贡献 (P_j V_j) 按缩放因子累加到 (O_i) 上。整个过程只需要保存当前块的 (S_j)大小是 ([B_q, B_k])(B_q)、(B_k) 是块大小不需要保存完整的 ([N, N]) 矩阵。2.2 分块计算的具体流程把上面的推导写成伪代码大概是这样的# 初始化 O zeros([N, d]) # 输出 l zeros([N]) # 指数和 m full([N], -inf) # 行最大值 # 外层遍历 Q 的块 for i in range(0, N, B_q): Q_i Q[i:iB_q] # 加载 Q 块到 SRAM O_i zeros([B_q, d]) l_i zeros([B_q]) m_i full([B_q], -inf) # 内层遍历 K, V 的块 for j in range(0, N, B_k): K_j K[j:jB_k] # 加载 K 块 V_j V[j:jB_k] # 加载 V 块 S_ij Q_i K_j.T # [B_q, B_k] m_ij max(S_ij, axis-1) # 当前块的行最大值 m_new maximum(m_i, m_ij) # 更新全局最大值 # 缩放因子 scale exp(m_i - m_new) P_ij exp(S_ij - m_new[:, None]) # 更新指数和与输出 l_i l_i * scale sum(P_ij, axis-1) O_i O_i * scale[:, None] P_ij V_j m_i m_new # 写回输出 O[i:iB_q] O_i / l_i[:, None]这段伪代码就是 FlashAttention 前向传播的核心。注意几个关键点(m_i) 和 (l_i) 是逐行维护的每个 Q 块对应一组每次更新都要重新缩放之前的累加结果这是在线 softmax 的精髓(S_{ij}) 只在 SRAM 里存在算完就丢不写回显存最终输出要除以 (l_i)因为累加过程中用的是未归一化的指数。2.3 为什么这样能省显存对比一下标准实现和 FlashAttention 的显存占用项目标准 AttentionFlashAttention中间矩阵 (S)([N, N]) 显存不存在中间矩阵 (P)([N, N]) 显存不存在块内 (S_{ij})无([B_q, B_k]) SRAM输出 (O)([N, d]) 显存([N, d]) 显存额外状态 (m, l)无([N]) 显存显存占用从 (O(N^2)) 降到了 (O(N))。当 (N8192) 时标准实现光中间矩阵就要几百 MB而 FlashAttention 只需要几十 KB 的 SRAM 来放块内矩阵。这就是“不存中间矩阵”的直接收益。但省显存只是表面更深层的收益是减少了显存 IO。标准实现要把 (S) 写回显存再读出来FlashAttention 全程在 SRAM 里操作IO 量大幅下降。在 memory-bound 的场景下这个收益比省显存更值钱。3. Triton 实战手写一个 FlashAttention Kernel3.1 为什么选 Triton要自己实现 FlashAttention理论上可以用 CUDA C 写但 CUDA 的门槛不低你得手动管理线程块、共享内存、寄存器分配、warp 调度还要处理各种边界情况。Triton 的出现大大降低了这个门槛。它用 Python 写 kernel编译器负责把高层描述映射到 GPU 的线程和存储层次上你只需要关注块级的计算逻辑。Triton 的几个关键概念program一个 kernel 实例对应一个块的计算blockTriton 里的数据块用tl.arange和tl.load加载tl.dot块级矩阵乘法编译器会映射到 Tensor Coretl.exp、tl.max、tl.sum块级逐元素和归约操作。用 Triton 写 FlashAttention代码量大概一百多行比 CUDA 版本短得多而且可读性好。下面我按自己的实现习惯一步步拆开讲。3.2 环境准备与依赖安装先把环境搭起来。Triton 对 Python 和 PyTorch 版本有要求我实测下来比较稳的组合是# 建议用 Python 3.10 或 3.11 pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cu118 pip install triton2.1.0如果你用的是更新的 PyTorchTriton 通常已经作为依赖被装上了可以直接import triton试试。检查一下版本import torch import triton import triton.language as tl print(torch.__version__) print(triton.__version__)注意Triton 版本和 PyTorch 版本、CUDA 版本之间有兼容性矩阵。如果import triton报错大概率是版本不匹配先对齐版本再往下走。我踩过的最坑的一次是 Triton 2.0 配 PyTorch 2.2编译 kernel 时直接段错误换成 2.1 就好了。另外如果你只是想验证功能CPU 上跑不了 Triton kernel必须有 NVIDIA GPU。显存建议 8GB 以上方便测试较长的序列。3.3 Kernel 的整体结构设计我的实现思路是把 kernel 分成几个部分确定 program 的划分每个 program 负责一个 Q 块(B_q) 行的输出。program 的数量是 (N / B_q)如果 batch 和 head 多就再乘上 batch 和 head 数。加载 Q 块把当前 program 对应的 Q 块加载到寄存器/SRAM。循环遍历 K、V 块对每个 K、V 块算得分、更新在线 softmax 状态、累加输出。归一化并写回最后把累加的输出除以指数和写回显存。这里有个设计选择Q 块和 K 块的大小怎么定。(B_q) 和 (B_k) 不能太大否则 SRAM 放不下也不能太小否则并行度不够、循环次数太多。经验值是 (B_q 64) 或 (128)(B_k 64) 或 (128)。具体要看 head 维度 (d)。如果 (d 64)(B_q 128)、(B_k 64) 是比较稳的组合。如果 (d 128)块大小要相应调小。3.4 关键代码逐段解析先看 kernel 的签名和初始化部分triton.jit def flash_attn_fwd( Q, K, V, O, L, M, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh, stride_om, stride_on, Z, H, N, scale, BQ: tl.constexpr, BK: tl.constexpr, D: tl.constexpr, ): # program id 分解 start_m tl.program_id(0) off_hz tl.program_id(1) off_z off_hz // H off_h off_hz % H # 计算 Q、K、V、O 的基地址 q_offset off_z * stride_qz off_h * stride_qh k_offset off_z * stride_kz off_h * stride_kh v_offset off_z * stride_vz off_h * stride_vh o_offset off_z * stride_oz off_h * stride_oh这里start_m是 Q 块的索引off_hz把 batch 和 head 合并成一个维度。这种二维 grid 的划分方式是为了让每个 program 独立处理一个 (batch, head, Q块) 的组合最大化并行度。接下来加载 Q 块# Q 块的行索引和列索引 offs_m start_m * BQ tl.arange(0, BQ) offs_d tl.arange(0, D) # 加载 Q 块形状 [BQ, D] q_ptrs Q q_offset offs_m[:, None] * stride_qm offs_d[None, :] * stride_qk q tl.load(q_ptrs, maskoffs_m[:, None] N, other0.0) q (q * scale).to(tl.float16)注意这里把scale通常是 (1/\sqrt{d})提前乘到 Q 上这样后面算 (QK^T) 时就不用再乘了。这是个常见的优化能省一次逐元素乘法。mask是为了处理序列长度不是块大小整数倍的情况越界的部分填 0。然后是初始化在线 softmax 的状态m_i tl.full([BQ], float(-inf), dtypetl.float32) l_i tl.zeros([BQ], dtypetl.float32) acc tl.zeros([BQ, D], dtypetl.float32)m_i是行最大值初始化为负无穷l_i是指数和初始化为 0acc是输出累加器用 float32 保证精度。主循环部分for start_n in range(0, N, BK): offs_n start_n tl.arange(0, BK) # 加载 K 块形状 [D, BK]转置后 k_ptrs K k_offset offs_n[None, :] * stride_kn offs_d[:, None] * stride_kk k tl.load(k_ptrs, maskoffs_n[None, :] N, other0.0) # 计算得分 S Q K^T形状 [BQ, BK] s tl.dot(q, k) # 在线 softmax 更新 m_ij tl.max(s, axis1) m_new tl.maximum(m_i, m_ij) p tl.exp(s - m_new[:, None]) alpha tl.exp(m_i - m_new) l_i l_i * alpha tl.sum(p, axis1) acc acc * alpha[:, None] # 加载 V 块形状 [BK, D] v_ptrs V v_offset offs_n[:, None] * stride_vk offs_d[None, :] * stride_vn v tl.load(v_ptrs, maskoffs_n[:, None] N, other0.0) # 累加输出 acc tl.dot(p.to(tl.float16), v) m_i m_new这段是整个 kernel 的核心。几个细节值得展开tl.dot(q, k)这里q是 ([BQ, D])k是 ([D, BK])结果是 ([BQ, BK])。Triton 会自动把tl.dot映射到 Tensor Core前提是数据类型是 float16 或 bfloat16且维度满足要求。m_ij tl.max(s, axis1)对每一行取最大值得到 ([BQ]) 的向量。m_new tl.maximum(m_i, m_ij)更新全局最大值。p tl.exp(s - m_new[:, None])这里减的是更新后的最大值保证数值稳定。alpha tl.exp(m_i - m_new)这是缩放因子用来把之前的累加结果按新的最大值重新缩放。acc acc * alpha[:, None]注意这里先缩放acc再加新的贡献。顺序不能反。acc tl.dot(p.to(tl.float16), v)把当前块的贡献累加进去。循环结束后归一化并写回acc acc / l_i[:, None] o_ptrs O o_offset offs_m[:, None] * stride_om offs_d[None, :] * stride_on tl.store(o_ptrs, acc.to(tl.float16), maskoffs_m[:, None] N)这里把acc除以l_i得到最终输出。注意l_i是 float32acc也是 float32除法在 float32 下做最后再转成 float16 写回。3.5 参数选择与性能调优块大小的选择对性能影响很大。我做过一组对比实验在 A100 上、序列长度 2048、head 维度 64 的情况下BQBK耗时 (ms)显存占用 (MB)64641.8242128641.35481281281.4156256641.5272可以看到 (BQ128, BK64) 是个不错的平衡点。(BQ) 太小并行度不够太大 SRAM 放不下(BK) 太大循环次数少但每次计算量大反而可能因为寄存器压力导致性能下降。另外几个调优点num_warps控制每个 program 的 warp 数量。默认是 4对于 (BQ128) 的情况可以试试 8。num_stages软件流水线的级数影响 K、V 加载和计算的 overlap。默认 3可以试 2 或 4。数据类型float16 比 bfloat16 在某些卡上快但 bfloat16 的动态范围更大数值稳定性更好。实测下来对于长序列bfloat16 更稳。提示调参的时候建议固定其他变量一次只改一个否则很难判断是哪个参数起了作用。我一般会写个小脚本把不同组合跑一遍记录耗时和显存然后画个表对比。4. 常见问题与排查技巧实录4.1 数值不稳定输出全是 NaN这是最常见的问题。原因通常有几个m_i初始化不对如果初始化为 0 而不是负无穷第一块的alpha exp(0 - m_new)可能溢出。scale没乘如果忘了把 (1/\sqrt{d}) 乘到 Q 上得分会很大exp直接溢出。l_i为 0如果某个 Q 块对应的所有 K 都被 mask 掉了l_i会是 0除法就出 NaN。排查方法在 kernel 里加tl.device_print打印中间值或者先用小规模数据比如 (N64)跑确认逻辑正确再放大。4.2 结果和 PyTorch 对不上如果输出不是 NaN但和 PyTorch 的scaled_dot_product_attention对不上通常是精度问题。FlashAttention 在累加过程中用的是 float32但tl.dot的输入是 float16这中间会有精度损失。一般来说相对误差在 (10^{-2}) 量级是正常的。如果误差很大检查mask 是否正确越界部分填 0 后exp(0 - m_new)不是 0会污染结果。正确的做法是填负无穷或者在算p之后把越界位置置 0。缩放因子的顺序acc acc * alpha必须在acc tl.dot(p, v)之前。l_i的更新l_i l_i * alpha tl.sum(p, axis1)顺序不能错。4.3 性能不如预期如果 kernel 能跑通但速度慢先确认是否用了 Tensor Coretl.dot的输入必须是 float16 或 bfloat16且维度是 16 的倍数。如果 (D) 不是 16 的倍数Tensor Core 用不上性能会差很多。是否 memory-bound用ncu或nsys看一下 kernel 的 occupancy 和 memory throughput。如果 memory throughput 接近峰值说明是 IO 瓶颈可以考虑增大块大小减少循环次数。是否 bank conflictSRAM 的 bank conflict 会拖慢速度。Triton 一般能自动处理但如果块大小选得不好可能会有冲突。4.4 常见问题速查表现象可能原因解决方法输出 NaNm_i初始化错误初始化为-inf输出 NaNl_i为 0检查 mask确保每行至少有一个有效 K输出 NaNscale未乘在加载 Q 后乘scale结果偏差大mask 填 0 污染越界位置填-inf或后处理置 0结果偏差大缩放顺序错误先缩放acc再加新贡献性能差未用 Tensor Core确保tl.dot输入是 fp16/bf16维度对齐性能差块大小不合适试 (BQ128, BK64)编译报错Triton 版本不匹配对齐 Triton 和 PyTorch 版本4.5 几个我踩过的坑第一个坑是mask 的处理。我一开始把越界位置填 0结果exp(0 - m_new)得到一个非零值污染了l_i和acc。正确的做法是填-inf这样exp(-inf - m_new) 0不会污染。或者在算完p之后用 mask 把越界位置显式置 0。第二个坑是alpha的计算时机。我一开始先更新m_i再算alpha结果alpha exp(m_new - m_new) 1缩放失效。正确的顺序是先用旧的m_i和新的m_ij算出m_new然后用m_i - m_new算alpha最后更新m_i m_new。第三个坑是tl.dot的维度要求。Triton 的tl.dot要求两个操作数的内维至少是 16且最好是 16 的倍数。如果 (D64)没问题如果 (D48)可能会报错或者性能很差。解决办法是 padding 到 64。5. 从原理到落地几个值得深挖的扩展方向5.1 反向传播怎么处理上面讲的是前向传播。训练的时候还需要反向传播而 FlashAttention 的反向传播比前向更复杂因为它需要重新计算前向的中间结果而不是存下来。这就是所谓的recomputation反向传播时重新跑一遍前向把 (S)、(P) 重新算出来用于计算梯度。这样显存里还是不需要存 ([N, N]) 的矩阵代价是多算一遍前向。反向传播的 kernel 通常比前向长不少涉及 (dQ)、(dK)、(dV) 三个梯度的计算。核心思路类似也是分块 在线累加。如果你要自己实现建议先把前向跑通再基于前向的代码改反向。5.2 和 PyTorch 官方实现的对比PyTorch 2.0 之后torch.nn.functional.scaled_dot_product_attention已经内置了 FlashAttention 的后端。你可以直接调用不用自己写 kernel。但自己写一遍的价值在于理解原理知道每一步在干什么遇到问题能定位定制化官方实现不一定支持你的特殊需求比如自定义 mask、特殊的缩放方式性能调优你可以针对自己的硬件和场景调参可能比通用实现更快。对比测试的时候建议用相同的输入、相同的精度跑多次取平均避免单次测量的噪声。5.3 在推理框架里的应用FlashAttention 在推理场景里同样重要尤其是长上下文推理。很多推理框架比如 vLLM、TensorRT-LLM都集成了 FlashAttention 的 kernel。如果你在做推理优化理解 FlashAttention 的分块逻辑能帮你更好地调优 KV cache 的管理和调度。另外现在有一些变体比如PagedAttention把 KV cache 分页管理配合 FlashAttention 的分块计算能进一步提升显存利用率和吞吐。这些变体的核心思想是一致的避免大矩阵的显存占用尽量在片上完成计算。5.4 一个实用的调试技巧最后分享一个我常用的调试方法先用小规模数据验证正确性再逐步放大。具体做法是用 (N64)、(D64)、单 batch 单 head 的数据跑和 PyTorch 的结果对比确认逻辑正确放大到 (N512)确认分块逻辑没问题再放大到 (N2048)确认性能和显存符合预期最后加上 batch 和 head 维度确认 grid 划分正确。每一步都打印中间结果和 PyTorch 的对应值对比。这样能把问题定位到具体的环节而不是一上来就跑大规模数据出了 NaN 都不知道哪里错了。我在实际项目里用这套方法基本能在半天内把一个 FlashAttention kernel 从零调到能跑通、结果正确。真正花时间的往往不是写代码而是调参和排查那些边界情况。希望这篇内容能帮你少走一些弯路把“不存中间矩阵”这件事真正搞明白。
阅读完成 · 觉得有帮助？
