Flash attention线性注意力机制考虑常规的注意力操作这是典型的关于线性注意力Linear Attention的推导起点。常规运算第一步计算 Q×KT矩阵形状(n×d)×(d×n)结果矩阵大小n×n。这个巨大的矩阵就是注意力分数矩阵。计算量生成这个 n×n 的矩阵需要n2次乘加运算n2个数每个数要 2d-1 次乘加运算。第二步将结果乘以 V矩阵形状(n×n)×(n×d)结果矩阵大小n×d计算量这也是一个矩阵乘法但前提是必须先把第一步算完。总复杂度在第一步生成 n×n矩阵的时候就已经引入了O(n2)的复杂度。这就是为什么普通 Attention 处理长文本时会崩溃——序列长度 n 变大一点计算量就会平方级爆炸。线性运算利用了乘法的结合律我们改变括号的位置第一步计算 KT×V矩阵形状(d×n)×(n×d)结果矩阵大小d×d。注意这个关键点结果是一个很小的矩阵它跟序列长度 n 没关系只跟模型隐藏层维度 d 有关d 是固定不变的比如 4096。计算量因为结果矩阵很小所以计算量也是固定的每次计算只需要做 d×n×d 实际上是 d×d×(2n-1) 次乘法与 n 是一次方成正比即 O(n)。第二步计算 Q×(上一步的结果)矩阵形状(n×d)×(d×d)结果矩阵大小n×d计算量同样这里计算量也是O(n)因为随着输入序列长度 n 变大计算量只会线性增加。这是线性的计算时间太棒了但更妙的是这看起来就像一个 RNN循环神经网络。这种“对偶性”使得我们既能高效地训练使用并行的、二次复杂度的形式又能高效地推理使用串行的、线性复杂度的形式。这里的线性 Attention把历史所有 Token 的信息压缩到了一个固定的矩阵 S 中大小为 d×d。每当新来一个 Tokenkt​ 和 vt​只需将它“累加”进 S 里即可。这个状态 SS 意味着什么意味着你不需要像传统 Transformer 那样为了生成下一个词必须把前面所有 10 万个 Token 的 Key 和 Value 都重新算一遍。你只需要维护这一个 SS 矩阵就行了。这种“有记忆状态”的递归形式就是 RNN 的核心特征。“线性 Attention”虽然快但它脑子不太好使记不住新东西也甩不掉旧垃圾。给它装上“记忆阀门”让它该记的记该忘的忘变得更聪明。 第一步认清老问题上一页提到的公式你可以把它想象成一个“永远在往杯子里倒水”的过程。杯子就是 St​记忆状态。每来一个新的字就相当于往杯子里倒一点水加上 ​。老问题的缺点致命伤水永远在加永远不减少。如果模型读了 100 万字杯子里装满了水最后面的字和前面的字混杂在一起模型就彻底糊涂了。这叫“记忆饱和”。 第二步Mamba-2 的解决方案公式是 Mamba-2 用的新手解释加入“遗忘阀”γt 这个希腊字母叫 Gamma这就是一个“遗忘阀”取值在 0 到 1 之间。原理每次要加入新东西之前先拿这个阀门把旧杯子里的水倒掉一点点比如说倒掉 10%。作用这样杯子就不会满出来了。模型可以记住最重要的信息慢慢遗忘掉不重要的老信息。 第三步Gated Delta Net 的更高级解法这是 PPT 的核心公式变复杂了新手解释加入“擦除刷子”这里引入了一个新东西叫 βtBeta0到1之间的值。我们换个比喻想象你在写课堂笔记。以前的方式线性 Attention老师每讲一句话你就直接写在笔记本最后面。时间久了本子写满了前面的内容和后面的内容混在一起根本找不到重点。Gated Delta Net 的升级这个新发明给了你一把“修正液/擦除刷子”。当你看完当前的字模型会判断“这个新信息重要吗”如果新信息非常重要βt≈1它不会直接写在最后而是会先找到笔记本里旧有的、意思相近的那句话也就是 ​指的方向把它擦掉/覆盖掉然后再把当前的关键信息写上去。如果新信息不重要βt≈0它就不写进去保持原样。MOE混合专家 首先看左侧Dense Model稠密模型/传统模型左边是传统的大模型比如早期的 GPT-3、Llama 2。架构逻辑输入进入“FFN Layer”前馈神经网络时必须经过同一个巨大的、唯一的“大脑”来计算。缺点FFN累死计算量大而且单个模型无法无限做大。如果把这个FFN变得特别巨大几百亿参数推理速度会变得极慢。 再来看右侧Sparse Model稀疏模型 / MoE这就是这张图的核心——MoE 架构。架构逻辑原来只有一个巨大的 FFN 层现在被拆成了多个小的 FFNFFN 1 ~ FFN 4。中间多了一个“Selector Layer”选择器/路由器。工作机制当词 x1​ (The) 进来时Selector 会判断这个字应该交给哪个专家比如箭头指向了 FFN 2。当词 x2 (Dog) 进来时Selector 判断它适合另一个专家比如箭头指向了 FFN 1 。MoE 可以让模型在“变大”的同时不增加计算量“通过 MoE我把模型的‘知识容量’扩大了 10 倍但是我每算一个字用的力气FLOPs / 电费 / 推理时间却和以前完全一样”在 1024 块显卡的训练场景下这种“稀疏化”意味着你可以把模型塞到7000 亿参数比如 GPT-4但每次推理只激活一小部分专家实际算力消耗相当于一个700 亿的模型。这就是大模型既能“懂得特别多”又能“回答得特别快”的根本原因。这张图告诉我们MoE 不仅“理论上变聪明了参数量变大算力不变”而且在“工程上分布式训练”它是天然为大规模 GPU 集群设计的。难点不是谁都能跑 MoE训练 MoE 需要极昂贵的高速互联网络把 1024 张卡串起来。没有这个基础设施MoE 跑不起来或者跑得比普通模型还慢。MoE 是“大厂特供”只有当你有了成百上千张卡做数据并行时你才有足够的算力余量去塞下那些“大模型参数”MoE 的算力优势才能发挥出来。训练过程像“走钢丝”想要让那个“路由器”不乱分配工作需要很高的技巧一旦平衡没调好几百万美金的训练就会白费。结合你的上下文你之前问了“用 1024 张 B100 训练 15T token”。要完成这个工程你不仅要有 1024 张卡你还必须拥有一个能够把卡片之间做到极高带宽比如 NVLink InfiniBand的超级网络以及一套能解决“路由坍塌”的复杂算法。这就是为什么全世界只有 OpenAI、Google、Meta、DeepSeek 等少数公司能做出万亿参数大模型的原因。MoE 的核心命脉——路由Routing机制1️⃣ 上半部分Top-k 路由当前主流方案这是目前几乎所有主流大模型都在用的动态路由。2️⃣ 下半部分哈希路由简单粗暴的基准线这是早期或作为对照实验的最原始的静态路由。稳定性问题训练 MoE 时如果路由算法出了点小毛病整个模型的 Loss损失值就会突然暴涨图里那个尖锐的反弹导致模型瞬间变傻。在大规模集群上一旦遇到这种情况工程师就得强行停止训练、找回快照、修改超参数再重新开机。每一次这样的折腾都要浪费掉数百万的算力成本。 第一部分为什么路由器会发疯罪魁祸首低精度如 bfloat16下的舍入误差。 第二部分底部解决方案FP32 路由器 Z-Loss既然低精度会导致这致命的 36% 偏差各大厂是怎么补救的怎么运作看中间的图输入一个词比如 x1​ The它会进入一个可学习的路由器Router通常是一个极小的神经网络。路由器会算出它跟每个专家的“匹配度”概率分数柱状图里。什么是 Top-k“Top-k”的意思是只激活分数排名最高的前 k 个专家来处理这个词。DeepSeek (7)每次激活 7 个专家。 重点解析 DeepSeek 的 k7通常 MoE 模型的专家总数可能高达 100 多个。DeepSeek 选择每次激活 7 个专家是一个非常高的数值。这意味着在保持总计算量FLOPs不变的情况下它把更多专家拉进来“共同会诊”让输出更接近普通稠密模型的效果极大提升了模型的上限。❌ 缺点训练不稳定。如果路由器“发疯”全把病人塞给同一个专家就会导致路由坍塌训练直接崩盘。工程师需要设计复杂的“负载均衡”辅助损失函数来修正它。怎么运作看下面的图没有神经网络没有学习能力。输入一个词直接扔进一个“哈希函数Hash Function”。这个函数就是一个固定的数学算式比如直接拿这个词的拼音首字母或者 ASCII 码做个换算。结果计算出的哈希值直接决定了这个词去哪个专家图里固定指向某一个专家。问题产生路由器Router在计算分配概率时需要计算一个Softmax找出排名前 k 的专家。路由器算出来的一组分数里最大的一个是128.5。毁灭性后果为了数值稳定Softmax 核心操作会统一减去最大值。本来正确的最大值是 128.5由于舍入误差变成了 128。这导致原本有一点点优势的那个专家瞬间失去了优势本来应该分配给最相关专家的概率是14.2%。因为 0.5 的误差概率硬生生被算成了9.1%。相差了整整 36%解法一硬件层面的精准定位只用 FP32 来算路由器PPT 明确提出Solution: Use Float 32 just for the expert router。大白话翻译其他几千亿的参数我们全用节省显存的bfloat16或者FP8唯独这个非常关键的路由器Router我们要用到精确度极高的FP32全精度单精度浮点数。解法二数学层面的限制器Z-Loss 辅助损失底部的公式 Lz(x)就是大名鼎鼎的Z-Loss。大白话解释路由器算出的分数不能太大。Z-Loss 就像给路由器的输出带上了一道“脚镣”刻意惩罚它算出过于极端的数值比如那个 128.5逼迫路由器的 logits 保持在一个合理的范围内。