Dropless 与分组矩阵乘

这一篇要回答的问题

token 为什么会被丢掉？是谁要求的？
MegaBlocks 到底改了什么？
一层 MoE 从头到尾，每一步张量是什么形状？烧多少算力？
那个「掩码」到底是怎么弄出来的？硬件需要会跳过空块吗？
为什么原来非得批量矩阵乘？现在这么算，创新在哪？
这个能力的名字到底叫 block-sparse 还是 grouped GEMM？
补齐到 128，跟 chunked prefill、PagedAttention 是一回事吗？
TPU 上是不是也这么干？
TPU 上 MXU 怎么被喂饱的？HBM → VMEM → VREG 的流水
三条路一个骨架：上界开在哪，「省」发生在哪
token 怎么见到它的专家：A2A / AG-RS / 不用 EP
V3 是真 dropless 还是近似的？

如果只记一句 token dropping 从来不是一个建模决定，它是矩阵乘法的形状约束漏进模型里的。 有人把那个约束在 kernel 里解决掉之后，「容量系数」这个超参数就自己消失了 —— 不是被调好了，是不存在了。

下面这一整篇，本质上就是把这句话拆开验一遍。

§1先把口径钉死：站在一张卡上数

所有「某个专家收到多少 token」的计数，都是某一张卡上的局部计数。不钉死这一点，后面每个数字都会差好几个量级。

专家并行（EP）的意思是：专家被拆开摆在不同的卡上。V3 的论文写得很明确 —— 256 个路由专家均匀部署在 64 张 GPU、8 个节点上，所以一张卡上住 4 个专家，一个节点上住 32 个。

那么「专家 #37 这一步收到多少个 token」，问的是： 住着 #37 的那张卡，从全网收到了多少行要交给 #37 算。它跟这张卡自己手里原本有几个 token 是两回事。

把这张卡这一步的工作量算出来

取 DeepEP 仓库自己给的口径（「Following V3's configuration, we tested with 8K tokens per batch, 7168 hidden dimensions, top 8 experts」）：

// 一张卡（rank）手里的 token 8,192 // EP 组一共 64 张卡 → 这一批全网的 token 8,192 × 64 = 524,288 // 每个 token 选 8 个路由专家 → 「token-专家」配对总数 524,288 × 8 = 4,194,304 // 均摊到 256 个专家 → 每个专家平均要算多少行 4,194,304 ÷ 256 = 16,384 // 一张卡上住 4 个专家 → 这张卡这一层要算多少行 16,384 × 4 = 65,536 // 自检：65,536 ÷ 8,192 = 8 = top-k ✓ // 注意这一步是「均摊」，不是守恒律 —— 见下面那条警告。

⚠️ 65,536 是均值，不是定值 —— 这一条很容易看漏

这张卡发出去的配对数才是定值： 8,192 × 8 = 65,536。 它手里就 8,192 个 token，每个选 8 个专家，跟路由怎么选毫无关系。

但它收进来多少行，是全网 64 张卡的路由结果共同决定的，完全可能不等于 65,536。 这一步哪几个专家火，火的那几个又恰好住在哪张卡上 —— 这些都是数据。热门专家扎堆的那张卡就是会多收。

全网守恒的只有总量：4,194,304 行分给 64 张卡。 和守恒，每一份都不守恒 —— 而且这里连「每张卡」这一份都不守恒，不只是「每个专家」那一份。

本文后面一律拿 65,536 当基准，是为了让所有百分比有一个稳定的分母；下面那张表四个数合计正好 65,536，也是我凑的。 真实系统里这个总数每一步都在动 —— 所以接收缓冲区是按一个申报的上界分配的（DeepEP 的 num_max_tokens_per_rank），不是按实际值。

记住 16,384 这个数，它是后面所有判断的分母 一个专家在一张卡上，一步大约收到一万六千行。 后面会看到「补齐到 128 的倍数」—— 128 之所以便宜，全靠它跟这个 16,384 比起来很小。换一个场景（比如推理 decode，一个专家可能只收到几十行），同一个 128 立刻就变贵了。

还有一个专家没算进来 —— 第 9 个，它故意不在这本账里 V3 的每个 token 实际上会经过 9 个专家：8 个路由专家＋ 1 个共享专家。 MaxText 的 deepseek3-671b.yml 写得很直白：

num_experts: 256 num_experts_per_tok: 8 shared_experts: 1 ← 每个 token 都要过它 base_emb_dim: 7168 base_moe_mlp_dim: 2048 base_num_decoder_layers: 61 first_num_dense_layers: 3 → 61 − 3 = 58 层 MoE，本文一直用的就是这个 58

但上面那笔账乘的是 8，不是 9 —— 这不是漏了，是因为 共享专家根本不走分组矩阵乘那条路。

MaxText 里这个 block 的主体就四行，两条路各走各的：

routed_experts = self.routed_moe(inputs, ...) // 排序 → grouped GEMM shared_experts = self.shared_experts(inputs, ...) // 一个普通 MlpBlock return routed_experts + shared_experts

注意第二行读的是 inputs —— 原始的那块 [T, h]，不是排过序的那个缓冲区。 没有 permute、没有复制、没有 group_sizes。

为什么共享专家可以「不用分组」 因为每个 token 都去。 没有「谁去谁不去」，就没有长短不一的段；没有变长的段，就没有任何需要 grouped GEMM 来解决的问题 —— 它退化成一个规规矩矩的稠密矩阵乘。

所以两本账要分开记：

左矩阵的行数：T × 8 —— 只有路由专家进这个数
这一层的总算力：按 9 算 —— 共享专家的中间维同样是 2048，跟一个路由专家等宽，所以它往上加 1/8 ＝ 12.5%

顺带纠一个很自然的直觉：「一个 token 被复制成 9 份」并不准确。 路由那 8 份在 permute 时确实会落成 8 行（MaxText 还有 moe_use_direct_token_gather 这个开关，用索引直接取、连这 8 份都不物化）； 而共享那一份一个字节都不复制，它读的就是原来那块激活。

真实的一步长什么样：不均匀

如果四个专家各收 16,384 行，这篇文档就不用写了 —— 形状本来就是齐的。麻烦在于它永远不齐。取一组有代表性的实际分布（合计仍是 65,536）：

这张卡上的专家	实收行数	相对均值
专家 e0 —— 这一步的热点	21,733	1.33×
专家 e1	15,004	0.92×
专家 e2	17,982	1.10×
专家 e3 —— 这一步的冷门	10,817	0.66×
合计	65,536	—

而且下一步就换一组数。 这四个数是路由算完才知道的，谁也没法提前告诉编译器。这就是第 3 步里说的「dense 的第一条好性质：形状静态可预测」在 MoE 上断掉的地方。

§2老办法：为了把形状凑齐，两头都赔

四步走完，你就知道 capacity factor 是从哪儿冒出来的了。

第 ① 步 · 重排（permute）

token 进来时是按句子顺序排的：第 0 行、第 1 行、第 2 行…… 而 e0 要处理的那些 token 散落在各处。

要做矩阵乘，同一个专家的行必须在内存里连续。所以第一步是物理搬家：按专家把行重新分组，复制出一个新的缓冲区。 这一步无论哪种方案都跑不掉，它不是问题所在。

第 ② 步 · 定一个盒子高度

问题出在这儿。下一步要用 batched matmul 一次算完四个专家 —— 而 batched matmul 要求这一批里每个矩阵的形状完全相同。

四个专家 21,733 / 15,004 / 17,982 / 10,817 行，高度不一样， 这个 kernel 直接用不了。于是倒过来要求模型： 每个专家一律按同一个高度处理。这个高度就是 expert capacity，写成 capacity_factor × 平均值。

为什么 batched matmul 非要形状一样—— 一句话的物理原因

batched matmul 说白了是「同一个矩阵乘法，重复 B 遍，每遍在内存里往前跳一个固定的步长」。它靠一组固定的 stride 来定位第 i 个矩阵在哪儿。

高度不一样 → 步长不固定 → 这套定位方式失效。 约束不在数学里，在寻址里。

MegaBlocks 论文图 3(A) 画的就是这件事：现有实现用 batched matmul 并行算所有专家，由此引入两条约束 —— “all experts are assigned the same number of tokens” 和 “all experts have the same shape”。

第 ③ 步 · 多的丢掉，少的补零

取 capacity_factor = 1.0，盒子高度 = 16,384：

丢掉 5,349 + 1,598 = 6,947 行，补零 1,380 + 5,567 = 6,947 行 —— 两个数一模一样。 这不是巧合：cf = 1.0 时盒子总容量恰好等于真实总行数，所以热门那边溢出多少，冷门那边就空出多少。你在同一张卡上、同一个时刻，一边扔掉真数据，一边拿零去凑算力。

第 ④ 步 · 现在形状齐了，可以 batched matmul 了

四个 16,384 × 7,168 的矩阵摞成一个三维张量，一次算完。代价刚才已经付过了。

那把盒子开大一点行不行

行。capacity_factor = 1.5，盒子 24,576 行 —— 最热的 e0 只有 21,733，装得下，一个都不丢。

代价是这张卡要算 24,576 × 4 = 98,304 行，而真数据只有 65,536 行 —— 比真实多算 50%。

真正难受的不是 50%，是你不知道该填多少这一步 1.5 够了，下一步不一定 —— 那四个数每一步都在变。
MegaBlocks 引 Tutel 的观察：有的模型容量系数要开到 11 才不丢，还有的模型这个值会在训练中途毫无征兆地飙上去。
所以它不是「调个参数就行」的事。它是一个你既不知道该设多少、设错了两个方向都赔的超参数 —— 设小了丢真数据，设大了烧算力和显存。

§3新办法：不设盒子，只补齐到 128 的倍数

同一张卡、同一组数字，重走一遍。

第 ① 步 · 重排 —— 一模一样

还是按专家把行分组搬到一起。这一步没有任何变化。

第 ② 步 · 不定盒子，只做对齐

e0 收了 21,733 行就摆 21,733 行，e3 收了 10,817 就摆 10,817。唯一的加工是：每一段往上补齐到 128 的整数倍。

专家	实收	补齐到	补了几行
e0	21,733	21,760	27
e1	15,004	15,104	100
e2	17,982	18,048	66
e3	10,817	10,880	63
合计	65,536	65,792	256

为什么「一个都不丢」是结构性的，不是运气好 补齐这个动作只会往上加行，永远不会减行。 没有任何一步在做「装不下就扔」的判断 —— 因为压根没有「装不下」这个概念。
这就是 dropless 的全部含义，没有别的玄机。 它的对立面不是「丢得少」，是「有没有一个上限」。

白算 256 行，占真实 65,536 行的 0.39%。上界也很好算：每个专家最多补 127 行，四个专家最多 508 行 = 0.78%。

第 ③ 步 · 高度还是不齐，那怎么乘？

这才是真正的技术内容。答案是：不再一个专家一个专家地乘。

把这张卡上四个专家的权重横着拼成一整块：每个专家的第一个矩阵是 7,168 × 2,048，四个拼起来就是 7,168 × 8,192。

现在设想拿 65,792 行 token 去乘这一整块：结果是 65,792 × 8,192 —— 每个 token 都跟四个专家算了一遍，四倍的浪费。

但你是知道谁该找谁的。 e0 那 21,760 行只需要拼接矩阵里 e0 那一竖条，e1 的行只要 e1 那条。所以结果矩阵里真正需要算的格子，排成一条阶梯，其余全是不用碰的。

关键的转变：从「同一个形状重复四遍」变成「告诉 kernel 算哪些格子」。 前者要求四遍的形状一致；后者只要求你能把「哪些格子」这件事描述清楚 —— 而每个格子的位置和大小，跟别的格子完全无关。整齐这个要求就是在这里消失的。

⚠️ 上面这张图画的是概念，不是内存 —— 这里必须说清楚 那个灰色的大矩形，在真实系统里从来没有被分配过，更没有「先全算一遍再拿掩码盖掉」这回事。
画成一个大矩形，是为了让「阶梯」这个结构看得见 —— 这也是 MegaBlocks 论文的语言（一个块稀疏矩阵，只有对角线上的块非零）。但「稀疏矩阵」是一个逻辑形状：物理上只存非零块，灰色部分不占一个字节。
而今天的 grouped GEMM 干脆连那根轴都不要了 —— 输出直接就是 65,792 × 2,048，专家号只活在权重张量上（§5 的图画的是这一版）。
两种写法的数值完全相同：把对角块压紧，就是后者。 详细的机制在 §5。

第 ④ 步 · 收回来

第二个矩阵乘反过来：拿刚才那个稀疏的中间结果，乘上四个专家竖着拼的第二个权重矩阵，得到一个稠密的 65,792 × 7,168。 每一行只会累加到自己那个专家的那一段，因为别的段结构上就是空的。

最后把补进去的那 256 行扔掉，按原顺序还原，乘上路由权重。完。

§4完整流水账：一层 MoE，每一步张量长什么样

前面讲的是道理，这一节是形状。从这张卡拿到 8,192 个 token 开始，到它把结果交回去为止，一步不跳。

步	做什么	出来的张量	形状
1	进这一层还是句子顺序	x	8,192 × 7,168
2	路由打分一个小矩阵乘，7,168 → 256	logits	8,192 × 256
3	取 top-8	idx / w	8,192 × 8
4	看这 8 个专家落在几个节点分组路由限死最多 4 个 → 决定往外发几份	—	≤ 4 份 / token
5	dispatch（all-to-all）收进来的是全网要交给我这 4 个专家算的行	recv_x counts	65,536 × 7,168 4
6	补齐：每段凑到 128 的倍数	recv_x	65,792 × 7,168
7	grouped GEMM ① gate 和 up 拼成一个 7,168 → 4,096	h	65,792 × 4,096
8	SwiGLU（逐元素）	h	65,792 × 2,048
9	grouped GEMM ② down：2,048 → 7,168	y	65,792 × 7,168
10	扔掉补的 256 行	y	65,536 × 7,168
11	combine（all-to-all 回程）每个 token 的 8 份结果按路由权重加权求和	out	8,192 × 7,168

第 5 步有一个容易漏掉的好处：排序是顺路做的 没有「先收下来、再排一次序」这个额外步骤。 DeepEP 的 dispatch 直接就把行按专家分好段吐出来，还顺手返回 num_recv_tokens_per_expert_list —— 每个专家实收多少。它的参数表里就有 expert_alignment（默认 1）， 补齐到 128 这件事也是在通信里顺手做的，不用事后再搬一趟内存。
所以「重排」在真实系统里不是一个独立开销 —— 它藏在 all-to-all 的落地写入里。

那这一步烧掉多少算力

一个专家对一行 token 要做三次矩阵乘（gate、up、down），每次都是 7,168 × 2,048 量级：

// 每一行 token 在一个专家里的 FLOP 2 × 7,168 × 2,048 × 3 = 88.08 MFLOP // 这张卡这一层，四种记账方式真实 65,536 行 = 5.772 TFLOP ← 分母补齐 65,792 行 = 5.795 TFLOP +0.39% cf=1.5 98,304 行 = 8.659 TFLOP +50.0% 不分组 65,792 行 × 4 专家 = 23.180 TFLOP 4.00× ← 见下 // 通信那一侧（dispatch，FP8） 65,536 × 7,168 × 1 B = 448 MiB 每卡每层

最后那行 23.18 TFLOP 是理解整件事的钥匙如果你把 65,792 行直接乘上「四个专家横着拼起来」的那块大权重，数学上完全成立，一次矩阵乘搞定，形状还特别整齐 —— 但你要付 4 倍的算力。
分组矩阵乘存在的全部意义，就是把这 4× 砍回 1×，同时不退回到「四个小矩阵乘分四次调用」。
专家越多这个倍数越大 —— 一张卡上住 8 个专家就是 8×。 这也解释了为什么 EP 度数越高（每卡专家越少），这件事越好办。

那这一步占多少显存

上面那笔账算的是算力。但 × 8 这个乘数不只出现在 FLOP 里 —— 它同样原样出现在字节里，而这一笔平时很少有人算。

把上面那张表里的大张量按 bf16 折成字节：

// 进这一层的时候，本卡的激活其实很小 x 8,192 × 7,168 × 2 B = 0.11 GiB // 排完序（或者 dispatch 落地）之后，行数变成 8 倍 recv_x 65,792 × 7,168 × 2 B = 0.88 GiB ← 步骤 6 h (gate+up) 65,792 × 4,096 × 2 B = 0.50 GiB ← 步骤 7 h (SwiGLU) 65,792 × 2,048 × 2 B = 0.25 GiB ← 步骤 8 y 65,792 × 7,168 × 2 B = 0.88 GiB ← 步骤 9

一个 0.11 GiB 的张量，进了 MoE 层就变成 0.88 GiB。 整整 8 倍，一个字节都没少 —— 因为 top-8 就是 top-8。

那它是「真复制」还是「一个 view」 是真复制。分组矩阵乘要的是同一个专家的行在物理上连着， view 给不了这个。
在 §14 的 A2A 那条路上这一点很显然 —— 那 65,536 行是从网络上收下来的，本来就得落地。但在「EP ＝ 1 ＋ FSDP」那条路上，这个缓冲区是本地 gather 出来的，很自然会指望它是个视图。它不是。

而且默认路径还不止一份。 MaxText 里默认是先 repeat 再 sort —— 等于连着开了两个「行数 × 8」的数组。

所以代码里有三个专门省这一块的东西，各省各的：

开关	省掉的是哪一份	代价
moe_use_direct _token_gather	中间那份 `repeat` —— 一次 gather 直接从「行数」跳到排好序的「行数 × 8」。docstring 的原话是 “Groups token activations by expert without materializing Top-K copies.”	走的是通用 JAX 反向，不是那个融合的 Mosaic gather-reduce kernel （后者写死了要求正好 8 个专家）—— 可能是拿显存换时间
use_custom _sort_vjp	反向要留的那一份。那个排序挂了自定义 vjp，前向存下来的 residual 只有 `sort_indices` —— 「行数 × 8」个整数，不是「行数 × 8 × 7,168」个 float。反向直接拿 `argsort` 反着排回去。	几乎没有 —— 反向多一次排序而已
num_moe _token_chunks	把峰值摊开。本卡 token 切成 N 块，同一时刻只有 1/N 的缓冲区活着	无 —— 而且它一招两用， §14 里同一个开关是拿来重叠通信的

再叠上 remat，峰值就是一层的量 上面那几个数是一层的瞬时值，不是全模型的。 大 MoE 训练一定开重算，所以这些中间张量算完就还，真正常驻的是每层的那个 checkpoint。

换一个真实配置算一遍 —— 顺便看两个特别容易踩的坑

上面用的是 V3 的口径。换一个不同形状的 295B 级 MoE （base_emb_dim 4,096、base_moe_mlp_dim 1,536、 192 个专家、top-8），在 TPU 上跑 max_target_length = 4,096、 per_device_batch_size = 12：

每 device 的 token 4,096 × 12 = 49,152 排序后的行数 49,152 × 8 = 393,216 排序缓冲区 393,216 × 4,096 × 2 B = 3.00 GiB 正好三个 G gate+up 393,216 × 3,072 × 2 B = 2.25 GiB 层输入 x 49,152 × 4,096 × 2 B = 0.375 GiB

两个坑，都会让你把这个数高估

把 hidden dim 拿错。 本文全篇用的 7,168 是 V3 的。这个模型是 4,096。同样 128K 行，7,168 下是 1.75 GiB，4,096 下正好 1.00 GiB —— 光这一项就差 1.75 倍。换模型必须换 h，不能沿用别家的口径。
把倍数估错。 393,216 是 393K，不是 320K；对着 128K 的锚点应该是 3 倍，不是 4 倍。

这两个错会相乘：

1.75 (h 拿错) × 1.33 (倍数估错) = 2.33 3.00 GiB × 2.33 = 7.0 GiB ← 一个看着很合理、但大了两倍多的答案

估算翻车通常不是因为公式错，是因为往公式里塞了别家的常数。 这也是反过来验算值得做的原因 —— 错的答案能被干净地分解成两个可指认的因子，说明你找到的是真错误，不是编了个说得通的理由。

别急着说「三个 G 而已」—— 在真实预算里它一点都不小 同一份配方实测：pdbs = 12 峰值 91.94 G，余量只剩 2.8 G，pdbs = 14 直接 OOM。
也就是说 —— 光这一个排序缓冲区（3.00 GiB），就比整张卡剩下的全部余量还大。 它当然早就算在那 91.94 里了，但这正说明：「行数 × 8」不是一个可以忽略的常数项，它是 batch 上限的直接组成部分。

§5kernel 怎么知道该算哪些块：三种记法，一件事

「掩码矩阵是怎么弄出来的」—— 答案是它通常根本不是一个矩阵，而是几个很小的一维数组。

在讲编码之前，先把三种布局下 LHS 和 RHS 各自长什么样摆出来 —— 很多说不清的地方，看一眼图就没了。

颜色是配对关系：蓝色那段 token 只跟蓝框那块权重 W0 相乘。
把三行的中间那一列（RHS）盖住对比看 —— 它一模一样。 四块权重，每块 7,168 × 2,048，从头到尾没被碰过。三种布局的差别全部发生在 LHS 和 OUT 的 M 轴上。
DeepGEMM 那句话说的就是这件事： “DeepGEMM groups only the M-axis, while N and K must remain fixed.” —— 所谓 grouped，只 group 了一根轴。

那 kernel 凭什么知道哪一段该配哪一块权重

这是看完上图必然会冒出来的问题，而且它比看上去重要 —— 因为答案解释了「元数据」为什么可以只有几百字节。

答案：靠顺序。配对关系根本没有被「记录」过。 约定死一条：LHS 里的段，严格按专家编号从小到大排。 于是「第 k 段配第 k 块权重」就是定义，不需要任何额外信息去描述它。
你只需要知道「第 k 段从第几行开始」—— 而这只是段长数组的前缀和。

把这条链完整走一遍，就能看清每一环各自负责什么：

环节	它保证了什么
dispatch（通信）	收到的行按专家号升序分段落地。顺序是在这一步建立的，不是事后排的。
段长数组长度 = 本卡专家数	每个专家实收多少行。DeepEP 直接返回它。
前缀和 → 段边界	把「第 k 段」翻译成「第 offset[k] 到 offset[k+1] 行」。 JAX 那份代码的 `group_offsets` 就是这个，文档字符串写得很直白： “group_offsets[i] indicates the row at which group [i] starts in the lhs matrix”。
tile → 段	kernel 逐个 tile 反查：我这一格的行落在哪个区间里，就用哪块权重。

所以「配对」这件事本身是免费的 四个整数（段长）→ 前缀和 → 五个整数（边界）。全部元数据就这么多。 它之所以不需要一张「哪一行属于谁」的大表，正是因为顺序已经把配对编码进去了。
反过来说，一旦顺序被打乱，这套就立刻崩掉 —— 这也是为什么 §4 第 5 步那个「按专家分好段落地」不是可有可无的优化，是前提。

再看一眼上面那张图：三种布局其实分两类 ① batched 和 ③ masked 的 LHS 都是三维的，第一根轴就是专家号 —— 配对关系明摆在形状里，压根不用问。代价是所有专家必须共用同一个高度。
② contiguous 是唯一把这根轴压扁了的 —— 压扁换来了「每段可以不一样高」，而丢掉的那根轴，就用「顺序 + 一个前缀和」补回来。
整个 dropless 的技术核心，浓缩起来就是这一笔交易。

知道了配对靠顺序，再来看具体怎么编码。要让 kernel 干活，它需要的信息只有一句： 第 i 个待算的块，落在哪些行上、因此该用哪块权重。 三套主流实现给出了三种写法，信息量完全一样。

会不会「先全算一遍，再拿掩码盖掉」

不会。而且这一点是整件事最容易想歪的地方。

直觉上很自然会这么想：LHS 源源不断喂进去，RHS 也源源不断喂进去，那前面那批 token 岂不是要跟四个专家都乘一遍？那不就得有个掩码把多余的盖掉？

思路要反过来：kernel 不是「按输入流」组织的，是「按输出格子」组织的 它手上是一张待办清单：输出有 515 个格子要填，第 i 个格子归谁、读哪段、用哪块权重。
它从来不是「让所有 LHS 流过所有 RHS，再挑结果」。
所以「跳过」根本不是一个动作 —— 那些组合从一开始就没上过清单。 没上清单的东西，不需要被跳过。

把 kernel 的骨架写出来，这件事就一目了然了：

// 每个计算单元认领一批 tile，各干各的 for i in 我认领到的 tile: e = group_ids[i] // 这一格归哪个专家 ← 全部的「掩码」就在这 m0 = m_tile_ids[i] * 128 // 这一格覆盖输出的哪 128 行 W = RHS[e] // ← 只取这一块。另外三块这一格从没读过 acc = 0 for k in range(0, 7168, tile_k): acc += LHS[m0 : m0+128, k : k+tile_k] @ W[k : k+tile_k, :] OUT[m0 : m0+128, :] = acc // 跨界的那一格改成累加（§11）

看那一行 W = RHS[e]。 「前面的 token 会不会跟四个专家都乘」这个担心，就消灭在这一行 —— 这一格的整个生命周期里，只有一块权重被读进来过。

所以「掩码」的真身是什么 它不是一张贴在结果上的布尔表，它是一张排班表。 作用发生在工作被分配之前，不是结果算出来之后。
而且它只有一维。 因为「一行属于哪个专家」是个单选题 —— 一行不可能同时属于两个专家。单选题的答案就是一个分段，分段只需要几个边界数字。

跟 splash attention 对照一下 —— 你的直觉方向是对的

这个套路和稀疏 attention 是同一家的： 先算清楚「哪些块真的需要算」，再只对那些块派活，其余的连排班都不排。 差别在于「哪些块」这件事本身有多复杂。

	splash / 稀疏 attention	分组矩阵乘
输出是什么形状	[查询, 键] 的方阵	[行, 特征] —— 没有「专家 × 专家」这种维度
稀疏模式的维数	二维：哪个 q 块要看哪个 kv 块	一维：哪一行属于哪个专家
因此需要什么	一个真正的块掩码结构（哪些块对存在）	几个边界整数就够了
边缘上的「半块」	比如因果掩码切过的那些块 —— 这些块要真的贴一个元素级掩码	跨了专家边界的那一格 —— TPU 的做法是拆成两次执行（§11）

一句话收束 两边都遵守同一条纪律：能在排班阶段解决的，绝不留到算完再修。
attention 的稀疏是二维的，所以它必须带一个块掩码； MoE 的稀疏是一维的，所以它退化成了「四个整数 + 一次前缀和」。
你会觉得需要掩码，是因为你脑子里画的是那张阶梯图（二维）—— 而实现层面它其实是一维的。

那 MegaBlocks 里到底有没有一个「掩码」—— 有，但它是输出规格，不是过滤器

论文的伪代码写成两步：先 sdd(x, w1, topology)，再 dsd(..., w2)。 sdd = sparse ← dense × dense，dsd = dense ← sparse × dense。

那个 topology 参数，就是你直觉里的「掩码」—— 它是三者里最接近的一个。但有两处关键差别：

它是块级的索引结构，不是元素级的布尔矩阵。 编码方式叫 Blocked-CSR-COO —— 一堆块坐标，不是一张跟结果同样大的表。
它是输出规格，不是事后过滤器。 它告诉 kernel「请生产这些块」，而不是「算完之后把这些扔掉」。方向完全相反 —— 这也是为什么灰色区域不占算力也不占内存。

到了 grouped GEMM 这一代，连这个都不需要了 —— 因为一维的分段用段边界就能完整描述，不必绕道稀疏矩阵。 这就是 §8 说的「同一件事的两种记法」，在这里能看到最具体的差别。

记法一 · 对齐 + 段长（DeepGEMM 的 contiguous 布局）

既然每段都已经对齐到 tile 高度，那一个 tile 里的行必然同属一个专家。于是「哪个 tile 归谁」就是一道除法，连查表都省了。

这才是「必须对齐」的真正原因 —— 不是因为对齐了算得快，是因为对齐之后， kernel 在一个 tile 的整个生命周期里只需要加载一份权重矩阵。如果一个 tile 横跨两个专家，它就得中途换权重，内层循环立刻变复杂。

记法二 · 掩码（DeepGEMM 的 masked 布局，decode 用）

形状固定成 [专家数, 最大行数, K]，另给一个长度等于专家数的小数组，说明每个专家前多少行是真的。 README 的原话是 “By providing a mask tensor, the kernel computes only the valid portions.”

所以「掩码」是一个长度为 4 的整数数组，不是一张 65,792 × 4 的布尔表。 它甚至不参与运算 —— kernel 读它来决定循环跑到哪儿停。

记法三 · 三张索引表（TPU 上的 megablox）

JAX 里那个 Pallas 实现走的是第三条路：输入就是一个 group_sizes（长度 = 专家数，放在设备上），由一段预处理算出三张表：

表	长度	含义
group_offsets	专家数 + 1	每个专家从第几行开始 —— 就是 `group_sizes` 的前缀和
group_ids	tile 数 + 专家数	第 i 个 tile 归哪个专家
m_tile_ids	tile 数 + 专家数	第 i 个 tile 落在输出的第几块

源码注释自己点破了血缘：“This metadata is similar to row offsets in a CSR matrix.” —— 正是 MegaBlocks 那套 Blocked-CSR 的 JAX 版本。 （这段代码在公开 JAX 的 pallas/ops/tpu/megablox 里，可以直接读。）

三种记法的共同点，才是本质 它们全都是几百字节的小数组，而且全都待在设备上。 CPU 从头到尾不需要知道任何一个专家收了多少行。
这一点是整件事能成立的地基 —— §7 会讲为什么它比 kernel 本身还关键。

把上面那段伪代码换成真的：grid 其实是三层

前面那段骨架为了讲清楚，写成了「一层 tile 循环＋一层 k 累加」。 真的 kernel 是三层。 把 TPU 上 megablox 那个 Pallas 实现的 grid 抄出来：

// Pallas 按行主序走 grid —— 最后一维变化最快 grid = (num_n, num_gm, num_k) ↑ ↑ ↑ 最外层中间层最内层

层	它在数什么	为什么在这个位置
`n` 最外	输出的列切片 —— 把 2,048 那个中间维切成几块	换得最慢，所以整个 LHS 会被走 `num_n` 遍
`gm` 中间	「专家」和「行块」被压成了一个一维索引	这就是关键的那一步 —— 见下
`k` 最内	在 7,168 上做累加	输出的累加器整个 `k` 循环都留在片上，不落地

那个 gm 轴怎么解回「哪个专家、哪些行」——两张查找表，就是记法三那两张：

gm_id_to_group_id[gm_id] → 这块 tile 该用哪个专家的权重 gm_id_to_m_offset[gm_id] → 这块 tile 盖住排序后的哪一段行 // 三个 index map 长这样（TPU megablox 源码，只改了排版） rhs_weight_index_map(n_id, gm_id, k_id) = (group_id, k_id, n_id) lhs_index_map (_ , gm_id, k_id) = (rows(gm_id), 0, k_id) out_index_map (n_id, gm_id, _ ) = (rows(gm_id), 0, n_id) ↑ 不含 k —— 所以累加器整轮不落地

所以「先 load 一个专家扫所有 token」还是「先 load 一个 token 扫所有专家」都不是。专家根本不是一层循环，它是一次查表。
循环走的是输出 tile；至于这块 tile 归哪个专家，是拿 tile 编号去 gm_id_to_group_id 里查出来的。

硬要说最接近「专家在外层」的，是这一条： gm 这一轴是按专家顺序铺开的，所以连着几个 tile 通常同属一个专家，权重换得很慢。 但那是数据的排列，不是循环的结构。

顺带钉死右矩阵的形状 —— 它是三维的 很容易把 RHS 想成「把 256 个专家横着拼成一个很宽的二维矩阵」。不是。
看 rhs_weight_index_map 返回的 (group_id, k_id, n_id) —— 第一个轴就是专家号。专家是一个独立的轴，不是被拼进列方向的。

这个区别不是抠字眼，它就是 §4 最后那行 23.18 TFLOP 的分水岭： 拼成宽矩阵，你每一行都得跟全部专家乘一遍，付 4×（或者专家数那么多倍）；留成独立的轴，一块 tile 一辈子只读一片权重。

最后一个细节，它一头连着 §7，一头连着 §11

num_gm —— 也就是「这一层一共有多少块 tile 要算」—— 是在设备上数出来的。源码里是一个 fori_loop 扫过所有专家，每个专家往总数里加一笔：

curr_num_gm = cdiv(aligned_group_size, tile_m) // 向上取整 num_gm += curr_num_gm // 在设备上累加

这两行小代码，把全文两条线接上了

「在设备上数出来」 —— 就是 §7 那条时间线里最后一公里的那个能力。tile 总数是数据， host 从头到尾不知道它是多少。少了这一条，你就得回去做一次 device-to-host 同步 —— 正是要躲开的那个东西。
那个向上取整 —— 就是 §11 那笔边界账的出处。每个专家都单独向上取整到整块 tile，所以段数越多、每段越短，浪费越大。它不在某个公式里，它就在这一行 cdiv 上。

顺带回答：硬件需要「能跳过空块」的能力吗

不需要。一点都不需要。 真正被算的每一个块，都是一次普普通通的稠密小矩阵乘。 MXU / tensor core 从头到尾不知道外面那个大矩阵是稀疏的 —— 它收到的就是一对稠密小矩阵，照常乘。

「跳过」发生在排班层，不是算术层。 那些空白块不是「算出来是零」，也不是「被硬件跳过了」—— 是 kernel 的 tile scheduler 压根没给它们派活。

要跟真正的硬件稀疏分清楚： 2:4 structured sparsity 那种才需要专门的 sparse tensor core、需要权重满足特定模式。 这里跟那个毫无关系，同一批算术单元、同一条指令。

唯一真正需要的新能力是这一条 —— kernel 得能在设备端读一小段元数据，并据此决定每个 tile 干什么。 GPU 上靠 persistent kernel 里的 tile scheduler； TPU 上靠 Pallas 的标量预取（把那三张索引表提前搬进标量内存，让 kernel 逐 tile 查）。 这是编程模型的能力，不是算术单元的能力。

§6trade-off：一张表看完

同一张卡、同一步、同一组实收行数（65,536 行真数据）。

方案	这张卡要算	丢掉	白算	要不要猜参数
盒子 · cf = 1.0	65,536	6,947	6,947 (10.6%)	要 —— 而且这次猜错了
盒子 · cf = 1.5	98,304	0	32,768 (多算 50%)	要 —— 这次够了，下一步不保证
分组矩阵乘 · 补齐 128	65,792	0	256 (多算 0.39%)	不要 —— 这个参数不存在

那付出的是什么

三条，都很实在：

现成的 batched matmul 用不了了，得有一个能吃「变长分段」的 kernel。 这是真活儿 —— 所以它是一篇论文，不是一个配置开关。
每一步都要现算一份「哪些块要算、在哪儿」的元数据，而且反向传播还要按转置的顺序再遍历一遍同一张图。 MegaBlocks 的办法是不搬数据、只多存一份按转置序排好的块偏移索引 —— 论文自己打的比方是「数据库的二级索引」。之所以划算，是因为块够大，元数据比真实数据便宜好几个量级。
你被锁在块粒度上。某个专家只收到 1 行，你照样得算满 128 行。 这条在训练时无关痛痒，在 decode 时会要命 —— 见 §10。

⭐ 真正的旋钮不是「丢不丢」，是块开多大 块越大，越接近密集矩阵乘的速度，但补零浪费越多；
块越小，浪费越少，但矩阵乘越碎、跑不满算力。
128 是这条曲线上的甜点 —— MegaBlocks 专门测了不同 tile 尺寸，结论是 128×128 「consistently on-par or better than other configurations」。
而这套东西对 MoE 特别管用，原因只有一条： 一个专家一步收到上万行，128 相对它很小。如果每个专家只有一百行，这招立刻失效。

§7「把不同长度的段接起来连着乘」——
这么朴素的想法，为什么以前没人这么干

先把话说明白：你的直觉是对的，这里数学上一点创新都没有。 真正的问题不是「想不到」，是「做出来打不过 bmm」。

先看被淘汰的那个方案：for 循环

「不同长度就分别算呗」——这个方案一直都存在，写出来五行：

// 谁都能写出来的版本 for e in range(4): seg = x[offset[e] : offset[e+1]] // 这个专家的那些行 out[offset[e] : offset[e+1]] = seg @ W[e]

它是对的，而且它很慢。三个原因，第一个是致命的：

	问题	为什么致命
①	那个 `offset` 在 GPU 上	它是路由算出来的。要让 Python 的 `for` 用上它，必须把它取回 CPU —— 一次 device-to-host 同步。每层两次、58 层，整条流水线被砍成 116 段，GPU 每次都得等 CPU。而且 CUDA graph 直接用不了，因为图要求形状固定。
②	kernel 启动次数	4 个专家 × 2 次 GEMM × 58 层 = 464 次 / 卡 / 步，反向还要翻倍。每次几微秒，加起来是实打实的一块。
③	单个 GEMM 喂不饱卡	把一块几百上千个计算单元的加速器，用一个只有一万多行的矩阵乘去填，尾巴上总有一批单元闲着。分四次，就闲四次。

所以 bmm 当年是一笔好交易，不是一个错误 它一次 launch 解决全部四个专家，CPU 全程不用知道任何形状， GPU 塞得满满的。代价只有一个：形状必须一样。
而在 2021–2022 年，这个代价看起来很便宜 —— 那时候 MoE 还是少数派架构，专家数少、多半是 top-1 或 top-2，「补零凑齐」的浪费也就没那么显眼。

再往回追一层：那个「形状必须一样」，最早是谁定的

§2 是从机械上讲 capacity factor 怎么冒出来的。但如果去翻 MoE 最早那批论文，会看到一件更有意思的事 —— 没有人「选择」过这个代价，它是写法自带的。

GShard（2020 年 6 月）把整个 MoE 层写成了一串 einsum：

// GShard 论文 Algorithm 2 里的核心两行 expert_outputs = einsum("EGCH,EHM->GECM", h, wo) outputs = einsum("GSEC,GECM->GSM", combine_weights, expert_outputs) ↑ E = 专家数 C = 「容量」—— 它是张量的一个轴

einsum 要求张量是矩形的。 矩形就意味着每个专家必须占一样多的行。 那个数需要一个名字，于是它被叫做 expert capacity。

所以「丢 token」从来不是一个算法设计 它是「张量得是方的」这件事的副产品。 而且论文里是把它当成 feature 写的 ——
“To ensure the load is balanced, we enforce that the number of tokens processed by one expert is below some uniform threshold, which we define as expert capacity.”

注意开头那半句：To ensure the load is balanced。 在原文的叙述里，容量是「保证均衡的手段」，不是「为了迁就形状而做的妥协」。溢出的 token 怎么办，同一段写得很干脆：门控输出退化成零向量，token 靠 residual connection 直接跳到下一层。

而 GShard 是长在 TPU ＋ XLA 上的。同一篇文章里还有一句更硬的话，出现在讲 SPMD 分区的那一节：

“XLA requires tensor shapes to be static.” 脚注：“The limited dynamism in the intermediate representation is often necessary to efficiently target accelerators.”

这不是 MoE 的特殊规定，这是整个编译器的底线。 MoE 只是恰好第一次撞上它。

一句话 capacity factor 是一个编译器约束，被提拔成了一个模型超参数。
之后所有围着它长出来的东西 —— 辅助均衡损失、容量系数调参、溢出走 residual —— 都是在修补这一个决定的下游后果。这也是为什么它后来那么难被拿掉：它已经不在 kernel 里了，它在模型定义里。

那后来到底变了什么

不是想法变了，是三件工程做成了。

形状留在设备上，CPU 不再参与。 段长、段偏移、掩码全是 device tensor，kernel 自己读（§5 那三种记法）。 DeepEP 描述 decode 路径时那句 “the layout is reused without CPU synchronization” 说的就是这件事。 这一条直接干掉了 for 循环方案的①，也是全部收益里最大的一块。
一次 launch 干完所有 tile。 kernel 用固定的 grid 启动，内部有个调度器循环认领「第 i 个 tile」，查表得知它属于哪个专家、要用哪块权重。 launch 次数从 464 回到跟 bmm 一样。
用一纸「契约」保住内层循环。 不管是 GPU 那边「每段对齐到 tile 高度」，还是 TPU 那边「允许一个 tile 被访问两次」（§11），目的都一样 —— 保证最内层还是一次纯粹的稠密小矩阵乘。这样几十年积累的 GEMM 内层优化一行都不用改。

还有一个不那么技术、但同样真实的原因

值不值得。写一个变长分组、还要打平 cuBLAS 的 kernel，是几个人月的活。 2022 年 MoE 是少数派，这笔投入很难论证 —— 所以它当年是一篇论文，而不是一个库里的默认选项。 到了今天 MoE 成了大模型的主流形态，同一件事的性价比彻底翻转。

硬件趋势还在同一个方向上推：加速器一代比一代宽， «一个矩阵乘要多大才能喂饱它» 的门槛一直在涨。 于是「每个专家单独调一次」这个方案，会随着硬件变强而相对变得更差；而「合成一次调用」的收益只增不减。

一条时间线：这个原语其实一直躺在货架上

上面讲的都是「为什么难」。但还有一件比「难」更尴尬的事实 —— grouped GEMM 不是为 MoE 发明的，它比 MoE 需要它早了两年就已经在库里了。

时间	发生了什么	说明什么
2020-06	GShard：MoE 层写成 einsum，容量 C 成了张量的一个轴	丢 token 的出身
2021-11	CUTLASS 2.8 带上了 grouped GEMM 示例（`examples/24_gemm_grouped`）	原语已经存在 —— 当时面向的是「一堆大小各不相同的小矩阵乘」这种通用场景，跟 MoE 毫无关系
2022-11	MegaBlocks 提出 dropless MoE —— 但走的是 block-sparse，不是 grouped GEMM	2022 年真能做出来的那条路，不是这一条
2023-09	一个很薄的 `grouped_gemm` CUTLASS 封装出现，成为 MegaBlocks 的可选依赖	到这时才有人把这个原语接到 MoE 上
2024-06	cuBLAS 12.5 把 grouped GEMM 做进官方 API	从「论文里的东西」变成「调一下就有」
2026-02	CUTLASS 4.4.0：grouped GEMM 示例开始支持纯设备端的 problem shape	最后一公里，刚补上没多久

这张表里最该看的是第二行和最后一行 原语 2021 年就在了，但在之后很长一段时间里，那个尺寸数组仍然得由 host 拼好、再拷到设备上。
对 MoE 来说这等于说：想调用它，先做一次 device-to-host 同步 —— 正好是你想躲开的那个东西。

所以「有这个原语」和「能用它做 dropless MoE」之间，隔着的从来不是矩阵乘，是那一次同步。 这也解释了第三行：MegaBlocks 到今天仍然把 block-sparse 那套（stanford-stk）列为必需依赖， grouped_gemm 只是一个可选 extra。同一个作者、两条路，先做出来的是 block-sparse 那条。

那么到底是什么变了 —— 三条，没有一条是矩阵乘本身

犯错的代价涨了。 上面那条注解说过，2021–22 年这个代价看着便宜。今天是 256 选 8、每张卡上放好几个专家，单专家负载抖得厉害（§1 那张表里热点是均值的 1.33 倍）。同样一个 1.5 的容量系数，现在是三分之一的算力。 浪费是跟着专家数一起长大的。
推理来了 —— 这一条最硬。 训练时丢一个 token，你还能把它说成正则化；服务时丢的是用户的 token，这话没法交代。 MoE 进了生产之后，dropless 就从一个性能问题 变成了一个正确性问题。
配套修好了。 上面那三件工程，再加上一个「按专家落地」的 all-to-all （§5「配对靠顺序」）。两头齐了，中间那个改动才显得小。

一句话收掉这一节：这件事不是被想出来的，是被允许的。 数学上它一天都没有难过；难的一直是「形状归谁管」。而那个答案，是编译器、通信库、kernel 三边一起改出来的 —— 没有哪一边单独能把它改掉。

硬件要求到底差在哪：一张表

	批量矩阵乘（bmm）	分组矩阵乘（grouped GEMM）
形状什么时候知道	启动前，CPU 知道	运行时，而且只有设备知道
第 i 个矩阵在哪	`base + i × 固定步长`	查一张偏移表
怎么分派工作	grid 按形状算好，一个 tile 一个 block	固定 grid，kernel 内部逐个认领 tile
要不要 CPU 同步	不要	不要（这是设计目标，不是附带）
算术单元	tensor core / MXU	完全一样，一条指令都没多
需要的新能力	—	设备端读元数据 + 间接寻址 + tile 调度器

反过来问一遍：那我就分四次算，不行吗

这是同一个问题的镜像，而且值得单独回答 —— 因为「分四次」听起来更简单，实际上不是。

先把答案放这儿 行。分四次能算对，而且在 prefill 这个尺度上，性能也只差个位数百分点。
但它并不更简单 —— 你以为省掉的那套机制，一样都省不掉。而做完那套机制之后，合成一次几乎不花钱，还严格更好。

① 「分四次」在你脑子里是这样，但 CPU 不知道每份多大

你想写的循环大概长这样：

for e in range(4): out[start[e] : end[e]] = lhs[start[e] : end[e]] @ W[e]

start[e] 和 end[e] 是数据，不是常数。 它们由这一步的路由结果决定，算在设备上。于是只剩两条路：

拷回 CPU 再发四次调用 —— 每层一次设备到主机的同步，流水线当场断掉。 这也是当年所有人绕开 for 循环的头号理由。 ⚠️ 但今天的 GPU 训练路径其实已经付了这次同步 —— 所以在那条路上「分四次」并不会多同步一次。下一小节的表格里有更新后的完整说法。
让每次调用自己去设备上读形状 —— 那你就得有：设备端的元数据、间接寻址、一个能按元数据认领 tile 的调度器。 这份清单跟 grouped matmul 需要的一模一样。

所以关键在这儿 「分四次」不是一个更便宜的方案，它是同一个方案做四遍。 难的那部分（把形状搬进设备的数据流）两边都躲不掉。一旦躲不掉，把四次并成一次基本是免费的 —— 循环边界从「一个专家的段」换成「所有专家的段」，就这一处改动。

唯一真正更简单的版本，是给每个专家一个固定大小的盒子 —— 形状变常数，四次调用干干净净。那就是 §2 的老办法，代价是丢 token。 这就是为什么这个问题的答案，跟 dropless 是同一个答案。

接着追一句：那 kernel 自己怎么分得开

这才是机制的核心。答案是：CPU 从头到尾只需要知道一个数 —— 「要派多少份活」。别的什么都不用知道。

因为 kernel launch 是 CPU 发出去的，grid 大小是 launch 接口的参数。除此之外，kernel 一旦跑起来，读显存是它的自由。所以真正的分界线在这儿：

	CPU 必须知道吗	为什么
一共要派多少份活	必须	这是 launch 的参数，而 launch 是 CPU 发的
第 i 份活该干什么	完全不用	那时候 kernel 已经在设备上跑了，自己查表就行

于是全部差别就落在第一行那个数上。而这里有个陷阱，值得单独拆开：

// 分四次：那个数是 m_e —— 一个变量 grid_e = ceil(m_e / 128) // 合一次：那个数是 m_total —— 也是一个变量！ grid = ceil(m_total / 128) (+ 边界余量) // 唯一真正静态的，是你自己申报的那个上界 num_max_tokens_per_rank ← DeepEP 按它分配接收缓冲区

⚠️ 这里我原先写错了，改过来我一开始写的是「合起来之后那个数是恒等式，等于 tokens × top-k」。不对。 那是发出去的配对数（§1 已修正）。 收进来多少行是数据，合起来也一样是数据。

所以真正的问题不是「合起来之后那个数变常数了吗」，而是 「那个数，CPU 到底知不知道」 —— 答案分三种情况，而且互不相同。

场景	CPU 知道每份多大吗	依据	分四次可行吗
GPU · 训练 / prefill	知道	DeepEP 的 `handle.num_recv_tokens_per_expert_list` 是个 Python list —— 说明已经拷回主机了。而 decode 那条路径的注释写着 “Reuse cached handle: skip layout recomputation and CPU sync”，反过来正好证明非缓存路径是要同步的。	可行
GPU · decode + CUDA graph	不知道	DeepGEMM 原话：masked 布局存在的理由就是 “when CUDA graph is enabled and the CPU is unaware of the number of tokens each expert receives”	不可行
TPU / XLA	永远不知道	jit 出来的是静态形状的程序。`group_sizes` 是设备数组，整个训练步不出编译边界 —— 没有任何一个时刻，主机有机会看到它	不可行

于是「为什么非得合成一次」的准确答案是这样

我上一小节说的「合起来 CPU 就不用参与了」， 在 GPU 训练路径上不成立 —— 那条路本来就有一次设备到主机的拷贝，而且是 dispatch 那一步付的，跟后面的 GEMM 分不分开无关。分四次并不会多一次同步，只会多三次 launch。

站得住的是这三条，而且理由各不相同：

TPU / XLA 上是结构性的。 静态形状编译，形状根本没有机会离开设备。「分四次」不是慢，是写不出来 —— 除非退回固定盒子，那就是丢 token。
GPU decode 上也是结构性的。 CUDA graph 一开，主机就被蒙住了。masked 布局就是为这个造的。
GPU 训练上只是性能性的。 尾巴（上面算过，约 2.4%）＋ 4 倍 launch。真实，但不是「不得不」。

而下面「kernel 内部靠什么区分」那一段完全不受影响 —— 顺序继承自 all-to-all、四个段长做前缀和，这一套在三种情况下一模一样。

那 kernel 内部到底靠什么「区分」

它不需要区分 —— 它是继承的。

token 是被 dispatch 的 all-to-all 按专家号送到这张卡上的。 换句话说，到达的先后顺序本身就已经把专家号编码进去了。 分组这件事不是矩阵乘 kernel 做的，是通信做的 —— kernel 只是继承了通信留下的那个顺序（这也是 §5「配对靠顺序」那一段说的事）。

所以剩下需要传递的信息，就只有每段多长：

group_sizes = [21733, 15004, 17982, 10817] // 4 个 int32，路由算出来的，在设备上 ↓ 前缀和（设备上，几十个周期） group_offsets = [0, 21733, 36737, 54719, 65536] // 5 个 int32 ↓ 查表第 i 份活覆盖第几行 → 落在哪两个边界之间 → 就是第几个专家

整个「掩码」的信息量就是这五个整数。 而最后那个 65,536 甚至不是算出来的，它是那个恒等式 —— 所以真正的自由度只有四个数。

这段前缀和跑在哪儿：TPU 那份代码可以直接读—— trip count 是 kernel 自己在 prologue 里算的，CPU 只给了「专家数」这个编译期常数

公开那份 Pallas GMM kernel 里，元数据是在 kernel 的序言里现算的（§12 讲过它落在 SMEM）：

// 静态上界：只跟「专家数」有关，是编译期常数 max_num_group = group_offset + size_group // 真实 trip count：kernel 自己在设备上跑一遍循环算出来 num_gm, _ = lax.fori_loop(0, max_num_group, outer_group_loop, (0, 0)) ↑ 循环体里读的是 lhs_group_sizes_ref[...] —— 路由的输出，在设备上 // 同一个循环顺手把两张表填进 SMEM metadata_ref.gm_id_to_group_id[tm_id] = group_id metadata_ref.gm_id_to_m_offset[tm_id] = curr_m_offset

注意这里的分工：CPU 给的是 max_num_group（＝专家数，常数），设备算的是 num_gm（＝真实要跑多少步）。 CPU 从头到尾没有见过任何一个 m_e。

GPU 那边把同一件事做得更彻底：常见做法是 persistent kernel —— grid 干脆等于计算单元的数量，跟问题尺寸完全无关，每个 block 起来之后自己去设备上的计数器领下一个 tile。 这样连「总行数」都不用告诉 CPU 了。

两句话 一、CPU 唯一需要知道的是「派多少份活」。合起来之后这个数是恒等式，分开之后是变量 —— 差别只有这一处。
二、kernel 内部不需要「区分」谁是谁，因为 all-to-all 已经按专家号排好序送过来了；它只需要四个段长做一次前缀和。

② 「拼在一起」这件事本身，既不花钱，也不是好处的来源

这里有个很容易混掉的地方，值得拆开：

	拼成一条（布局决策）	合成一次调用（调度决策）
要额外花钱吗	不用。dispatch 的 all-to-all 本来就是按专家把 token 送过来的， contiguous 布局是它落地时的自然形状，没有人执行过一次「拼接」	不用，就是循环边界换个写法
好处来自这里吗	基本不是	是
反证	decode 用的 masked 布局根本没有拼接 —— 形状是 [专家数, 最大行数, K]，四个专家各占一层，中间全是空的。但它仍然是一次 kernel 调用。（见 §10）

所以你问的那个「好处」，归「一个 kernel」，不归「拼一起」。 拼起来只是因为数据本来就是那么躺着的，不拼反而要多做一次搬运。

③ 合成一次，实打实赚的是这一块：尾巴只剩一条

每个 kernel 收尾的时候，最后一波派出去的活填不满全部计算单元，剩下的单元只能空转等它结束。四次调用 = 四条这样的尾巴。

// 数量级估算，模型：一个计算单元同时算一个输出 tile 输出 tile 总数 = (65,536 ÷ 128) × (2,048 ÷ 128) = 512 × 16 = 8,192 每条尾巴平均空转 ≈ 半波 ≈ 66 个 tile 的位置 // 按 132 个计算单元估合成一次 1 条尾巴 ≈ 66 分成四次 4 条尾巴 ≈ 264 差 198 ÷ 8,192 = 约 2.4%

2.4% 是真的，但也就是 2.4% —— 这是个粗估，真实值取决于 tile 尺寸和占用率，我把模型写在上面了，你可以自己代别的数。

真正要注意的是它怎么长大：这个数跟专家数成正比，跟每个专家的行数成反比。 这张卡上 4 个专家、每个一万多行，所以只有 2.4%。换成一张卡放 32 个专家，就是 32 条尾巴；换成 decode、每个专家只剩一百多行，那连第一波都填不满 —— 就不是百分之几的事了。

三句话收掉 一、难的那部分（形状留在设备上）分不分都躲不掉，所以「分四次」并不更简单。
二、真正带来好处的是「一次调用」，不是「拼成一条」 —— 拼是因为数据本来就那么躺着。
三、合起来赚的是「尾巴只剩一条」，这张卡上约 2.4%；专家越多、每个专家越小，这个数涨得越快。
你的直觉「token 只跟自己的专家发生关系」完全正确 —— §12 那条流水里，不匹配的权重连搬都没搬进片上过。合成一次调用没有改变这一点，它只是让这些互不相干的乘法 排在同一条流水线上跑完，中间不停机。

一句话回答「有什么创新」 创新不在矩阵乘法里，在谁在什么时候知道形状这件事上。 把「形状」从 CPU 的控制流里搬进 GPU 的数据流里 —— 就这一下，「变长」从一个必须先解决掉的障碍，变成了一个普通的输入参数。

§8这个能力到底叫什么：block-sparse 还是 grouped GEMM

同一件事，两代人用了两个名字。今天该用后一个。

把 §3 第 ③ 步那个能力抽象出来，它是一句话：

一次 kernel 调用，解决一串 M 各不相同、N 和 K 相同的矩阵乘法。

这个东西的通用名字就是 grouped GEMM（分组矩阵乘）。 MegaBlocks 在 2022 年没有用这个名字 —— 它用的是 block-sparse matmul：把这一串矩阵摆成一个块对角矩阵，然后说「我只算对角线上那些块」。

两种说法描述的是同一个计算。 区别在于你把「哪些块要算」这件事，是编码成一个稀疏矩阵的拓扑，还是编码成一个「每段多高」的长度列表。

	MegaBlocks（2022）	今天的实现
叫法	block-sparse matmul，dMoE	grouped GEMM
怎么描述「算哪些」	稀疏矩阵拓扑（Blocked-CSR-COO 编码）	一个长度列表每个专家收到多少行
代表	MegaBlocks / stk	CUTLASS grouped GEMM、DeepGEMM

DeepGEMM 的文档把这层意思写得最直白：

“Unlike traditional grouped GEMMs in CUTLASS, DeepGEMM groups only the M-axis, while N and K must remain fixed. This design is tailored for scenarios where experts in an MoE model share the same shape.”

所以 MegaBlocks 的历史地位该怎么讲 它的贡献不是「发明了分组矩阵乘」。 它的贡献是第一次把「MoE 可以完全不丢 token」做成一个能跑满硬件的实现，并且用一组对照实验证明了这件事值得做（见 §16 的三个数）。
但接口层面，今天大家用的不是它那套 block-sparse 编码，而是更直白的 grouped GEMM。 讲课时可以说「MegaBlocks 把这条路走通了」，不要说「今天的 dropless 都是 MegaBlocks 那套」——不是。

§9补齐到 128，跟 chunked prefill 是一回事吗

同一族思想，但更贴切的对照是 PagedAttention 的 page，不是 chunked prefill。

共同的那句话是：

一个变长的东西，要交给只认固定形状的下游去处理，就得先量化成固定大小的块。
代价是最多浪费一个块；收益是下游从「处理不了」变成「高效处理」。

	切的是什么	块大小定在哪	为了迁就谁	最多浪费
专家分段补齐	矩阵的行数	GEMM 的 tile 高度	矩阵乘 kernel	一个块（127 行）
PagedAttention	KV cache 的长度	一个 page 的 token 数	显存分配器	一个 page
chunked prefill	一次前向的工作量	一个 chunk 的 token 数	调度器	—— 见下

为什么 PagedAttention 更像。 两者都是空间上的量化：一段连续的东西长度不定，切成固定大小的格子来放，尾巴上那格没填满就是浪费，浪费的上界就是一个格子。 「专家收到 21,733 行 → 占 170 个 128 行的块，最后一块只用了 101 行」跟「一个序列 21,733 个 token → 占 N 个 page，最后一页没填满」，是同构的。

chunked prefill 差在哪。 它切的是时间轴上的工作量，目的是调度 —— 让长 prefill 不要一口气霸占住卡，好跟 decode 交错、把延迟毛刺压下去。它的收益是延迟和公平性，不是「让一个 kernel 能跑」。分块本身不产生填充浪费（最后一个 chunk 短就短了，照算），它的代价在别处：KV 要反复读、attention 被切成几段。

一句话分辨 补齐 128 和 PagedAttention 是「为了放得下 / 算得动」而量化，浪费在填充；
chunked prefill 是「为了排得开」而量化，代价在重复读取。
都是「变长 → 定长」，但迁就的对象和赔的东西不一样。

§10为什么 decode 要换一套：128 在那儿就贵了

§1 埋的那个伏笔在这里兑现 —— 128 便宜与否，取决于分母。

训练和 prefill 时，一个专家一步收上万行，补 127 行是 0.8%。 decode 时分母塌下来了。

按同样的口径推一遍（假设每张卡同时服务 B 条序列， decode 每步每条序列出一个 token）：

// 全网这一步的 token B × 64 // × top-8 ÷ 256 个专家 → 每个专家收到 B × 64 × 8 ÷ 256 = 2B // 训练那一步代进去自检：B = 8,192 → 16,384 ✓ 跟 §1 对上 // B = 512 → 每个专家 1,024 行　补齐 128 → 正好整除，不浪费 // B = 64 → 每个专家 128 行　补齐 128 → 正好一个块 // B = 16 → 每个专家 32 行　补齐 128 → 算 4 倍的量 // B = 4 → 每个专家 8 行　补齐 128 → 算 16 倍的量

同一个 128，在训练时是 0.8% 的税，在小批量 decode 时是几倍的浪费。

所以 DeepGEMM 给 decode 准备的是另一套布局。文档原话：

“During the inference decoding phase, when CUDA graph is enabled and the CPU is unaware of the number of tokens each expert receives, we support masked grouped GEMMs. By providing a mask tensor, the kernel computes only the valid portions.”

思路反过来了：不再把每段补齐到块边界，而是开一个固定大小的缓冲区，再用一个 mask 告诉 kernel 哪些行是真的。 这样既满足了 CUDA graph 要求形状固定，又不用为对齐付整块的钱。

两套布局，两个约束 contiguous 布局（训练前向 / prefill）：段长可变、补齐到块 —— 因为行数很多，对齐很便宜。
masked 布局（decode）：形状固定、用 mask 标有效 —— 因为行数很少，对齐会翻好几倍。
两者都不丢 token。「dropless」这件事跟用哪套布局无关。

§11TPU 那边不补齐 —— 同一个问题的另一种解法

前面十节都默认「补齐到块」是唯一出路。不是的。 JAX 里那个 TPU 实现走了另一条，值得单独看，因为它正好补上了 MegaBlocks 论文留的那个口子。

MegaBlocks 在 §5.2 说了两句话：一是「我们把每组补齐到 128 的倍数」，二是「这个约束是可以去掉的，只要支持边缘的不完整块」—— 当时没做。JAX 的 megablox 做了。

它是怎么绕过补齐的

核心就一句：允许一个 tile 被访问两次。

假设 tile 高度 128，而专家 e1 从第 21,733 行开始 —— 这一行落在第 169 个 tile 的中间。那么这个 tile 会被排两次班：

第一次算它属于 e0 的那部分（第 21,632–21,732 行），用 e0 的权重
第二次算它属于 e1 的那部分（第 21,733 行往后），用 e1 的权重
两次结果累加进同一个输出块 —— kernel 里靠「这个 tile 的专家号跟上一个 tile 不一样吗」来决定要不要清零累加器

同一条边界，两种付法：A 多出 27 行零，B 多执行 1 个 tile。 B 里那两次执行不是浪费在「算零」上 —— 它们算的全是真数据，只是同一块输出被写了两遍（第二遍累加）。 kernel 靠「这个 tile 的专家号跟上一个一样吗」来决定要不要先清零累加器。

代价因此换了一种形式，而且源码注释直接把上界写出来了：

// 要执行的 tile 总数，夹在这两个数之间 tiles_m ≤ 实际执行的 tile 数 ≤ tiles_m + 专家数 − 1 // 代进我们这张卡（65,536 行、tile 高 128、4 个专家） tiles_m = 65,536 ÷ 128 = 512 四段起点 0 / 21,733 / 36,737 / 54,719 后三个都不在 128 的边界上 → 3 次额外访问实际执行 515 个 tile = 等效 65,920 行 = 多算 0.59% // 对照 GPU 那条路（补齐到 65,792 行） = 多算 0.39%

两条路，同一个量级，赔的东西不一样

	补齐（DeepGEMM 的 contiguous）	重复访问（JAX megablox）
段长要不要动	要，往上凑到块的整数倍	不动，原样就是实收行数
浪费的形式	补进去的零行	跨界 tile 的重复执行
上界	每专家 < 一个块	总共 < 专家数 − 1 个块
内层循环	最简单：一个 tile 一份权重	要处理「块内换专家」和累加器清零
这个例子里	+0.39%	+0.59%

谁更好取决于专家数和 tile 高度的比例，不存在通用答案。 专家多、每个专家行数少的时候，「补齐」的上界（专家数 × 127 行）会迅速变难看，「重复访问」的上界（专家数 × 1 个 tile）其实同阶 —— 真正的差别在内层循环的复杂度，以及要不要在通信阶段动数据布局。

一个很好的旁证 那份 JAX 代码的注释里，把这套元数据形容成 “similar to row offsets in a CSR matrix”。
而 MegaBlocks 论文的编码就叫 Blocked-CSR-COO。同一个血统，隔了三年，在另一种加速器上、用另一种语言重写了一遍，还顺手把当年没做的那件事做了。

§12再往下一层：MXU 是怎么被喂饱的
—— HBM → VMEM → VREG 的那条流水

前面十一节，「tile」一直是最小的砖。但 tile 不是硬件概念，是软件概念。 真正的硬件里有四层各自认死了形状的东西，一个 tile 得逐层换装才能走到乘法器面前。这一节把那条流水拆开 —— 用的是能公开读到的那份 TPU kernel 源码 （Apache-2.0，MaxText/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py，文件头写明 forked from openxla/tokamax）。

先说清楚：这是第二代 kernel，跟 §5 §11 讲的不是同一份 §5 和 §11 讲的是公开 JAX 里 pallas/ops/tpu/megablox 那一版（三张索引表、跨界 tile 执行两次）。这一节读的是更新的一版，思路一脉相承，但把边界粒度从「128 行」细化到了「一条 sublane」—— 那正好是硬件真正的最小单位。两版都在跑，不冲突。

四层，每层认一种形状

这条链上最反直觉的一点：它们的单位互不整除，也互不让步。 一个数据块每下一层，都要重新按下一层的形状裁一次。

层	它认什么形状	谁负责把东西搬进来	这个 kernel 在这一层放什么
HBM 片外主存	随便，但按 DMA 描述符成批搬	DMA 引擎（描述符由标量单元发）	整个 lhs[65536, 7168]、 rhs[4, 7168, 2048]、输出
VMEM 片上 SRAM	末两维必须是 8 的倍数 × 128 的倍数（或正好等于整维）	DMA，由 Pallas 的流水线自动排	当前几个 tile + 累加器 + 一个跨界残片
VREG 向量寄存器	8 sublane × 128 lane，每格一个 32 位值	向量 load / store 指令	正在参与这一次乘法的那一小片
MXU 脉动阵列	固定边长的方阵，v7 上是 256 × 256	—— 数据从两边流进去	一次吃 [tile_m, k] × [k, 256]

这些数字 kernel 自己不写死 —— 它当场去问硬件 源码里到处是 pltpu.get_tpu_info()： .mxu_column_size（MXU 边长）、.num_lanes（128）、 .vmem_capacity_bytes（VMEM 容量）、 .get_sublane_tiling(dtype)（一条 sublane 装几个逻辑行）。
所以同一份 kernel 在 v5p / v6e / v7 上会自动切出不同的 tile。下文凡是给具体数字的地方，都是代入 v7（MXU 256）之后的结果，不是源码里的常量。

整条链上，「哪个专家」这件事只在最顶上那一层出现过一次。 标量单元从 SMEM 读一个组号，把它变成一个 HBM 地址。从 DMA 往下的每一层 —— VMEM、VREG、MXU —— 都完全不知道 MoE 的存在，它们看到的只是一对稠密小矩阵。

第 0 层 · 元数据必须比数据先到

整套机制的起点是 Pallas 的标量预取：

// 外层 pallas_call grid_spec = PrefetchScalarGridSpec(num_scalar_prefetch=2, ...) // scratch 里给元数据留的位置 —— 注意 memory space 是 SMEM 不是 VMEM gm_id_to_group_id = pltpu.SMEM((max_num_gm,), jnp.int32) gm_id_to_m_offset = pltpu.SMEM((max_num_gm + 1,), jnp.int32)

为什么非得是 SMEM。因为读它的不是向量单元，是标量单元，而标量单元干的活是算地址。公开 JAX 文档把这条规矩写得很直白： “any data used to perform control-flow decisions should be placed in SMEM”。

然后这段就是全篇的题眼 —— 三个 index map，各自决定「这一步该从 HBM 的哪里搬」：

// 权重该取哪一块？读一次 SMEM 就知道了 def rhs_weight_index_map(n_id, gm_id, k_id): group_id = metadata_ref.gm_id_to_group_id[gm_id] ← 全部的「选专家」就在这一行 return (group_id, k_id, n_id) // 这一步该搬 lhs 的哪些行？ def lhs_index_map(_, gm_id, k_id): m_start = metadata_ref.gm_id_to_m_offset[gm_id] m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1] row_start = m_start // size_lhs_sublane ← 注意除的是 sublane，不是 128 row_end = cdiv(m_end, size_lhs_sublane) return (pl.ds(row_start, row_end - row_start), 0, k_id)

上一节那个问题，在这里有了物理答案 你问「怎么保证第 k 段只跟第 k 个专家相乘」—— 答案是一次 SMEM 标量读，结果被拿去当 DMA 描述符的源地址。
不匹配的那些组合，权重从来没有被搬进 VMEM 过。不是算了再丢，是那条 DMA 压根没被发出去。「掩码」在这个层面上的物理形态，就是一个 int32。

第 1 层 · HBM → VMEM：双缓冲，权重开三层

外层把三个张量全声明成 memory_space=pltpu.HBM —— 这在 Pallas 里的意思是「别帮我自动搬，我自己排流水」。真正的搬运交给 pltpu.emit_pipeline，grid 是三层 (num_n, num_gm, num_k)。

缓冲深度是不对称的，而且能从 VMEM 预算函数里直接读出来：

// _gmm_vmem_estimate() 的实际构成 lhs_vmem = 2 * tile_m * tile_k * bytes // 双缓冲 rhs_vmem = 3 * tile_k * tile_n * bytes // 三缓冲 ← 只有它是 3 acc_vmem = tile_m * tile_n * 4 // 累加器，fp32 out_vmem = 2 * tile_m * tile_n * bytes // 双缓冲 partial_out = size_lhs_sublane * tile_n * bytes // 跨界残片，下面讲 // 而 block spec 上写得明明白白 rhs_weight_spec = pl.BlockSpec(..., pipeline_mode=pl.Buffered(buffer_count=3))

为什么只有权重开三层？源码没写理由，我的推断是形状： RHS tile 是 tile_k × tile_n，而 tile_n 至少是 MXU 边长的两倍（v7 上 ≥ 512）， LHS tile 才 128 × tile_k —— 权重那一块通常大好几倍， DMA 要跑更久，得多一层缓冲才追得上计算。 这是我从两边的尺寸推的，不是源码里的说法。

tile 多大，是解出来的不是拍的

vmem_limit = get_tpu_info().vmem_capacity_bytes * 0.9 // 先把 tile_n 往下缩，直到装得进 VMEM while estimate(tile_n, tile_k) > vmem_limit and tile_n > tile_n_limit: tile_n = ... // 切得更碎 // 那个下限是硬的，源码注释写明了原因： // "To avoid stalling MXU, we add some buffer room where // tile_n cannot go smaller than 2x of mxu_column_size." tile_n_limit = mxu_column_size * 2 // v7 上 = 512

这条下限值得停一下。为什么 tile_n 不能只有一个 MXU 那么宽？因为下一节要讲的重叠 —— 只剩一个 256 宽的面板时，就没有「下一块」可以拿来跟当前这块的收尾动作重叠了。 流水线需要至少两节车厢。

第 2 层 · VMEM → VREG → MXU：按 256 宽切列面板

进到 kernel 内部，第一句就是 VMEM → VREG：

tiled_lhs = tiled_lhs_ref.reshape(-1, tile_k)[...] ← 这个 [...] 就是向量 load

然后是真正喂 MXU 的那两层循环。外层步长正好是 MXU 的边长：

mxu_size = get_tpu_info().mxu_column_size // v7 = 256 for start_n in range(0, rhs_tile_n, mxu_size): // 一次一个 256 宽的列面板 acc_n = zeros((tile_m, 256)) for start_k in ...: acc_n += jnp.matmul( tiled_lhs[:, k0:k1], tiled_rhs[k0:k1, start_n:start_n+256], preferred_element_type=jnp.float32) // 累加恒为 fp32 acc_list.append(acc_n)

为什么 n 在外、k 在内 —— 源码把理由完整写出来了 “Without n outer loop, result of quantized matmul becomes available only at the last iteration of the loop. This means [tile_m, tile_n] value needs to be stored until the last iteration. By adding n outer loop, result of [tile_m, mxu_size] becomes available at the end of every k inner loop, which can be used to pipeline subsequent VPU or VST ops with MXU ops for the next [tile_m, mxu_size].”

翻译成人话：如果 k 在外、n 在内，那要等整个 [tile_m, tile_n] 全算完才有第一个能用的结果，中间这一大块只能干等着占寄存器。把 n 提到外层之后，每跑完一轮 k，就立刻攒够一个 256 宽的完整结果 —— 这块可以马上交给向量单元做后处理、交给存储单元写回， 而 MXU 同时已经在啃下一个 256 宽的面板了。

这就是「vreg 层的流水」最具体的样子 —— 循环顺序不是风格问题，它决定了 MXU、VPU、存储三个单元能不能同时忙。

累加器住在 VMEM，不在寄存器里

K 轴要跑 7,168 长，一个 tile 装不下，所以得跨步累加。但累加器 [tile_m, tile_n] 用 fp32 存 —— 代进 v7：128 × 512 × 4 B = 256 KB，寄存器堆装不下这个量级。于是它被显式声明成一块 VMEM scratch：

acc_ref = pltpu.VMEM((tile_m, acc_cols), acc_dtype) // 于是 K 轴每一步的收尾动作，取决于它在 K 上的位置 if not is_first_k_step: acc += acc_ref[...] // 中间步：读回来接着加 if is_last_k_step: ...写出去... // 末步：不必再写 acc_ref else: acc_ref[...] = acc // 否则存回去

源码为此写了四个 matmul 变体 （first_last / first / last / 普通），运行时用 lax.cond 按 k 的位置挑一个。 目的很朴素：首步别做那次无意义的读，末步别做那次无意义的写。在一个跑几十万次的内层循环里，这两次访存值得单独写四份代码。

回到 MoE：那条参差的边界，在硬件这一层长什么样

现在可以回答一个前面绕过去的问题了 —— §11 说 TPU 靠「重复访问一个 128 行的 tile」来对付跨界。但 128 是软件挑的 tile 高度，硬件真正的最小单位是什么？

是 一条 sublane。这一版 kernel 把 LHS 重新看成三维：

// [m, k] 被 reshape 成 [m / sublane, sublane, k] lhs_in = lhs_ref.reshape(-1, size_lhs_sublane, k) // 于是动态起点的粒度 = 一条 sublane size_lhs_sublane = get_tpu_info().get_sublane_tiling(lhs.dtype) // 32 位类型：8 行（就是 VREG 的 8 sublane） // 更窄的类型会把多个逻辑行压进同一条物理 sublane，这个数更大

所以边界的浪费从「不到 128 行」缩到了「不到 8 行」（窄类型下略多）。 §11 那张图里那个「多执行一个 128 行 tile」的代价，在这一版里小了一个数量级。

你直觉里的那个「掩码」，在这里真的出现了 —— 但它掩的不是专家 它掩的是「同一条物理 sublane 里，不属于本组的那几行」。 而且它只在 K 轴的最后一步才施加：

m_offset = m_start - m_start % size_lhs_sublane // 退到 sublane 边界 m_start_local = m_start - m_offset m_end_local = m_end - m_offset iota = broadcasted_iota(int32, acc.shape, 0) // 沿行方向数 0,1,2,... mask = (m_start_local <= iota) & (iota < m_end_local) acc = where(mask, acc, 0) // 不属于本组的行清零

注意它的形状：这是一个一维的行区间判断，不是一张二维布尔表。 而且它作用在 acc 上（结果），范围只有一条 sublane 的两端 —— 跟「四个专家」没有任何关系。专家那件事，在四层之前就已经由 SMEM 里那个 int32 解决掉了。

两个组共用一条 sublane 怎么办：partial_out

既然边界可以落在 sublane 内部，就会出现「一条物理 sublane 里坐着两个组的行」。两个组分属两次执行、用两份权重，但它们要写进同一条 sublane。源码用一块 [size_lhs_sublane, tile_n] 的 scratch 接住这个残片，注释里那张小图画得比什么都清楚：

// 假设一条 sublane 装 4 行，数字是组号，| 是 sublane 边界 | 0 0 1 2 | 2 2 2 2 | 3 3 4 4 | ↑ ↑ │ └─ 组 2 正好填满整条 → partial_out 清零，免得污染组 3 └─ 组 1 只占 1 行，填不满 → 把这条残片存进 partial_out，下一步（组 2）读回来累加到同一条 sublane 上

这就是 §11 说的「跨界 tile 执行两次、结果累加」在最细粒度上的实现 —— 只不过「两次」发生在一条 sublane 上，不是一个 128 行的 tile 上。

顺带一个只有在这一层才看得见的细节：MXU 边长决定了量化怎么做—— 反量化放在乘法前还是乘法后，是 mxu_column_size 说了算

FP8 权重要配一个 scale。乘完再乘 scale（省事）还是先还原成宽类型再乘（费事）？源码里的判据只有一行：

should_dequantize_before_matmul = (quant_block_size < mxu_size)

意思是：如果量化块比 MXU 还窄，就必须在 VMEM 里先反量化。 源码给的理由是 “to avoid small contracting dimensions” —— 因为「乘完再乘 scale」要求同一个 scale 覆盖整段收缩维；如果 scale 每 128 个元素换一次，而 MXU 一口吃 256，那就只能把 matmul 切成 128 长的小段 —— 脉动阵列立刻只装了一半。

这是一个很好的例子：一个看起来纯属数值格式的选择（量化块多大），实际由脉动阵列的物理边长决定。 也解释了 §3 里 128 这个数在 GPU 那边的地位，在 TPU 这边为什么会变成 256。

这一节的一句话 「哪个专家」这个问题，在标量单元上就被解决掉了；从 DMA 往下的所有硬件，跑的都是一次普普通通的稠密矩阵乘。
整条流水里唯一跟 MoE 有关的东西，是 SMEM 里那两个 int32 数组，以及边界那条 sublane 上不到 8 行的清零。其余全是常规 matmul 的功夫 —— 双缓冲、三缓冲、按 MXU 边长切列面板、fp32 累加器住 VMEM、循环顺序让三个单元并行。

§13三条路，同一个骨架
—— 上界开在哪，「省」发生在哪

走到这儿，前面十二节其实一直在同一个模式上打转。 把这个模式抽出来，剩下所有的差别都只是参数。

骨架就三句话

形状按一个上界开死 —— 让编译器 / CUDA graph 满意
真实长度以数据的形式进到 kernel 里 —— 不穿过 CPU
真正的「省」发生在 kernel 内部 —— 上界只决定形状，不决定算力

第三句是全部的重点。 §2 那个老办法之所以非要丢 token，病根就在它把第三句做反了 —— 它让上界直接进了算力账。上界进算力账，上界就只能开小；开小，就必须丢。

关键的那条不对称：上界只花内存，不花算力

上界开在哪不重要，重要的是它进不进算力账。 B 里那条灰虚线可以随便加宽 —— 它只是一块预留内存。真正被执行的，永远是实收的那些行（§12 里那个「trip count 由设备端现算」就是这句话的实现）。

	上界是什么	超出会怎样	进算力账吗
老办法 capacity factor	每个专家一个盒子	token 被丢掉 —— 这是个质量事件，不是资源事件	进
contiguous 训练 / prefill	一张卡的接收缓冲区 `num_max_tokens_per_rank`	缓冲区不够，属于配置错误，必须避免	不进
masked decode	每个专家一个等大的槽 max_m	同上	不进（见下）

那 decode 为什么还是贵 —— 贵的不是 max_m masked 布局的 max_m 同样不进算力账 —— DeepGEMM 的说法是 “By providing a mask tensor, the kernel computes only the valid portions.”
decode 贵在另一个地方：每个专家至少要占一个完整的 tile。 训练时一个专家收一万六千行，128 的下限微不足道； decode 时一个专家可能只收三十几行，同一个 128 立刻变成 4 倍浪费。 贵的是「粒度下限」，不是「上界」（§10 算过这笔账）。

上界开在哪一层，决定它能贴多紧

这是老办法和 contiguous 之间一条不太被提起、但很实在的差别：

老办法把上界开在「每个专家」头上。 单个专家的负载起伏很大 —— §1 那张表里，这一步的热点就是均值的 1.33 倍。所以盒子得留 1.2 到 1.5 倍才敢用。
contiguous 把上界开在「一整张卡的总量」上。 四个专家的起伏会互相抵消一部分，所以总量的相对波动，小于任何单个专家 —— 这是求和的一般性质（若四份独立，相对标准差按 1/√4 收缩）。 真实路由是有相关性的，不会这么理想，但方向是对的。
masked 又退回「每个专家」。 所以 decode 这边拿不到这个红利 —— 这跟上一条的粒度问题是两回事，但方向一致。

三条路，一张表看完

	左矩阵形状	上界开在哪	变长信息长什么样	谁编译 / 捕获	「省」发生在哪
TPU / XLA 训练	[M上界, K] 一整条	一张卡的总量	group_sizes 设备数组 → 前缀和 → 段边界	整步一个 XLA 程序	kernel 内部：grid 步数由设备端现算，尾巴不进循环
GPU · contiguous 训练 / prefill	[M实收, K] 一整条	接收缓冲区	host 已知 M ＋每行的段号	不捕获（eager）	M 本来就是实收 —— 没有需要省的
GPU · masked decode	[E, max_m, K] 分层，不拼接	每个专家一个槽	一个 mask 数组（每层前多少行是真的）	CUDA graph 捕在 max_m 上	kernel 内部：只算 valid portion

最后一句，也是这一整篇最该带走的 注意第三行：decode 那条路根本没有拼接。 左矩阵是四层分开的槽，中间全是空的。但它仍然是一次 kernel 调用，仍然一个 token 都不丢。

所以「拼成一条」从来不是必须的 —— 必须的是上面那三句骨架。 「grouped matmul」这个名字容易让人以为重点在「拼」。重点其实在：把形状从 CPU 的控制流里搬进设备的数据流里，并且让上界只花内存、不花算力。 做到这两条，「变长」就从一个必须先解决掉的障碍，退化成一个普通的输入参数。

§14换一个轴：token 是怎么见到它的专家的
—— A2A / AG-RS / 干脆不用 EP

前面十三节全都在回答同一个问题： 这些行已经躺在这张卡上了，接下来怎么乘。 这一节换一个正交的轴 —— 这些行当初是怎么到这张卡上来的。 它跟 §13 那三条路不是同一件事，别叠在一起看。

先把一个特别容易混的术语钉死 AG-RS（all-gather ＋ reduce-scatter）搬的是 token，不是权重。
这个词在 MoE 语境里指的是「all-to-all 的替代品」，而 all-to-all 搬的从来就是 token，所以它的替代品搬的也是 token。三个互相独立的实现口径一致：

Megatron-Core 的开关叫 --moe-token-dispatcher-type，两个取值就是 allgather 和 alltoall —— 两个都是 token dispatcher
TensorRT-LLM 的文档写得最直白： “AllGather brings all tokens to every rank. AlltoAll sends each token only to the ranks that own its routed experts.”
MaxText 里这条路叫 ring of experts，代码注释： “The ring-of-experts strategy first duplicates the inputs to all expert shards, and then routes within each shard.”

所以「把专家权重全 all-gather 到每张卡」不是 AG-RS。 那件事真实存在、而且在 TPU 上是对的选择，但它的名字是 「EP ＝ 1 ＋ FSDP」，或者干脆就叫「不用专家并行」。两条路的唯一共同点，只是都不做 all-to-all。

三条路，机制上差在哪

	谁在动	一张卡看到什么	怎么把多余的去掉
A2A expert parallel	token	只收到选了本地专家的那些行	不需要去 —— 收到的每一行都要算
AG-RS ring of experts	token（复制给所有人）	全部 token，一个不落	本地路由一遍，mask 掉不属于自己专家的行，算完 reduce-scatter 回去
EP ＝ 1 ＋ FSDP TPU 上的实际选择	权重	只有自己原本那批 token ＋全部专家的权重	不需要去 —— 全部专家都在本地，谁都不用找

通信量：三笔账，每卡每层

还是 §1 那个口径 —— h = 7168、每卡 8,192 个 token、 top-8、256 个专家、 64 张卡。

// ① A2A —— 每个 token 只走一趟，走到它该去的地方 dispatch 8,192 × 8 = 65,536 行 × 7168 × 1 B (fp8) = 448 MiB combine 65,536 行 × 7168 × 2 B (bf16) = 896 MiB 合计 ≈ 1.31 GiB // 关键性质：跟 EP 度数无关。EP 开到 8 还是 64，这个数不变。 // ② AG-RS —— 每个 token 复制给所有人 all-gather (64−1) × 8,192 × 7168 × 2 B = 6.89 GiB reduce-scatter 同上 = 6.89 GiB 合计 ≈ 13.8 GiB // 关键性质：正比于 EP 度数。EP 越大越亏。 // ③ EP=1 + FSDP —— token 一个字节都不动，动的是权重 256 专家 × 3 个矩阵 × 7168 × 2048 = 112.7 亿参数 × 2 B (bf16) = 21.0 GiB 每层每卡 // 关键性质：跟 token 数无关 —— batch 越大越摊得开。

①②之间的交叉点可以直接解出来 两边都是「行数 × h × 字节」，把 h 和每卡 token 数约掉就剩：

A2A : k × (1 + 2) = 8 × 3 = 24 AG-RS : (EP − 1) × (2 + 2) = 4 × (EP − 1) 相等 ⇒ EP − 1 = 6 ⇒ EP ≈ 7

也就是说 EP 超过七八路，AG-RS 就打不过 all-to-all 了。 这跟 Megatron 文档那句话正好对上：allgather 这条「achieves better performance and efficiency when only tensor parallelism is used」—— 也就是 EP 很小的时候。

那 TPU 为什么反而选字节最多的第三条

因为便宜的不是字节，是「字节能不能藏起来」。

	这笔通信依赖什么	能不能提前发
FSDP all-gather 权重	只依赖走到第几层 —— 编译期就知道	能。第 L 层在算的时候，第 L+1 层的权重已经在路上了
MoE 的 all-to-all	依赖这一层 router 刚算出来的结果	不能预取。它卡在 router 和 GEMM 中间 —— 但能用别的办法盖住，见下面那条

再叠上拓扑：ICI 是 3D torus， 64 路 all-to-all 要横穿整个环面，没有 NVLink 那种全连接域；而 all-gather 是规则的近邻通信，在 TPU 上还能整体卸载到 SparseCore，跟 TensorCore 的矩阵乘重叠。

实测：这不是推演，是量过的 同一个 295B MoE，v7 上把 ici_expert_parallelism 从 1 打开：

EP = 2（64 芯片） −39.6% EP = 4（半 batch，16 芯片） −71.36%

v5p 那边官方 DeepSeek3-671B 的 256 芯片配方， ici_expert_parallelism 直接写 1 —— 192 个专家全靠 FSDP 沿 embed 维切开，all-gather 走 SparseCore。

所以「97% 参数在专家里，所以 EP 是主旋钮」这句话， 对 GPU 成立，对 TPU 不成立。 这是把一种硬件的经验直接搬到另一种上的典型翻车。

那 A2A 就真的藏不住吗 —— 藏得住，但不是靠预取

把上面那个「不能」收窄一点 A2A 不能预取，但可以用另一个 micro-batch 的计算去盖住它。
DeepSeek-V3 的 DualPipe 做的就是这件事：双向流水线，官方仓库的说法是 “achieves full overlap of forward and backward computation-communication phases”。代价也写在报告里 —— 专门划出 20 个 SM 跑手写 PTX 的通信 kernel。

所以真正的差别不是「能不能藏」，是「藏起来要花多少工程量」：

FSDP all-gather —— 改一行（预取下一层），编译器自己就会做
MoE all-to-all —— 要重排整个 pipeline 调度，还得给通信留出计算单元

MaxText 那边同一招的小号版本叫 num_moe_token_chunks：把本卡 token 切成 N 块，让第 c+1 块的通信跟第 c 块的 GMM 重叠。

先别拿字节数当选路依据 —— 包括我上面那三笔账

把上面①和③摆平，能解出一个「每卡要多少 token，FSDP 才在字节上打平」。

// ① A2A : T × k × h × 3 （fp8 dispatch 1 B ＋ bf16 combine 2 B） // ③ FSDP : P × 2 （P = 一层的全部专家参数 = 112.7 亿）相等 ⇒ T = 2P ÷ (3 k h) = 2×112.7亿 ÷ (3 × 8 × 7168) ≈ 131,000 token / 卡

而 TPU 那份配方只有每卡 32,768 个 token 也就是说：在 TPU 上，FSDP 按字节算依然亏四倍。它照样赢了。

所以上面那三笔账算得没错，但不能拿它选路。 如果谁只看字节数，会得出「两个平台都该用 EP」的结论 —— 而这跟实测正好相反。把这条写在这儿，就是防止这张表被那样用。

字节数不是依据，那依据是什么？两条，上面讲了一条，这里补另一条。

同一个 collective，在两种 fabric 上代价天差地别。 （就是上一小节 —— NVLink 全连接域里 all-to-all 每一对都有直连，接近最优；3D torus 上它多跳又拥塞。而 all-gather 反过来。）
内存：EP 是「切了就不用还原」的切分，FSDP 是「切了还得还原」的切分。 这一条上面没说，但它才是 GPU 那边的决定性理由。

	算这一层时，卡上要放多少专家权重	V3 一层的量
EP 切了不还原	永远只有本地那 4 个专家 —— 从头到尾不摊开	336 MiB
FSDP 切了必须还原	整层全部 256 个专家，而且要为预取留双缓冲	21 GiB × 2 ≈ 42 GiB

H100 一共 80 GB。42 GiB 摊在那儿，激活基本没地方放了 —— 这就是 GPU 侧「不能只用 FSDP」最硬的那条。 TPU v7 每 device 有 94.7 GiB，同一件事才勉强做得下。

推理那边更是没得选 FSDP 的成本跟 batch 无关。 decode 一次就一两百个 token，你却要为这一层搬 21 GiB 的权重进来 —— 完全说不通。
所以 vLLM / SGLang / TensorRT-LLM 的 MoE 推理路径全都是 EP ＋ A2A，没有一个走 FSDP。

一句话的判据

FSDP 的成本 ∝ 专家参数量，跟 batch 无关 A2A 的成本 ∝ 每卡 token 数，跟 EP 度数无关

token 多就偏 FSDP，token 少就偏 EP。 训练 token 多＋环面偏爱 all-gather ＋ HBM 够大 ⇒ TPU 敢用 FSDP； GPU 训练 token 也多，但卡在内存和「NVLink 让 A2A 本来就便宜」这两条上； 推理 token 少，两个平台都只能 EP。

代价：EP ＝ 1 会把段数顶上去

不用 EP 意味着所有专家都在本地。于是 §11 那个「tile 被访问两次」的边界开销，乘数从「几个专家」变成「一两百个专家」：

// §11 的界：tiles_m ≤ 实际 tile 数 ≤ tiles_m + 段数 − 1 边界重访开销（最坏） ≈ (段数 − 1) ÷ (总行数 ÷ tile_m) // 每卡 8,192 token（GPU 那个口径）、256 专家、tile_m = 512 总行数 65,536 ÷ 512 = 128 块，段数 256 → 最坏 +199% ← 顶不住 // 每卡 32,768 token（TPU 那份配方）、192 专家、tile_m = 512 总行数 262,144 ÷ 512 = 512 块，段数 192 → 最坏 +37%，均值约 +19%

所以那份配方为什么要把每卡 token 数拉大八倍 不是为了吞吐好看，是为了让 EP ＝ 1 这条路成立。 段数被专家数钉死了，你唯一能动的是分母 —— 每个专家分到的行数够多，边界那点开销才摊得掉。 大 batch 在这里不是调优，是前提。

回头看，这又是 §13 那句话的一个实例：开销不会消失，只会换账本。 A2A 把它放在通信量上，AG-RS 放在「通信量随 EP 增长」上， EP ＝ 1 放在段数、也就是 tile 粒度上。三条路，三本账，你挑的其实是「让哪一本最好平」。

§15V3 到底是真 dropless 还是近似的

这个问题要拆成两问，答案不一样。

问一：token 层面，丢不丢？

不丢，而且是结构性的不丢。整条链路上没有任何一处存在「容量上限」这个概念 —— 没有上限，就没有「超出」，也就没有「丢弃」这个动作可执行。

三处证据，都可公开核验：

来源	说了什么
V3 技术报告标题就叫 No Token-Dropping	训练全程不丢任何 token（归因于负载均衡策略有效）；推理也不丢（靠专门的部署策略）。只给结果，没说机制。
DeepEP dispatch/combine 通信库	`dispatch` 的参数里有 `expert_alignment`（默认 1），返回的 handle 带 `num_recv_tokens_per_expert_list` —— 每个专家实收多少，是运行时才知道的一个变长列表，不是一个固定盒子。全库没有 capacity 这个概念。
DeepGEMM 矩阵乘库	contiguous 布局明写：训练前向和 prefill 时 “each expert may process a varying number of tokens”，拼成一个张量，每段对齐到 GEMM 的 M block size。同样没有 capacity。

这三处拼起来，结构跟 §3 走的那一遍完全一致：变长分段 → 对齐补齐 → 分组矩阵乘。

问二：算力层面，有没有浪费？

有，但性质完全不同。

盒子方案的浪费是「按容量对齐」 —— 每个专家都补到同一个高度，浪费跟失衡程度成正比，可以到几十个百分点。
分组方案的浪费是「按块对齐」 —— 每个专家最多补一个块，浪费有一个跟失衡无关的硬上界。

所以「dropless」指的是前者那件事被解决了，不是「一点算力都不浪费」。 这两件事经常被混为一谈。MegaBlocks 论文自己在 §5.2 也说得很清楚： kernel 要求每个专家的 token 数是块大小的倍数，所以「we pad each group of tokens with zeros to the nearest multiple of 128」—— 而且它还补了一句，这个约束是可以去掉的，只要支持边缘的不完整块，只是当时没做。

⚠️ 这里有一句不能说过头 可以说：V3 训练和推理都不丢 token（报告明写）； DeepSeek 自己开源的算子栈里没有 capacity 这个概念，它按每个专家的实收行数分段（DeepEP + DeepGEMM 的文档可核）。
不能说：「V3 预训练那次跑的就是 DeepGEMM 这套 kernel」。
DeepGEMM 和 DeepEP 都是 2025 年才开源的，两个 README 里 没有任何一句把它们钉到 V3 那次预训练上。这是很强的旁证，不是直接证据。
差别很大：「靠均衡好到没触发上限」和「压根没有上限」，在系统实现上是两回事。

顺带：推理侧不丢，靠的是另一套东西

训练时靠动态偏置把负载掰匀；上线之后没有训练循环了，靠的是搬专家 —— 把观测到的高负载专家复制一份额外部署（prefill 阶段设 32 个冗余专家，每张卡除原本的 8 个再多带 1 个），热点根据线上统计检测，每隔一段时间重新调整一次。

值得记的是这个动作的性质：它明确是在节点内部重排，避免增加跨节点 all-to-all 的开销。又一次 —— 先认拓扑，再谈优化。

§16MegaBlocks 论文的三个数，以及它们能证明什么

丢 token 是真的伤质量。同一组对照里，容量系数取 1（会丢）的 MoE 只把验证损失降了 0.15；完全不丢的那版降了 0.26 —— 1.73 倍。
靠「把盒子开大」来不丢，很贵。原文：为了不丢 token， MoE 层的数学运算量涨了 两倍多。
而且你不知道该开多大。论文引 Tutel 的观察：有的模型容量系数要开到 11 才不丢，还有的模型这个值会在训练中途毫无征兆地飙上去。

正面结果（同一套设置）：相比用 padding 硬凑的做法，端到端训练快 1.38× / 2.0× / 4.35×（三个规模）；相比同质量的稠密 Transformer 快 1.8×–2.4×。 显存那边同样明显：padding 方案把最大微批量压掉了 2× / 4× / 8×。

用之前先知道它的适用范围 —— 这不是 671B 的实验

维度	MegaBlocks 那组实验	DeepSeek V3
专家数	64	256 + 1 共享
每 token 选几个	top-1	top-8
模型规模	Transformer-Small 量级	671 B
均衡机制	常规辅助损失	动态偏置 + 极小辅助损失

差别最大的是 top-1。 top-1 意味着每个 token 只押一个专家，这个专家满了就整个丢掉； top-8 天然摊得开，同样的失衡下丢得没那么惨。
所以这三个数能证明的是「容量系数这个旋钮本身是个麻烦」，不能直接换算成「V3 如果丢 token 会掉多少分」。

§17回到主线

专题一从头到尾在讲一件事：硬件的形状会渗进模型的定义里。

分组路由那一节是最赤裸的一例 —— topk_group = 4 跟模型质量毫无关系，它的存在只是为了让 all-to-all 少跨机。硬件划线，模型让步。

这一篇是那件事的镜像。 capacity factor 曾经也长得像一个模型超参数，人们认真地给它调参、写进论文的消融表。而它其实只是 batched matmul 的形状约束换了身衣服。 有人回头把 kernel 那一侧修好之后，这个「超参数」就自己脱落了。

留给学员的一句 看到一个说不出物理意义的超参数，先问一句：它到底是模型需要的，还是某个 kernel 的限制伪装成的？ 前者要调，后者该拆。