Amazon SageMaker AI 上的混合专家内核工程
aws ·
本文详细介绍了在 Amazon SageMaker AI 上针对混合专家（MoE）模型的内核工程优化，重点在于提升大型语言模型推理的性能。作者描述了三项关键技术：Swap-AB 转置，用于将矩阵维度与 NVIDIA Hopper 架构的 WGMMA 硬件要求对齐；后端路由调整，使用 Triton 而非 CUTLASS 进行 FP8 量化；以及消除不必要的 Triton JIT 重新编译以降低延迟。这些优化解决了 MoE 架构中的特定瓶颈，例如解码过程中的小批次大小和可变的 token 数量。实施这些更改后，在 H200 GPU 上实现了高达 10% 的吞吐量提升和 9% 的首个 token 生成时间（TTFT）改善。其中几项增强功能已合并到 vLLM 开源社区中。这项工作表明，深入的硬件特定调优可以显著提高稀疏激活模型的效率。
aws