SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations
本文针对 MoE 模型推理中的两大瓶颈——细粒度专家带来的激活显存膨胀与 IO 开销增大,以及稀疏专家导致 Grouped GEMM 中的 padding 浪费——提出 SonicMoE 优化方案。核心技术包括:(1) 内存高效的前向/反向计算算法,最小化激活缓存;(2) GPU 内核级 IO 与计算重叠优化,适用于所有 MoE 架构;(3) 新颖的 token rounding 策略减少 padding 浪费。该工作来自 FlashAttention 作者 Tri Dao 团队,具有很强的工程实践价值。
INSIGHTSonicMoE 的核心创新在于从 GPU 内核层面解决 MoE 的效率问题。当前 MoE 趋势是更多更小的专家(细粒度)和更高稀疏度,这虽然提升模型质量但严重降低硬件利用率。SonicMoE 通过 IO-computation overlap 内核设计和 tile-aware token rounding 直接在算子层面提升效率。对搜索/推荐系统部署 MoE 模型的启示:选择推理框架时应关注其 MoE 内核实现质量,细粒度 MoE 的性能瓶颈往往不在计算量而在 IO。