Kimi-K2-Thinking Inference on NPU

【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer

概述

本样例基于CANN平台对Kimi-K2-Thinking模型进行迁移，可在华为 Atlas A3 集群上运行，并行策略和性能优化点详细介绍可参见NPU Kimi-K2-Thinking推理优化实践。

硬件要求

产品型号：Atlas A3 系列

操作系统：Linux ARM

镜像版本：cann8.5_pt2.6.0_kimi_k2_aarch_image:v0.1

驱动版本：Ascend HDK 25.2.0

npu-smi info 检查Ascend NPU固件和驱动是否正确安装。如果已安装，通过命令npu-smi info确认版本是否为 25.2.0。如果未安装或者版本不是 25.2.0，请先下载固件和驱动包，然后根据指导自行安装。

快速启动

下载源码

在各个节点上执行如下命令下载 cann-recipes-infer 源码。

mkdir -p /home/code; cd /home/code/ git clone git@gitcode.com:cann/cann-recipes-infer.git cd cann-recipes-infer

下载数据集

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

mkdir -p dataset/InfiniteBench

下载权重

下载Kimi-K2-Thinking原始权重，并上传到Atlas A3各节点某个固定的路径下，比如/data/models/Kimi-K2-Thinking。

获取 docker 镜像

从ARM镜像地址中下载 docker 镜像，然后上传到A3服务器的每个节点上，并通过命令导入镜像docker load -i cann8.5_pt2.6.0_kimi_k2_aarch_image_v0.1.tar。

拉起 docker 容器

在各个节点上通过如下脚本拉起容器，默认容器名为 cann_recipes_infer。注意：需要将权重路径和源码路径挂载到容器里。

docker run -u root -itd --name cann_recipes_infer --ulimit nproc=65535:65535 --ipc=host \ --device=/dev/davinci0 --device=/dev/davinci1 \ --device=/dev/davinci2 --device=/dev/davinci3 \ --device=/dev/davinci4 --device=/dev/davinci5 \ --device=/dev/davinci6 --device=/dev/davinci7 \ --device=/dev/davinci8 --device=/dev/davinci9 \ --device=/dev/davinci10 --device=/dev/davinci11 \ --device=/dev/davinci12 --device=/dev/davinci13 \ --device=/dev/davinci14 --device=/dev/davinci15 \ --device=/dev/davinci_manager --device=/dev/devmm_svm \ --device=/dev/hisi_hdc \ -v /home/:/home \ -v /data:/data \ -v /etc/localtime:/etc/localtime \ -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ -v /etc/ascend_install.info:/etc/ascend_install.info -v /var/log/npu/:/usr/slog \ -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /sys/fs/cgroup:/sys/fs/cgroup:ro \ -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/sbin:/usr/local/sbin \ -v /etc/hccn.conf:/etc/hccn.conf -v /root/.pip:/root/.pip -v /etc/hosts:/etc/hosts \ -v /usr/bin/hostname:/usr/bin/hostname \ --net=host \ --shm-size=128g \ --privileged \ cann8.5_pt2.6.0_kimi_k2_aarch_image:v0.1 /bin/bash

在各个节点上通过如下命令进入容器：

docker attach cann_recipes_infer cd /home/code/cann-recipes-infer/models/kimi-k2-thinking

修改代码

• 在各个节点上修改 set_env.sh 文件中的IPs。

  export IPs=('xxx.xxx.xxx.xxx' 'xxx.xxx.xxx.xxx') # 所有节点的IP，确保第1个IP是master，多个节点的ip通过空格分开

• 在各个节点上修改 config/ 路径下需要执行的yaml文件中的model_path路径。关于YAML文件中的更多配置说明可参见YAML参数描述。

  model_path: "/data/models/Kimi-K2-Thinking"

• 在各个节点上修改 infer.sh 文件中的YAML_FILE_NAME，指定为上一步需要执行的yaml文件名。默认的yaml路径为32卡推理。

  # prefill export YAML_FILE_NAME=kimi_k2_thinking_rank_64_64ep_prefill_benchmark.yaml # decode export YAML_FILE_NAME=kimi_k2_thinking_rank_128_128ep_decode_benchmark.yaml

  Note: 本样例Prefill支持8-192卡，Decode支持8-192卡，可分别在config下的yaml文件中修改world_size配置。

拉起多卡推理

在各个节点上同步执行如下命令即可拉起多卡推理任务。

bash infer.sh

附录

FAQ

• HCCL_BUFFSIZE不足问题：如果报错日志中出现关键字"HCCL_BUFFSIZE is too SMALL, ..., NEEDED_HCCL_BUFFSIZE..., HCCL_BUFFSIZE=200MB, ..."，可通过配置环境变量export HCCL_BUFFSIZE=实际需要的大小解决，所有Rank上的该环境变量需保持一致。HCCL_BUFFSIZE参数介绍可参考昇腾资料中的详细描述。
• 自定义算子导入失败：如果报错日志中出现类似关键字"'_OpNamespace' 'custom' object has no attribute"，可参考自定义算子指南编译所需算子。

【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer