紧接着昨天那篇PyTorch Blog的内容MetaShuffling：Meta的Fused MoE kernel工程方案，更激进的Kernel优化和尽量避免Padding，我把fbgemm开源的moe grouped gemm kernel(https://github.com/pytorch/FB...)拷贝了一下，fix了一个小bug，然后在H100(Hopper)和SGLang的Grouped GEMM Triton Kernel对比了一下正确性和性能，在正确性没问题的情况下，性能可以提升挺多的。细节放在这里了：https://github.com/sgl-projec... 。结论就是fbgemm可以在MoE模型上相比于SGLang的grouped gemm实现较大的性能提升，这个kernel可以在fp16/bf16和fp8 per-tensor quant的情况下直接应用到sglang的ep-moe的grouped gemm kernel中提升性能。不过TP模式下的Triton Fused MoE不是最直接的Grouped GEMM，需要有小伙伴按照这个Triton优化的技巧去改kernel才可以。
FBGEMM GroupedGEMM 基准测试结果
当使用 triton==3.2.0 运行基准测试时,会出现以下警告:我们无法使用 warp专用化, 但持久化kernel和 TMA load/store 仍然可用。
我把fbgemm开源的moe grouped gemm kernel(https://github.com/pytorch/FBGEMM/tree/main/fbgemm_gpu/experimental/gen_ai)拷贝了一下，fix了一个小bug，然后在H100(Hopper)和SGLang的Grouped GEMM Triton Kernel对比了一下正确性和性能，在正确性没问题的情况下，性能可以提升挺多的。细节放在这里了：https://github.com/sgl-project/sglang/pull/6924 。结论就是fbgemm可以在MoE模型上相比于SGLang的grouped gemm实现较大的性能提升，这个kernel可以在fp16/bf16和fp8 per-tensor quant的情况下直接应用到sglang的ep-moe的grouped gemm kernel中提升性能。不过TP模式下的Triton Fused MoE不是最直接的Grouped GEMM，需要有小伙伴按照这个Triton优化的技巧去改kernel才可以。Qwen2-57B-A14B-Instruct BF16 W8A8 TP4
python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4
grouped-gemm-performance:
batch_size FBGEMM Grouped GEMM BF16 SGLang Grouped GEMM BF16
0 1.0 0.032352 0.022272
1 2.0 0.032096 0.022080
2 4.0 0.032640 0.021984
3 8.0 0.031840 0.021472
4 16.0 0.030832 0.021536
5 32.0 0.032192 0.021632
6 64.0 0.393504 0.595008
7 128.0 0.393872 0.598048
8 256.0 0.394848 0.589760
9 512.0 0.397488 0.605888
10 1024.0 0.401248 0.581952
11 2048.0 0.407232 0.559232
12 4096.0 0.416368 0.717936Qwen2-57B-A14B-Instruct FP8 W8A8 TP4
python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4
grouped-gemm-performance:
batch_size FBGEMM Grouped GEMM BF16 SGLang Grouped GEMM BF16
0 1.0 0.032352 0.022272
1 2.0 0.032096 0.022080
2 4.0 0.032640 0.021984
3 8.0 0.031840 0.021472
4 16.0 0.030832 0.021536
5 32.0 0.032192 0.021632
6 64.0 0.393504 0.595008
7 128.0 0.393872 0.598048
8 256.0 0.394848 0.589760
9 512.0 0.397488 0.605888
10 1024.0 0.401248 0.581952
11 2048.0 0.407232 0.559232
12 4096.0 0.416368 0.717936meta-llama/Llama-4-Scout-17B-16E-Instruct FP16 TP8
python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4
grouped-gemm-performance:
batch_size FBGEMM Grouped GEMM BF16 SGLang Grouped GEMM BF16
0 1.0 0.032352 0.022272
1 2.0 0.032096 0.022080
2 4.0 0.032640 0.021984
3 8.0 0.031840 0.021472
4 16.0 0.030832 0.021536
5 32.0 0.032192 0.021632
6 64.0 0.393504 0.595008
7 128.0 0.393872 0.598048
8 256.0 0.394848 0.589760
9 512.0 0.397488 0.605888
10 1024.0 0.401248 0.581952
11 2048.0 0.407232 0.559232
12 4096.0 0.416368 0.717936meta-llama/Llama-4-Scout-17B-16E-Instruct FP8 TP8
python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4
grouped-gemm-performance:
batch_size FBGEMM Grouped GEMM BF16 SGLang Grouped GEMM BF16
0 1.0 0.032352 0.022272
1 2.0 0.032096 0.022080
2 4.0 0.032640 0.021984
3 8.0 0.031840 0.021472
4 16.0 0.030832 0.021536
5 32.0 0.032192 0.021632
6 64.0 0.393504 0.595008
7 128.0 0.393872 0.598048
8 256.0 0.394848 0.589760
9 512.0 0.397488 0.605888
10 1024.0 0.401248 0.581952
11 2048.0 0.407232 0.559232
12 4096.0 0.416368 0.717936结论是FBGEMM相比SGLang的分组GEMM实现在MoE模型上可以实现显著的性能提升。这个kernel可以直接应用到SGLang的EP-MoE分组GEMM kernel中,以提升fp16/bf16和per-tensor量化fp8条件下的性能。
局限性
当前的局限性在于如果不编译Meta特定的Triton版本,warp专用kernel似乎无法使用。此外,这个kernel目前仅支持fp16/bf16和per-tensor量化的fp8w8a8 - 要与DeepSeek兼容还需要进一步修改。
END
作者：BBuf
来源：GiantPandaLLM
推荐阅读
- Soft 最近邻插值 SNI 对齐特征 +ESD 增强空间保留，效率精度双突破
- 图解Vllm V1系列4：加载模型权重(load_model)
- 减少 55% 代码行数、 CPU 和 GPU 峰值内存，加速训练推理！
- 用大模型实现 VHDL 代码解释、生成与优化，加速高性能微处理器设计！
- MetaShuffling：Meta 激进优化 Fused MoE kernel，规避Padding
欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
