返回课堂

完整深入材料 · 保留原页与全部停点。回到相关主讲

全部块走完,写回的是同一个注意力输出

N51/b0 · 第一块:累计保留

同一合并规则 → 三块执行示意HBM:尚未处理的块仍保留片上工作区固定 Q 块K/V 1 已处理K/V 2 待处理K/V 3 待处理本轮 K/V 1累计量 m、l、u局部表释放,累计保留还不能写最终输出
第一块:累计保留

把刚才两块的合并规则推广到更多块,这次示意分三块,没有另给一套数值。先固定右侧的 Q 块,把第一块 K/V 读进片上完成局部计算。现在图中第一块标为已处理:临时分数表用完可以释放,累计框中的 m、l、u 留下来。左边第二、三块仍待处理,没有被删除,因此这时不能写最终输出。

教师备课:指图与模拟问答

指图与推进 · 不朗读指固定Q、已处理的第1块与累计框,再指仍灰色的第2/3块;不把局部累计称为最终输出。

N51/b1 · 第二块:接着合并

同一合并规则 → 三块执行示意HBM:尚未处理的块仍保留片上工作区固定 Q 块K/V 1 已处理K/V 2 已处理K/V 3 待处理本轮 K/V 2累计量 m、l、u换块,统一尺度后更新还不能写最终输出
第二块:接着合并

下一步读第二块。固定的 Q 没变,新的 K/V 块参与打分和汇总;必要时统一数值尺度,再更新同一个累计框。现在第一、二块都已处理,但第三块还在等待,所以依然不能收束。我们释放的是不用再保存的局部大表,不是已经积累的贡献。

教师备课:指图与模拟问答

指图与推进 · 不朗读第1/2块标为已处理;指“本轮K/V2”与累计框,再指第3块仍待处理。

N51/b2 · 第三块:完整输出

同一合并规则 → 三块执行示意HBM:尚未处理的块仍保留片上工作区固定 Q 块K/V 1 已处理K/V 2 已处理K/V 3 已处理全部块完成累计量 m、l、u所有允许位置均已参与该 Q 块输出:u / l写回输出输出 O:这个 Q 块
第三块:完整输出

再把第三块也走完,现在三块都标成已处理,才得到这个query块完整的输出。请注意最后这一步:分块不是只算一个局部答案,也不是把第一块算完就提前相信它;所有应当参与的位置都进入了结果。我们只是不把完整的大分数表和概率表反复写到显存里。

教师备课:指图与模拟问答

指图与推进 · 不朗读第3块完成,所有块变绿后才出现归一化与写回输出。

N51/b3 · 反向重算与代价

同一合并规则 → 三块执行示意HBM:尚未处理的块仍保留片上工作区固定 Q 块K/V 1 已处理K/V 2 已处理K/V 3 已处理全部块完成累计量 m、l、u所有允许位置均已参与该 Q 块输出:u / l写回输出输出 O:这个 Q 块反向按块重算 → 少保存、少搬运大中间表FlashAttention · 稠密算术仍有 T²(T = 序列长度)
反向重算与代价

这就是FlashAttention的核心执行思路。训练反向需要一些中间结果时,还可以按块重算,增加一部分算术,换取少保存、少搬运。它计算的是同一类精确attention目标,不表示不同浮点实现逐位相同,也没有把稠密注意力的平方量级算术消掉。到这里,预算、存储和搬运这几个现实问题都能落到具体动作上了。现在让训练暂时结束,把参数保存下来,看看同一个模型面对用户请求时,哪些状态还在,哪些状态要重新产生。

进一步解释

训练反向可按块重算所需中间量,以额外部分算术换取更少保存和搬运。数学目标保持同类精确 attention,但浮点运算顺序变化不保证逐 bit 相同。本页的固定Q扫描KV是教学视角,不规定某份内核代码的循环次序。

教师备课:指图与模拟问答

指图与推进 · 不朗读保留全部完成范围,加入反向重算与T²边界。

补充推导、边界与来源

本页把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边界。

本页证据与条件

FlashAttention 机制;教学数值与执行示意,不是性能实测

进入主题正文