从内存访问到稀疏激活:两种不同的效率问题
课堂先在硬件段提出计算、存储与搬运三个约束,再在规模化训练中走完“取一块、算一块、累计、继续”。这里解释为何累计仍得到完整 attention,并把 MoE 作为另一种容量与每 token 执行量的取舍。
FlashAttention 与 MoE 都与大模型效率有关,但改变的对象不同。前者重新安排同一注意力计算的数据移动和中间存储;后者改变一个 token 经过哪些参数模块。把两者统一叫“让模型变快”会丢掉适用条件,也无法判断代价落在哪里。
大表的代价不只是乘加
令一段序列有 T 个位置,Q、K、V 是各位置的查询、键和值。稠密注意力需要形成各 query 与 key 的比较关系,常写为 softmax(QKᵀ/√d)V;因果 mask 将未来位置的权重置零。直接分步实现可能把 T×T 的分数和概率中间量写入 HBM,再读回后续计算。HBM 是 GPU 的大容量高带宽显存;片上 SRAM 容量较小,访问更快。一次计算是否有效率,既取决于算术量,也取决于数据访问和硬件利用率。
FlashAttention 原论文的核心是分块执行并采用在线 softmax:把可容纳的 Q/K/V 小块搬进片上快速存储,完成块计算,更新少量统计与输出,然后继续下一块。分块不意味着只取部分 key,也不意味着取消因果约束。Pallas/Splash 材料用同样的归一化问题说明底层程序如何安排块。Triton 与 Pallas 是编写设备计算程序的工具,Splash 是具体 attention 实现;不要把 Triton 编程语言与 NVIDIA Triton Inference Server 当成同一个产品。
为什么必须保留分子与分母
N50 的当前算例使用已经取过指数的正权重 w=(2,1,4),value 为 v=(10,20,30)。完整输出为:
$$o=\frac{2\times10+1\times20+4\times30}{2+1+4}=\frac{160}{7}\approx22.857.$$
块 A 包含前两个位置,贡献分子40、分母3,局部输出40/3;块 B 包含最后一个位置,贡献分子120、分母4,局部输出30。合并时加分子、加分母,全部块完成后才相除:(40+120)/(3+4)=160/7。两个局部输出直接相加不对,简单平均也只得(40/3+30)/2=65/3≈21.667。一般需要用各块分母给局部输出加权,而不是按块数取平均。
N51把同一合并规则推广到三块示意:每处理一块更新累计分子、分母和数值尺度,处理完全部允许的K/V块再写回输出。三块图是在推广规则,不是把上述三个位置私自改成新的数值算例。固定一个Q块扫描K/V是教学视角,不声称逐行复现某个kernel的循环次序。
历史独立算例保留用于回查:旧权重(2,1,3)会得到130/6,两个块分母恰好均为3,局部平均偶然正确。这是不同输入,不是当前160/7算例的中间步骤,也不能据它推出一般可平均。
在线 softmax 怎样保持数值稳定
分数 s 直接取指数可能溢出。softmax 可先减去共同最大值 m:分子分母都乘同一个因子,最终比值不变。对已处理位置,保存 m、l=Σ exp(sᵢ−m) 和 u=Σ exp(sᵢ−m)vᵢ。新块有最大值 m_b、归一化量 l_b 与分子 u_b,合并时设 m′=max(m,m_b),然后:
$$l'=e^{m-m'}l+e^{m_b-m'}l_b,$$ $$u'=e^{m-m'}u+e^{m_b-m'}u_b,\qquad o'=u'/l'.$$
公式里的两个缩放系数把旧量和新量换到同一尺度。运行最大值更新时只调整分子,或只调整分母,都会破坏比值。
用当前 logits=(ln2,0,ln4) 验证。第一块的最大值为ln2,平移后的权重是1和1/2,因此l=3/2、u=20。第二块带来更大的ln4,旧分子分母乘exp(ln2−ln4)=1/2;新块平移后的权重为1。合并l′=(1/2)(3/2)+1=7/4,u′=(1/2)20+30=40,输出40/(7/4)=160/7,与完整加权输出相同。真实 kernel 以向量、多个 query 行和 tile 执行,此例只展开一个标量输出。
mask 也参与有效位置定义。完全被遮罩的块不应贡献概率或分母;实现可跳过明确全空的块,但必须处理尚无任何有效位置的初始化,避免把负无穷相减产生 NaN。能跳过某些块,不代表任意稠密注意力都变成线性算术。
反向重算为什么可能划算
反向传播需要相关中间量。FlashAttention 不保存完整概率矩阵,而使用保留的输入、输出和归一化统计,按块重算反向需要的部分。这以额外部分算术换取更少 HBM 访问和中间矩阵存储。它保持同一注意力数学目标,但不同运算顺序和浮点误差意味着不保证逐 bit 一致。是否更快要在实际形状、序列长度、精度和设备上测量;课堂或 CPU 玩具不能证明某个 GPU 加速倍数。不同代际工作还优化线程分工、流水重叠和硬件指令,不能把跨硬件性能图视为单一算法的普适倍数。
MoE 改变的是激活路径
混合专家 MoE 在部分层放置多个可训练专家,用路由器选择少数专家处理当前 token,再把输出加权组合。课堂构造的共享部分为 0.2B,8 个专家各 0.1B,总参数 1B;每 token 选两个时,按本例汇总口径激活约 0.4B。这个口径将其他计入参数含在共享部分,不对应一个特定真实 checkpoint。
总参数决定需要容纳多少权重;激活路径影响每 token 实际执行量。未选专家仍需存储,也可能被同批其他 token 使用。专家并行把不同专家放在不同设备,token 表示经分发后在目标专家计算,再合并返回;all-to-all 等通信和专家负载不均会进入关键路径。辅助均衡目标、容量限制、分组专家及通信实现会影响效果。专家是训练形成的模块,不保证各自对应人类命名的领域。稠密训练近似 C≈6ND 不能直接拿 MoE 总参数替换 N 而忽略路由与激活结构。
CPU/GPU 异构为什么在这里出现
KTransformers探索利用 CPU 与 GPU 的不同优势执行大模型,尤其是权重容量很大而每 token 只激活部分专家的 MoE。CPU 内存可容纳更多专家权重,适当的 CPU kernel、GPU attention 和调度可分担工作。这是容量、带宽、计算与传输的联合设计,不能说“未激活参数不用存”或“CPU 内存就是免费显存”。
实际性能还取决于 CPU 指令能力、内存带宽、NUMA 拓扑、线程绑定、专家选择、批量和设备传输。量化可能进一步减少权重容量,同时引入误差和反量化成本;必须注明精度并评价任务质量。A08 中的具体硬件压缩和延迟实验保留原实验身份,本课不把它们移植成当前机器的实测承诺。
课堂从N49内存层次进入N50分子分母,再由N51全部块与反向重算完成闭环。MoE和异构承载从N43/N51说明深入。原论文、课程材料的具体版本及采用位置见 来源清单。教学计算与硬件性能实测分别报告。