NPU Kimi-K2-Thinking推理优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer11月6日Moonshot AI发布了最新的Kimi-K2-Thinking思考模型擅长深度推理打榜性能再创新高。其模型结构和Kimi-K2保持一致原生支持MOE W4A16 Attention BF16混合精度模式兼顾推理性能与效果。本文主要介绍基于A3集群的Kimi-K2-Thinking模型的推理优化由于其结构与Kimi-K2基本一致基础优化点可平滑继承。针对该模型涉及到的MOE W4A16量化本实践已0day适配支持并提供了A3集群的推荐部署策略未来将持续进行性能优化。HighlightsCANN已0day支持Kimi-K2-Thinking 256K序列推理部署模型推理代码已开源同时也适配了主流开源推理框架SGLang本实践支持Kimi-K2-Thinking原生量化模式MOE GroupedMatmul采用W4A16计算Attention保留BF16精度模式兼顾推理性能与效果配套的GroupedMatmul算子实现已全面开源针对小batch size、长序列生产等典型推理场景Decode MLA已实现Flash Decode加速机制有效降低时延并提升算力利用率配套的FusedInferAttentionScore算子实现已全面开源本实践的最小部署单元为单机可支持4K序列长度的推理任务同时支持多机大EP部署模式通过多机集群协同部署可满足256K超长序列的推理需求Outline并行策略融合KernelMTPFuture Plan并行策略Atlas A3推荐部署策略如下图所示Prefill使用M个节点部署Decode使用N个节点部署每个节点包含8卡推荐根据资源数量、SLA等约束M和N在1~24内动态调整。Prefill并行策略考虑到长序列场景Prefill Attention选用Context Parallel(CP)并行多个rank均摊长序列的计算单rank的计算量和activation内存都较小TTFT较为可控用户体验更好。MoE模块则沿用DeepSeek-V3.1的EP并行兼顾吞吐与时延。Decode并行策略Decode阶段依旧沿用DeepSeek V3.1的部署策略选用Attention DP MoE EP部署。特别地由于O_proj和LM_Head权重内存较大且在Decode阶段表现为明显的访存瓶颈本实践选用局部TP并行。同时为了降低设备内存占用Embedding层同样使用TP切分。为了尽可能地减小TP并行带来的通信开销TP域控制在单机内。融合Kernel整网计算流如下图所示本实践使用了MLAProlog、FusedInferAttentionScore、MoeDistributeDispatch、MoeDistributeCombine、GroupedMatmul等融合Kernel可供其他同类模型在昇腾平台的高效落地提供实践参考。MTP相较于DeepSeek V3.1Kimi-K2-Thinking的Attention Head数量从128精简至64使得Attention计算量显著降低更利好MTP场景。LLM推理Decode阶段通常为访存瓶颈MTP可通过少量计算代价来缓解访存压力而 Kimi-K2-Thinking的Attention本身计算负载更轻因此在 MTP模式下更难触及计算瓶颈性能更优。本实践对应的代码已支持多头 MTP 功能若具备训练完成的MTP权重可直接启用该功能以实现推理加速。Future Plan量化目前支持MOE W4A16推理未来可针对Attention开发低比特量化版本进一步降低系统内存占用与时延长序列性能优化长序列场景下的TTFT仍有优化空间后续将探索更亲和的部署方案并针对性优化融合 Kernel 性能全面提升长序列推理的响应速度与运行效率【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考