news 2026/5/9 14:26:29

ATB自注意力C++示例

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
ATB自注意力C++示例

加速库SelfAttentionOperation C++ Demo

【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost

介绍

该目录下为加速库SelfAttentionOperation C++调用示例。

使用说明

  • 首先source 对应的CANN和nnal包的安装路径

    1. source [cann安装路径]/set_env.sh 默认：source /usr/local/Ascend/ascend-toolkit/set_env.sh
    2. source [nnal安装路径]/set_env.sh
      默认：source /usr/local/Ascend/nnal/atb/set_env.sh
      ①. 如果使用加速库源码编译，source [加速库源码路径]/output/atb/set_env.sh
      例如： source ./ascend-transformer-boost/output/atb/set_env.sh

  • 运行demo

    bash build.sh

    注意：

    • 使用cxx_abi=0（默认）时，设置D_GLIBCXX_USE_CXX11_ABI为0，即：

      g++ -D_GLIBCXX_USE_CXX11_ABI=0 -I ...

    • 使用cxx_abi=1时，更改D_GLIBCXX_USE_CXX11_ABI为1，即：

      g++ -D_GLIBCXX_USE_CXX11_ABI=1 -I ...

    • 生成的二进制文件***_demo可以额外传入一个int参数作为deviceId，默认为0，如：

      ./self_attention_encoder_demo 0

    • 提供的build脚本仅用于编译和运行self_attention_encoder_demo.cpp，如需编译其他demo，需要替换“self_attention_encoder_demo”为对应的cpp文件名

额外说明

示例中生成的数据不代表实际场景，如需数据生成参考请查看根目录下的python用例目录： tests/apitest/opstest/python/operations/self_attention/

产品支持情况

SelfAttention在Atlas A2/A3系列仅支持部分场景，且Encoder场景在Atlas 推理系列产品上调用与Atlas A2/A3有所区别

场景说明

提供demo分别对应不同产品/场景，具体场景区别可以参见官网，编译运行时需要对应更改build脚本：

self_attention_encoder_demo.cpp
  • 场景：FA Encoder基础场景，分开传入key，CacheK，value和CacheV
  • 默认编译脚本可编译运行
  • 该demo仅支持在Atlas A2/A3上运行
  • demo中使用全量的上三角mask演示

参数设置：

成员名称取值
headNum32
kvHeadNum32
qkScale1/sqrt(128)
calcTypeENCODER
kernelTypeKERNELTYPE_DEFAULT
maskTypeMASK_TYPE_NORM

数据规格：

tensor名字数据类型数据格式维度信息cpu/npu
queryfloat16nd[160,2048]npu
keyfloat16nd[160,2048]npu
valuefloat16nd[160,2048]npu
cacheKfloat16nz[1, 10, 1024, 2048]npu
cacheVfloat16nz[1, 10, 1024, 2048]npu
attentionMaskfloat16nd[10, 1024, 1024]npu
tokenOffsetint32nd[10]cpu
seqLenint32nd[10]cpu
layerIdint32nd[1]npu
Output
outputfloat16nd[160, 2048]npu
  • q，k，v第一维度为总词元长度，即sum(seqlen)
  • q，k，v第二维度headNum，headSize合轴，实际为headNum(32) $\times$ headSize(128)
self_attention_encoder_inference_demo.cpp
  • 场景：fa encoder基础场景在Atlas推理系列上的实现，分开传入key，CacheK，value和CacheV
  • 该demo仅支持在Atlas推理系列上运行

参数设置：

成员名称取值
headNum16
kvHeadNum16
qkScale1/sqrt(128)
calcTypeENCODER
kernelTypeKERNELTYPE_DEFAULT
maskTypeMASK_TYPE_UNDEFINED

数据规格：

tensor名字数据类型数据格式维度信息cpu/npu
queryfloat16nd[16, 256]npu
keyfloat16nd[16, 256]npu
valuefloat16nd[16, 256]npu
cacheKfloat16nz[1, 1, 16, 256, 16]npu
cacheVfloat16nz[1, 1, 16, 256, 16]npu
tokenOffsetint32nd[1]cpu
seqLenint32nd[1]cpu
layerIdint32nd[1]npu
Output
outputfloat16nd[16, 256]npu
self_attention_pa_encoder_demo.cpp
  • 场景：FA使用PA Encoder的场景，使用FA输入，只需传入key，value
    • 传入不同的headNum，kvHeadNum且headNum可以被kvHeadNum整除时开启GQA（Grouped Query Attention）。
  • 该demo仅支持在Atlas A2/A3系列上运行
  • demo中使用全量的上三角mask演示

参数设置：

成员名称取值
headNum32
kvHeadNum16
qkScale1/sqrt(128)
calcTypePA_ENCODER
kernelTypeKERNELTYPE_DEFAULT
maskTypeMASK_TYPE_NORM

数据规格：

tensor名字数据类型数据格式维度信息cpu/npu
queryfloat16nd[3072, 32, 128]npu
keyfloat16nd[3072, 16, 128]npu
valuefloat16nd[3072, 16, 128]npu
attentionMaskfloat16nd[4, 1024, 1024]npu
seqLenint32nd[4]cpu
Output
outputfloat16nd[3072, 32, 128]npu
self_attention_pa_encoder_qwen_demo.cpp
  • 场景：FA使用PA Encoder的场景，使用FA输入，只需传入key，value
    • 传入不同的headNum，kvHeadNum且headNum可以被kvHeadNum整除时，开启GQA（Grouped Query Attention）。
  • 该demo仅支持在Atlas A2/A3系列上运行
  • demo中为了应对长序列的情况，使用压缩上三角mask演示

参数设置：

成员名称取值
headNum5
kvHeadNum1
qkScale1/sqrt(128)
isTriuMask1
calcTypePA_ENCODER
kernelTypeKERNELTYPE_DEFAULT
maskTypeMASK_TYPE_NORM_COMPRESS

数据规格：

tensor名字数据类型数据格式维度信息cpu/npu
queryfloat16nd[1024, 5, 128]npu
keyfloat16nd[1024, 1, 128]npu
valuefloat16nd[1024, 1, 128]npu
attentionMaskfloat16nd[128, 128]npu
seqLenint32nd[1]cpu
Output
outputfloat16nd[1024, 5, 128]npu
self_attention_prefix_encoder_demo.cpp
  • 场景：FA使用Prefix Encoder的场景，传入PA的依据blockTables存放的key，value
    • 此场景支持q和kv不等长，但是要求： $$\forall i \lt len(seqLen)，kvSeqLen[i] - seqLen[i] = 0 \ (mod \ 128) $$
  • 该demo仅支持在Atlas A2/A3系列上运行
  • demo中使用Alibi上三角mask叠加bias slopes演示

参数设置：

成员名称取值
headNum32
kvHeadNum8
qkScale1/sqrt(128)
isTriuMask1
calcTypePREFIX_ENCODER
kernelTypeKERNELTYPE_HIGH_PRECISION
maskTypeMASK_TYPE_ALIBI_COMPRESS

数据规格：

tensor名字数据类型数据格式维度信息cpu/npu
queryfloat16nd[96, 32, 128]npu
keyfloat16nd[480, 128, 8, 128]npu
valuefloat16nd[480, 128, 8, 128]npu
blockTablesfloat16nd[4, 4]npu
maskfloat16nd[32, 96, 128]npu
seqLenint32nd[4]cpu
kvSeqLenint32nd[4]cpu
slopesfloat32nd[128]npu
Output
outputfloat16nd[96, 32, 128]npu

【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

版权声明: 本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若内容造成侵权/违法违规/事实不符，请联系邮箱：809451989@qq.com进行投诉反馈，一经查实，立即删除！
网站建设 2026/5/9 14:25:56

CANN驱动获取设备逻辑ID

dcmi_get_device_logic_id 【免费下载链接】driver 本项目是CANN提供的驱动模块，实现基础驱动和资源管理及调度等功能，使能昇腾芯片。 项目地址: https://gitcode.com/cann/driver 函数原型 int dcmi_get_device_logic_id(int *device_logic_id,…

作者头像 李华
网站建设 2026/5/9 14:23:03

CANN/pto-isa内存操作（GM ＜-＞ Tile）

Memory (GM <-> Tile) 【免费下载链接】pto-isa Parallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations a…

作者头像 李华
网站建设 2026/5/9 14:21:38

探索率如何驱动多智能体Q学习中的自发合作行为

1. 项目概述&#xff1a;当AI学会“社交”&#xff0c;探索与合作的微妙平衡最近在复现一个多智能体强化学习的实验时&#xff0c;我遇到了一个非常有趣的现象&#xff1a;几个原本各自为战、只追求自身奖励最大化的AI智能体&#xff0c;在没有任何显式通信或中央协调的情况下&…

作者头像 李华
网站建设 2026/5/9 14:20:41

CANN/metadef动态算子自动映射

AutoMappingByOpFnDynamic 【免费下载链接】metadef Ascend Metadata Definition 项目地址: https://gitcode.com/cann/metadef 函数功能 动态输入/输出算子的自动映射回调函数。 函数原型 Status AutoMappingByOpFnDynamic(const ge::Operator &op_src, ge::Ope…

作者头像 李华
网站建设 2026/5/9 14:19:33

无监督模型评估新范式：可靠性-完整性曲线与AURCC指标详解

1. 项目概述&#xff1a;当模型评估遇上“无米之炊”在机器学习项目的落地过程中&#xff0c;我们总会遇到一个经典的“鸡生蛋还是蛋生鸡”的困境&#xff1a;要评估一个模型的性能好坏&#xff0c;你需要有高质量的标注数据&#xff1b;但很多时候&#xff0c;获取这些标注数据…

作者头像 李华
网站建设 2026/5/9 14:18:04

CANN社区新建仓与仓开放操作指引

CANN社区新建仓与仓开放操作指引 【免费下载链接】community 本项目是CANN开源社区的核心管理仓库&#xff0c;包含社区的治理章程、治理组织、通用操作指引及流程规范等基础信息 项目地址: https://gitcode.com/cann/community 本指引结合TSC、PMC会议时间&#xff0c;为…

作者头像 李华