返回课堂

算得完、装得下、搬得动

本页问题

显存保存大对象,片上工作区接收当前计算需要的数据。直接分阶段 attention 可以把 S、P 等完整中间表写回后再读;这使计算吞吐、容量与搬运带宽成为不同约束。后续分块执行改变的是保存与读写安排,仍需完成规定的读取目标。

RC-H3/b0 · 给数据找到物理位置

大容量显存 HBM片上工作区Q / K读取 Q/K计算分数 S小工作区,逐次取数据
给数据找到物理位置

左侧是高带宽显存 HBM,保存较大的输入、参数和中间结果;右侧是靠近计算的片上工作区,容量小得多。当前要算 Q 与 K 的匹配,先沿蓝色箭头把数据读过来。数据已经放在显存里,并不等于计算单元已经拿到它。

进一步解释

HBM 是 High Bandwidth Memory;本图只区分显存与片上工作区的职责,不给某型号芯片的容量或速度。

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

指图与推进 · 不朗读指左侧 HBM、右侧片上工作区,再沿读取 Q/K 路径。

本态来源与条件

教学构造与机制示意;非模型实测

RC-H3/b1 · 算完还要写回和再读

大容量显存 HBM片上工作区Q / K读取 Q/K计算分数 S写回 S中间分数 S下一步再读 S下一步使用 S小工作区,逐次取数据
算完还要写回和再读

算出分数表 S 以后,一种直接实现把它写回 HBM;下一阶段再读出来做归一化。于是一次预测除了乘加,还包含保存、写回和再次读取。红线表示写回,蓝线表示读入,方向不同但处理的是同一个中间对象 S。

进一步解释

S=QKᵀ(可带任务规定的缩放与 mask),后续 P=softmax(S),再 O=PV。这里保留代表性的 S 写回与重读,完整三阶段读写和分块比较在 RC-T3 展开。

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

指图与推进 · 不朗读沿Q/K读取→计算S→写回S→下一步再读S这一条代表路径;其余阶段留待分块比较。

本态来源与条件

教学构造与机制示意;非模型实测

RC-H3/b2 · 把三个约束带到训练

大容量显存 HBM片上工作区Q / K读取 Q/K计算分数 S写回 S中间分数 S下一步再读 S下一步使用 S小工作区,逐次取数据容量:要保存多少带宽:来回搬多少算力:乘加有多少
把三个约束带到训练

现在可以区分三张账:容量看要保存多少,带宽看单位时间要搬多少,算力看乘加要做多少。完整大表能否少写回,后面还会回到这幅图。先问这些昂贵运算为何值得反复做:MAE 用原图构造目标,普通文字又怎样提供训练目标?

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

指图与推进 · 不朗读分别指计算、容量和搬运三个位置,再指大表问号。

本态来源与条件

教学构造与机制示意;非模型实测

接下来

先看训练不断在做什么题,再回来减少大表搬运。

补充推导、边界与来源

进一步阅读

条件与来源

以下从计算对象、条件与原始证据继续展开。

原 N23:算得快之外,还要装得下、搬得动

直接 attention 实现会生成分数矩阵,写入显存,再读出做 softmax,随后使用归一化权重汇总 V。HBM 容量、片上工作区容量、计算吞吐和数据搬运带宽是不同约束。这里先定位问题;N49–N51 回到同样的内存层级,用分块与在线归一化解释如何减少保存和搬运,不提前把名称当答案。

来源:FlashAttention v2,§3.1 IO 分析。

相关的完整阅读材料保留原有公式、引用和原始图解。

本页证据与条件

教学构造与机制示意;非模型实测

进入主题正文