在 MoE 推理链路里，真正决定整机吞吐上限的往往不是 Attention，而是那个看起来不起眼的稀疏专家计算模块。我最近在业务里深入打磨一个低精度 GMM 算子（Grouped Matrix Multiplication，分组矩阵乘），从最初能跑通、到精度对齐、再到把性能拉到接近设备极限，前后折腾了好几轮。这篇会把整个过程完整盘一遍：为什么稀疏专家网络绕不开这个算子、低精度输入要如何做数值保真、kernel 调度该怎么设计，以及最后一批容易踩的暗坑。对正在搭 MoE 推理引擎、或者在做量化算子优化的人来说，里面很多结论可以直接拿去做方案参考。
这不是一篇纯理论推导，文中所有参数、数据和方法都来自实际测试。我尽量把取舍背后的原因讲清楚，而不是只贴结论，这样大家可以根据自己的模型规模和设备特性做调整。
1. GMM 算子是什么，为什么绕不开
1.1 MoE 中的矩阵乘法不是一杯匀浆
MoE 模型前向时，每个 token 经过 router 分配，被送往少数几个专家。这里的关键结构是不同专家拥有各自独立的权重矩阵，而且每个专家实际承接的 token 数量还各不相同。如果简单地把整批 token 当作一个大矩阵 X，把专家 FFN 权重拼成一个大矩阵 W，理论上可以一次普通矩阵乘算完，但那样所有权重就变成共享一套参数，和 MoE 的定义完全冲突，还会激活大量闲置参数，导致显存和计算量成倍上涨。所以实际计算只能按专家分组来做，每组一个小矩阵乘，最后把结果合并起来才是整层的输出。
这就是 GMM。实际推理引擎里，它通常表现为：先把 token 按路由 ID 排个序，统计每个专家分到的 token 数，再让一个内核按组执行分批矩阵乘。把这种计算收敛成一个算子的好处非常明显：一次启动、统一管理显存分配，分组信息直接写在内核参数里，避免循环执行几十次独立小矩阵乘带来的时间碎片和 kernel launch 开销。尤其当专家数量上百时，少一次循环就是少一次通信间隙，整体延迟能差出一个量级。
1.2 为什么要低精度，而不是用 FP32 硬扛
先看一组典型参数：假设模型有 32 个专家，每个专家的 FFN 隐藏层是 K=4096、N=4096，一次推理总共处理 2048 个 token。每个专家平均分到 64 个 token，那么单个专家的矩阵乘是 64×4096 @ 4096×4096，大约 2.1 GFLOPs，全部专家加起来接近 67 GFLOPs。这个计算量不算小，但如果只有一个 token 被分到某个专家，那 M_e 就等于 1，计算强度会跌到个位数，算子大概率变成访存受限，而不是计算受限。低精度权重能把权重体积直接砍半，再配合 INT8 或 FP8 的矩阵指令，才能同时把带宽压力和算力压力降下来。
低精度也不是没有代价。量化误差会让数值结果出现偏差，进而影响模型端 logits、路由选择以及最终生成质量。所以低精度 GMM 算子设计里的“精度”不是一个可选项，而是最核心的约束条件。做这个算子，本质上是在算力、带宽、精度三者之间找动态平衡。先跑通 FP32 基线，再把低精度版本对齐到这个基线上，这是后续所有调优工作的前提。
2. 数学定义与瓶颈判断
2.1 标准 GMM 表达形式
GMM 算子的数学描述很简单，就是一组按分组索引区隔的矩阵乘拼在一起。设输入矩阵为 X，每一行是一个 token 的 K 维向量。E 个专家，第 e 个专家有一个形状为 [N_e, K] 的权重矩阵 W_e，承接的行集合是 S_e，该组 token 数量为 M_e。算子输出 Y 满足：
Y_e = X[S_e] @ W_e^T, e = 0, ..., E-1
最终输出 Y 需要按照 S_e 的行顺序重新排回 token 顺序。这里有一个关键细节：真实推理时 X 通常已经按专家 ID 排好序，所以 S_e 就是一段连续区间，不需要散列索引，只需要一个 offsets 前缀和以及 sizes 数组。这个设计能省掉一次索引访问，也让 kernel 里循环加载更规整。算子真正要管理的元数据只有两个小数组，复杂度都集中在矩阵乘内部的数据搬运和调度上。
2.2 先算算术强度，别一上来就写 CUDA
我踩过的第一个坑就是没算算术强度就直接套用大矩阵乘的优化策略。GMM 和普通矩阵乘最大的不同在于每个分组的 M_e 差异巨大。M_e 很小的时候，访存开销会超过计算开销，这时候堆再多的低精度计算单元也没用，因为算力都在等数据回来。
算术强度可以按下面的公式估算：
AI_e = (2 * M_e * N_e * K) / (M_e * K + N_e * K + M_e * N_e)
分子是计算量，分母是必须从显存搬运的字节数，如果统一按 FP16 算，再乘以 2 就是实际字节数。以 N=K=4096 为例，M_e=64 时算术强度大约在 60 FLOP/byte 左右；M_e=8 时只有约 8 FLOP/byte。对大多数 GPU 来说，FP16 下的计算访存比大约在 40 到 80 之间。这么一算，小 M_e 场景基本就是访存瓶颈。改成 INT8 权重后，权重字节直接减半，算术强度几乎翻倍，整个算子会从偏访存型被拉到更均衡的最优区间。
所以低精度在这里不是“锦上添花”，而是让稀疏场景重新回到可高效利用算力的关键手段。后续做 kernel 优化时，我习惯把算术强度表格先列出来，再决定优先优化调度还是优化访存。
2.3 低精度数值误差从哪里来
低精度主要引入三类误差：权重量化误差、激活量化误差、中间累加误差。权重量化误差跟量化策略直接相关，如果一个输出通道的数值波动范围很大，却用同一个尺度去量化，误差就会被放大；激活量化误差随 token 的动态范围变化，同一批数据里不同 token 的激活幅度可能差出十倍；中间累加误差在 INT8 场景必须用 FP32 累加来消除，如果偷懒用 FP16 累加，K 一大误差就会快速累积。
我的处理方式是：权重用 per-channel scale，激活用 per-token scale，算子内部一律 FP32 累加，只有喂给低精度计算指令的那一刻才把数据降下去。这样输出的数值精度可以逼近 FP16 基线，却只付出约一半的存储和更高的指令吞吐。对于一个推理引擎来说，这种“外部高精度语义、内部低精度执行”的设计是最稳的。
3. 核心设计细节与实现过程
3.1 算子接口怎么设计
设计一个可落地的 GMM 算子，接口需要同时表达“分组矩阵乘法”和“低精度”两件事。我最终使用的接口大致长这样：
gmm(input, weights, scales_x, scales_w, expert_offsets, expert_sizes, output, M, K, N, E)
其中参数含义如下：
- input: [M, K] 的输入 token，已经按专家 ID 排好序
- weights: [E, N, K] 的低精度权重，通常是 INT8 或 FP8
- scales_x: [M] 的 per-token 激活缩放
- scales_w: [E, N] 的 per-expert 输出通道权重缩放
- expert_offsets: [E+1] 的分组起止偏移
- expert_sizes: [E] 每个专家承接的行数
把 scales 直接传进算子，而不是在核外先乘好，是为了把量化尺度融合留在 kernel 内部，减少一次中间张量写回。这样也能和后续的残差、激活函数一起做融合，省掉几次额外的显存往返。接口一旦定了，后面优化 kernel 的时候基本不用再改上层代码。
3.2 量化策略与精度保真
权重对称量化是最稳的起点。对每个输出通道独立统计绝对最大值，除以 127 得到 scale，然后四舍五入到 INT8。伪代码大概是这样：
def quantize_weight(w): # w shape: [N, K] scale = w.abs().amax(dim=1, keepdim=True) / 127.0 w_q = torch.clamp(torch.round(w / scale), -127, 127).to(torch.int8) return w_q, scale.squeeze(1)这里有一个容易被忽略的点：对称量化牺牲了零点，但如果权重分布不是严重不对称，对称量化带来的误差可以忽略，而且实现上少处理一个 zero point，内核里计算路径会干净很多。权重分布经过绝对最大值归一化之后，一般不会出现特别重的尾部分布，所以效果可控。
激活量化要复杂一些。Attention 输出的数值分布会随输入内容变化，全局固定 scale 很容易在某些极端 prompt 下爆炸，所以我对激活采用 per-token 动态量化，也就是每个 token 单独计算量化因子。实现方式就是加载一行时取绝对最大值后除一次，成本很低，收益却非常大。我在精度验证中观察到：权重误差是系统性的，会影响所有样本；激活误差更依赖具体 token，在 batch 内表现为偶发飘移。所以权重侧我做得更保守，激活侧的量化则允许在线实时计算。
3.3 内核调度：静态还是动态
GMM 内核调度通常有两种思路：一种是按专家分配线程块，另一种是动态工作块调度。按专家分配逻辑最简单，每个 block 负责一个专家，循环处理该专家的若干行 tile。缺点是专家 token 数差异很大，如果某个专家只有 2 个 token，它的 block 几乎瞬间干完然后闲置，另一个专家却还有几百个 token，SM 资源白白浪费。
我最终采用的是动态工作块：把每个专家按固定 BLOCK_M 拆成若干小任务，再把所有这些小任务排成一个任务队列，由各 SM 上的 block 到队列里领取任务。核心逻辑类似：
// 伪代码：所有 block 共享同一个任务计数器 while (true) { int task = atomic_add(&task_counter, 1); if (task >= task_num) break; int e = task_tbl[task].expert; int mt = task_tbl[task].m_tile; compute_expert_tile(e, mt); }真实实现里还要注意取单之后的内存栅栏，否则任务会被重复执行。这个动态调度方案对负载不均衡非常敏感，但收益也最直接。任务粒度我选的是 32 行，正好等于一个 block 内一行 tile 的大小，领单后不需要再开一次循环，调度开销和负载均衡效果都处于可接受范围。
3.4 权重内存布局与大 K 分块
低精度 GMM 要跑出理想效果，权重的内存布局很关键。我建议把全部专家权重统一排成 [E, N, K] 的连续块，并且让 N 在外层、K 在连续内存方向。这样计算 Y[m, :] 时，权重可以按照 K 维连续加载，在核内高效点积。INT8 权重在这里还有额外优势：加载一个 32×32 的权重 tile 只需要 1KB 带宽，是 FP16 的一半、FP32 的四分之一。省下来的带宽悄悄提高了 L2 缓存的有用数据命中率，专家数量多时这个优势会被放大到整层延迟上。
分块参数我最终锁定在 BLOCK_M=32、BLOCK_N=64、BLOCK_K=32。这个组合对中等 K 值最稳定，不会因为 BLOCK_K 太大导致共享内存开销过剩，也不会因为 BLOCK_N 太小导致每个 SM 资源利用不充分。如果遇到特别大的 K，比如 K=8192，我会把 BLOCK_K 适当调大，换取更长的卷计算流水线。
4. 性能调优的几条主线
4.1 任务粒度与负载均衡
第一版 GMM 用的是静态调度，逻辑简单，但实测下来慢得离谱，因为 32 个专家里最大的组有 256 个 token，最小的组只有 2 个 token。每个专家分配固定块后，小专家白白占着 SM，大专家又让其他 block 等很久。改成动态调度后，问题基本消失。
任务包太大，负载均衡会退化成静态调度；任务包太小，原子加的调度开销又会吃掉收益。我用 32 行作为基本单位，因为这一层正好和一个 block 的 tile 大小一致，从一个任务包里只能取出一片连续区间，不会因为跨专家操作而破坏权重读取的局部性。实测整层效率从 43% 提升到 78% 左右，这个提升基本全部来自负载均衡的改善。
4.2 小 K 场景的访存优化
有时候专家权重 K 维度并不大，但 N 维却很大，比如 K=768、N=3072。这种情况下，单次矩阵乘的 K 维度太小，低精度指令的流水线很难打满。更好的做法是同时并行多个专家的 tile，让一个 block 的流水线交替加载不同专家的数据。等待某一组权重从显存读回的时候，另一组的数据已经在路上，计算单元不会长期空转。
我实际使用时会用手工把两个连续的小专家合并到同一个 block 内交替执行，也就是所谓的组间交错。实现上多写一层索引，但效果非常明显。单独跑那个小 K 场景时，合并后的吞吐大概提升了 1.5 倍。如果你也想尝试，建议先在小规模模型上验证一下，因为组间交错会让调试变复杂，不是所有场景都值得。
4.3 融合 epilogue 与中间激活
MoE 层通常不止一个 GMM，经典的专家 FFN 结构是 up 投影加 gate 投影，再经过激活函数后进入 down 投影。如果 up 和 down 各写各的，中间就会多出一轮完整的显存读写。我在低精度 GMM 里把 gate 和 up 两个分支的矩阵乘拆成两个并行 GMM 执行，但 output 端共享同一个 epilogue：先完成两个 acc 的相加，再走激活函数，最后直接送进下一个 GMM 的输入 tile。中间激活的读写就减少了一半以上。
这里有个实践建议：不要一开始就试图把两个 GMM 完全融合成一个 kernel，除非你的 block 内并行度足够。推理引擎里更常见的做法是先保持两个低精度 GMM 内核，但让第二个内核直接从 L2 缓存读第一批输出，而不是从显存读。这样可以获得接近融合的缓存命中率，同时避免一个巨型 kernel 的调试复杂度。
4.4 实际观测数据
我在一个模拟 MoE 网络上做的基准测试结果如下：64 个专家，K=4096，N=2048，总 token 数 1024，平均每个专家分到 16 个 token，单卡推理场景。三个版本分别是独立的 FP16 矩阵乘循环、FP16 GMM、INT8 GMM，耗时如下：
| 版本 | 耗时 | 说明 |
|---|---|---|
| FP16 独立矩阵乘循环 | 约 2.9 ms | 启动开销大，带宽浪费明显 |
| FP16 GMM | 约 1.4 ms | 单次内核启动，负载均衡后提升显著 |
| INT8 GMM | 约 0.8 ms | 权重带宽减半，延迟进一步下降 |
对应到模型效果上，INT8 权重与 FP16 基线之间的路由匹配率维持在 99.3%，输出 logits 的余弦相似度在 0.996 以上。如果只是做线上推理，这个精度损失完全可接受；如果要做量化感知训练蒸馏，也还有余量继续调 scale。
5. 常见问题与调试实录
5.1 量化后输出波动大
现象是整体生成质量漂移，或者同一个 batch 内部分样本明显不正常。大概率是激活量化 scale 太小，把某些 token 的重要数值截成了边界值。解决办法是把 per-token 的 scale 计算放到内核里实时算，不要用统计出来的全局常数。同时检查是否出现异常 token，对这种 outlier 加一个 clip 阈值，避免个别极端值拉高所有人 scale 导致量化分辨率下降。
5.2 个别专家特别慢
这通常是静态调度惹的祸。把专家块拆成任务包、改成动态领取后，能解决大部分负载不均问题。如果仍有个别专家慢，看一眼是不是该专家权重不在 L2 缓存里，每次都 miss。可以在加载阶段做位置重排，把经常一起出现的专家权重放到相邻显存区域，提升读取局部性。
5.3 INT8 数值溢出
INT8 乘法的单个乘积上限是 127×127，但中间累加必须用 FP32 或 INT32，否则结果必炸。很多框架为了方便可能默认走了低精度累加，这时就会出现突发性 NaN。调试阶段一定要把累加类型单独打印出来确认。低精度算子的第一版先别急着优化性能，先把累加精度校正干净，否则后面所有性能分析都会被无效数据干扰。
5.4 低精度后性能不升反降
这种情况多半不是低精度本身的问题，而是算子没有吃到低精度红利。常见原因有两个：一是 M_e 太小，整个算子是纯访存型，低精度权重体积小了，但 block 调度依然低效；二是 BLOCK_K 太小，低精度计算单元不能连续工作。建议先做算术强度诊断，如果算出的强度低于目标设备的基本比率，就应该先去改调度和布局，而不是继续堆指令。
6. 收尾的一点经验
回到最初的问题：低精度 GMM 算子难在哪里？它不难在“用 INT8 做矩阵乘”，而是难在“让每次矩阵乘都在正确的精度尺度下运行”。从 per-channel 权重量化到 per-token 激活量化，再到动态任务调度、组间交错、融合 epilogue，每一步单独看都不复杂，但组合起来才能让这个算子真正在 MoE 稀疏计算里站住脚。
我最后想分享一个更容易被忽略的细节：凡是做算子优化，先落一个能正确跑通的低精度版本，再考虑性能。精度和性能不是前后顺序，而是在同一个版本里要一起校验的两条线。把精度基线锁定在 FP32 参考实现上，每次改一个调度策略都跑一遍全量校验，后面排查问题能省掉一半时间。这个原则，在我做这个 GMM 算子的过程中救了好几次。
