相同的读取目标,能否少保存、少搬运?
本页问题
固定一个 Q 块,逐次处理所有允许的 K/V 块。每块计算后释放局部分数,保留可合并分子 u、分母 l 与稳定尺度 m;全部块完成后才输出 u/l。这样的执行减少完整 S/P 表写回与再读,不删除合法位置,也不消除稠密配对的二次算术。
RC-T3/b0 · 先定位重复落显存的大表
回到前面的硬件图。直接实现先由 Q/K 算出 S,再将 S 读回做归一化得到 P,最后用 P 汇总 V。这里的问题不是这些关系不该算,而是完整 S 和 P 在阶段之间反复写回显存、再被读出来。我们保持同一读取目标,改变执行顺序。
教师备课:指图与模拟问答
指图与推进 · 不朗读沿三行读/写指S、P、O;确认HBM与片上位置不变。
本态来源与条件
教学构造与机制示意;非模型实测
RC-T3/b1 · 第一块结束,只留下可续算信息
现在固定一个 Q 块,只把第一块允许的 K/V 送入片上工作区。局部分数算完即可释放,但留下累计权重 l、累计加权内容 u 和用于稳定计算的尺度 m。其余块仍在显存里等待,不能把第一块的局部平均当作最终输出。
进一步解释
l 为归一化分母的累计,u 为加权 value 的累计;m 是稳定指数计算的最大值尺度。当新的最大值变大时,旧 u 和 l 必须一起重缩放。三块流程是执行示意,不配成下述两块手算的数值结果。
教师备课:指图与模拟问答
指图与推进 · 不朗读指固定Q、已处理块1、仍待处理块2/3与累计框。
本态来源与条件
教学构造与机制示意;非模型实测
RC-T3/b2 · 后续块进入同一累计
第二块加入时,旧贡献继续保留,再把新块的权重与内容合进同一个累计。两块贡献有不同的总权重,直接平均两份局部答案通常不对;必须连分母一起合并。绿色区表示已经处理,灰色区提醒我们还有允许的位置尚未加入。
进一步解释
完整手算 w=(2,1,4)、v=(10,20,30):块1 u₁=40,l₁=3;块2 u₂=120,l₂=4;合并输出 160/7≈22.86。稳定递推与向量公式在本页深入正文保留。
教师备课:指图与模拟问答
指图与推进 · 不朗读指已完成块1/2与待处理块3,箭头接到同一累计框。
本态来源与条件
教学构造与机制示意;非模型实测
RC-T3/b3 · 全部允许位置参与,才写回输出
第三块完成后,当前 Q 块所允许的全部 K/V 都参加了共同归一化,才用 u 除以 l 得到完整输出并写回。减少的是完整中间表的保存与搬运;反向还可以重算部分量来节省空间。稠密配对的平方算术仍在。现在停止更新,看看训练好的参数如何处理一份请求。
进一步解释
T 是序列长度。FlashAttention 减少中间量物化与 IO,训练反向可以重算;仍覆盖全部允许 key,并没有把一般稠密注意力所有二次算术变成线性。
教师备课:指图与模拟问答
指图与推进 · 不朗读确认所有块完成,再指输出写回;最后指重算和T²边界。
本态来源与条件
教学构造与机制示意;非模型实测
接下来
停止更新后,新请求仍要产生当前状态。
补充推导、边界与来源
进一步阅读
条件与来源
以下从计算对象、条件与原始证据继续展开。
原 N23:算得快之外,还要装得下、搬得动
直接 attention 实现会生成分数矩阵,写入显存,再读出做 softmax,随后使用归一化权重汇总 V。HBM 容量、片上工作区容量、计算吞吐和数据搬运带宽是不同约束。这里先定位问题;N49–N51 回到同样的内存层级,用分块与在线归一化解释如何减少保存和搬运,不提前把名称当答案。
来源:FlashAttention v2,§3.1 IO 分析。
相关的完整阅读材料保留原有公式、引用和原始图解。
原 N49:同一张注意力大表,为什么要搬来搬去?
直接分步 attention 实现可能物化完整分数矩阵 S=QKᵀ,再写入 HBM;读出 S 做 softmax 得到概率 P,再写回;之后读取 P 与 V 相乘。蓝色箭头为读入片上、红色为写回。课堂路径代表会物化大中间表的一类直接实现,不表示所有现有库都如此。
T为序列长度。稠密关系按 query 与允许的 key 比较,分块不能随意漏掉本应参与的位置。片上工作区放不下整表时,可分块计算并保留可合并累计;难点是 softmax 仍需跨所有有效位置共同归一。
FlashAttention 的IO模型与范围;论文 §2、§3.1。
原 N50:分块以后,分子和分母都要一起接着算
原 N50 保留的校准算例:正权重w=(2,1,4),value=(10,20,30),前两项为块1、第三项为块2。
$$u_1=2\times10+1\times20=40,\quad l_1=2+1=3;\qquad u_2=4\times30=120,\quad l_2=4.$$
$$o=\frac{u_1+u_2}{l_1+l_2}=\frac{160}{7}\approx22.86.$$
直接平均两个局部均值是$(40/3+30)/2=65/3\approx21.67$,与正确输出不同,因为两块分母不等。
稳定递推。 logits=(ln2,0,ln4)。首块用m=ln2缩放,得到u=20、l=1.5;第二块的最大值ln4更大,旧分子和分母一起乘1/2,第二块贡献u_b=30、l_b=1。因此新u=40、l=1.75,输出40/1.75=160/7。一般公式为$m'=\max(m,m_b)$,$u'=e^{m-m'}u+e^{m_b-m'}u_b$、$l'=e^{m-m'}l+e^{m_b-m'}l_b$。
完整向量实现与因果mask见在线softmax正文。旧权重(2,1,3)与130/6只作为历史算例,不是原 N50 材料当前数据。
原 N51:全部块走完,写回的是同一个注意力输出
原 N51 材料把N50两块数值验证推广为三块过程示意,不把三块流程配成160/7这一数值结果。固定一个Q块,依次将所有允许的K/V块装入片上工作区,更新最大值m、归一化分母l与分子u。局部中间表释放后,累计量仍保留。第3块完成、所有允许位置都参与后,才计算u/l并写回该Q块输出。
训练反向可以按块重算所需中间量,以额外部分算术换取少保存、少搬运。FlashAttention 保持同类精确attention数学目标,浮点求值顺序变化不保证逐bit一致;稠密算术仍含$T^2$,T为序列长度。固定Q扫描KV是课堂执行视角,不声称逐行复现Algorithm1循环次序。
FlashAttention §3.1 / Algorithm1 / Theorem1;完整在线稳定递推、反向重算、MoE与KTransformers边界。
本页证据与条件
教学构造与机制示意;非模型实测