基于 CANN 的 Kimi-K2-Thinking 高性能推理部署指南Atlas A3 集群实战与优化解析【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本文是基于 cann-recipes-infer 仓库中 Kimi-K2-Thinking 样例编写的技术指南围绕在华为 Atlas A3 集群上完成 Kimi-K2-Thinking 思考模型的迁移、环境搭建、权重准备、离线/在线推理执行以及 MLA Prolog、多流并行、superkernel 等 NPU 优化特性的配置与原理展开。读完本文你将掌握该样例从零到跑通的完整操作路径并理解其并行策略与融合算子加速机制可直接复用到其他基于 DeepSeek 结构的 MoE 思考模型的昇腾部署实践中。概述为什么关注 Kimi-K2-Thinking 在 NPU 上的推理Kimi-K2-Thinking 是 Moonshot AI 发布的深度推理思考模型模型结构与 Kimi-K2 保持一致原生支持MOE W4A16 Attention BF16 混合精度模式兼顾推理性能与效果。本样例基于 CANN 平台对 Kimi-K2-Thinking 进行迁移可在华为Atlas A3 系列产品上运行模型推理代码位于 models/kimi_k2_thinking。从仓库内配套的优化实践文档NPU Kimi-K2-Thinking 推理优化实践可知该样例具备以下关键能力0day 适配支持CANN 已 0day 支持 Kimi-K2-Thinking 256K 序列推理部署推理代码随仓库开源。原生量化模式MOE GroupedMatmul 采用 W4A16 计算Attention 保留 BF16 精度。Flash Decode 加速针对小 batch、长序列生产场景Decode 阶段 MLA 已实现 Flash Decode 加速机制降低时延并提升算力利用率。灵活部署规模最小部署单元为单机可支持 4K 序列长度推理同时支持多机大 EP 部署模式通过多机集群协同可满足 256K 超长序列推理需求。由于其结构与 Kimi-K2 基本一致DeepSeek-R1 Decode 优化实践中的基础优化点可平滑继承因此本样例是理解昇腾 NPU 上 MoE 类思考模型推理加速的典型范本。支持的产品型号与软硬件版本本样例的编译执行依赖完整的 CANN 软件栈具体版本约束如下组件版本要求硬件Atlas A3 系列产品CANN 软件CANN 9.0.0含开发套件包与二进制算子包Ascend Extension for PyTorchtorch_npuv26.0.0PyTorch2.8.0Python仅支持3.11其中 torch_npu 安装包对应torch_npu-2.8.0.post4-cp311-cp311-manylinux_2_28_${arch}.whl版本${arch}表示 CPU 架构aarch64、x86_64。各 Python 依赖库的具体版本可在 models/kimi_k2_thinking/requirements.txt 中查看例如 transformers 4.53.3、datasets 3.6.0、compressed-tensors 0.6.0、tiktoken 0.9.0 等其中compressed-tensors与 MOE W4A16 量化推理相关。环境准备1. 安装 CANN 软件包从昇腾社区官方下载渠道获取Ascend-cann-toolkit_${version}_linux-${arch}.run与Ascend-cann-A3-ops_${version}_linux-${arch}.run两个软件包并参考 CANN 官方安装文档完成安装${version}表示 CANN 包版本号例如9.0.0${arch}表示 CPU 架构例如aarch64、x86_64。2. 安装 Ascend Extension for PyTorchtorch_nputorch_npu 是支撑 PyTorch 框架运行在 NPU 上的适配插件。下载对应版本的torch_npu-2.8.0.post4-cp311-cp311-manylinux_2_28_${arch}.whl安装包后参考 torch_npu 二进制包安装文档完成安装。安装完成后建议验证import torch_npu正常且能够识别 NPU 设备。3. 下载项目源码并安装 Python 依赖# 下载项目源码以 master 分支为例 git clone https://gitcode.com/cann/cann-recipes-infer.git # 安装依赖的 python 库仅支持 python 3.11 cd cann-recipes-infer pip3 install -r ./models/kimi_k2_thinking/requirements.txt4. 配置样例运行所需环境信息修改 executor/scripts/set_env.sh 中的如下字段IPs配置所有节点的 IP按照 rank id 排序多个节点的 IP 通过空格分开例如(xxx.xxx.xxx.xxx xxx.xxx.xxx.xxx)。首节点为 master 节点。cann_pathCANN 软件包安装路径例如/usr/local/Ascend/ascend-toolkit/latest。在线推理场景下如果 prefill 与 decode 角色共享同一批主机可在调用infer.sh前通过设置ASCEND_RT_VISIBLE_DEVICES按角色隔离 NPU见 executor/scripts/set_env.sh 中PREFILL_IPS/DECODE_IPS的注释说明。说明HCCL 相关配置如HCCL_SOCKET_IFNAME、HCCL_OP_EXPANSION_MODE可参考昇腾集合通信环境变量文档并在 executor/scripts/function.sh 中自定义配置。权重准备从 Hugging Face 官方模型仓库下载Kimi-K2-Thinking 原始权重并上传到 Atlas A3 各节点某个固定的路径下例如/data/models/Kimi-K2-Thinking。该路径将在下一步通过 YAML 文件中的model_path参数指定。推理执行1. 配置推理执行参数推理的配置分为两部分YAML 模型配置文件以及 executor/scripts/infer.sh 启动脚本参数。修改 YAML 文件参数以 models/kimi_k2_thinking/config/kimi_k2_thinking.yaml 为例需将model_config下的model_path修改为权重实际存储路径例如/data/models/Kimi-K2-Thinking。该文件完整内容如下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 only data_config: dataset: InfiniteBench input_truncated_len: 65536 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: 32YAML 中的通用参数model_config、data_config、parallel_config、scheduler_config各字段含义详见 YAML 参数描述。结合上方样例几个关键字段说明如下exe_mode执行模式可选eager、npugraph_ex、ge_graph。样例默认使用npugraph_ex。enable_static_kernel仅npugraph_ex模式下生效将算子编译为静态 kernel 以降低调度开销。enable_weight_nz启用权重 NZ 格式布局优化访存。input_truncated_len输入序列截断长度样例配置为 65536配合 256K 超长序列推理场景。除框架统一配置之外样例还额外支持以下特性放置在 YAML 文件model_config的custom_params字段下参数名类型默认值含义enable_mla_prologboolfalse启用 MLA prolog 优化用融合算子处理 Attention 的前置计算提高吞吐。enable_multi_streamsboolfalse在图模式下启用 MoE 共享专家多流并行计算提升吞吐。enable_superkernelboolfalse在ge_graph模式下启用 superkernel 加速将多个算子融合为大核以提高执行效率npugraph_ex模式不支持。在 models/kimi_k2_thinking/config 目录下已提供较优性能的 YAML 样例供参考kimi_k2_thinking.yaml32 卡、CP32 的离线推理配置enable_mla_prolog与enable_multi_streams均默认开启。pd/prefill.yaml 与 pd/decode.yaml在线 PDPrefill/Decode分离推理配置。prefill 侧重 CP 并行cp_size: 32摊薄长序列计算decode 侧重较大 batchbatch_size: 64、max_new_tokens: 4096与局部 TP 切分。配置 infer.sh 启动参数统一入口脚本位于 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从 executor/scripts/infer.sh 源码可以看到脚本内部会先解析命令行参数覆盖默认值再依次 source executor/scripts/function.sh 与 executor/scripts/set_env.sh并通过 executor/scripts/validate_infer_args.py 校验模型目录、模式与 YAML 路径的合法性最后调用launch $MODE启动对应推理流程。在线模式下 YAML 路径的默认规则为models/${MODEL}/config/${MODEL}_pd/prefill.yaml对应本样例即kimi_k2_thinking_pd/仓库中实际以pd/目录组织可通过--p-yaml-name显式指定。在线模式 IP 等更多配置参见 executor 设计文档 §5.1 启动方式。2. 准备输入 prompt本样例默认使用InfiniteBench 数据集进行长序列推理。需要从 InfiniteBench 数据集的公开发布源下载长序列输入数据集longbook_qa_eng并上传到各个节点上新建的dataset/InfiniteBench路径下mkdir -p dataset/InfiniteBench若需要使用内置 prompt 或自定义 prompt可以修改 YAML 文件中的dataset参数为dataset: default并在 dataset/default_prompt.json 文件中自定义 prompt 输入。3. 执行统一推理脚本使用方式一命令行传参# 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脚本会打印各参数的使用示例含 offline、online 以及 mooncake KV 传输等用法。使用方式二直接修改脚本默认值后执行编辑 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说明如果是多机环境需要在每个节点上执行。推理日志和结果保存在models/kimi_k2_thinking/res/路径下。优化特性源码级解析样例中的三项custom_params优化特性并非空壳开关在 models/kimi_k2_thinking/models/modeling_deepseek.py 中有完整的调用链实现理解它们有助于针对不同推理场景做出正确的参数选择。MLA Prolog融合 Attention 前置计算enable_mla_prolog对应源码中的 MLA Prolog 优化。从源码看Attention 的 Q/K/V 投影、RoPE 位置编码等前置计算被封装为一次torch_npu.npu_mla_prolog_v3融合算子调用见 modeling_deepseek.py 中mla_prolog_input_args的构造与调用将多个离散算子融合为单算子执行减少 kernel 启动与中间张量搬运开销从而提升吞吐。源码中还校验了kv_cache_quant_mode int8时必须开启enable_mla_prolog否则直接抛出ValueError说明该特性与 KV Cache 量化C8存在依赖关系。Flash DecodeFusedInferAttentionScore与 MLA Prolog 配合的是 Decode 阶段的注意力算子。源码中通过self.fa_ops_prefill.npu_fused_infer_attention_score(...)与self.fa_ops_decode.npu_fused_infer_attention_score(...)分别处理 prefill 与 decode 两个阶段的融合注意力计算。Decode 场景下基于 MLA 的 Flash Decode 加速机制可有效降低时延并提升算力利用率这也是长序列生产场景下推荐开启enable_mla_prolog的原因。MoE 共享专家多流并行enable_multi_streams在图模式下启用 MoE 共享专家多流并行计算。源码中enable_gegraph_and_multistream与enable_npugraphex_and_multistream两个标志分别对应ge_graph与npugraph_ex两种执行模式下的多流开关且 decode 阶段的多流并行仅在enable_multi_streams and not is_prefill时生效对应wait_event(enable_multi_streams, self.npu_events, 1)的事件同步逻辑。同时源码显示多流会通过stream-fusion1选项注入算子编译配置将共享专家计算拆分到多条 NPU 流上并行执行以提升吞吐。需要注意的是该特性依赖图模式npugraph_ex或ge_grapheager模式不支持。Superkernel算子大核融合enable_superkernel仅在ge_graph模式下生效npugraph_ex模式不支持。源码中以superkernel_scope(self.enable_superkernel and not is_prefill, label, option)的形式包裹共享专家计算区域将多个相邻算子融合为一个 superkernel 大核执行减少核间调度与中间结果落盘开销。若在源码校验逻辑中指定了exe_mode不支持该特性会抛出明确的ValueError提示见 modeling_deepseek.py 中dynamo_feat相关校验。从源码结构看三项特性均与执行模式exe_mode强耦合eager模式不支持下述图级优化npugraph_ex支持多流但不支持 superkernelge_graph则三者皆可。因此建议严格按样例默认配置npugraph_exenable_mla_prologenable_multi_streams运行仅在切换到ge_graph时再考虑开启enable_superkernel。并行策略从单机到 256K 超长序列NPU Kimi-K2-Thinking 推理优化实践给出了 Atlas A3 上的推荐部署策略Prefill 使用 M 个节点、Decode 使用 N 个节点每个节点包含 8 卡M 与 N 可根据资源数量与 SLA 约束在 1~24 内动态调整。Prefill 并行策略长序列场景下 Attention 选用 Context ParallelCP并行多个 rank 均摊长序列计算单 rank 计算量与 activation 内存都较小TTFT 可控MoE 模块沿用 DeepSeek-V3.1 的 EP 并行兼顾吞吐与时延。对应kimi_k2_thinking.yaml中cp_size: 32、moe_tp_size: 1的配置。Decode 并行策略沿用 Attention DP MoE EP 部署由于O_proj与LM_Head权重内存较大且 Decode 阶段呈明显访存瓶颈采用局部 TP 并行o_proj_tp_size: 8、dense_tp_size: 4为降低设备内存占用Embedding 层同样使用 TP 切分embed_tp_size: 16、lmhead_tp_size: 16并将 TP 域控制在单机内以减小通信开销。对应 pd/decode.yaml 中的配置。整网计算流中使用了 MLAProlog、FusedInferAttentionScore、MoeDistributeDispatch、MoeDistributeCombine、GroupedMatmul 等融合 Kernel可供其他同类模型在昇腾平台的高效落地提供实践参考MTP 多 Token 预测支持相较于 DeepSeek V3.1Kimi-K2-Thinking 的 Attention Head 数量从 128 精简至 64Attention 计算量显著降低更利好 MTPMulti-Token Prediction场景。LLM 推理 Decode 阶段通常为访存瓶颈MTP 可通过少量计算代价缓解访存压力而 Kimi-K2-Thinking 的 Attention 本身计算负载更轻因此在 MTP 模式下更难触及计算瓶颈性能更优。样例代码已支持多头 MTP 功能若具备训练完成的 MTP 权重可直接启用该功能以实现推理加速YAML 中next_n参数与此相关默认0表示关闭。未来演进方向量化目前支持 MOE W4A16 推理未来可针对 Attention 开发低比特量化版本进一步降低系统内存占用与时延。长序列性能优化长序列场景下的 TTFT 仍有优化空间后续将探索更亲和的部署方案并针对性优化融合 Kernel 性能全面提升长序列推理的响应速度与运行效率。总结通过本文你可以完整掌握 Kimi-K2-Thinking 在 Atlas A3 集群上的推理部署流程从 CANN 9.0.0 与 torch_npu v26.0.0 环境搭建、权重准备到离线/在线两种模式的 YAML 配置与infer.sh启动再到 MLA Prolog、多流并行、superkernel 等优化开关的源码级理解。该样例不仅是 Kimi-K2-Thinking 的落地参考其 CPEP 并行策略与融合算子体系对 DeepSeek 系列及其他 MoE 思考模型在 CANN 平台上的高效推理同样具有借鉴意义。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 基于 CANN 的 Kimi-K2-Thinking 高性能推理部署指南：Atlas A3 集群实战与优化解析
基于 CANN 的 Kimi-K2-Thinking 高性能推理部署指南Atlas A3 集群实战与优化解析【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本文是基于 cann-recipes-infer 仓库中 Kimi-K2-Thinking 样例编写的技术指南围绕在华为 Atlas A3 集群上完成 Kimi-K2-Thinking 思考模型的迁移、环境搭建、权重准备、离线/在线推理执行以及 MLA Prolog、多流并行、superkernel 等 NPU 优化特性的配置与原理展开。读完本文你将掌握该样例从零到跑通的完整操作路径并理解其并行策略与融合算子加速机制可直接复用到其他基于 DeepSeek 结构的 MoE 思考模型的昇腾部署实践中。概述为什么关注 Kimi-K2-Thinking 在 NPU 上的推理Kimi-K2-Thinking 是 Moonshot AI 发布的深度推理思考模型模型结构与 Kimi-K2 保持一致原生支持MOE W4A16 Attention BF16 混合精度模式兼顾推理性能与效果。本样例基于 CANN 平台对 Kimi-K2-Thinking 进行迁移可在华为Atlas A3 系列产品上运行模型推理代码位于 models/kimi_k2_thinking。从仓库内配套的优化实践文档NPU Kimi-K2-Thinking 推理优化实践可知该样例具备以下关键能力0day 适配支持CANN 已 0day 支持 Kimi-K2-Thinking 256K 序列推理部署推理代码随仓库开源。原生量化模式MOE GroupedMatmul 采用 W4A16 计算Attention 保留 BF16 精度。Flash Decode 加速针对小 batch、长序列生产场景Decode 阶段 MLA 已实现 Flash Decode 加速机制降低时延并提升算力利用率。灵活部署规模最小部署单元为单机可支持 4K 序列长度推理同时支持多机大 EP 部署模式通过多机集群协同可满足 256K 超长序列推理需求。由于其结构与 Kimi-K2 基本一致DeepSeek-R1 Decode 优化实践中的基础优化点可平滑继承因此本样例是理解昇腾 NPU 上 MoE 类思考模型推理加速的典型范本。支持的产品型号与软硬件版本本样例的编译执行依赖完整的 CANN 软件栈具体版本约束如下组件版本要求硬件Atlas A3 系列产品CANN 软件CANN 9.0.0含开发套件包与二进制算子包Ascend Extension for PyTorchtorch_npuv26.0.0PyTorch2.8.0Python仅支持3.11其中 torch_npu 安装包对应torch_npu-2.8.0.post4-cp311-cp311-manylinux_2_28_${arch}.whl版本${arch}表示 CPU 架构aarch64、x86_64。各 Python 依赖库的具体版本可在 models/kimi_k2_thinking/requirements.txt 中查看例如 transformers 4.53.3、datasets 3.6.0、compressed-tensors 0.6.0、tiktoken 0.9.0 等其中compressed-tensors与 MOE W4A16 量化推理相关。环境准备1. 安装 CANN 软件包从昇腾社区官方下载渠道获取Ascend-cann-toolkit_${version}_linux-${arch}.run与Ascend-cann-A3-ops_${version}_linux-${arch}.run两个软件包并参考 CANN 官方安装文档完成安装${version}表示 CANN 包版本号例如9.0.0${arch}表示 CPU 架构例如aarch64、x86_64。2. 安装 Ascend Extension for PyTorchtorch_nputorch_npu 是支撑 PyTorch 框架运行在 NPU 上的适配插件。下载对应版本的torch_npu-2.8.0.post4-cp311-cp311-manylinux_2_28_${arch}.whl安装包后参考 torch_npu 二进制包安装文档完成安装。安装完成后建议验证import torch_npu正常且能够识别 NPU 设备。3. 下载项目源码并安装 Python 依赖# 下载项目源码以 master 分支为例 git clone https://gitcode.com/cann/cann-recipes-infer.git # 安装依赖的 python 库仅支持 python 3.11 cd cann-recipes-infer pip3 install -r ./models/kimi_k2_thinking/requirements.txt4. 配置样例运行所需环境信息修改 executor/scripts/set_env.sh 中的如下字段IPs配置所有节点的 IP按照 rank id 排序多个节点的 IP 通过空格分开例如(xxx.xxx.xxx.xxx xxx.xxx.xxx.xxx)。首节点为 master 节点。cann_pathCANN 软件包安装路径例如/usr/local/Ascend/ascend-toolkit/latest。在线推理场景下如果 prefill 与 decode 角色共享同一批主机可在调用infer.sh前通过设置ASCEND_RT_VISIBLE_DEVICES按角色隔离 NPU见 executor/scripts/set_env.sh 中PREFILL_IPS/DECODE_IPS的注释说明。说明HCCL 相关配置如HCCL_SOCKET_IFNAME、HCCL_OP_EXPANSION_MODE可参考昇腾集合通信环境变量文档并在 executor/scripts/function.sh 中自定义配置。权重准备从 Hugging Face 官方模型仓库下载Kimi-K2-Thinking 原始权重并上传到 Atlas A3 各节点某个固定的路径下例如/data/models/Kimi-K2-Thinking。该路径将在下一步通过 YAML 文件中的model_path参数指定。推理执行1. 配置推理执行参数推理的配置分为两部分YAML 模型配置文件以及 executor/scripts/infer.sh 启动脚本参数。修改 YAML 文件参数以 models/kimi_k2_thinking/config/kimi_k2_thinking.yaml 为例需将model_config下的model_path修改为权重实际存储路径例如/data/models/Kimi-K2-Thinking。该文件完整内容如下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 only data_config: dataset: InfiniteBench input_truncated_len: 65536 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: 32YAML 中的通用参数model_config、data_config、parallel_config、scheduler_config各字段含义详见 YAML 参数描述。结合上方样例几个关键字段说明如下exe_mode执行模式可选eager、npugraph_ex、ge_graph。样例默认使用npugraph_ex。enable_static_kernel仅npugraph_ex模式下生效将算子编译为静态 kernel 以降低调度开销。enable_weight_nz启用权重 NZ 格式布局优化访存。input_truncated_len输入序列截断长度样例配置为 65536配合 256K 超长序列推理场景。除框架统一配置之外样例还额外支持以下特性放置在 YAML 文件model_config的custom_params字段下参数名类型默认值含义enable_mla_prologboolfalse启用 MLA prolog 优化用融合算子处理 Attention 的前置计算提高吞吐。enable_multi_streamsboolfalse在图模式下启用 MoE 共享专家多流并行计算提升吞吐。enable_superkernelboolfalse在ge_graph模式下启用 superkernel 加速将多个算子融合为大核以提高执行效率npugraph_ex模式不支持。在 models/kimi_k2_thinking/config 目录下已提供较优性能的 YAML 样例供参考kimi_k2_thinking.yaml32 卡、CP32 的离线推理配置enable_mla_prolog与enable_multi_streams均默认开启。pd/prefill.yaml 与 pd/decode.yaml在线 PDPrefill/Decode分离推理配置。prefill 侧重 CP 并行cp_size: 32摊薄长序列计算decode 侧重较大 batchbatch_size: 64、max_new_tokens: 4096与局部 TP 切分。配置 infer.sh 启动参数统一入口脚本位于 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从 executor/scripts/infer.sh 源码可以看到脚本内部会先解析命令行参数覆盖默认值再依次 source executor/scripts/function.sh 与 executor/scripts/set_env.sh并通过 executor/scripts/validate_infer_args.py 校验模型目录、模式与 YAML 路径的合法性最后调用launch $MODE启动对应推理流程。在线模式下 YAML 路径的默认规则为models/${MODEL}/config/${MODEL}_pd/prefill.yaml对应本样例即kimi_k2_thinking_pd/仓库中实际以pd/目录组织可通过--p-yaml-name显式指定。在线模式 IP 等更多配置参见 executor 设计文档 §5.1 启动方式。2. 准备输入 prompt本样例默认使用InfiniteBench 数据集进行长序列推理。需要从 InfiniteBench 数据集的公开发布源下载长序列输入数据集longbook_qa_eng并上传到各个节点上新建的dataset/InfiniteBench路径下mkdir -p dataset/InfiniteBench若需要使用内置 prompt 或自定义 prompt可以修改 YAML 文件中的dataset参数为dataset: default并在 dataset/default_prompt.json 文件中自定义 prompt 输入。3. 执行统一推理脚本使用方式一命令行传参# 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脚本会打印各参数的使用示例含 offline、online 以及 mooncake KV 传输等用法。使用方式二直接修改脚本默认值后执行编辑 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说明如果是多机环境需要在每个节点上执行。推理日志和结果保存在models/kimi_k2_thinking/res/路径下。优化特性源码级解析样例中的三项custom_params优化特性并非空壳开关在 models/kimi_k2_thinking/models/modeling_deepseek.py 中有完整的调用链实现理解它们有助于针对不同推理场景做出正确的参数选择。MLA Prolog融合 Attention 前置计算enable_mla_prolog对应源码中的 MLA Prolog 优化。从源码看Attention 的 Q/K/V 投影、RoPE 位置编码等前置计算被封装为一次torch_npu.npu_mla_prolog_v3融合算子调用见 modeling_deepseek.py 中mla_prolog_input_args的构造与调用将多个离散算子融合为单算子执行减少 kernel 启动与中间张量搬运开销从而提升吞吐。源码中还校验了kv_cache_quant_mode int8时必须开启enable_mla_prolog否则直接抛出ValueError说明该特性与 KV Cache 量化C8存在依赖关系。Flash DecodeFusedInferAttentionScore与 MLA Prolog 配合的是 Decode 阶段的注意力算子。源码中通过self.fa_ops_prefill.npu_fused_infer_attention_score(...)与self.fa_ops_decode.npu_fused_infer_attention_score(...)分别处理 prefill 与 decode 两个阶段的融合注意力计算。Decode 场景下基于 MLA 的 Flash Decode 加速机制可有效降低时延并提升算力利用率这也是长序列生产场景下推荐开启enable_mla_prolog的原因。MoE 共享专家多流并行enable_multi_streams在图模式下启用 MoE 共享专家多流并行计算。源码中enable_gegraph_and_multistream与enable_npugraphex_and_multistream两个标志分别对应ge_graph与npugraph_ex两种执行模式下的多流开关且 decode 阶段的多流并行仅在enable_multi_streams and not is_prefill时生效对应wait_event(enable_multi_streams, self.npu_events, 1)的事件同步逻辑。同时源码显示多流会通过stream-fusion1选项注入算子编译配置将共享专家计算拆分到多条 NPU 流上并行执行以提升吞吐。需要注意的是该特性依赖图模式npugraph_ex或ge_grapheager模式不支持。Superkernel算子大核融合enable_superkernel仅在ge_graph模式下生效npugraph_ex模式不支持。源码中以superkernel_scope(self.enable_superkernel and not is_prefill, label, option)的形式包裹共享专家计算区域将多个相邻算子融合为一个 superkernel 大核执行减少核间调度与中间结果落盘开销。若在源码校验逻辑中指定了exe_mode不支持该特性会抛出明确的ValueError提示见 modeling_deepseek.py 中dynamo_feat相关校验。从源码结构看三项特性均与执行模式exe_mode强耦合eager模式不支持下述图级优化npugraph_ex支持多流但不支持 superkernelge_graph则三者皆可。因此建议严格按样例默认配置npugraph_exenable_mla_prologenable_multi_streams运行仅在切换到ge_graph时再考虑开启enable_superkernel。并行策略从单机到 256K 超长序列NPU Kimi-K2-Thinking 推理优化实践给出了 Atlas A3 上的推荐部署策略Prefill 使用 M 个节点、Decode 使用 N 个节点每个节点包含 8 卡M 与 N 可根据资源数量与 SLA 约束在 1~24 内动态调整。Prefill 并行策略长序列场景下 Attention 选用 Context ParallelCP并行多个 rank 均摊长序列计算单 rank 计算量与 activation 内存都较小TTFT 可控MoE 模块沿用 DeepSeek-V3.1 的 EP 并行兼顾吞吐与时延。对应kimi_k2_thinking.yaml中cp_size: 32、moe_tp_size: 1的配置。Decode 并行策略沿用 Attention DP MoE EP 部署由于O_proj与LM_Head权重内存较大且 Decode 阶段呈明显访存瓶颈采用局部 TP 并行o_proj_tp_size: 8、dense_tp_size: 4为降低设备内存占用Embedding 层同样使用 TP 切分embed_tp_size: 16、lmhead_tp_size: 16并将 TP 域控制在单机内以减小通信开销。对应 pd/decode.yaml 中的配置。整网计算流中使用了 MLAProlog、FusedInferAttentionScore、MoeDistributeDispatch、MoeDistributeCombine、GroupedMatmul 等融合 Kernel可供其他同类模型在昇腾平台的高效落地提供实践参考MTP 多 Token 预测支持相较于 DeepSeek V3.1Kimi-K2-Thinking 的 Attention Head 数量从 128 精简至 64Attention 计算量显著降低更利好 MTPMulti-Token Prediction场景。LLM 推理 Decode 阶段通常为访存瓶颈MTP 可通过少量计算代价缓解访存压力而 Kimi-K2-Thinking 的 Attention 本身计算负载更轻因此在 MTP 模式下更难触及计算瓶颈性能更优。样例代码已支持多头 MTP 功能若具备训练完成的 MTP 权重可直接启用该功能以实现推理加速YAML 中next_n参数与此相关默认0表示关闭。未来演进方向量化目前支持 MOE W4A16 推理未来可针对 Attention 开发低比特量化版本进一步降低系统内存占用与时延。长序列性能优化长序列场景下的 TTFT 仍有优化空间后续将探索更亲和的部署方案并针对性优化融合 Kernel 性能全面提升长序列推理的响应速度与运行效率。总结通过本文你可以完整掌握 Kimi-K2-Thinking 在 Atlas A3 集群上的推理部署流程从 CANN 9.0.0 与 torch_npu v26.0.0 环境搭建、权重准备到离线/在线两种模式的 YAML 配置与infer.sh启动再到 MLA Prolog、多流并行、superkernel 等优化开关的源码级理解。该样例不仅是 Kimi-K2-Thinking 的落地参考其 CPEP 并行策略与融合算子体系对 DeepSeek 系列及其他 MoE 思考模型在 CANN 平台上的高效推理同样具有借鉴意义。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考