企业官网建设流程全解析

上个月帮组里一位做推荐的同事排查训练速度问题，模型是接近万亿参数的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，激活参数只有37B，256个路由专家加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 · Σ_{i=1}^{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常见两种格式：E4M3（1位符号、4位指数、3位尾数，动态范围大，适合前向）和E5M2（1位符号、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要访问“别人的专家”时，通信就不可避免。

流程可以拆成四步：

1. dispatch：每个token先算出门控结果，知道自己要去哪些专家；
2. 跨设备All-to-All通信：把token的hidden表示发给目标专家所在的GPU；
3. 专家计算：目标GPU上的专家做GEMM，得到局部结果；
4. 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架构上常见的wgmma（warpgroup-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的M（token数量）还都不一样。如果逐个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 先把四块短板补齐

我带过不少算法同学，发现能顺利把这类源码吃下来的人，通常都先补齐了四块知识：

1. GPU内存层级（HBM、L2、Shared Memory、寄存器）和CUDA线程模型——不一定要会写kernel，但要能看懂一张内存层次图，知道数据从哪里来、到哪里去；
2. GEMM的经典实现演进：从最朴素的三重循环，到共享内存分块，再到Tensor Core和wgmma。可以跑几个CUTLASS示例感受一下每一层优化带来的收益；
3. 分布式训练常用通信原语：AllReduce、AllToAll、Broadcast，以及各自的大致时间模型，搞明白什么情况下通信会盖过计算；
4. 性能分析工具：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推理里的内存布局和调度策略，那就下一期见。