DeepSeek-R1或Kimi-K2模型在NPU实现高性能推理

概述

DeepSeek-R1和Kimi-K2都是2025年开源的大语言模型，二者结构类似，代码可以复用。本样例基于Deepseek开源代码进行迁移，并完成对应的优化适配。

支持的产品型号

Atlas A3 系列产品

环境准备

1. 安装CANN软件包。

   本样例的编译执行依赖CANN开发套件包与CANN二进制算子包，支持的CANN软件版本为CANN 9.0.0。

   请从软件包下载地址下载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 PyTorch（torch_npu）。

   Ascend Extension for PyTorch（torch_npu）为支撑PyTorch框架运行在NPU上的适配插件，本样例支持的Ascend Extension for PyTorch版本为v26.0.0，PyTorch版本为2.8.0。

   请从软件包下载地址下载torch_npu-2.8.0.post4-cp311-cp311-manylinux_2_28_${arch}.whl安装包，并参考torch_npu安装文档进行安装。

   • ${arch}表示CPU架构，如aarch64、x86_64。

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/deepseek_r1/requirements.txt
   

4. 配置样例运行所需环境信息。

   修改executor/scripts/set_env.sh中的如下字段:

   • IPs：配置所有节点的IP，按照rank id排序，多个节点的ip通过空格分开，例如：('xxx.xxx.xxx.xxx' 'xxx.xxx.xxx.xxx')。
   • cann_path: CANN软件包安装路径，例如/usr/local/Ascend/ascend-toolkit/latest。

   说明：HCCL相关配置，如：HCCL_SOCKET_IFNAME、HCCL_OP_EXPANSION_MODE，可以参考集合通信文档并在executor/scripts/function.sh中自定义配置。

权重准备

请根据所使用的模型类型自行下载原始权重到本地路径，例如/data/models/origin/。

Deepseek-R1与Kimi-K2的原始权重下载地址如下：

• Deepseek-R1权重
• Kimi-K2权重

权重转换

在各个节点上使用weight_convert.sh 脚本完成fp8到bfloat16/int8权重转换。

入参介绍：input_fp8_hf_path：原始fp8权重路径；output_hf_path：转换后输出的权重路径；quant_mode：量化模式

如果权重转换的运行环境为NPU，需要先执行：

cann_path=/usr/local/Ascend/ascend-toolkit/latest  # cann包安装路径
source ${cann_path}/bin/setenv.bash

权重转换拉起示例：

# 转换为bfloat16权重，适用于DeepSeek-R1和Kimi-K2。
bash utils/weight_convert.sh --input_fp8_hf_path /data/models/origin/DeepSeek-R1-FP8 --output_hf_path /data/models/origin/DeepSeek-R1-Bfloat16 --quant_mode bfloat16

# 转换为W8A8C16权重，适用于DeepSeek-R1和Kimi-K2。
bash utils/weight_convert.sh --input_fp8_hf_path /data/models/origin/DeepSeek-R1-FP8 --output_hf_path /data/models/origin/DeepSeek-R1-W8A8C16 --quant_mode w8a8c16

# 转换为W8A8C8权重，仅适用于DeepSeek-R1。
bash utils/weight_convert.sh --input_fp8_hf_path /data/models/origin/DeepSeek-R1-FP8 --output_hf_path /data/models/origin/DeepSeek-R1-W8A8C8 --quant_mode w8a8c8

注意：仅DeepSeek-R1支持转W8A8C8权重。

推理执行

1. 配置推理执行需要加载的权重文件以及YAML文件。

   • 修改YAML文件参数。

     本文以models/deepseek_r1/config/decode_r1_rank_16_16ep_a8w8.yaml文件为例，修改其中的model_path参数，将其设置为权重转换阶段准备好的权重文件存储路径，例如/data/models/origin/DeepSeek-R1-W8A8。 在models/deepseek_r1/config目录下已提供了较优性能的YAML样例供您参考，您可以根据模型类型、集群规模以及量化类型选择对应的YAML文件，具体参数说明如下：

     • YAML文件中的配置说明可见YAML参数描述。

     • 除框架统一配置之外，还额外支持以下特性，放置在 YAML 文件 model_config 的 custom_params 字段下：

       参数名	类型	默认值	含义
       enable_mla_prolog	bool	false	启用 MLA prolog 优化，用融合算子处理 Attention 的前置计算，提高吞吐。
       enable_multi_streams	bool	false	在图模式下，启用 MoE 共享专家多流并行计算，提升吞吐。
       enable_superkernel	bool	false	在图模式下，启用 superkernel 加速，将多个算子融合为大核以提高执行效率。
       moe_chunk_max_len	int	65536	MoE 算子 token 分发时的最大 chunk 长度，用于避免长序列场景下 OOM。
       micro_batch_mode	int	0	控制 prefill 阶段的 micro batch 执行模式：0 表示关闭；1 表示按 DP/EP 维度切分 micro batch；2 表示按 SP/TP/EP 维度切分 micro batch。
       enable_o_proj_alltoall	bool	false	启用 attention 中 o_proj 的 all-to-all 通信优化。

   • 配置executor/scripts/infer.sh脚本中的参数。

     离线推理模式下，将--yaml设置为config文件夹下YAML文件名称，例如decode_r1_rank_16_16ep_a8w8.yaml。 在线推理模式下，将--mode设置为online，--pd-role设置为prefill或decode，可通过--p-yaml-name和--d-yaml-name指定prefill/decode的YAML文件。

2. 准备输入prompt。

   • 使用内置prompt。

     本样例已在dataset/default_prompt.json中内置了输入prompt，若您直接使用内置prompt，本步骤可直接跳过。

     当然，您也可以在dataset/default_prompt.json文件中自定义prompt输入。

   • 使用长序列prompt。

     本样例默认使用内置prompt，若您需要使用长序列prompt，可以选择LongBench数据集或者InfiniteBench数据集。需要执行以下操作：

     1. 修改YAML文件中的dataset参数，将其修改为dataset: "LongBench"/dataset: "InfiniteBench"，使用LongBench数据集或InfiniteBench数据集作为长序列prompt。

     2. 若您选择LongBench数据集，且机器无法联网，需要您从huggingface手动下载数据集至dataset/LongBench目录下，LongBench文件夹需手工创建，目录中包含LongBench.py和data目录，并需要在LongBench.py中修改数据集加载路径；若您的机器可正常联网，样例执行过程中会自动在线读取LongBench数据集，您无需手工下载。

     3. 若您选择InfiniteBench数据集，需要从链接中下载长序列输入数据集longbook_qa_eng，并上传到各个节点上新建的路径 dataset/InfiniteBench下。

     mkdir -p dataset/InfiniteBench
     

     说明：

     • 在使用LongBench/InfiniteBench数据集或其他自定义数据集时，默认执行文本摘要任务，可在cann-recipes-infer/executor/utils/data_utils.py的build_dataset_input函数里修改默认的system prompt。
     • 长序列请求执行中若出现out of memory问题，可参见附录中的长序列请求out of memory问题处理。

3. 执行统一推理脚本。

   统一入口脚本位于 executor/scripts/infer.sh，通过以下参数控制启动：

   参数	含义	取值示例
   --model	模型目录名，对应 models/ 下的子目录	deepseek_r1
   --mode	推理模式	offline（离线推理）/ online（在线PD分离推理）
   --yaml	离线模式：yaml 文件名	decode_r1_rank_16_16ep_a8w8.yaml
   --pd-role	在线模式：PD 部署角色	prefill / decode
   --p-yaml-name	可选，在线模式：prefill yaml 文件名，不传则默认 deepseek_r1_pd/prefill.yaml	deepseek_r1_pd/prefill.yaml
   --d-yaml-name	可选，在线模式：decode yaml 文件名，不传则默认 deepseek_r1_pd/decode.yaml	deepseek_r1_pd/decode.yaml

   在线模式 IP 等更多配置参见 executor 设计文档 §5.1 启动方式。

   使用方式一：命令行传参

   # offline 模式
   bash executor/scripts/infer.sh --model deepseek_r1 --yaml decode_r1_rank_16_16ep_a8w8.yaml
   # online 模式
   bash executor/scripts/infer.sh --model deepseek_r1 --mode online --pd-role prefill
   # online 模式（指定 prefill/decode yaml）
   bash executor/scripts/infer.sh --model deepseek_r1 --mode online --pd-role prefill --p-yaml-name deepseek_r1_pd/prefill.yaml --d-yaml-name deepseek_r1_pd/decode.yaml
   

   如需查看参数说明，可以执行 bash executor/scripts/infer.sh --help。

   使用方式二：直接修改脚本默认值后执行 编辑 executor/scripts/infer.sh，按照需求修改 MODEL / MODE / YAML_FILE / PD_ROLE / P_YAML_NAME / D_YAML_NAME 等参数的默认值，例如：

   MODEL=deepseek_r1
   MODE=offline
   YAML_FILE=decode_r1_rank_16_16ep_a8w8.yaml
   

   保存后直接执行：

   bash executor/scripts/infer.sh
   

   说明：

   • 如果是多机环境，需要在每个节点上执行。
   • 推理日志和结果保存在 models/deepseek_r1/res/ 路径下。

优化点参考

• 本样例prefill阶段采用的详细优化点介绍可参见基于Atlas A3集群的DeepSeek-R1模型prefill阶段推理性能优化实践。

• 本样例decode阶段采用的详细优化点介绍及性能Benchmark可参见基于Atlas A3集群的DeepSeek-R1模型decode阶段推理性能优化。

Benchmark

基于Atlas A3，本实践使用config/decode_r1_rank_128_128ep_a8w8c8_mtp_benchmark.yaml作为运行配置文件，对DeepSeek-R1 W8A8C8 量化版本进行了性能Benchmark测试。

Quant Mode	Global Batch Size	Seq Length	Chips	TPOT (ms)	Throughput (tokens/p/s)
W8A8C8	6144	4096	64	44.9	2138

注：性能数据基于 MTP1 与 perfect eplb 配置采集，平均 1 个 draft token 中 accept token 为 0.7 个。

附录

常见问题处理

长序列请求out of memory问题处理

长序列请求可能导致device内存out of memory，尤其是在prefill阶段:

• Attention的Softmax操作通常为float32计算，其内存大小为batch_size * num_heads * q_s * kv_s * (2Bytes + 4Bytes)。

• MoE的Routing分发，可能存在极端负载不均，导致个别卡上的grouped_matmul算子占用较大内存。

为了缓解这两处峰值带来的OOM问题，可分别采用以下方法：

• 使能Paged Attention进行内存管理并调用Flash Attention融合算子，算子内会切块计算Attention，避免了q_s * kv_s的峰值内存产生。

• Prefill内存通常与batch_size大小成正比，当decode需要推理的global batch size过大时，prefill可能会由于OOM而无法在一轮推理中处理完所有的batch，因此我们可进行多次小batch串行推理，从而降低峰值内存。

  当前executor-core的prefill默认采用packed sequence执行，可通过配置YAML中的max_prefill_tokens限制单次prefill batch的总prompt token数，从而控制prefill阶段的峰值内存。

• 为了缓解MoE负载不均带来的峰值内存，我们可进行Chunk MoE推理，即在MoE切Chunk串行推理，降低极端场景下的峰值内存，可通过YAML中的moe_chunk_max_len开关设置chunk的大小。当前该开关只针对prefill生效，开启后，由于MoE部分将串行计算各chunk，会对prefill的性能产生相应的影响。