CANN融合算子库实现

概述

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

此项目是基于昇腾硬件Atlas A2/Atlas A3/Ascend 950PR/Ascend 950DT的融合算子库。当前项目包含的算子说明详见同级docs目录。

目录结构

├── pypto_python # pypto算子代码目录 │ │ ├── docs # 自定义算子文档 │ │ ├── impl # 自定义算子计算流代码目录 │ │ │ ├── compressed_flash_attention_pypto.py # compressed_flash_attention算子kernel实现 │ │ │ ├── compressor_pypto.py # compressor算子kernel实现 │ │ │ ├── hc_pre_pypto.py # hc_pre算子kernel实现 │ │ │ ├── lightning_indexer_prolog_quant_pypto.py # lightning_indexer_prolog_quant算子kernel实现 │ │ │ ├── mla_prolog_pypto.py # mla_prolog算子kernel实现 │ │ │ ├── mla_prolog_quant_pypto.py # mla_prolog_quant算子kernel实现 │ │ │ ├── sliding_window_attention_pypto.py # sliding_window_attention算子kernel实现 │ │ │ ├── sparse_compress_flash_attention_pypto.py # sparse_compress_flash_attention算子kernel实现 │ │ ├── example # 自定义算子测试代码目录 │ │ │ ├── test_compressed_flash_attention_pypto.py # compressed_flash_attention算子测试样例 │ │ │ ├── test_compressor_pypto.py # compressor算子测试样例 │ │ │ ├── test_hc_pre_pypto.py # hc_pre算子测试样例 │ │ │ ├── test_lightning_indexer_prolog_quant.py # lightning_indexer_prolog_quant算子测试样例 │ │ │ ├── test_mla_prolog_pypto.py # mla_prolog算子测试样例 │ │ │ ├── test_mla_prolog_quant_pypto.py # mla_prolog_quant算子测试样例 │ │ │ ├── test_sliding_window_attention_pypto.py # sliding_window_attention算子测试样例 │ │ │ ├── test_sparse_compressed_flash_attention_pypto.py # sparse_compressed_flash_attention算子测试样例

PyPto自定义算子开发资料：PyPto文档

环境准备

硬件要求

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

下载源码

可以选择在宿主机或者容器内下载源码，如果在容器内下载，应在主机挂载在容器的目录下下载；在宿主机内下载则无此约束。 执行如下命令即可下载 cann-recipes-infer 源码。

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

获取 docker 镜像

从ARM镜像地址中下载 docker 镜像，然后上传到需要A3服务器每个节点上，并通过命令导入镜像docker load -i cann9.0.pt2.8.0_ds_pypto_aarch_image.tar

拉起 docker 容器

容器拉起脚本如下，默认容器名为 cann_recipes_infer_pypto：

docker run -u root -itd --name cann_recipes_infer_pypto --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 \ cann9.0.pt2.8.0_ds_pypto_aarch_image:v0.2 /bin/bash

进入容器：

docker attach cann_recipes_infer_pypto

设置环境变量

source /usr/local/Ascend/ascend-toolkit/latest/bin/setenv.bash

编译执行

PyPTO算子工程编译安装：

PyPTO已切换为Python代码实现（镜像中已安装PyPTO的whl包）。

若涉及PyPTO框架源码修改，可下载PyPto开源仓 https://gitcode.com/cann/pypto ，参考ReadMe “快速入门” 编译PyPTO软件包并安装。

示例算子执行：

在pypto_python/examples目录通过执行以下脚本执行示例算子

cd /home/code/cann-recipes-infer/ops/pypto_python/example python3 test_hc_pre_pypto.py

DeepSeek-V4 整网集成样例执行

算子已支持集成到DeepSeek-V4整网，样例执行过程如下：

权重和数据集准备

DeepSeek-V4模型和数据集准备，请参考模型权重和数据集准备中相关章节

代码修改适配

网络执行前需对配置做一些调整，参考修改代码章节进行适配

修改网络配置和环境配置

当前网络脚本中，在各个节点上修改models/deepseek-v4/config/ 路径下需要执行的yaml文件中model_config配置项，配置过程如下：

• 增加 enable_pypto: True配置将pypto算子集成到网络中
• 修改 enable_limit_core: False配置将limit_core配置关闭

model_config: enable_limit_core: False enable_pypto: True

推理执行

参考拉起多卡推理章节。

执行结束后，出现model run success，则表示推理执行成功。

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