CANN 平台 Kimi-K2-Thinking 推理优化实践W4A16 混合精度、PD 分离并行策略与融合 Kernel 全解析【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本篇技术指南聚焦 CANN 开源仓库中 Kimi-K2-Thinking 模型的端到端推理优化实践模型原生支持 MOE W4A16 Attention BF16 混合精度代码在 CANN 上 0day 适配并支持 256K 超长序列部署。文章将完整讲解 Atlas A3 集群的 PD 分离部署策略Prefill CP EP、Decode Attention DP MoE EP 局部 TP、MLA Prolog / FusedInferAttentionScore / GroupedMatmul 等融合 Kernel 的落地方式、多头 MTP 加速原理并结合仓库中真实的 YAML 配置与推理脚本给出可复现的部署执行方案。背景Kimi-K2-Thinking 与 CANN 的 0day 适配Kimi-K2-Thinking 是 Moonshot AI 发布的深度推理思考模型其模型结构与 Kimi-K2 保持一致原生支持MOE W4A16 Attention BF16的混合精度模式MoE 专家模块的 GroupedMatmul 以 W4A16 量化计算Attention 保留 BF16 精度兼顾推理性能与效果。由于结构与 Kimi-K2 基本一致DeepSeek-R1 系列沉淀的基础优化点MLA 低秩压缩融合算子、EP 并行、MLP 线性层合并、torchair 图模式、SuperKernel 等可以平滑继承到本模型。围绕该模型特有的 MOE W4A16 量化本实践已 0day 完成适配并给出了 Atlas A3 集群的推荐部署策略。实践要点Highlights如下0day 序列能力CANN 已 0day 支持 Kimi-K2-Thinking 256K 序列推理部署模型推理代码已开源同时适配了主流开源推理框架 SGLang。原生量化模式支持 Kimi-K2-Thinking 原生量化MOE GroupedMatmul 采用 W4A16 计算、Attention 保留 BF16配套的 GroupedMatmul 算子实现已全面开源。Flash Decode 加速针对小 batch size、长序列生产等典型推理场景Decode 阶段 MLA 已实现 Flash Decode 加速机制有效降低时延并提升算力利用率配套的 FusedInferAttentionScore 算子实现已全面开源。部署规模弹性最小部署单元为单机8 卡可支持 4K 序列长度的推理任务同时支持多机大 EP 部署模式通过多机集群协同部署满足 256K 超长序列的推理需求。并行策略Atlas A3 集群 PD 分离部署在 Atlas A3 集群上推荐采用Prefill / Decode 分离PD 分离的部署策略Prefill 使用 M 个节点部署Decode 使用 N 个节点部署每个节点包含 8 卡。推荐根据资源数量、SLA 等约束将 M 和 N 在 1~24 范围内动态调整实现 Prefill 吞吐与 Decode 时延的独立扩缩容。对应到仓库中的实际配置models/kimi_k2_thinking/config/pd/目录下分别提供了 prefill.yaml 与 decode.yaml 两套 PD 分离配置并在在线推理模式下通过--pd-role指定角色加载对应 YAML。Prefill 并行策略CP EP 双管齐下考虑到长序列场景Prefill 阶段的 Attention 选用Context ParallelCP并行多个 rank 均摊长序列的计算单 rank 的计算量和 activation 内存都较小TTFT 较为可控用户体验更好。MoE 模块则沿用 DeepSeek-V3.1 的Expert ParallelEP并行兼顾吞吐与时延。在 prefill.yaml 中cp_size: 32表示 32 个 rank 全量参与 Context Parallel 切分world_size: 32attn_tp_size: 1表示 Attention 不做 Tensor Parallelmoe_tp_size: 1表示 MoE 采用纯 EP 而非 TP同时o_proj_tp_size: 1、dense_tp_size: 1表明 Prefill 阶段不启用局部 TP。Decode 并行策略Attention DP MoE EP 局部 TPDecode 阶段依旧沿用 DeepSeek-V3.1 的部署策略选用Attention DP MoE EP部署。特别地由于O_proj 和 LM_Head 权重内存较大且在 Decode 阶段表现为明显的访存瓶颈本实践对这两部分选用局部 TP 并行同时为了降低设备内存占用Embedding 层同样使用 TP 切分。为了尽可能地减小 TP 并行带来的通信开销TP 域被控制在单机内。在 decode.yaml 中可以看到这种差异化切分的具体体现parallel_config: world_size: 32 cp_size: 1 # Decode 阶段不做 Context Parallel attn_tp_size: 1 # Attention 使用 DP o_proj_tp_size: 8 # O_proj 局部 TP切 8 份 dense_tp_size: 4 # Dense 部分局部 TP moe_tp_size: 1 # MoE 纯 EP embed_tp_size: 16 # Embedding TP 切 16 份降低设备内存 lmhead_tp_size: 16 # LM Head TP 切 16 份缓解访存瓶颈除 PD 分离配置外仓库还提供了单 YAML 的离线/统一配置 kimi_k2_thinking.yamlcp_size: 32、o_proj_tp_size: 8、dense_tp_size: 4可按集群规模选择对应的 YAML 文件。融合 Kernel整网计算流的关键加速器整网计算流如下图所示本实践使用了MLAProlog、FusedInferAttentionScore、MoeDistributeDispatch、MoeDistributeCombine、GroupedMatmul等融合 Kernel可为其他同类模型在昇腾平台的高效落地提供实践参考各融合 Kernel 的作用与对应实现可归纳如下与 DeepSeek-R1 基础优化点一脉相承融合 Kernel作用使能方式MLAProlog融合 Attention 前置计算Q/K/V 线性层、RoPE、RMSNorm 及 KV Cache 更新enable_mla_prolog: trueFusedInferAttentionScore融合 MLA 注意力计算Decode 阶段实现 Flash Decode 加速配合 MLA Prolog 使用MoeDistributeDispatch / CombineEP 并行下多卡间的高性能通信路由dispatch 与 combineMoE EP 并行默认使能GroupedMatmul一次调用同时处理多个专家含 W4A16 量化的矩阵计算提高计算与搬运效率MoE 层默认使能MLA Prolognpu_mla_prolog_v3在源码 modeling_deepseek.py 的 Attention 实现中通过构造mla_prolog_input_args并调用torch_npu.npu_mla_prolog_v3完成 Q、K、V 低秩投影、RoPE、KV Cache 更新等前置计算的一体化融合减少算子下发次数。它由 YAML 的custom_params.enable_mla_prolog控制默认false启用后可显著提高吞吐。Flash Decode针对小 batch、长序列场景FusedInferAttentionScore 的 Flash Decode 机制降低了 Decode 时延并提升算力利用率。GroupedMatmulW4A16MoE 专家计算采用 W4A16 量化模式权重以 INT4 存储、计算时反量化为 FP16/BF16配合 GroupedMatmul 批量处理多专家是兼顾显存占用与精度的核心手段。MoE 多流并行custom_params.enable_multi_streams默认false在图模式下启用 MoE 共享专家的多流并行计算进一步提升吞吐源码中通过stream-fusion1等 option 控制流融合行为。除上述特性外模型配置还支持enable_static_kernelnpugraph_ex模式下将算子编译为静态 kernel、enable_weight_nz权重 NZ 格式等开关详见后文 YAML 配置章节。MTP多头多 token 预测加速相较于 DeepSeek-V3.1Kimi-K2-Thinking 的Attention Head 数量从 128 精简至 64使得 Attention 计算量显著降低更利好 MTPMulti-Token Prediction场景。LLM 推理 Decode 阶段通常为访存瓶颈MTP 可通过少量计算代价来缓解访存压力而 Kimi-K2-Thinking 的 Attention 本身计算负载更轻因此在 MTP 模式下更难触及计算瓶颈性能更优。本实践对应的代码已支持多头 MTP功能若具备训练完成的 MTP 权重可直接启用该功能以实现推理加速在 YAML 的model_config中设置next_n参数使能 MTP如next_n: 1、next_n: 2源码中支持next_n取0/1/2/3超过 3 会报错。源码 modeling_deepseek.py 中实现了DeepseekV3ModelMTPLayer、DeepseekV3ModelMTP等 MTP 专用模块MTP 头位于主模型最后一层之后model.layers.num_hidden_layers i并有独立的mtp_unique_weight_mapping权重映射逻辑负责加载 MTP 头权重。工程落地环境准备、YAML 配置与推理执行环境准备根据 模型 README样例运行环境要求如下CANN 软件包依赖 CANN 开发套件包与 CANN 二进制算子包支持版本为CANN 9.0.0Ascend-cann-toolkit_${version}_linux-${arch}.run与Ascend-cann-A3-ops_${version}_linux-${arch}.run。torch_npu版本v26.0.0PyTorch 版本2.8.0安装torch_npu-2.8.0.post4-...whl。Python仅支持 python 3.11。下载项目源码并安装依赖依赖清单见 requirements.txt包含 torch 2.8.0、transformers 4.53.3、datasets 3.6.0、compressed-tensors 0.6.0、accelerate 1.0.1 等git clone https://gitcode.com/cann/cann-recipes-infer.git cd cann-recipes-infer pip3 install -r ./models/kimi_k2_thinking/requirements.txt随后修改 executor/scripts/set_env.sh 中的IPs所有节点 IP按 rank id 排序空格分隔与cann_pathCANN 软件包安装路径如/usr/local/Ascend/ascend-toolkit/latest。HCCL 相关配置如HCCL_SOCKET_IFNAME、HCCL_OP_EXPANSION_MODE可参考集合通信文档并在 executor/scripts/function.sh 中自定义其中HCCL_OP_EXPANSION_MODEAIV可借助 Vector Core 加速通信执行。权重方面请下载 Kimi-K2-Thinking 原始权重并上传到 Atlas A3 各节点固定路径如/data/models/Kimi-K2-Thinking并在 YAML 的model_path中指定。YAML 配置详解以 kimi_k2_thinking.yaml 为例配置分为model_config、data_config、parallel_config、scheduler_config四部分统一参数说明见 YAML 参数描述。model_config: model_name: kimi_k2_thinking model_path: /data/models/Kimi-K2-Thinking with_ckpt: True exe_mode: npugraph_ex # [eager, npugraph_ex, ge_graph] next_n: 0 enable_profiler: False force_eplb: False enable_static_kernel: True # npugraph_ex only: compile operators to static kernels enable_cache_compile: False enable_weight_nz: True custom_params: enable_mla_prolog: True enable_multi_streams: True enable_superkernel: False # ge_graph onlyexe_mode执行模式支持eager单算子下发、npugraph_exNPU Graph 图模式本样例默认与ge_graph对应 torchair max-autotune 图模式。next_nMTP 头数量取0/1/2/30表示关闭 MTP。force_eplb是否强制专家负载均衡Expert Parallel Load Balance。enable_weight_nz权重是否以 NZ 格式存储。除框架统一配置外模型还额外支持custom_params字段下的特性开关参数说明如下参数名类型默认值含义enable_mla_prologboolfalse启用 MLA prolog 优化用融合算子处理 Attention 的前置计算提高吞吐。enable_multi_streamsboolfalse在图模式下启用 MoE 共享专家多流并行计算提升吞吐。enable_superkernelboolfalse在ge_graph模式下启用 superkernel 加速将多个算子融合为大核以提高执行效率npugraph_ex模式不支持。说明若kv_cache_quant_mode为 C8INT8 KV Cache 量化则enable_mla_prolog必须置为true源码中有对应校验。并行与调度配置parallel_config: world_size: 32 cp_size: 32 attn_tp_size: 1 o_proj_tp_size: 8 dense_tp_size: 4 moe_tp_size: 1 embed_tp_size: 16 lmhead_tp_size: 16 scheduler_config: block_size: 128 max_prefill_tokens: 65536 max_new_tokens: 256 batch_size: 32world_size参与推理的总卡数32 卡对应 4 节点 × 8 卡。cp_sizeContext Parallel 切分大小Prefill 配置为 32Decode 配置为 1。attn_tp_size/o_proj_tp_size/dense_tp_size/moe_tp_size/embed_tp_size/lmhead_tp_size分别控制 Attention、O_proj、Dense、MoE、Embedding、LM Head 的 TP 切分数其中 Decode 阶段通过o_proj_tp_size: 8、embed_tp_size: 16、lmhead_tp_size: 16缓解 O_proj、Embedding 与 LM_Head 的访存瓶颈。block_sizeKV Cache 块大小128 token/块max_prefill_tokens单次 Prefill 最大 token 数max_new_tokens最大生成 token 数batch_size批大小。Decode 配置中还包含num_reserved_decode_tokens: 64用于为 Decode 阶段预留 KV Cache token 空间。推理执行统一入口脚本位于 executor/scripts/infer.sh支持以下参数参数含义取值示例--model模型目录名对应models/下的子目录kimi_k2_thinking--mode推理模式offline离线推理/online在线 PD 分离推理--yaml离线模式yaml 文件名kimi_k2_thinking.yaml--pd-role在线模式PD 部署角色prefill/decode--p-yaml-name可选在线模式prefill yaml 文件名不传则默认pd/prefill.yamlpd/prefill.yaml--d-yaml-name可选在线模式decode yaml 文件名不传则默认pd/decode.yamlpd/decode.yaml使用方式一命令行传参# offline 模式 bash executor/scripts/infer.sh --model kimi_k2_thinking --yaml kimi_k2_thinking.yaml # online 模式 bash executor/scripts/infer.sh --model kimi_k2_thinking --mode online --pd-role prefill # online 模式指定 prefill/decode yaml bash executor/scripts/infer.sh --model kimi_k2_thinking --mode online --pd-role prefill --p-yaml-name pd/prefill.yaml --d-yaml-name pd/decode.yaml如需查看参数说明可以执行bash executor/scripts/infer.sh --help在线模式 IP 等更多配置参见 executor 设计文档 §5.1 启动方式。使用方式二直接修改脚本默认值后执行编辑 executor/scripts/infer.sh修改MODEL/MODE/YAML_FILE/PD_ROLE/P_YAML_NAME/D_YAML_NAME等参数默认值例如MODELkimi_k2_thinking MODEoffline YAML_FILEkimi_k2_thinking.yaml保存后直接执行bash executor/scripts/infer.sh。输入数据准备样例默认使用 InfiniteBench 数据集longbook_qa_eng进行长序列推理需将其上传到各节点dataset/InfiniteBench目录若需使用内置或自定义 prompt可将 YAML 中dataset改为default并在 dataset/default_prompt.json 中自定义输入。多机环境下需在每个节点上分别执行推理脚本推理日志和结果保存在models/kimi_k2_thinking/res/路径下。未来规划Future Plan量化演进目前支持 MOE W4A16 推理未来可针对 Attention 开发低比特量化版本进一步降低系统内存占用与时延。长序列性能优化长序列场景下的 TTFT 仍有优化空间后续将探索更亲和的部署方案并针对性优化融合 Kernel 性能全面提升长序列推理的响应速度与运行效率。小结Kimi-K2-Thinking 推理样例在 CANN 上完整展示了原生混合精度适配 PD 分离并行 融合 Kernel MTP 加速的组合式优化方法论Prefill 侧以 CP 摊薄长序列计算、以 EP 保障吞吐Decode 侧以 Attention DP MoE EP 局部 TP 化解 O_proj/LM_Head 访存瓶颈W4A16 GroupedMatmul 与 MLA Prolog / Flash Decode 等融合算子则从算子层面压榨算力。相关代码、YAML 配置与算子实现均可在当前仓库中直接查看与复现也可作为其他同类 MoE 长序列模型在昇腾平台落地的参考范本。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考