MoE 之年的总结和 MoE 推理优化的一些认识

0x0. 前言

首先我会简单回顾下 2024 年的学习收获，然后我会聊一聊我在 SGLang 中度过的几个月业余开源开发的经历。无论是最近火遍全球的 DeepSeek V3/R1，还是在 2024 年各家厂商发布的重量级 MoE 模型都将 MoE 架构重新带上了历史舞台，所以我个人把 2024 年定义为 MoE 之年。因此，最后我会讨论一下我对 MoE 模型中的 Fused MoE 核心算子的推理优化的理解，我会对比目前的一些业界开源方案以及更有前景的解法。

0x1. 学习和收获

在 2024 年，我的前 3 个季度是很平淡的，就是普通的写代码完成任务以及被各种新的模型发布消息轰炸，导致对几乎每周发布的新模型都无感了。这段时间我个人最大的收获应该是熟练使用了 Cursor 帮我完成各种重复的简单工作。

大概 9 月的时候我开始以做工程的角度切入 LLM 推理框架，首先研究的就是 VLLM 和 SGLang，看这些框架本身还是比较痛苦，因为感觉没有触发感兴趣的点。直到后面接触到 MoE 模型中的 Fused MoE 算子实现调研优化，我才找到了我这一年以来的兴趣点。后面也从这一点出发成为了 SGLang 的开发者，并且在 SGLang 的开源贡献过程中收获了很多，不仅是技术，也包括从 SGLang Team 成员获得的情绪价值。认识了 @Lianmin Zheng @Yineng Zhang @Chenyang Zhao 等 SGLang Team 成员，特别是 @Yineng Zhang 在我做开源贡献时提供的专业技术帮助和情绪价值(hhh。我会在下面的 0x2 节聊一下我这几个月参与 SGLang 框架的从零开始的开源开发经历。

SGLang 贡献图，和大佬合影

此外，我也继续坚持在 GiantPandaCV 公众号创作了一年，原创文章大概保持了 1 个月 2 篇左右，频率这么低的原因是因为自己半开摆了。然后做了一些 PyTorch Blog 以及 CUDA MODE 课程的笔记，算是业余充电补充知识了（CUDA-MODE 学习笔记我也是算子那个一个月 2 篇里面凑数的，看得出来是真的有点摆了，今年会在 GiantPandaCV 尽量多更一点笔记）。

我的 github 的两个学习笔记仓库更新频率在 2024 后半年降到了很低，但是 star 数却一直在涨，感谢关注我的网友们。我也争取新年继续更新一下，特别是优质博客链接应该多更新一下。

0x2. SGlang 几个月的开源开发经历

在推理框架这块的背景下，我个人对算子开发和模型 profile 的这一套逻辑稍微熟悉一点，然后我之前是没有大模型推理框架开发经验的。我这里先聊一下，我是如何给 SGLang 做开源贡献的，大概是 9 月中旬开始参与。我选定的就是我熟悉的方向即算子和算子/模型性能优化。

首先是做了一些基础的 bug 修复，添加 fused moe 测试相关的工作，以及完善了一下 fused moe triton 算子的 auto tuning 脚本等，还有就是讨论了一下在低端 GPU 上如何恰当的使用 chunked prefill 的问题：

接着做的贡献来自于我 profile 模型的一个发现，我发现 VLLM 和 SGLang 在离线推理一个相同的 qwen2-7b 模型时通过 Nsight System 拿到的结果显示 sglang 的 rmsnorm 是明显更快的。所以我在 https://github.com/sgl-projec... 这里基于 Triton Benchmark 写了一个 rmsnorm kernel 的 micro benchmark 对比脚本，可以证明在各个情况下 SGLang 使用的 FlashInfer 的 rmsnorm 总是比 vllm 的实现更快。我这个发现和 benchmark 脚本后续也被 vllm 采用了：https://github.com/vllm-proje... 。在此基础上继续做了一下 write_req_to_token_pool_triton 这个算子的 benchmark 和初步优化，然后就是在 GTX 4090, h100 和 h200 上针对目前常用的 MoE 模型例如 Mixtral 8x7B/8x22B 和 Qwen2-57B-A14B 等跑了跑 fused moe triton 算子 Tuning，都是一些比较简单的零碎开源工作。

再之后就是最近这一个月了，参与到了 DeepSeek V3 的优化中，提升了 Fused MoE 模块中 moe_align_kernel 这个 cuda 算子的性能。https://github.com/sgl-projec... & https://github.com/sgl-projec... ，不过 kernel 仍有优化空间。我写的这个 kernel 也被 vllm 采用了：https://github.com/vllm-proje... （笑

接着就是针对和 DeepSeek V3 同期发布的 MiniMax-01-Text 模型的 Lightning Attention 的 Prefill 和 Decode 做了 Benchmark 和正确性测试，以及针对 Lightning Attention 的 Decode 做了几次 Triton 和 CUDA 优化，具体可以看使用 NCU 和 Cursor Claude-sonnet-3.5 写出高效 cuda 算子的正确姿势，同时也参与到了 sgl-kernel 库的维护和 review。

最后，在春节假期的时候，也就是最近一周，review 一个 contributor 提出的针对 deepseek V3 的 fused gate 优化中提出把 TensorRT-LLM 的 Fused MoE 算子单独拆分以供 SGLang 来定制和调用，https://github.com/sgl-projec... 。然后我就负责了这个工作，踩了几个大坑把 TensorRT-LLM 的 Fused MoE 算子单独拆分并跑通了。这个工作也是为了给 DeepSeek V3 的优化做进一步铺垫。初步的代码在这里：https://github.com/BBuf/tenso... ，这里提一下我踩的一些坑：

我是在 H100 上从 TensorRT-LLM 的最新分支把 FusedMoE 单独移出来，所以要跑通 Hopper 下的 Cutlass MoE GEMM 需要一些特殊的编译选项和架构，例如使用 sm90_a 而不是 sm_90 架构，并且需要打开 Hopper 下的一些 NVCC Flags。这段代码配置在这里，调试了很久：https://github.com/BBuf/tenso...
FusedMoE 涉及到的代码写得算是比较独立，只涉及到 common 和 cutlass MoE GEMM 两个比较独立的模块，所以代码本身是比较好抽取的，需要按照源码的目录放到固定位置，然后对于cpp/tensorrt_llm/kernels/mixtureOfExperts下面的 MoE 核心实现，由于不需要在这里考虑 LORA 所以逐步去手动删掉了和 LORA 实现相关的代码，让代码更加 clean。这里有一个坑是这些文件中 data type 使用了 TensorRT 中的 nvinfer1 namespace 下的 dtype，所以需要我们在自己的 docker 中apt install tensort然后把 TensorRT 的 include 和 lib 放到正确位置。https://github.com/BBuf/tenso...
由于 Fused MoE 的 Cutlass Grouped GEMM 依赖固定版本的 CutLass，所以需要使用 submodule 的方式引一个 CutLass 到 3rdparty 中：https://github.com/BBuf/tenso...
接下来的一个坑是需要运行 MoE Grouped GEMM 的实例化 Python 脚本来为各种 dtype 和各种 Tile Config 的 Grouped GEMM 进行实例化，否则运行的时候会提示 GEMM 的符号找不到。Python 脚本也是 copy 过来跑就可以：https://github.com/BBuf/tenso... ，需要注意的一点是需要把 PYTHONPATH 设置到https://github.com/NVIDIA/cut...目录下才可以正常运行。注意，这个脚本生生的.cu也需要加到setup.py的source才可以。
注意到以上问题之后基本就可以编译通过这个移植的程序了，接着就是参考一下TensorRT-LLM/cpp/tests/kernels/mixtureOfExpertsTest.cu调用 CutlassMoeFCRunner 的方式来编写一个接口并写一下测试代码来测试了。https://github.com/BBuf/tenso... & https://github.com/BBuf/tenso...
在测试过程中结果对不上，debug 了数天，第一个坑是 renormalize 的参数没有正确传递，https://github.com/BBuf/tenso... 这里应该传递tensorrt_llm::kernels::MOEExpertScaleNormalizationMode::RENORMALIZE,但是我根据测试程序来写传递的是 None，导致 topk softmax 之后的结果没对应上，这个问题通过程序对拍了 2 半天解决了。
接下来的一个问题是使用 PyTorch 写的测试程序的 expert 的权重默认是 Row Major 的，但是这个 Cutlass GEMM 的版本要求的权重是 Col Major 的，debug 了好几天，一直在纠结为什么 gemm 的输入数据和其它超参数完全一样但是输出就不对，直到和 @Yineng Zhang 交流了几次才想到有可能是数据的 Stride 可能不对，结果真的是踩了这个坑，修复之后 fp32 精度下的测试就过掉了。https://github.com/BBuf/tenso...
最后一个坑是在 BF16 的时候发现测试没过，定位到原因是 TensorRT LLM 会把 gating 先 cast 成 fp32 再做 topk softmax，我的测试程序里面的这里有点问题，修复之后就可以了。https://github.com/BBuf/tenso...

踩这些坑花了几乎一个假期，中途一度颓废开摆，不过还好在假期结束前 debug 出来了。在这之后应该会基于这个版本做一些接口封装或者 feature 定制化工作。

0x3. Fused MoE 推理方案盘点

这一节我来盘点一下 2024 年我们经历的 Fused MoE 推理算子的演进，或者说各家开源推理框架的方案。

VLLM/SGLang Triton Fused MoE

这是 VLLM/SGLang 开源之初支持 MoE 模型的方案，直接使用了 AnyScale 开发的 Triton Fused MoE 算子，后面又在这个算子上做了一些优化包括开发 Tuning 策略，FP8 PerTensor, FP8 Blockwise，INT8 Fused MoE 等等。目前处于功能丰富，但限制也明显，毕竟用 Triton 写的，相比于 CutLass 实现的 Grouped GEMM 会有性能上的削弱，并且由于 Triton Grouped GEMM 要求把每个 expert 的 token 数 padding 到矩阵 Tile config['M']的倍数引入的moe_align_block_size kernel 也是一个开销无法忽略的算子，另外 Triton Grouped GEMM1 和 silu 等激活函数无法融合也持续增大了整个算子的 overhead。关于这个 Triton 算子实现，硅基流动的 zhuping 在 Triton 中国做分享的时候做过一个 Slides 总结得比较好，我把这几页截图一下：

以 Token 为中心还是以 Expert 为中心

上面最后两张 Slides 提到我们即可以先确定 Experts，也可以先确定 Tokens。VLLM/SGlang 目前的版本就是先确定 Experts 然后再确定 Tokens。LmDeploy 中提供了一个先确定 Tokens 再确定 Expert 的 Fused MoE 算子，https://github.com/InternLM/l... 。不过缺少 Auto Tuning，量化相关的各种支持，感兴趣的读者可以考虑自行探索一下。

SplitK Grouped GEMM Fused MoE

来自：https://github.com/pytorch-la...

我在 SGLang 中尝试过这个算子，为其新增了 FP8 的支持和 Auto Tuning，但是在 Qwen2-57B TP4 w8a8 情况下没有看到明显的收益，如下图所示。https://github.com/sgl-projec...

由于 SplitK GEMM 会引入tl.atomic_add，而 Triton 的atomic_add不支持 BF16，如果要求用 dtype 为 bf16 来推理模型这里就会报错，这也是一个限制。

SGLang Expert Parallel Fused MoE

这个我单独写过了一篇博客：SGLang 的 Expert Parallel 特性解读

文章的最后我也提到了它目前的优势和限制。

TensorRT-LLM Fused MoE

上面提到的 Fused MoE 变种实现全都是基于 Triton 的，限制也比较多。TensorRT-LLM 也开源了一个 Fused MoE 的实现，是纯 CUDA 写的，它不仅支持 gating 的融合，也支持把 GEMM1 后面的 Activation 通过 Cutlass Epilogue fuse 起来，它的 grouped gemm 也是通过 cutlass 来定制的，并且有特定的 tuning 策略。此外，它的实现也可以完美对接 Expert Parallel，这应该是当前开源版本中最优的一个实现。然后由于这个实现由 c++编写并且嵌入到 TensorRT-LLM 内部，外部如果想使用目前比较困难。在上面提到了我们正在做一个移植工作将其独立为一个单独调用的库，并且初步成功并验证了正确性，相信之后在开源社区的努力下可以更广泛的使用这个实现。请关注 SGLang Project 的进展。