Attention融合算子Experimental使用说明

前提条件

1. 环境部署：调用算子之前，请先参考环境部署完成基础环境搭建，其中安装依赖章节的版本约束以如下内容为准：

   • python >= 3.10.0
   • torch >= 2.7.0
   • gcc >= 9.0.0
   • torch_npu >=2.7.0

2. 安装torch与torch_npu包

   根据实际环境，下载对应torch包并安装：torch-${torch_version}+cpu-${python_version}-linux_${arch}.whl 下载链接为：官网地址

   安装命令如下：

   pip install torch-${torch_version}+cpu-${python_version}-linux_${arch}.whl
   

   根据实际环境，安装对应torch-npu包：torch_npu-${torch_version}-${python_version}-linux_${arch}.whl 下载链接为：官网地址

   安装命令如下：

   pip install torch_npu-${torch_version}-${python_version}-linux_${arch}.whl
   

   • ${torch_version}：表示torch包版本号。
   • ${python_version}：表示python版本号。
   • ${arch}：表示CPU架构，如aarch64、x86_64。

算子列表

项目提供的所有算子列表如下：

算子名	接口文档	算子实现		torch调用	图模式调用		算子执行硬件单元	说明
		kernel	host		Ascend IR	aclgraph
compressor	文档	√	√	√	×	√	AI Core	将每4或128个token的 KV cache 压缩成一个，然后每个token与这些压缩的 KV cache进行 DSA 计算。 算子torch接口调用依赖torch_ops_extension，具体安装方法见安装指导。
kv_quant_sparse_attn_sharedkv	文档	√	√	√	×	√	AI Core	支持量化模式的Sliding Window Attention、Compressed Attention以及Sparse Compressed Attention计算。 算子torch接口调用依赖torch_ops_extension，具体安装方法见安装指导。
kv_quant_sparse_attn_sharedkv_metadata	文档	√	√	√	×	√	AI Cpu	该算子为kv_quant_sparse_attn_sharedkv算子提供分核结果。 算子torch接口调用依赖torch_ops_extension，具体安装方法见安装指导。
quant_lightning_indexer	文档	√	√	√	×	√	AI Core	该算子是推理场景下，稀疏Attention前处理的计算，选出关键的稀疏token，并对输入query和key进行量化实现存8算8，获取最大收益。 算子torch接口调用依赖torch_ops_extension，具体安装方法见安装指导。
quant_lightning_indexer_metadata	文档	√	√	√	×	√	AI Cpu	该算子为quant_lightning_indexer算子提供分核结果。 算子torch接口调用依赖torch_ops_extension，具体安装方法见安装指导。
sparse_attn_sharedkv	文档	√	√	√	×	√	AI Core	该算子支持非量化的Sliding Window Attention、Compressed Attention以及Sparse Compressed Attention计算。 算子torch接口调用依赖torch_ops_extension，具体安装方法见安装指导。
sparse_attn_sharedkv_metadata	文档	√	√	√	×	√	AI Cpu	该算子为sparse_attn_sharedkv算子提供分核结果。 算子torch接口调用依赖torch_ops_extension，具体安装方法见安装指导。
blitz_sparse_attention	/	/	/	/	/	/	/	算子具体使用方法见README文档。
fused_infer_attention_score	/	/	/	/	/	/	/	算子具体使用方法见README文档。
typhoon_mla	/	/	/	/	/	/	/	算子具体使用方法见README文档。

自定义算子编译

1. 编译自定义算子包

   进入项目根目录，执行如下编译命令：

   bash build.sh --pkg --experimental --soc=${soc_version} --ops=${op_list}
   # 如要使用DeepSeek-V4，910c环境编译命令示例如下：
   # bash build.sh --pkg --experimental --soc=ascend910_93 --ops=compressor,quant_lightning_indexer,quant_lightning_indexer_metadata,sparse_attn_sharedkv,sparse_attn_sharedkv_metadata
   
   # 如要使用DeepSeek-V4，ascend950环境编译命令示例如下：
   # bash build.sh --pkg --experimental --soc=ascend950 --ops=compressor,quant_lightning_indexer,quant_lightning_indexer_metadata,kv_quant_sparse_attn_sharedkv,kv_quant_sparse_attn_sharedkv_metadata
   

   • --soc：${soc_version}表示NPU型号。Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件使用"ascend910b"，Atlas A3 训练系列产品/Atlas A3 推理系列产品使用"ascend910_93"。
   • --ops：自定义算子名称，多个自定义算子通过,分割。

   若提示如下信息，说明编译成功。

   Self-extractable archive "cann-ops-transformer-custom_linux.${arch}.run" successfully created.
   

   编译成功后，run包存放于项目根目录的build_out目录下。

2. 安装自定义算子包

   ./cann-ops-transformer-${vendor_name}_linux-${arch}.run --install-path=${install_path}
   

   自定义单算子包安装时若不指定参数--install-path，将默认安装于${ASCEND_HOME_PATH}/opp/vendors目录下。${ASCEND_HOME_PATH}已通过环境变量配置，表示CANN ops包安装路径。如指定--install-path=${install_path}，则自定义算子包将安装于${install_path}/vendors目录下。注意自定义算子包不支持卸载。安装完成后，需要激活自定义算子包环境变量或导出环境变量，具体方式如下：

   # experimental自定义单算子包自定义安装路径
   source ${install_path}/vendors/custom_transformer/bin/set_env.bash
   # experimental自定义单算子包默认安装路径${ASCEND_HOME_PATH}/opp/vendors
   # export LD_LIBRARY_PATH=${ASCEND_HOME_PATH}/opp/vendors/custom_transformer/op_api/lib/:$LD_LIBRARY_PATH
   

torch_ops_extension包编译安装（可选）

根据算子列表章节介绍，如需使用torch接口进行算子调用，则需前置安装torch_ops_extension包。

torch_ops_extension包为自定义算子提供了torch.ops的拓展接口，具体编译安装方法请见torch_ops_extension安装方法。

自定义算子执行

项目中各算子通过pytest验证各算子的功能是否正常，各算子的pytest调用方法如下表。

算子名	pytest示例
		compressor	示例
kv_quant_sparse_attn_sharedkv	示例
quant_lightning_indexer	示例
sparse_attn_sharedkv	示例