1. 先别急着读源码：MoE 的数学模型里藏着“调度”的影子
1.1 路由公式到底在算什么：概率、采样还是排序？
很多同学第一次接触 MoE，看到的是这样的公式：
score_i = Wg_i · x p_i = softmax(score)_i output = Σ_{i ∈ top-k} p_i · FFN_i(x)于是很自然地把它理解成“按概率加权求和”。这个理解在数学上没错，但在系统实现里非常容易误导人。推理阶段的 MoE 几乎不会做随机采样，而是严格做 top-k 硬选择：先对 N 个专家打分，然后取分数最高的 k 个，只让这 k 个专家参与计算。换句话说，p_i不是“这个 token 有 30% 概率被专家 A 处理”，而是“调度系统决定把 token 派给专家 A，权重是 0.3”。
这个视角转换非常重要。一旦你接受了“路由结果是调度指令而不是概率分布”，你才会理解为什么 MoE 在 GPU 上的实现需要 token 重排、分组 GEMM、动态 block 调度这些听起来很工程化的东西。稀疏性不是模型自带的魔法，而是路由结果天然把计算切成了“只访问少数专家”的碎片化计算。
为什么是 top-k 而不是“全部加权”？两层原因。第一是计算量，N 个专家全部过一遍，等于退回稠密 FFN，MoE 的参数扩展优势瞬间消失。第二是梯度质量，被小概率选中的专家很难累积到稳定梯度，训练后期几乎学不到东西。所以 top-k 是一个典型的“算法和工程共同决定”的超参数：它同时影响模型收敛质量和推理延迟，调参时不能只对着 loss 曲线看。
1.2 路由崩溃与负载均衡：辅助损失的数学动机
自由路由有一个臭名昭著的问题：赢家通吃。如果门控网络完全自由发展，少数学会通用表达方式的专家会被疯狂调用，其余专家饿死，整体模型退化成只有几个专家干活。这就是路由崩溃（routing collapse）。
经典解法是辅助负载均衡损失，常见形式是：
L_aux = α · N · Σ_{i=1}^{N} f_i · P_i f_i = 该批次中被路由到专家 i 的 token 占比 P_i = 门控网络对专家 i 输出概率的平均值这个公式妙在：当所有专家被均匀使用时，f_i ≈ 1/N、P_i ≈ 1/N，乘积和接近最小；一旦某个专家被过度使用，或者门控对它输出过高概率，乘积就会变大，形成惩罚。α 通常取 0.01 量级，太大会让“均衡”压倒“任务效果”，模型训练完变成均匀但没有专长的“平庸模型”。
另一个配套机制是专家容量（expert capacity）。容量定义了一个专家在一个迭代步内最多能处理多少 token：
C = ceil(k / N × batch_tokens × c)c 是容量因子，常见取 1.25。举个例子：假设 N=256、k=8、batch_tokens=2048，那么理想分配是每专家 64 个 token，容量因子 1.25 则放宽到 80。超出容量的 token 会被直接丢弃（注意，不是排队）。这个设计把每个专家的最坏工作量钉死，才能在推理时提前规划显存和计算资源。代价是丢弃 token 会破坏序列连续性，所以现代实现更多用放大容量的 no-drop 变体，或者引入重路由。
1.3 总参数与激活参数：区分两个完全不同量级的数字
MoE 模型最容易被算法同学忽略的，是“参数规模”和“计算规模”不再是同一个东西。总参数决定模型文件大小、显存占用、分布式切分方式；激活参数决定单个 token 前向真正触碰的权重数量，直接影响推理解码延迟。
传统稠密 Transformer 两者几乎相等。MoE 把 FFN 复制成 N 份，但每个 token 只走 k 份，于是激活参数远小于总参数。业界公开的大规模 MoE 里，总参几百 B、激活参数几十 B 是很常见的量级比，激活占比通常在 5% 到 10% 之间。这意味着你可以把模型做得很大，但每一次前向计算的成本只相当于一个小得多稠密模型。这套“参数冗余换能力、稀疏激活换速度”的交换，是整个 MoE 系统设计的最上层前提，后面所有 GEMM 优化都建立在这个前提之上。
2. MegaMoE 的演进方向：细粒度专家、共享专家与稀疏边界
2.1 先搭一个教学用的“参照配置”
由于 V4 的细节未必是所有人都熟悉的公开信息，我先把这套分析建立在“开放的架构样例”上。下面给出的配置是为了让讨论有抓手，不代表官方数字，但它符合大规模 MoE 一路迭代下来的主流方向：专家数变大、单个专家变细、顶层再叠一个共享专家。
假设一个 MegaMoE 风格的参照配置：
| 属性 | 参考值 | 说明 |
|---|---|---|
| 总参数量 | 900B 级别 | 文件存储和切分的依据 |
| 激活参数量 | 30B 级别 | 单 token 前向实际计算的权重 |
| 路由专家数 | 上千个 | 细粒度专家，技能分离度高 |
| 共享专家数 | 1 个 | 每个 token 无条件通过 |
| top-k | 6 到 10 | 典型稀疏激活比例 |
这个配置下，单 token 激活的路由专家占全部路由专家比例大约在 1% 到 3% 区间，加上共享专家，稀疏性非常明显。记住这个比例，后面计算 GEMM 的形状和预估带宽时会反复用到。
2.2 细粒度专家的好处与代价：为什么“小专家”反而更强
如果把一个传统的大 FFN 专家拆成多个窄专家，每个专家的隐层宽度变小、数量变多，这带来两个直接好处。第一，技能分离更干净：大专家内部隐层宽，不同 token 的特征混在一起容易互相干扰；小专家只负责一个窄的技能面，路由更容易精准匹配。第二，并行粒度更好：专家数量多了，可以更均匀地分散到不同 GPU 或不同计算节点上，减少热点专家导致的负载倾斜。
代价也很现实：专家数量变多，路由矩阵Wg变大，top-k 排序开销上升，通信模式更复杂，kernel 调度的碎片化问题更严重。细粒度专家不是免费的午餐，它把一部分模型能力收益转移到了系统侧的优化压力上。这也是为什么你会看到专门为 MoE 写的内核库，而不是直接调通用矩阵乘库。
2.3 共享专家：给所有 token 一条“公共通道”
在路由专家之外，MegaMoE 这类架构通常还会放少量共享专家，每个 token 无论如何都必须通过。原因非常朴素：有些能力是所有 token 都需要的，比如句法结构、常用搭配、全局上下文线索。如果这些通用能力也去抢占有限的路由专家名额，路由带宽被浪费，专家容量也被挤占。
共享专家相当于一条“公共知识通道”，把高频通用计算从路由专家的职责里剥离出来。实现上，共享专家的宽度通常大于单个路由专家，但又远小于一个传统大 FFN。参与每个 token 的计算，意味着它在激活参数里始终占一份固定成本。当你看到激活参数占比时，不能简单用“top-k / 专家总数”去估，还要把共享专家这一份算进去。
2.4 路由稳定性在系统层的另一种含义
专家变多之后，门控打分的维度变大，score 的方差随之上升，softmax 更容易输出接近 one-hot 的分布，梯度更难传播。MegaMoE 的常见做法是对打分做归一化，比如在Wg · x上除以一个可学习温度，等效于把分数拉到余弦相似度的尺度上。这能让 logits 的方差更平稳，路由决策更稳定。
从系统角度看，路由稳定性还意味着另一个关键性质：不同 batch 之间的专家负载波动变小。如果路由结果忽而把所有 token 都丢给 32 号专家，忽而又均衡分散，那么分组 GEMM 的组大小就在剧烈变化，缓存的利用率、block 调度的效率都会下降。所以“路由稳定”不只是模型收敛问题，它同时是 kernel 性能问题。这是一条算法和 Infra 深度耦合的典型证据。
2.5 预填充与解码的稀疏形态会两极分化
同一个 MoE 模型，在预填充（prefill）阶段和解码（decode）阶段的计算形态几乎完全不同。预填充时一个序列往往有上千个 token，batch 内 token 总数很大，top-k 选择之后几乎所有专家都能分到足够多的 token。这时候 MoE 层实际上退化成“对每个专家跑一个不小的稠密 GEMM”，整卡算力容易吃满。
解码阶段就惨了：一次只生成一个 token，batch 内通常只有几十甚至几个请求，每个 token 只命中 k 个专家，每个专家只分到几个 token。这是典型的“稀疏小 GEMM”困境：计算量不大，但访问的专家权重却不少，吞吐瓶颈落在访存和 kernel 调度上。读任何 MoE 内核源码之前，先分清你关心的是预填充吞吐还是解码延迟，因为这两个场景对内核的优化方向几乎是相反的。
3. 从数学模型到 GEMM 问题：一次计算形态的“翻译”
3.1 把 MoE 层翻译成分组矩阵乘法
每个专家本质上就是两层线性变换加激活：
h_e = act(X_e · W1_e + b1_e) out_e = h_e · W2_e + b2_e其中 X_e 是路由给专家 e 的 token 子矩阵。如果把所有专家看成一个整体，就会发现整个 MoE 层是一个“分组 GEMM”（Grouped GEMM）：不同专家拥有不同的权重矩阵，被路由到同一专家的 token 块只与该专家的权重做乘法。
门控本身也是一个小的 GEMM：score = X · Wg^T，输出只有 N 行（专家数）。但它后面紧跟着 top-k 排序和 token 重排，把散落在原始序列里的 token 按照目标专家重新排列成连续块。这一步在源码里叫 permute 或 gather，往往是最容易被忽略的隐形开销。
3.2 为什么不能简单写成“N 个小 GEMM 的 for 循环”
最容易想到的实现方式，就是对每个专家分别调用一次矩阵乘法。小规模实验这么做完全没问题，但规模化之后会撞上三堵墙。
第一是 kernel 启动开销。一次 kernel 启动有微秒级固定成本，假设模型有 100 层、每层几百个专家，一次 forward 下来启动次数可能上千，纯启动开销就能吃掉大量时间。第二是并行度不足。解码阶段每个专家只有几个 token，单看一个小 GEMM 的 M 维度只有个位数，无法填满 GPU 上的所有 SM，算力利用率惨不忍睹。第三是内存碎片。为每个专家临时分配小形状的张量，会让显存分配器疲于奔命，碎片化严重时甚至触发不必要的显存搬运。
于是工程上出现两个更聪明的方向：一是先把 token 按专家排列成连续大块，再用批量 GEMM（Batched GEMM）一把算完；二是保留分组结构，在同一个 kernel 内部遍历所有专家，这就是分组 GEMM。它们都把“按专家切开再逐个算”的启停开销压缩到最小。
3.3 为什么通用矩阵乘库不够用：FP8、动态形状和融合需求
有人会问：CUDA 生态里已经有 cuBLAS 和 CUTLASS，为什么还要为 MoE 专门写内核？答案是通用库在三个维度上无法兼顾。
- 精度控制：FP8 是大模型的主流方向，E4M3 表示范围很窄，直接做矩阵乘容易溢出或丢精度。通用库对 FP8 的支持通常只覆盖标准形状，对 block 级缩放、混合精度累加这些 MoE 需要的控制点不给足自由度。
- 动态形状：MoE 的分组大小每层每 batch 都在变，通用 grouped GEMM 接口虽然存在，但对“每个组的 M 都不同、且动态变化”的场景优化并不充分，性能和可编程性都不理想。
- 算子融合：把分组 GEMM、缩放、激活、token 重排融合进一个 kernel，可以减少中间张量的全局内存往返。在大模型场景里，少一次整张量读写就能省下可观延迟，通用库做不到这种跨算子优化。
3.4 先算一笔账：什么时候 GEMM 会成为瓶颈
写内核前先估算一下工作量和瓶颈类型。一次 GEMM 的浮点运算量大约是：
FLOPs ≈ 2 × M × N × K假设解码阶段某专家分到 16 个 token，专家隐层宽度为 8192，输入维度为 4096，那么M=16、N=8192、K=4096，FLOPs 只有约 1.07 TFLOPs。听起来不小，但现代 GPU 的稠密 FP8 算力是几百 TFLOPS，这个专家上的实际计算只占 GPU 峰值算力不到 1%。真正的成本花在把权重从显存搬到计算单元上：8192 列的权重矩阵读一遍要消耗大量访存带宽。
这个估算告诉我们，解码阶段的 MoE 瓶颈落在访存与调度，而不是算力；预填充阶段 token 多、专家分到的块大，才可能逼近算力上限。这两个结论直接影响内核设计：解码侧要重点优化访存流水线和启动开销，预填充侧则更关注矩阵分块和指令发射效率。
4. DeepGEMM 源码手把手拆解：FP8、TMA 与分组调度的落点
标题既然点到了“源码”，这部分就是重点。读这类内核源码不需要逐行背，抓住三个模块就够：启动器与分组参数、FP8 缩放逻辑、主循环里的访存与矩阵乘指令。下面我会按源码阅读视角拆解四个关键位置。为了便于说明，我会给出一些等价伪代码——真实项目里的命名和细节以你手上的版本为准，但核心机制是共通的。
4.1 启动器与分组参数：一切从“组列表”开始
内核启动器不会只接受一个简单的M/N/K，它接收的是分组的描述：
group_ids : 每个分组对应的专家 id offsets : 每个分组在原始张量中的行/列偏移 M_list : 每个分组的实际 token 数启动器拿到这些参数后，把整个 CUDA grid 按专家切块，每个 block 负责一个专家的一个 tile。为了减少空闲，通常还会在启动前对M_list排序，把小分组放到 grid 后段统一处理。读源码时，先找启动器里对分组做归一化、补零或对齐的那几行，那里是整个 kernel 布局逻辑的入口。真正理解分组参数怎么被翻译成 GPU block 的坐标，比看主循环的算数更容易建立整体图景。
4.2 FP8 缩放：块级缩放与反量化时机
FP8 矩阵乘最容易犯的错误，是把量化当成一次性预处理。实际上，缩放必须伴随计算走完全程。一个等价的核心循环是：
// 伪代码：展示源码等价逻辑 for (int k = 0; k < K; k += 16) { half a_bf16 = load_A_block(...); half b_bf16 = load_B_block(...); fp8 a_fp8 = fp8_round(a_bf16 / scale_A); // 进入内核前已量化 fp8 b_fp8 = fp8_round(b_bf16 / scale_B); acc += mma_fp8(a_fp8, b_fp8); // 累加器保持 fp32 } out = acc * scale_A * scale_B; // 最后一次反量化关键设计点是scale_A和scale_B通常是块级缩放，整个 K 维度只用一个或一组缩放值，而不是每个元素一个。块级缩放省内存、省指令，但块内动态范围受最大值限制。源码里经常会同时维护 per-block 和 per-group 两套 scale，供不同场景选择。另一个重点是累加器必须用 FP32，FP8 乘法本身只有 8bit 精度，如果累加也用低精度，误差会在 K=4096 的累加链上滚成不可接受的结果。读源码时盯住acc的声明类型和最后一次scale乘法写在哪一行，基本就把握住了精度控制策略。
4.3 TMA 与软件流水线：让访存时间藏进计算时间
现代大模型的 kernel 瓶颈经常不在算力而在访存。在 Hopper 及后续架构上，TMA（Tensor Memory Accelerator）可以把全局内存到共享内存的数据搬运交由专用硬件完成，再用异步操作把搬运和计算重叠起来。DeepGEMM 这类内核的主循环基本都是“双缓冲或多缓冲 + 生产者消费者”模式：
- 生产者 warp：用 TMA 把下一个 tile 的 A、B 矩阵从全局内存拷到共享内存；
- 消费者 warp：对当前 tile 执行矩阵乘累加指令；
- 共享内存拆成多个 buffer 轮流使用，搬运和计算重叠，让 HBM 带宽尽量保持繁忙。
读源码的快速路径是：找主循环里的异步拷贝指令，看它后面是否立刻接等待和屏障，再判断它是同步等待还是异步推进。异步推进做得越彻底，流水线覆盖率越高，kernel 的实测吞吐和理论峰值的差距就越小。
4.4 分组调度的尾块处理：性能陷阱都在“尾巴”上
分组 GEMM 最麻烦的不是正常 tile，而是每个分组最后剩下的尾巴。比如某专家分到 37 个 token，tile 尺寸是 64，那就要补到 64，产生 27 个“假 token”。如果每个专家都这么补，计算浪费可能是百分之几十。
源码里的对策通常有两类：一是把M_list排序，让小分组集中到 grid 后段，统一走一个针对小 M 优化的内核路径；二是在最后一列 tile 上缩小实际计算范围，避免无意义的填充乘法。解码阶段更常用“M 维度合并”的招数：把多个 token 都很少的专家合并成一个大的组，计算时共享同一块 tile，写回时按各自偏移分开。这个技巧效率很高，因为它同时解决了启动开销和填充浪费两个问题。
5. 算法同学的上手路径：从迷你 MoE 到真正读懂源码
5.1 最小复现路线：不要一上来就啃巨型模型
我给算法同学的建议非常直接：不要试图直接啃 671B 级别的源码和配置，先在单卡上复现一个迷你 MoE。比如 8 个专家、隐层维度 1024、batch=32，先用 PyTorch 的常规矩阵乘法把逻辑跑通，再手动替换成分组 GEMM 的简单实现，最后再尝试接触 FP8 内核。
下面是这个迷你 MoE 前向的朴素 PyTorch 实现，它就是后面所有优化的基准：
import torch import torch.nn.functional as F def moe_forward(x, gate_w, experts_w1, experts_w2, top_k=2): """ x: [T, d] gate_w: [N, d] experts_w1: [N, d, d_ff] experts_w2: [N, d_ff, d] """ score = x @ gate_w.T # [T, N] topk_idx = torch.topk(score, k=top_k, dim=-1).indices topk_w = torch.gather(F.softmax(score, dim=-1), 1, topk_idx) out = torch.zeros_like(x) for e in range(experts_w1.shape[0]): mask = (topk_idx == e).any(dim=-1) if mask.any(): h = F.gelu(x[mask] @ experts_w1[e].T) out[mask] += h @ experts_w2[e].T return out这段代码的好处是逻辑清晰，坏处是性能极差，但恰恰是这种“朴素正确”的版本，最适合用来验证后续所有 kernel 优化的数值正确性。拿它当基准，你对分组 GEMM 的每一步改造都能立刻知道对不对。
5.2 源码阅读顺序建议
跑通迷你 MoE 之后，再读真正的内核源码就有目标了。我习惯的顺序是：
- 先读 README 和启动器入口，搞清楚 kernel 接收什么布局的参数，尤其是分组、缩放相关的结构体；
- 再读控制流，确认每个专家、每个 group 是怎么映射到 block 的；
- 接着读 FP8 缩放部分，理解 scale 从哪来、反量化在哪做；
- 最后才看主循环里的 TMA 和矩阵乘指令，这一步最容易迷路，要对照架构手册慢慢看。
阅读时准备一张纸，画出“全局内存 → 共享内存 → 寄存器 → 累加器 → 全局内存”的数据流动路径，比盯着源码逐行脑补有效得多。内核代码的难点从来不是某一行指令，而是整条数据流的节奏。
5.3 常用验证工具与调试手段
没有工具辅助的源码阅读是空中楼阁。我推荐三件套：
- 用 Nsight Compute 看每个 kernel 的吞吐、访存命中和 stall 原因；
- 用 torch profiler 看各阶段耗时占比，确认瓶颈到底在 GEMM 还是在前面的 permute；
- 写数值对比脚本，在 FP8 内核和 FP32 参考实现之间做最大误差检查。
检查项通常包括：分组排列是否正确、scale 有没有重复乘或漏乘、尾块填充有没有污染输出、top-k 索引与专家权重对不对得上。这些 bug 在普通神经网络训练里几乎不会出现，但在 MoE 内核里全是常见病。
5.4 三个最容易踩的坑
- 容量因子设太极端：为了指标好看把容量因子压到 1.0，结果 token 大量被丢，训练损失曲线看似下降，实际在线效果崩盘。容量因子要结合真实路由负载分布去调，不能只凭感觉。
- FP8 量化范围误判：统计的激活最大值来自离线样本，线上突然出现异常 token 超出最大值，scale 一旦溢出，整条累加链全部污染。要加 clip 层和一个高水位监控，线上排查速度会快很多。
- 只优化 GEMM 不优化 permute：MoE 层的 token 重排和结果回写同样消耗显存带宽，经常出现 GEMM 跑得飞快，前后 permute 却占了 30% 时间的情况。整体优化时要连重排逻辑一起看。
6. 写在最后的个人经验
6.1 我最大的冲击：路由结果是调度指令，不是概率
这次拆解给我最大的冲击，不是公式多复杂、内核多精妙，而是一个看起来很小的视角差异：MoE 的路由输出，在系统实现里被当成调度指令而非概率分布。算法同学习惯把模型看作一堆张量运算，而 Infra 的视角则是把这些运算翻译成访存、调度和指令发射。同一个模型，两种视角下看到的瓶颈完全不同。算法侧优化往往在调超参、换激活函数、改注意力结构；系统侧优化则是在改数据排布、调缩放时机、优化 kernel 启动方式。
6.2 最有效的练习：让负载形状决定内核设计
如果你也想真正跨过“算法懂 Infra”这道坎，我的建议是把周末时间花在迷你 MoE 上。先跑通朴素实现，然后手动把其中一个专家组改成分组 GEMM，观察 8 个专家在 kernel 内部按组切分、每组不等长却仍然能保持较高吞吐时，那种“稀疏模型为什么快”的体感会非常强烈。到这一步再回头看完整源码，你会发现它其实就是把你手写实现里最丑陋的部分——动态形状、尾块浪费、缩放时机——逐一打磨干净的结果。教程里的参数永远是针对特定场景调出来的，你自己的 batch 分布和 top-k 分布长什么样，最终决定应该改哪一行代码。
这个内容后续可以往两个方向扩展：一是深入到多 GPU 场景下专家并行的通信优化，二是把 FP8 量化的数值误差分析方法补进来。先把今天这条主链路吃透，后面会顺很多。
