做算法的人学 Infra最怕的就是被一堆源码和术语劝退。但如果你真的想把一个大模型跑明白绕不开 GPU 底层发生了什么。今天这篇是“算法同学学 Infra 系列”的第三篇专门讲 DeepSeek V4 里最核心的 MegaMoE 架构以及把它跑起来的发动机——DeepGEMM 源码。这篇文章会从数学模型一步步走到 CUDA 级别的实现聊清楚为什么 MegaMoE 能兼顾超大参数和高效推理也看看 DeepGEMM 这类底层算子库究竟做了哪些事。适合那些对 MoE 只有模糊概念、想深入研究 Infra 细节、或者准备做高性能推理优化的朋友。我自己推动过大模型的训练和上线过去也一直在跟显存瓶颈、路由失衡、算子开销这些东西打交道。看完 V4 这套设计之后最强烈的感受是它的竞争力不只是模型本身强而是从数学设计到底层算子都打通了这才是 Infra 工程真正体现价值的地方。下面我会按自己的理解从整体设计、数学模型、DeepGEMM 源码、实操细节和避坑经验这五个维度来拆解尽量保证你看完能对 MegaMoE 有一个立体的认知。1. 先搞清楚我们在聊什么MegaMoE 到底解决了什么问题1.1 从稠密模型到稀疏模型的必然演进先聊几句背景。2023 年之前我们训练一个千亿参数模型通常按稠密的方式来做也就是说每个 token 进来的前向推理必须计算模型里所有的参数。比如一个 1750 亿参数的传统模型无论输入是什么全套参数都要参与运算这就导致计算量非常奢侈。后来大家发现一个现象实际输入的 token 有很多是高度相似的或者至少它们只需要激活一小部分能力就能得到不错的表示。就像一个团队里不是所有专家都要处理每个任务只有相关的少数专家站出来就够了。于是稀疏混合专家架构MoE开始流行这类模型把一个巨大的前馈网络拆分成多个并行的专家子网络每次只激活其中一部分。DeepSeek V3 已经用了 MoE 思路V4 的 MegaMoE 则是把这个思路做到更极致。MegaMoE 的核心变化在于专家数量更多、激活策略更精细同时配合了大规模的专家并行和底层算子优化。它的好处很直接参数总量很大但每个 token 实际参与计算的参数显著减少训练和推理成本按比例下降。这就是为什么它能撑起更大规模模型的同时在单卡部署上仍保持可控的显存占用。1.2 MegaMoE 的“Mega”体现在哪里MegaMoE 里的 “Mega” 不只是一个营销词。我理解它包含三个层面第一专家规模变得非常大。V4 延续了主流 MoE 的做法把 Transformer 里的 FFN 模块替换成若干专家但这些专家的数量比上一代有明显的提升同时每个专家被设计成比较规整的矩阵结构方便底层算子做并行优化。第二稀疏性被提升了。每个 token 不一定固定路由到固定数量的专家而是可以根据任务难度动态调整。简单 token 可能只要两三个专家复杂 token 会激活更多专家。这种动态稀疏的做法使得模型在容量和计算成本之间取得了更好的平衡。第三系统层面为“大”做了充分的配套设计。模型并行、专家并行、显存调度不再像以前那样靠后添加而是从一开始就内建在训练和推理框架里。这部分正是算法同学最容易忽略、但实际影响最大的环节。1.3 一个算法工程师的直觉为什么要研究 Infra 才能读懂 MegaMoE我见过不少算法同学看模型结构图能理解 MoE 在做什么但一遇到性能报告就茫然了。本质原因是模型的“数学好”和“跑得快”之间隔着一整层 Infra 工程。MegaMoE 这种结构如果只是用 PyTorch 按最直白的方式写出来效率会惨不忍睹因为专家分布在不同 GPU 上跨卡通信会拖垮整体吞吐。就拿路由来说模型决定把 token 分配给哪些专家但如果底层没有高效的 token 到专家映射机制这种灵活性反而会带来巨大的调度开销。DeepGEMM 这种底层算子库就是在解决这类问题它深刻理解 MegaMoE 的数学结构然后针对 GPU 的硬件特性做了极致优化让理想中的稀疏高效真正落到现实的算力上。所以我一直觉得懂一点 Infra再回头看模型设计你会发现完全不同的层次。2. MegaMoE 的数学模型拆解路由、负载均衡与稀疏激活2.1 稀疏门控的数学表达先给一个经典的 MoE 层数学框架。假设输入 token 的隐藏表示为 x通过一个门控网络Gating Network来计算该 token 和各个专家之间的匹配分数。通常门控是一个线性层加 softmaxG(x) Softmax(W_g · x)其中 W_g 是门控权重矩阵输出形状是所有专家数量。这个输出的含义很简单每个专家在这个 token 上的“发言权”有多大。比如有 64 个专家输入 x 是 4096 维向量那么 W_g 就是 64×4096 的矩阵。但这里有一个关键问题如果对 64 个专家都做完整的专家 FFN 计算MoE 的优势就消失了因为计算量跟稠密模型没有区别。所以 MegaMoE 采用稀疏门控策略只挑分数最高的 top-k 个专家其余专家的分数直接忽略TopK(G(x)) 保留前 k 个最大分数其余置为 -∞ 或置为 0在 V4 的实现里k 的取值通常比较小比如 6 到 8 之间但专家总数可能达到一两百个。这样做的好处是在模型容量不缩水的情况下计算量只跟 k 成正比而跟总专家数没有直接关系。2.2 路由机制 top-k 和容量因子的配合top-k 路由看起来很简单但实际工程里有一个问题不同 token 的 top-k 分数分布差异很大有的 token 可能前 6 个专家分数很接近有的 token 则集中在某一个专家上。如果不加约束热门专家会被大量 token 涌入形成路由失衡。所以 MegaMoE 用了一个容量因子capacity factor来限制每个专家最多能处理多少 token。容量因子 C 的定义是C (token 总数 / 专家总数) × CF其中 CF 是人为设定的一个超参数一般取 1.25 到 2.0。这个公式的意思是理论上让每个专家处理均分的 token但预留一些余量防止波动。如果某个专家的 token 数量超过了容量上限超出的部分会被丢弃或重新路由。从数学角度看容量因子是对稀疏路由的一种软约束。CF 设置太小会导致 token 被丢弃丢失信息设置太大则失去了稀疏的意义。在实际调试中CF 是一个需要反复做的实验通常会结合后续的负载均衡损失一起调。2.3 辅助损失让门控学会公平分配为了让门控网络本身学出均衡的路由策略MegaMoE 还引入了辅助负荷均衡损失auxiliary loss。最常用的设计思想是想让每个专家被分配的 token 比例尽量接近均匀同时每个专家的平均门控分数也要尽量均衡。简化表示就是计算所有专家被路由到的概率分布然后和均匀分布做比较常用 KL 散度或者 L2 距离作为惩罚项加到总损失里。这样门控网络会在训练中学会避免“偏科”让专家之间的利用率趋于平衡减少某些专家过载而另一些专家闲置的情况。不过辅助损失也不能设太大否则会矫枉过正。它可能让门控不再关注 token 的真实语义而是一味追求平均分配导致模型效果下降。V4 这种大规模 MoE 里对这个损失做幂次调整或自适应加权是常见的做法。2.4 从数学到实现模型并行与专家并行数学模型定义清楚之后怎么把它拆分到多台 GPU 上就是 Infra 的活了。MegaMoE 在训练和推理阶段一般会同时使用两种并行方式首先是模型并行也就是把 Transformer 的不同层拆分到不同设备其次是专家并行把不同的专家分布到不同设备上。专家并行有一个经典问题如果 token 被路由到的专家不在本设备就需要跨设备通信把 token 的隐藏表示传到目标设备上去。这个通信过程叫 All-to-All。它就像一个大型中转站每个设备把自己管辖的 token 按目标专家分桶发给对应的设备同时接收其他设备发来的数据。在 MegaMoE 的框架里专家并行的调度和 DeepGEMM 的调用是紧密耦合的。每个设备收到的 token 需要被整理成连续的 batch然后一次性喂给本地的专家做矩阵乘。如果数据布局混乱GEMM 就没法高效运行。这也是为什么算法层面的路由策略会直接影响底层算子的发挥——你选择的 top-k 和容量因子最终都要通过 All-to-All 和 GEMM 来兑现成实际的吞吐。3. DeepGEMM把数学跑起来的底层引擎3.1 矩阵乘法的本质从线性代数到 GPU 的 FMAMegaMoE 的数学结构最终会落到一个核心操作上矩阵乘法。专家模块的本质就是一个大矩阵token 的隐藏表示和这个矩阵相乘得到输出。DeepGEMM 就是专门为这类矩阵乘法设计的高性能 GPU 算子库。矩阵乘法的数学定义很简单C[m, n] sum_k A[m, k] × B[k, n]实际在 GPU 上执行时底层是由很多小的计算单元并行完成的。比如 A 是 4096×4096 的矩阵B 是 4096×4096 的矩阵乘积 C 有 1600 多万个元素。如果让一个计算单元顺序算完要循环几十亿次乘加这显然不现实。GPU 的做法是把矩阵切成小块分给成千上万个线程同时计算。DeepGEMM 在传统优化之上还利用了 Nvidia Hopper 架构特有的 Tensor Core 和 TMATensor Memory Accelerator机制。Tensor Core 可以一次完成多个矩阵元素的乘累加操作相当于把普通 GPU 里的标量 FMA 升级成了小规模矩阵的 FMA。3.2 FP8 量化两个精度的混合骗术DeepGEMM 最能拿出手的是它对 FP8 数据类型的支持。FP8 是 8 位浮点数相比 FP16 精度低很多但计算速度更快、占用显存更少。DeepGEMM 的策略不是无脑使用 FP8而是把输入 A 和权重 B 分开量化到不同精度模式。它支持两种 FP8 格式E4M3 和 E5M2。E4M3 的尾数位数更多精度稍高适合用于需要准确表示的激活值E5M2 的指数范围更大适合用于权重这种变化范围大的数据。DeepGEMM 的做法是让激活值用 E4M3权重用 E5M2这样可以在损失极小精度的情况下把计算总量压缩一半。当然FP8 的误差问题也需要处理。DeepGEMM 提供了 scale 参数来做缩放通俗地说就是把数值先乘一个因子放大到 FP8 能表达的范围计算完再乘回原来的尺度把精度损失控制在可接受范围内。3.3 DeepGEMM 为什么快从 CUDA Core 到 Tensor Core 的迁移传统 GEMM 实现会把计算任务大量放在 CUDA Core 上即普通的 GPU 计算单元。虽然 CUDA Core 数量很多但每个核心的计算能力有限在执行大型矩阵乘时非常吃力。DeepGEMM 的做法是让计算主要在 Tensor Core 上完成而 CUDA Core 主要做数据搬运和预处理。这里有个很重要的概念叫 warpgroup level MMA即一个 warp 组共同完成一次大块矩阵乘的部分结果。在 Hopper 架构上它通过新型的 WGMMA 指令warpgroup matrix multiply accumulate来完成。WGMMA 指令可以把一个比较大的矩阵块直接加载到 Tensor Core 附近的寄存器或共享内存里然后一次性完成乘法累积极大减少了指令分发的次数。配合 TMA 机制DeepGEMM 可以异步地把数据从全局内存搬到共享内存。这就像工厂里提前把原材料运到工位旁边计算单元一开工就能立刻拿到数据而不是等待运输从而避免算力空闲。这种“计算与搬运重叠”的设计让 DeepGEMM 在运行 MegaMoE 的专家计算时可以做到接近硬件极限的吞吐。3.4 从简化代码看懂 DeepGEMM 的流程直接看完整源码会把很多人劝退所以我简化成几个核心步骤来描述它的流程这样你能理解框架再看源码时就会轻松很多。第一步拿到输入激活矩阵 A 和专家权重矩阵 B判断当前批次应该用哪个 FP8 精度并读取对应的 scale 参数。第二步通过 TMA 指令把 A 和 B 的矩阵块异步搬运到共享内存中这一步不占用计算单元。第三步计算块小矩阵乘法用的是 WGMMA 指令让 Tensor Core 执行实际上的一次性大矩阵乘累加。第四步把结果矩阵 C 的块按行对应关系写回全局内存。如果存在分组 GEMM 场景还要根据专家 ID 找到对应权重块再执行类似的操作。这就解释了为什么 DeepGEMM 可以同时服务于稠密计算和分组场景。其实分组 GEMM 就是很多个小 GEMM 的集合DeepGEMM 通过索引和线程块调度让这些小矩阵乘也能被并行高效处理不需要反复启动 kernel。这是它在 MoE 推理场景里非常关键的优势。4. 实操拿 DeepGEMM 源码做一次手把手的拆解4.1 源码目录结构与入口如果你去翻 DeepGEMM 的开源代码会发现它的结构非常克制。核心目录一般只有少数几个文件但每个文件都极其精炼。建议新手不要上来就看最深层的 CUDA kernel 实现而是从统一入口开始。入口通常是一个处理层host 端代码负责把用户调用转换成 GPU kernel 的启动。这里你会看到对 M、N、K 维度的解析以及 FP8 的布局转换。M 可以理解为 batch size 和序列长度的乘积N 是输出通道数K 是输入特征维度。DeepGEMM 对 M 特别敏感因为 M 的大小决定了数据加载时是否值得用 TMA 做异步搬运。真正核心的 kernel 代码会包含模板参数用于控制是否启用 TMA、是否处理 Grouped GEMM、是否熔合激活函数等。看明白这些模板开关基本就能理解 DeepGEMM 的设计边界在哪。4.2 TMA 与 WGMMA 的配合细节我花了好几天才真正理解 TMA 和 WGMMA 是怎么配合的。你可以把 TMA 想成一个擅长搬东西的工人而 WGMMA 是计算流水线上的核心机床。TMA 负责把共享内存里需要的数据按特定布局放好WGMMA 直接从共享内存读取并计算。DeepGEMM 会预先申请一块共享内存作为 buffer然后通过 TMA 把 A 和 B 的块轮流搬进去。每次搬完一块数据WGMMA 就立刻算这块数据的一部分结果同时 TMA 开始搬下一块。这种双缓冲机制能让计算单元始终有活干而不是干等数据。实际操作中需要特别注意的是内存地址对齐和同步屏障的位置。如果同步太早计算单元会等待不必要的数据搬运完成如果同步太晚可能会读到脏数据。DeepGEMM 的源码里这些地方都有非常精细的同步控制这也是它跟普通实现拉开差距的地方。4.3 跑通一个可调试的工程三个关键配置如果你想在本地跑通 DeepGEMM我建议先别急着启动大型模型而是编译运行它的测试程序对比 CPU 参考实现和 GPU 计算结果。在这个过程中有三个关键配置特别容易出问题第一个是 FP8 的 scale 设置如果 scale 不合适结果误差会非常大第二个是共享内存的大小TMA 和 WGMMA 都会占用共享内存一旦超限 kernel 会启动失败第三个是编译选项里的架构设置必须是 sm_90a 才能支持 TMA 和 WGMMA如果用老架构编译根本跑不起来。我实际操作时把精度标准设置为误差相对值小于 1e-2 就可以用于大多数推理任务。如果误差过大优先检查 scale 初值然后看数据有没有溢出 FP8 范围。4.4 FP8 精度控制与误差排查FP8 的精度控制永远是绕不开的坑。DeepGEMM 的每个结果都带一个 scale这是量化过程的核心。简单说FP8 能表示的数范围有限必须先把原始浮点数缩放到这个范围内再量化。假设原始激活数值分布在 -64 到 64 之间E4M3 的最大值大约是 448理论上可以容纳但尾数只有 3 位精度损失很大。这时候可以通过乘以一个小于 1 的 scale 把数值压小用更多的有效尾数位去表达小数部分。这个 scale 通常是在初始收集时统计出来的。如果验证时不匹配我的排查顺序是先确认权重和激活是否都已经是量化后的 FP8 数据再看 scale 是否被正确传递到反向量化环节最后查 K 维度的累加顺序是否导致浮点误差累积。多数情况下都不是 DeepGEMM 的问题而是我们自己的数据预处理兜错了。5. 常见问题与排查技巧算法视角的 Infra 避坑5.1 显存爆了到底是谁的锅很多人在跑 MegaMoE 推理时遇到显存溢出第一反应是模型太大了要换小模型。但实际显存消耗分为三块权重本身占用的显存、KV cache 占用的显存、以及前向计算时中间激活占用的显存。权重显存是固定的但中间激活在 MoE 里可能非常大因为路由之后多个专家同时计算它们的中间结果会短暂同时存在。用 DeepGEMM 这类底层算子其实可以在一定程度上降低中间显存占用因为它的融合计算能当场算完当场写回的减少临时张量。但如果你用的是直白的 PyTorch 实现每个专家的输出都会保留一份完整张量显存自然容易爆。我建议先做显存剖析统计每一个环节的峰值内存再判断是权重、KV cache 还是中间激活主导。很多时候只要调整专家并行策略或融合计算就能省出一大块显存不需要更换更小的模型。5.2 路由不均衡导致的性能退化MegaMoE 在训练时通过辅助损失来保证路由均衡但推理时也有可能遇到路由倾斜。如果一批 token 集中涌向少数专家这些专家所在的 GPU 会被打满而其他专家的设备却空闲整体吞吐反而可能不如稠密模型。这时候单看理论稀疏度是没有意义的因为瓶颈转移到了跨卡通信和热点专家上。我做性能分析时习惯统计每个专家的负载分布如果方差过大就要考虑调整容量因子或者给热点专家做进一步的权重切分把它的矩阵按层拆到多个设备上。如果设备支持多副本推理把热点专家复制多份也是一种有效手段。这需要路由表参与调度会增加一部分工程复杂度但效果立竿见影。5.3 DeepGEMM 计算结果与参考实现不一致怎么办有一类经典问题DeepGEMM 的输出和 PyTorch 参考实现的输出对不上。我的经验是不要立刻怀疑算子有 bugFP8 量化本身就会带来误差需要先建立合理阈值。判断误差的合理标准是相对误差在 1e-2 级别通常可以正常参与后续计算如果到了 0.1 级那就要检查。常见原因有三个scale 设置不合适、分组 GEMM 的索引映射错位、或者某个维度没有对齐到 TMA 要求的对齐尺寸。建议先从单专家小矩阵开始测试对比 DeepGEMM 输出和 PyTorch 输出在相同输入下的差异逐步扩大规模。如果小型测试通过而大型测试失败大概率是 data layout 变动导致索引丢失仔细查看分组索引的重建逻辑就好。5.4 一点工程心得我跑过很多次 MoE 模型的上线有一个感触把算法模型读懂和把底层算子调通其实是相辅相成的。你越理解 MegaMoE 的稀疏设计初衷就越知道 DeepGEMM 里的分组计算为什么要把专家索引提前绑定在 kernel 参数里而不是在运行时临时判断。这类底层优化都要提前把“可变”的东西尽量变成“固定”的编译期信息以减少运行时的分支判断。DeepGEMM 的模板参数把 M、N、K、是否启用 TMA、是否启用 grouped GEMM 这些选择全部放到了编译期原因就在这。对算法同学来说理解这种“编译期展开”的思路对于理解任何高性能算子库都有很大帮助。6. 从模型到算子一次完整的 Infra 视角思考6.1 一张图理解 MegaMoE 的完整生命周期从输入 token 到最终输出MegaMoE 经历了一条非常清晰的链路Embedding 层把 token 变成向量然后经过多层 Transformer 模块其中 MoE 层的门控网络决定每个 token 去哪些专家专家并行系统负责把 token 调度到正确的设备上最后设备上的 DeepGEMM 算子执行矩阵乘并返回结果。这个过程看似复杂但核心只有两个路由调度和矩阵计算。路由调度决定“谁来做”矩阵计算决定“怎么做”。DeepGEMM 解决的正是第二个问题里最耗时的那部分。理解了这条链路以后再看任何推理框架的 profiling 数据你都不会再一头雾水。6.2 给算法同学的下一步建议如果你也想从模型工程师往 Infra 方向跨一步我的建议是先跑通一个简单的 MoE 推理实验记录耗时分布找出矩阵乘占比和通信占比。然后去读 DeepGEMM 的 README对照着你记录的耗时数据去理解它优化的是哪一个环节。接着可以尝试修改 FP8 的 scale 初始值观察精度和速度的变化亲手感受一下精度成本与算力收益的权衡。最后再看 Grouped GEMM 的实现理解它是如何把不同的专家矩阵放在同一个 kernel 里跑的。这一系列操作下来你会发现很多以前觉得很抽象的技术名词都变得具体了TMA 就是高效搬数据WGMMA 就是批量矩阵乘FP8 就是低精度换速度。它们没有多玄妙只是在工程上把硬件能力用到极致。我做 Infra 这么多年最深的体会是真正拉开模型应用差距的往往不是某个花哨的网络模块而是那些看起来不起眼却能决定吞吐上限的底层细节。算法同学如果愿意花一点时间弄懂 DeepGEMM 这类库的运作方式在优化模型和排查瓶颈时会有完全不一样的视野。这套思路不仅适用于 DeepSeek V4也适用于之后每一个要落地的大模型系统。