摘要
注意力层的执行时间不仅取决于乘加次数，也取决于中间结果写到哪里、被读取几次。本文围绕分块计算，给出在线 Softmax 的独立推导、可运行算例和访存预算。所讨论的分块与重计算方法来自 Tri Dao、Daniel Y. Fu、Stefano Ermon、Atri Rudra 和 Christopher Ré 提出的 FlashAttention，本文在该研究基础上组织工程解释与教学实现。
01 / 先计算中间结果有多大
设单个注意力头的查询、键和值均有 N 行，头维度为 d。暂不加入掩码与 dropout，标准缩放点积注意力为：
S = QKᵀ / √d
P = softmax(S) # 按行归一化
O = PV按三个独立算子实现时，得分矩阵 S 写入显存，再读出求 P，P 又被写入并读出参与下一次乘法。大容量 HBM 承载输入与输出；容量较小的片上存储更适合临时计算。分块让一个局部得分算完后立即参与输出累加，从而避免把完整 S、P 保存在 HBM。输入和最终输出仍需访问显存。
用一个教学预算说明数量级：取单头 N = 4096、d = 64，假设所有表内元素均占 2 字节，暂不计训练梯度和工作区。MiB 按 2²⁰ 字节计算。
| 对象 | 元素数 | 单份容量 |
|---|---|---|
| Q、K、V 或 O | 4096 × 64 | 每个 0.5 MiB |
| 完整 S 或 P | 4096 × 4096 | 每个 32 MiB |
| 一个得分块 | 64 × 64 | 8 KiB |
若 S 与 P 各发生一次完整写入和读取，仅这两份中间矩阵就对应 128 MiB 的逻辑访问量。这是张量层面的算例，实际 HBM 流量还受缓存、融合和精度选择影响。分块避免 N×N 中间结果物化，但稠密注意力仍要处理所有查询与键的配对，计算量仍为 O(N²d)。
02 / 三个状态连接所有数据块
不能分别计算每个块的 Softmax 后直接相加，因为每块的分母不同。可从 Milakov 与 Gimelshein 的在线归一化方法出发，为每一行维护三个状态：已处理得分的最大值 m、缩放后的指数和 l，以及尚未归一化的加权向量 u。
m = max(sⱼ)
l = Σ exp(sⱼ − m)
u = Σ exp(sⱼ − m) vⱼ
o = u / l上式求和覆盖已经处理的键。新块 B 到达后，旧状态必须先换到同一个指数基准，再加入新贡献：
m′ = max(m, max(sⱼ in B))
α = exp(m − m′)
pⱼ = exp(sⱼ − m′), j in B
l′ = αl + Σ pⱼ
u′ = αu + Σ pⱼvⱼ这一步的依据是 exp(s − m) · exp(m − m′) = exp(s − m′)。旧分子与旧分母同时乘 α，比例不变；所有指数的自变量不大于零，避免直接计算大正数的指数。初始化为 m = −∞, l = 0, u = 0，处理完全部块后仅需一次除法。
03 / 可运行的单行教学实现
以下代码仅依赖 Python 标准库。scores 是预先算好的某一行 QKᵀ / √d，values 是对应的值向量。为突出递推，使用四个有限得分、二维值向量和块大小 2，不加入掩码与 dropout。这段 CPU 程序核对代数关系，不模拟 GPU 的性能。
from math import exp
def tiled_row(scores, values, tile=2):
m, total = float("-inf"), 0.0
acc = [0.0] * len(values[0])
for start in range(0, len(scores), tile):
block = scores[start:start + tile]
new_m = max(m, max(block))
alpha = exp(m - new_m)
weights = [exp(s - new_m) for s in block]
total = alpha * total + sum(weights)
acc = [
alpha * old + sum(
w * values[start + j][c]
for j, w in enumerate(weights)
)
for c, old in enumerate(acc)
]
m = new_m
return [value / total for value in acc]
scores = [1.0, 2.0, 0.0, 3.0]
values = [[1., 0.], [0., 1.], [2., 1.], [1., 3.]]
weights = [exp(s - max(scores)) for s in scores]
reference = [
sum(w * v[c] for w, v in zip(weights, values)) / sum(weights)
for c in range(2)
]
result = tiled_row(scores, values)
error = max(abs(a - b) for a, b in zip(result, reference))
print([round(x, 6) for x in result])
assert error < 1e-12可以沿着状态检查计算过程。第一块的得分为 1、2，得到 m = 2、l ≈ 1.367879，加权向量约为 [0.367879, 1]。第二块含得分 3，最大值更新为 3；先将已有分母与向量都乘 exp(−1)，再加入两个新权重。旧块不必重新读取，每次迭代保留的数值含义也始终一致。
本例执行结果为 [0.795176, 2.200684]，与一次性归一化的参考计算在上述容差内一致。分块得到数学定义相同的精确注意力，不代表浮点结果逐位相同：求和顺序、指数近似与累加精度都可能改变末尾数位。NVIDIA 对浮点运算结合律的说明解释了这一差别。
04 / 反向传播为什么选择重算
训练需要用注意力概率求梯度。若前向保存完整 P，反向能直接读取它，却要承担长序列的显存占用。FlashAttention 保存输入、输出及每行归一化统计量，在反向阶段重新生成局部得分与概率，再累计 Q、K、V 的梯度。具体推导见原论文附录 B。
重算增加一部分运算，交换的是更少的大矩阵读写。因此应同时检查前向、反向时间和峰值显存，不能只凭乘加次数判断快慢。实现若选择保存 L = m + log(l)，可用 exp(s − L) 重建局部概率；它与保存 m、l 表达相同的归一化信息。
05 / 块大小决定怎样使用片上资源
一个查询块包含 Bᵣ 行，一个键值块包含 B𝚌 行。扩大查询块有机会让更多查询复用同一次读取的 K、V，但得分块与输出累加器也随之增长。实际片上预算还包括寄存器、共享内存、缓冲区与累加精度，不能只看表中的 8 KiB。
块过小会增加循环与归一化更新次数；块过大可能使可同时驻留的线程块减少，甚至引入额外数据溢出搬运。依据 CUDA 的占用率说明，本文建议把块配置视为实验参数：固定输入后比较执行时间、寄存器和共享内存用量，再判断复用收益是否超过资源压力。
06 / 怎样建立可比较的实验
实验先固定批量、头数、查询与键的长度、头维度、dtype、数据布局、缩放系数和因果掩码；同一轮比较使用相同输入。推理关闭 dropout，训练对齐其概率及前反向范围。检查输出与梯度的误差容限后，再开始计时。
预热要覆盖编译、分配和首次内核启动。GPU 异步执行时，应使用设备事件，或在主机计时边界显式同步；输入准备是否计入需事先一致。PyTorch benchmark 工具提供预热、重复采样和加速器同步支持。记录多次测量的中位数与波动，并单独报告峰值分配显存。
还要记录基线实际调用的内核：PyTorch SDPA可能选择融合实现，不能把接口名等同于朴素三算子实现。最终报告应包含设备、软件版本、完整形状和选用内核。
参考文献
- Dao et al. — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness，2022：分块、访存分析及反向重计算。
- Milakov & Gimelshein — Online normalizer calculation for softmax，2018：在线归一化与分块状态合并。
- Dao-AILab / flash-attention：官方实现与接口说明。
- NVIDIA — CUDA C++ Best Practices Guide：存储、占用率和浮点数值行为。
- PyTorch — Benchmark Utils；Scaled Dot Product Attention：计时工具与基线选择。