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) 级别的。以 (N=8192)、单精度浮点（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 最精妙的部分——在线 softmax（online 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:i+B_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:j+B_k] # 加载 K 块 V_j = V[j:j+B_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:i+B_q] = O_i / l_i[:, None]这段伪代码就是 FlashAttention 前向传播的核心。注意几个关键点：
- (m_i) 和 (l_i) 是逐行维护的，每个 Q 块对应一组；
- 每次更新都要重新缩放之前的累加结果，这是在线 softmax 的精髓；
- (S_{ij}) 只在 SRAM 里存在，算完就丢，不写回显存；
- 最终输出要除以 (l_i)，因为累加过程中用的是未归一化的指数。
2.3 为什么这样能省显存
对比一下标准实现和 FlashAttention 的显存占用：
| 项目 | 标准 Attention | FlashAttention |
|---|---|---|
| 中间矩阵 (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))。当 (N=8192) 时，标准实现光中间矩阵就要几百 MB，而 FlashAttention 只需要几十 KB 的 SRAM 来放块内矩阵。这就是“不存中间矩阵”的直接收益。
但省显存只是表面，更深层的收益是减少了显存 IO。标准实现要把 (S) 写回显存再读出来，FlashAttention 全程在 SRAM 里操作，IO 量大幅下降。在 memory-bound 的场景下，这个收益比省显存更值钱。
3. Triton 实战：手写一个 FlashAttention Kernel
3.1 为什么选 Triton
要自己实现 FlashAttention，理论上可以用 CUDA C++ 写，但 CUDA 的门槛不低：你得手动管理线程块、共享内存、寄存器分配、warp 调度，还要处理各种边界情况。Triton 的出现大大降低了这个门槛。它用 Python 写 kernel，编译器负责把高层描述映射到 GPU 的线程和存储层次上，你只需要关注块级的计算逻辑。
Triton 的几个关键概念：
- program：一个 kernel 实例，对应一个块的计算；
- block：Triton 里的数据块，用
tl.arange和tl.load加载； tl.dot：块级矩阵乘法，编译器会映射到 Tensor Core；tl.exp、tl.max、tl.sum：块级逐元素和归约操作。
用 Triton 写 FlashAttention，代码量大概一百多行，比 CUDA 版本短得多，而且可读性好。下面我按自己的实现习惯，一步步拆开讲。
3.2 环境准备与依赖安装
先把环境搭起来。Triton 对 Python 和 PyTorch 版本有要求，我实测下来比较稳的组合是：
# 建议用 Python 3.10 或 3.11 pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu118 pip install triton==2.1.0如果你用的是更新的 PyTorch，Triton 通常已经作为依赖被装上了，可以直接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, mask=offs_m[:, None] < N, other=0.0) q = (q * scale).to(tl.float16)注意这里把scale（通常是 (1/\sqrt{d})）提前乘到 Q 上，这样后面算 (QK^T) 时就不用再乘了。这是个常见的优化，能省一次逐元素乘法。mask是为了处理序列长度不是块大小整数倍的情况，越界的部分填 0。
然后是初始化在线 softmax 的状态：
m_i = tl.full([BQ], float("-inf"), dtype=tl.float32) l_i = tl.zeros([BQ], dtype=tl.float32) acc = tl.zeros([BQ, D], dtype=tl.float32)m_i是行最大值，初始化为负无穷；l_i是指数和，初始化为 0；acc是输出累加器，用 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, mask=offs_n[None, :] < N, other=0.0) # 计算得分 S = Q @ K^T，形状 [BQ, BK] s = tl.dot(q, k) # 在线 softmax 更新 m_ij = tl.max(s, axis=1) 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, axis=1) 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, mask=offs_n[:, None] < N, other=0.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, axis=1)：对每一行取最大值，得到 ([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), mask=offs_m[:, None] < N)这里把acc除以l_i得到最终输出。注意l_i是 float32，acc也是 float32，除法在 float32 下做，最后再转成 float16 写回。
3.5 参数选择与性能调优
块大小的选择对性能影响很大。我做过一组对比实验，在 A100 上、序列长度 2048、head 维度 64 的情况下：
| BQ | BK | 耗时 (ms) | 显存占用 (MB) |
|---|---|---|---|
| 64 | 64 | 1.82 | 42 |
| 128 | 64 | 1.35 | 48 |
| 128 | 128 | 1.41 | 56 |
| 256 | 64 | 1.52 | 72 |
可以看到 (BQ=128, BK=64) 是个不错的平衡点。(BQ) 太小并行度不够，太大 SRAM 放不下；(BK) 太大循环次数少但每次计算量大，反而可能因为寄存器压力导致性能下降。
另外几个调优点：
num_warps：控制每个 program 的 warp 数量。默认是 4，对于 (BQ=128) 的情况，可以试试 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打印中间值，或者先用小规模数据（比如 (N=64)）跑，确认逻辑正确再放大。
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, axis=1)，顺序不能错。
4.3 性能不如预期
如果 kernel 能跑通但速度慢，先确认：
- 是否用了 Tensor Core：
tl.dot的输入必须是 float16 或 bfloat16，且维度是 16 的倍数。如果 (D) 不是 16 的倍数，Tensor Core 用不上，性能会差很多。 - 是否 memory-bound：用
ncu或nsys看一下 kernel 的 occupancy 和 memory throughput。如果 memory throughput 接近峰值，说明是 IO 瓶颈，可以考虑增大块大小减少循环次数。 - 是否 bank conflict：SRAM 的 bank conflict 会拖慢速度。Triton 一般能自动处理，但如果块大小选得不好，可能会有冲突。
4.4 常见问题速查表
| 现象 | 可能原因 | 解决方法 |
|---|---|---|
| 输出 NaN | m_i初始化错误 | 初始化为-inf |
| 输出 NaN | l_i为 0 | 检查 mask，确保每行至少有一个有效 K |
| 输出 NaN | scale未乘 | 在加载 Q 后乘scale |
| 结果偏差大 | mask 填 0 污染 | 越界位置填-inf或后处理置 0 |
| 结果偏差大 | 缩放顺序错误 | 先缩放acc再加新贡献 |
| 性能差 | 未用 Tensor Core | 确保tl.dot输入是 fp16/bf16，维度对齐 |
| 性能差 | 块大小不合适 | 试 (BQ=128, BK=64) |
| 编译报错 | 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 的倍数。如果 (D=64)，没问题；如果 (D=48)，可能会报错或者性能很差。解决办法是 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 一个实用的调试技巧
最后分享一个我常用的调试方法：先用小规模数据验证正确性，再逐步放大。具体做法是：
- 用 (N=64)、(D=64)、单 batch 单 head 的数据跑，和 PyTorch 的结果对比，确认逻辑正确；
- 放大到 (N=512)，确认分块逻辑没问题；
- 再放大到 (N=2048)，确认性能和显存符合预期；
- 最后加上 batch 和 head 维度，确认 grid 划分正确。
每一步都打印中间结果，和 PyTorch 的对应值对比。这样能把问题定位到具体的环节，而不是一上来就跑大规模数据，出了 NaN 都不知道哪里错了。
我在实际项目里用这套方法，基本能在半天内把一个 FlashAttention kernel 从零调到能跑通、结果正确。真正花时间的往往不是写代码，而是调参和排查那些边界情况。希望这篇内容能帮你少走一些弯路，把“不存中间矩阵”这件事真正搞明白。
