上个月帮组里一位做推荐的同事排查训练速度问题模型是接近万亿参数的MoE架构看着显存利用率挺高但整机吞吐就是上不去。我让他先跑一遍nsys的profile结果kernel时间线里密密麻麻全是小GEMM的启动每个只有几十微秒中间穿插着一堆等待。他当时的表情我印象很深“我以为瓶颈在模型计算原来全程卡在底层调度上。”后来我们把小专家合并、改成grouped GEMM之后吞吐翻了将近一倍。这件事跟DeepSeek V4 MegaMoE放在一起看特别有意思。最近社区里关于V4 MegaMoE和DeepGEMM的讨论非常多各种群都在传开源仓库的消息。但作为一个一直关注大规模模型训练的算法工程师我反而觉得最有价值的不是“参数又变大了”这种新闻而是它背后那条“数学模型→GPU实现→源码落地”的完整链路。这篇是“算法同学学Infra”系列的第三篇我想换个角度不聊八卦带着你把MegaMoE的数学模型拆一遍然后一路看到DeepGEMM的源码组织方式搞清楚一个token从门控打分到专家计算中间到底发生了什么。适合谁看算法出身、最近在搞大规模MoE训练或推理、想系统补一补系统知识底子的同学。你已经知道MoE大概是什么样但还不清楚为什么一到分布式就各种意外。看完这篇你能建立一条完整的认知链路模型设计如何映射到硬件源码里的关键机制都在解决什么问题以及你下一步该学什么。1. 算法同学为什么会被MegaMoE“卡住”——从一次性能排查说起1.1 先把MoE的基础框架摆清楚如果你之前做过MoE模型应该清楚它的核心就是“稀疏激活”模型总参数量非常大但每个token不会走完整网络只通过一个路由器从几百个专家里挑几个出来计算。DeepSeek V3的公开数据就是一个典型例子总参数671B激活参数只有37B256个路由专家加1个共享专家每个token只激活8个专家。所谓MegaMoE我理解就是在同一个方向上继续放大专家数量更多、粒度更细、路由更复杂同时用更激进的工程手段把每一分算力都压出来。为什么方向是这条路核心原因是纯稠密模型继续扩大规模的边际收益越来越差而稀疏激活理论上能用更少的计算量干更多的活。但“稀疏”这两个字恰恰是工程噩梦的起点。稠密模型的计算形状非常规整一个batch过来矩阵乘法维度固定GPU利用率可以长时间保持在高位。MoE模型完全不是这样每个token去哪个专家是运行期才决定的专家之间的token数量天然不均匀GEMM的形状也各不相同。算法同学写forward时很爽一行torch.mm就完事可一旦跑大规模训练或推理所有问题都会在这层“稀疏”上暴露出来。1.2 那次排查教会我的事瓶颈在小GEMM的启动开销回到开头的故事。当时那位同事的模型其实计算量和显存都没到极限GPU利用率只有30%。用nsys拉时间线问题清楚得让人意外大量小GEMM的kernel launch把流水线打碎了。单个小GEMM的执行时间只有几十微秒但kernel启动的固定开销、显存访问的延迟、以及每次切换算子带来的流水线排空全都摊在了极短的计算时间上。打个比方你有一辆超跑但在市区里每开500米就要停车等红灯重新起步再好的发动机也快不起来。稠密GEMM相当于高速巡航MoE里的小GEMM则是频繁起步。业界对这个问题的解法其实已经明确第一把相同形状的小GEMM合并成grouped GEMM一次启动处理一批问题第二在推理时通过连续批处理让连续token尽量去往连续的专家提升缓存命中。两边都指向同一个底层能力——你得理解GPU是怎么调度矩阵计算的。这正好是Infra的范畴也是系列文章想带你建立的核心视角。1.3 “Mega”到底巨在哪规模增长会加剧哪些问题很多算法同学把“Mega”简单理解成“参数更多了”但这只是表象。真正可怕的是随之而来的几个非线性增长问题。专家数量增加负载不均衡会更严重热门的专家可能收到几十倍于冷门专家的token这部分专家成为瓶颈其他专家却在空转。token分配更散通信模式从“少数设备之间做AllReduce”变成“全员All-to-All”通信成本随设备数量快速上涨。小GEMM的数量指数级增加kernel启动开销和显存访存的碎片化会一点点吃掉规模红利。所以DeepGEMM这种基础设施才会被重视——它解决的不是“某个模型跑得快”而是“一类稀疏架构的通用性能底座”。这也是我要一路拆到源码的原因不看底层很多上层决策就只能靠猜。接下来先从数学模型说起因为源码的每一个设计最终都能追溯到某个数学需求。2. MegaMoE的数学模型拆解路由、稀疏激活与损失函数2.1 路由token是怎么找到专家的MoE最核心的公式其实就几行。第一门控网络给每个专家打分p_i Softmax(W_r · x_t)_i第二取前k个高分专家得到集合top_k第三把选中专家的输出按门控概率加权合并并加上共享专家的输出y E_shared(x_t) Σ_{i∈top_k} p_i · E_i(x_t)这里的细节值得展开。p_i是经过softmax的概率它同时承担两个角色一是“选择谁”二是合并时每个专家输出占多少权重。共享专家的作用是提供一条稳定的通用表示通路避免完全依赖少量路由专家导致信息瓶颈。还有一个容易被忽略的地方门控打分时V3以及很多现代MoE会在logits上加入可学习的噪声项让路由在训练初期不那么“自信”给冷门专家更多被选中的机会。这是模型侧缓解“马太效应”的一种手段但它只是延缓问题真正兜底还是要靠下一节的辅助损失。从数学结构看MoE的好处非常直观你把671B参数分成几百个专家每个token只用其中8个那计算量大约是“8/数百”的比例远小于同尺寸稠密模型。这就是稀疏激活在算力效率上的核心价值。2.2 辅助损失如何让“打工人”不扎堆如果路由完全没有约束训练到最后大概率是一两个专家累死、其余专家摸鱼。为了把负载拉平MoE论文里几乎标配一个辅助损失L_aux α · E · Σ_{i1}^{E} f_i · P_i其中f_i是实际被路由到专家i的token比例P_i是该专家的平均门控概率E是专家数量α是权重系数。每个字母都值得解释一下。如果专家i很热门f_i会很大同时门控给它的平均概率P_i也高那么乘积就大损失就上升训练时梯度会倾向于把概率挪给冷门专家。乘以E是为了让损失尺度不随专家数量变化不然专家越多这个辅助损失的绝对数值就越大。在DeepSeek V3的技术报告里除了这种经典的Aux Loss还额外引入了鼓励专家分配更均匀的互补损失并通过“无token丢弃”的工程设计保证训练稳定性。到了MegaMoE阶段专家规模更大负载不平衡问题会更尖锐辅助损失的设计会更讲究但核心思想仍然是把这个“f_i·P_i”的因果变量压平。理解这一点你再看训练曲线里那些loss突变就不会只盯着主loss了。2.3 从FP32到FP8精度管理是“省一切”的前提模型侧聊完必须聊精度。MegaMoE规模摆在那里如果权重全部用FP32或FP16存放显存和带宽都会出问题。FP8就是目前大规模MoE都在走的方向。FP8常见两种格式E4M31位符号、4位指数、3位尾数动态范围大适合前向和E5M21位符号、5位指数、2位尾数精度低但范围更大适合梯度这类动态范围变化大的数据。为什么需要逐token、逐group放缩本质原因是FP8的尾数太少直接用全局scale会出现极端值被削平、小数值被吞掉的问题所以得给不同的token、不同group单独维护一个缩放系数。DeepGEMM整个存在的意义用一句话概括就是这样一个量化计算过程C ≈ (A_quant / scale_A) · (B_quant / scale_B) · scale_C它把FP8的量化误差控制在可用范围同时把计算吞吐拉上去。算法同学需要建立的关键认知是FP8不是“直接省一半显存”这么简单它是一个从模型设计、量化策略到kernel实现都要多层配合的系统工程。下面开始进入硬件视角。3. 从数学公式到GPU硬件算法与Infra之间的那条鸿沟3.1 你写的torch.mm在GPU上到底发生了什么我见过不少算法同学把GPU当成一个“会加速矩阵乘法的黑盒子”这不怪大家PyTorch封装得太好了。但一旦你要优化大规模MoE就得知道黑盒里有什么。GPU上有三级存储概念全局显存HBM最慢共享内存Shared Memory在SM内部寄存器最快。GEMM在GPU上的高效实现核心思路是分块把需要相乘的A、B矩阵切片加载到共享内存尽可能复用然后交给Tensor Core执行乘加运算。Tensor Core是NVIDIA自Volta开始引入的专用计算单元支持低精度混合计算FP8 GEMM的高吞吐基本都靠它。这里有一个对算法同学极其重要的概念算术强度也就是“单位字节访存对应多少次计算”。稠密GEMM如果矩阵够大计算量远大于访存量属于计算密集GPU很擅长。而MoE场景里几万个token被切成几百个小组每个小组的GEMM矩阵都很小访存开销和kernel启动开销占比一下子涨上去算术强度掉得非常厉害GPU自然就“闲”下来了。我常用的生活类比是大矩阵GEMM像自助餐厅客人一次拿一堆食物回去慢慢吃数据加载一次、计算很多次MoE小GEMM像外卖小哥每单都要跑回餐厅取餐每次计算都要重新加载数据骑手再快也跑不出高吞吐。3.2 专家并行一切并发问题的起点训练或推理MegaMoE把几百个专家放在一张卡上是不可行的必须分布到多张GPU。常见做法是专家并行每个GPU上放一部分专家。这样当token要访问“别人的专家”时通信就不可避免。流程可以拆成四步dispatch每个token先算出门控结果知道自己要去哪些专家跨设备All-to-All通信把token的hidden表示发给目标专家所在的GPU专家计算目标GPU上的专家做GEMM得到局部结果combine把结果聚合回来做加权求和继续下一层。这个All-to-All的通信量可以算一笔账。假设一批4096个token、hidden size 7168、每个token激活8个专家那么每个token要往8个目标位置发送7168 × 2字节BF16的数据约14KB。单层单方向的总数据量大约就是4096 × 8 × 14KB ≈ 450MB。MoE层一多这个数字乘上几十层再算上前向和反向通信开销非常可观。这也是为什么工程上如此重视通信算子的优化通信要和计算尽量重叠负载也要尽量均衡否则一个节点慢了整批token都得等它。算法同学平时感知不到这层是因为框架帮你藏起来了但一旦性能出问题最先露馅的就是通信。3.3 显存账本MegaMoE的存储开销不能只看参数很多算法同学以为“显存够不够”只取决于模型权重实际上训练和推理的显存开销都是一个大账本。权重本身671B参数如果按FP8存储大约336GB。即使铺到64张卡每张卡也要承担5GB以上的权重。但这只是冰山一角。推理场景还需要KV Cache并发请求越多占的额外显存越大训练场景更夸张AdamW优化器要为每个参数额外保存momentum和variance每个参数动辄多出8字节。671B × 8字节大约是5.4TB已经超过很多小集群的单卡显存总和。所以FP8量化、激活重计算、梯度检查点这些技术本质上都是在给显存账本做精细化管理。算法同学要建立的直觉是模型越大显存和数据搬运越可能成为瓶颈纯计算次数反而不是。MegaMoE这种几百专家的结构更是在挑战存储和带宽的极限。带着这个视角你去看DeepGEMM源码就会有完全不同的代入感——那不是在炫技是每个字节都得抠着用的必然结果。4. DeepGEMM源码走读一个FP8 GEMM Kernel是如何被组织起来的4.1 先搞清楚DeepGEMM在整条链路里的位置DeepGEMM是DeepSeek开源的一个FP8通用矩阵乘法库定位很简单给训练和推理两条线提供高性能的FP8 GEMM覆盖两类场景——普通稠密GEMM和MoE的场景下那类Grouped GEMM。为什么不用cuBLAS因为cuBLAS的FP8性能在不同shape下波动很大而且它是通用库没有针对“一堆形状不同的小矩阵”这种MoE场景做专门优化。自研的好处是可以把kernel参数、调度策略、JIT编译都完全掌控在自己手里专门服务自家模型的计算特征。我读源码的第一感受是这个库非常“头文件驱动”核心kernel大量使用模板和inline尽量在编译期决定寄存器分配和循环展开。它没有做一个巨大复杂的框架而是把注意力集中在GEMM本身的调度上。这种“小而锐”的风格恰恰是它能在边缘场景里榨出性能的原因。读这类源码我的建议是不要从文件开头逐行读。先看README、看测试用例看入口函数暴露的API然后直接跳到核心kernel的主循环。接下来我就按这个路径走。4.2 核心主循环加载、计算、存储的三步曲一个高性能GEMM kernel的骨架本质上就三件事把数据搬上芯片、让Tensor Core算、把结果写回显存。DeepGEMM的主循环也逃不出这个框架但每一部分都做了精细的调度。以Hopper架构上常见的wgmmawarpgroup-level MMA为例主循环大致长这样对每个输出Tile 把A的一个Tile加载到共享内存producer生产者负责 把B的一个Tile加载到共享内存 当共享内存数据就绪 执行wgmma计算consumer消费者负责 把结果累加到寄存器 按需应用缩放因子 写回全局显存关键在于producer和consumer是分离的一部分warp专门从HBM往共享内存搬数据另一部分warp专心用Tensor Core做计算两边通过共享内存和同步机制衔接让“搬下一块数据”和“算当前这块数据”尽可能重叠。这就是流水线思想的雏形后面第五章我会再展开。为什么要把这一步理解到位因为所有挂在嘴上的优化最后都要落到这个循环里Tile怎么切、数据放哪个存储层级、谁来搬、谁来算、算完怎么存。一旦你看懂了主循环一个GEMM kernel里90%的代码就都顺了。4.3 Grouped GEMM把一堆形状各异的“小矩阵”揉在一起MoE场景最特殊的地方是它同时需要计算很多个专家GEMM每个GEMM的Mtoken数量还都不一样。如果逐个kernel计算CPU端的调度开销和GPU端的启动开销会吃掉大部分性能。Grouped GEMM就是为此设计的一次kernel启动同时处理很多个GEMM问题。实现上通常是这样调用方把所有GEMM的A、B矩阵都准备好再传一个“问题描述”数组里面记录每个GEMM的M、N、K以及各自的指针偏移。kernel内部循环遍历这些problem或者用一组独立的warpgroup分别处理不同的problem从而把多次启动合并成一次。还有一个小细节值得注意在MoE计算中先对token做TopK分派如果我们在物理上把去往同一个专家的一批token在显存里排成连续区域那么专家计算就变成了一系列连续的小GEMM。DeepGEMM的实现会尽量利用这种布局来提升缓存命中率减少随机访存。说得直白一点数据排布决定了kernel的寿命。很多看起来“同样的算法”只是把数据排布微调了一下性能就能翻倍原因也在这里。5. 源码中不容易一眼看懂的细节缩放因子、双缓冲与边界处理5.1 FP8缩放因子为什么不能“一个scale用到底”FP8的问题很简单尾数位数太少动态范围不够。如果我们对一整块权重只用一个缩放因子一旦出现几个异常大或异常小的值量化误差就会被整体拉高。所以现代的FP8方案基本都是“分段缩放”——每个token有自己的scale每个group甚至每个block都有自己的scale。落实到DeepGEMM工程层面缩放因子不是一个附加的独立计算而是被精确地编排进GEMM流程里的。对A矩阵的scale可以在加载阶段就乘进去对B矩阵的scale也可以提前乘累加过程尽量保持FP32精度最后再统一应用C相关的scale。所以你在代码里看到的往往不是一行“乘以scale”而是一套贯穿加载、计算、写回的缩放体系。我个人觉得这是阅读FP8 GEMM源码时要优先搞懂的点量化误差的管理一半在量化算法本身另一半在kernel怎么组织这些scale系数。你如果只是把torch里的量化API换一换很可能看到性能不变甚至变差原因就在于scale的编排没有跟着一起优化。5.2 双缓冲把“等待”变成“重叠”GEMM访存的延迟相对计算来说大得多如果不做缓冲GPU就会一会儿等数据、一会儿猛算利用率忽高忽低。双缓冲的意思是在共享内存里准备两套缓冲区一套给当前这一轮计算使用另一套提前预取下一轮要用到的数据。这样访存和计算就可以并行相当于一边吃饭一边做下一道菜。对于MoE的Grouped GEMM来说双缓冲尤其关键。单个小GEMM的计算量本来就不大如果每次都要等数据从HBM到共享内存再开始算等待时间对短计算来说占比太高。通过双缓冲甚至多级缓冲把访存时间尽量藏到上一轮的计算里kernel才能持续把Tensor Core喂饱。读源码时可以重点观察shared memory数组有没有声明成“两块”类似__shared__ float sA[2][BLOCK_K][BLOCK_M]的形态以及同步原语出现的节奏。理解了双缓冲你就理解了为什么某些kernel能把GPU利用率从30%拉到90%。5.3 边界当M不是tile的整数倍时MoE里每个专家的token数量是由路由动态决定的这个数字几乎不可能正好是tile大小比如128的整数倍。所以GEMM kernel必须处理“最后一块不够满”的情况。常见做法是predicate或mask正常路径用满tile计算边界路径只对有效行做乘加其余部分填零或跳过。边界处理写得不好会产生两个问题一是性能回退比如整个正常路径被拖进低效分支二是正确性问题比如tail里的垃圾数据被卷进累加。源码里看到if (row M)这样的判断或者落到PTX的wgmma带mask参数都属于这类逻辑。我自己读这类kernel的体会是边界代码虽然看起来只有几行但往往是判断一个kernel是Demo级别还是生产级别的分水岭。真实运行中这些边界分支可能被触发数百万次任何一个小分支跳转的代价都会被放大。算法同学如果以后自己写kernel一定记得把tail case当成一等公民来设计。6. 给算法同学的Infra学习路径与实战建议6.1 先把四块短板补齐我带过不少算法同学发现能顺利把这类源码吃下来的人通常都先补齐了四块知识GPU内存层级HBM、L2、Shared Memory、寄存器和CUDA线程模型——不一定要会写kernel但要能看懂一张内存层次图知道数据从哪里来、到哪里去GEMM的经典实现演进从最朴素的三重循环到共享内存分块再到Tensor Core和wgmma。可以跑几个CUTLASS示例感受一下每一层优化带来的收益分布式训练常用通信原语AllReduce、AllToAll、Broadcast以及各自的大致时间模型搞明白什么情况下通信会盖过计算性能分析工具nsys、ncu、dcgmi至少会用nsys看kernel时间线和通信时间线。这不等于要求算法同学变成系统工程师而是给自己打开“黑盒”的能力。以后遇到性能问题你能第一眼判断瓶颈大概在计算、访存还是通信层。6.2 四个可以动手的小实验理论说再多不如跑两个实验。我建议按下面的顺序做每一步的产出都是下一阶段的基础。实验1用Triton或CUDA写一个最简单的SGEMM然后逐步加Shared Memory、双缓冲观察不同版本在Profiler里的性能变化。做完你会对GEMM优化的每一步有体感。实验2拿一个小型MoE模型几百M参数就够加不加aux loss各训练一遍把每个专家收到的token数量画成直方图。你会直观看到“负载不均衡”长什么样以及辅助损失是怎么把它拉平的。实验3在一个多卡环境里写一个All-to-All通信带宽测试把token规模变化和通信耗时画成曲线。前面那个450MB量级的账是死的跑一次你就知道它为什么是瓶颈。实验4把DeepGEMM跑起来在你自己的Transformer模块里把某个FP16 GEMM替换成FP8 GEMM对比速度和精度损失。别只盯着benchmark数字重点观察kernel参数block大小、stage数对性能的影响。这四个实验是我认为从“听说过”到“摸得着”的最短路径。我自己团队里的新同学完成这套之后再去读复杂源码基本都能抓住主线。6.3 我踩过的一些坑和心得最后说几个我真实踩过、也经常看别人踩的坑。第一个坑只看FLOPs不看算术强度。之前我们优化过一个GEMM理论计算量下降了15%实际耗时反而涨了不少原因就是shape变小之后访存占比更高。算法的“计算量”和硬件的“耗时”之间隔着一个访存带宽忽略它就会得出完全相反的结论。第二个坑分布式训练里只盯着单卡利用率。有时某张卡利用率很低不是它计算不行而是在等通信。把通信时间线和计算时间线放在一起看别把单卡指标当成全部真相。第三个坑读源码时从第一行开始线性读。GEMM kernel的template参数、defines和辅助函数能把你淹死。正确的路径是“README→API入口→主循环→细节”这也是我这篇源码走读采用的顺序。拿我自己来说读完DeepGEMM最大的收获不是“我也能写高性能GEMM了”而是终于在自己脑子里搭起了一条完整的因果链门控网络打分→TopK稀疏选择→token在设备间搬运→grouped小GEMM→结果聚合。以前这些环节在我的认知里是分散的现在每次写MoE模型我都会下意识地想一下底层kernel会怎么消化这个结构。如果这篇对你有帮助建议你今晚就做一件事把你正在跑的最大的那个模型丢进nsys里跑一次看看占用时间前五的kernel或通信原语分别是什么。大概率你会回来读第二遍源码。系列第4篇我可能会继续沿着这条链路往下走聊聊KV Cache在MegaMoE推理里的内存布局和调度策略那就下一期见。