Glen Li's Gallery
站点导航
模型工程 WIKI
← 返回书架人工智能 · 交互式 Wiki
模型工程/分块 Attention 与反传资料与边界
训练系统与性能进阶计算演示

EXACT MECHANICS / ATTENTION KERNELS

分块 Attention：重建概率，累计梯度

前向用少量统计量完成归一化，反向再按块重建概率。先看一行输出如何形成，再追踪多个 query 对同一个 K/V 的梯度贡献。

01 / RUNNING STATE

三个累计状态足够计算一行输出

m′=max⁡(m,max⁡j∈Bsj),α=em−m′,ℓ′=αℓ+∑j∈Besj−m′,a′=αa+∑j∈Besj−m′vj,o′=a′/ℓ′.\begin{aligned}m'&=\max(m,\max_{j\in B}s_j),\quad\alpha=e^{m-m'},\\\ell'&=\alpha\ell+\sum_{j\in B}e^{s_j-m'},\\a'&=\alpha a+\sum_{j\in B}e^{s_j-m'}v_j,\quad o'=a'/\ell'.\end{aligned}
m 为当前最大分数，ℓ 为相对此最大值的指数和，a 为向量加权和；B 是新块。初始 m=−∞、ℓ=0、a=0，首块旧状态缩放视为 0。

全行减去同一最大值不会改变 softmax。如果新块含更大的分数，旧 ℓ 和 a 需要乘 e^(m−m′)，才能把新旧块放到同一尺度。只对每个块单独 softmax 再平均，通常得不到全行结果。

02 / BLOCK LEDGER

按块推进，看旧累计值怎样重新缩放

六个分数与二维值向量

双精度在线计算

原分数 [1,2,−1,3,0,2.5]。平移会改变运行最大值，不改变数学上的概率与输出。

已读块数
当前 m
当前 ℓ

当前前缀输出：

整行参考输出：

与整行参考的最大绝对差：

块本块分数块最大值运行最大值旧状态缩放累计分母

读完所有块后才应与整行结果一致；中间状态只覆盖已读取位置。最大误差反映数值比较，不是模型质量。当前示例不含 mask、dropout、batch 或多 query。

整行的值向量与参考概率
位置V全行概率
03 / GPU EXECUTION

FlashAttention 在这个关系上做了哪些工程工作？

01 / TILE

分块读取 Q/K/V

把块送到片上存储，计算局部分数，避免把全部分数写回 HBM。

02 / MERGE

更新各行累计状态

维护最大值、分母与输出，局部块处理完后复用空间。

03 / BACKWARD

重算局部概率

反向使用保存的少量统计与 Q/K/V 重建需要的块，计算梯度。

04 / MEASURE

核对数值与成本

比较参考输出与梯度、峰值内存、访存和运行耗时。

FlashAttention 的“精确”描述标准注意力目标；浮点求和顺序不同，结果不必逐位相同。它减少物化的中间矩阵和内存流量，密集全注意力的交互数量仍含序列长度平方项。稀疏 Attention 则还改变可访问的位置集合。

真实实现还需设计块尺寸、线程组织、寄存器占用、Tensor Core 计算、mask 和反向重算。更少 HBM 访问不等于任意形状都同样加速，必须在目标硬件、长度与 batch 下测量。全 mask 行与 dropout 的约定也需与参考实现匹配。

Attention 反传公式稳定归一化与浮点误差访存、计算与等待
04 / RECONSTRUCT & DIFFERENTIATE

没有保存 P，反向怎样得到 dQ、dK、dV？

前向完成后保留 Q/K/V、输出 O 和每行的 log-sum-exp（LSE）。当前块重新计算分数，再减去整行 LSE，即可重建该块对应的全局概率。下式使用无 dropout 的标准 Attention；上游梯度 G=∂L/∂O 固定。

λi=log⁡∑jeSij,Pij=eSij−λi,Di=Gi⋅Oi,∂Sij=Pij(Gi⋅Vj−Di),∂Qi+=∂SijKj/d,∂Kj+=∂SijQi/d,∂Vj+=PijGi.\begin{aligned} \lambda_i&=\log\sum_j e^{S_{ij}},\\ P_{ij}&=e^{S_{ij}-\lambda_i},\qquad D_i=G_i\cdot O_i,\\ \partial S_{ij}&=P_{ij}(G_i\cdot V_j-D_i),\\ \partial Q_i&\mathrel{+}=\partial S_{ij}K_j/\sqrt d,\\ \partial K_j&\mathrel{+}=\partial S_{ij}Q_i/\sqrt d,\\ \partial V_j&\mathrel{+}=P_{ij}G_i. \end{aligned}
S=QKᵀ/√d，λ 是整行统计，不能换成块内归一化。∂S 表示损失对分数的导数；多个 query 块对相同 K/V 的贡献必须相加。

五个位置的逐块反向

二维 Q/K/V · L=Σ O⊙G · 无 dropout
反向块进度
前向输出最大差
最终梯度最大差

块QueryKey重建的概率行

前向保存的逐行统计与输出
行LSE输出 OD=G·O
目前累计的梯度与完整参考
对象累计梯度完整矩阵参考
全部输入矩阵与空行约定
Q = [[.2,-.3],[.5,.1],[-.4,.7],[.6,-.2],[.1,.8]]
K = [[.3,.5],[-.2,.4],[.7,-.1],[.1,-.6],[-.5,.2]]
V = [[1,0],[0,2],[-1,1],[2,-1],[.5,.3]]
G = [[.2,-.1],[.3,.5],[-.4,.2],[.1,-.3],[.6,.4]]

全部屏蔽的 query 行在本例显式定义 O=0、相关导数为零，LSE 显示为 −∞。实现单独处理该行，避免计算 −∞−(−∞)。这是一项明确的边界约定，不把无定义的普通 softmax 直接当成零。

05 / STORAGE & WORK

省下的保存量，与增加的重算分别记账

本例对比“保留完整 P”与“保留一行一个 LSE”的路径。Q/K/V/O 在两条路径中共同保留；上游 G 和 dQ/dK/dV 另计。临时概率块按 query 块×key 块计容量，但真实反向还会产生分数、导数、矩阵乘 workspace 等临时张量，因此这个容量不是物理内存峰值。

5×5 概率表有 25 个元素，LSE 有 5 个；默认 2×2 概率块最多 4 个。前向对因果可见位置计算 15 次 QK 点积，反向重建再计算 15 次。JavaScript 交互保存的逐步快照和完整参考仅用于展示，未包含在算法账本里。

完整矩阵反向公式编码器梯度缓存运行独立验证与实测报告
继续看上下文并行：跨卡 K/V 与梯度归还

适用范围与一手资料

入门为一行六分数；反向为五位置二维 Q/K/V、固定上游 G，无 dropout。分块重建与完整参考及 PyTorch autograd 核对；全屏蔽行显式取零。账本排除界面快照，不实现 GPU kernel。

公式与机制的原始依据 · 4 项
知识目录
搜索知识库