Study note

Prime Flash MoE: Faster MoE Kernels Optimized for Blackwell

Properties

Type: Blogs
Status: 待读
Domain: AI / ML
Category: LLM 推理与服务
Source: primeintellect.ai
Vault note: library/articles/ai_ml/Prime-Flash-MoE-Faster-MoE-Kernels-Optimized-for-Blackwell-65056e3198917319.md

Summary

围绕路由完成后的 MoE SwiGLU feed-forward，设计 fused 与 split 两套 Blackwell kernel。文章讲解如何用 TMA tensor map 交错 gate/up 权重并 gather token，利用 TMEM、`tcgen05`、共享内存复用、stage barrier 和流水线完成两次 GEMM，同时减少中间激活写回 HBM；还覆盖 split-K partial reduction、PDL 交接和 MXFP8 scale/quantization 路径。

Highlights

文章没有把“融合越多越快”当作定律，而是定量比较省掉激活 materialization 与增加 split-K partial-output reduction 的代价，并据工作集大小选择 fused 或 split pipeline。作者在 B200 的给定 benchmark 中报告相对 PyTorch grouped GEMM 最高约 2.4 倍加速，适合完成 TMA、CuTe/CUTLASS、Blackwell Tensor Core 和 MoE 基础后精读。

Notes

面向 B200 的 BF16/MXFP8 MoE 前馈 CUDA kernel 深度解析。

Comments

Loading comments...