Ascend Transformer Boost RopeOperation C++ 调用示例详解：从环境配置到源码校验

Ascend Transformer Boost RopeOperation C 调用示例详解从环境配置到源码校验【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost旋转位置编码Rotary Position EmbeddingRoPE是大语言模型LLM中至关重要的位置编码方案用于在不引入额外参数的前提下为注意力机制注入序列位置信息。本文以 CANN ascend-transformer-boost 加速库中 RopeOperation C Demo 为核心完整讲解 RopeOperation 的环境配置、编译运行、参数语义、输入输出张量规格并结合仓库源码与测试用例剖析其底层校验逻辑。读完本文你将掌握在 Ascend 昇腾设备上通过 C 接口调用 RopeOperation 的完整实战方案并能够根据实际模型规模如 Qwen 系列灵活调整 demo 代码。一、RopeOperation 在 Transformer 加速库中的定位ascend-transformer-boost 是 CANN 提供的一款基于华为 Ascend AI 处理器的 Transformer 加速库面向 Transformer 定制化场景提供高性能融合算子。旋转位置编码算子 RopeOperation 是其中面向推理与训练场景的常用算子之一它对 queryQ和 keyK按位置施加旋转使注意力计算能够感知 token 的相对位置。从仓库的算子实现看RopeOperation 的宿主侧实现位于 src/ops/ops_infer/rope/rope_operation.cpp对应的算子参数结构体定义在 include/atb/infer_op_params.h//! //! \brief 旋转位置编码。hiddenSizeQ必须是hiddenSizeK的整数倍且满足hiddenSizeQ headDim * headNum。 //! struct RopeParam { //! \brief rope旋转系数对半旋转是2支持配置2、4或headDim / 2。 int32_t rotaryCoeff 4; //! \brief 训练用参数支持配置0或1 int32_t cosFormat 0; //! //! \brief 预留参数 //! uint8_t rsv[8] {0}; };RopeParam仅包含两个业务参数rotaryCoeff旋转系数。对半旋转取 2也可配置为 4 或headDim / 2默认值为 4cosFormat训练用参数支持 0 或 1默认值为 0rsv[8]预留参数必须保持为 0。示例目录 example/op_demo/rope 提供了 4 个 C demo 文件对应不同模型规模与输入场景是快速上手 RopeOperation 的最佳入口。二、环境准备加载 CANN 与 NNAL 运行环境运行 RopeOperation demo 前必须首先加载 CANN 与 NNALAscend Transformer Boost 加速库两套安装包的运行环境按顺序执行# 1. 加载 CANN 安装路径下的环境变量 source [CANN installation path]/set_env.sh # 默认路径 source /usr/local/Ascend/ascend-toolkit/set_env.sh # 2. 加载 NNAL 安装路径下的环境变量 source [NNAL installation path]/set_env.sh # 默认路径 source /usr/local/Ascend/nnal/atb/set_env.sh需要特别说明的是如果直接从加速库源码构建应改为 source 加速库源码编译产物的环境脚本例如source ./ascend-transformer-boost/output/atb/set_env.sh从仓库结构看该编译产物目录对应构建脚本scripts/set_env.sh、scripts/install.sh等所生成的output/atb目录其中包含了 ATB 的头文件与动态库-latb、-lasccendcl。环境变量加载完成后ATB_HOME_PATH与ASCEND_HOME_PATH会被指向正确的安装/编译产物位置供编译脚本使用。三、编译与运行 Demo示例目录下已提供构建脚本 example/op_demo/rope/build.sh进入目录后直接执行bash build.sh该脚本的核心逻辑为先通过 Python 探测当前环境中 PyTorch 的 CXX11 ABI 编译开关再据此设置_GLIBCXX_USE_CXX11_ABI宏随后用g编译并直接运行rope_democxx_abi$(python3 -c try: import torch print(1 if torch.compiled_with_cxx11_abi() else 0) except ImportError: print(1) ) g -D_GLIBCXX_USE_CXX11_ABI$cxx_abi -I ${ATB_HOME_PATH}/include -I ${ASCEND_HOME_PATH}/include \ -L ${ATB_HOME_PATH}/lib -L ${ASCEND_HOME_PATH}/lib64 \ rope_demo.cpp ../demo_util.h -l atb -l ascendcl -o rope_demo ./rope_demo编译时需注意cxx_abi 与_GLIBCXX_USE_CXX11_ABI的一致性当使用cxx_abi0默认时需要设置-D_GLIBCXX_USE_CXX11_ABI0g -D_GLIBCXX_USE_CXX11_ABI0 -I ...当使用cxx_abi1时需改为g -D_GLIBCXX_USE_CXX11_ABI1 -I ...这一点至关重要ATB 库本身是按某个 ABI 编译的如果 demo 编译时的 ABI 开关与库不一致链接或运行时会出现 std::string / std::vector 布局不兼容的问题。示例中的 build.sh 已通过探测 torch 的编译开关自动适配因此推荐直接使用该脚本。构建其他 demo 的方法构建脚本仅用于编译并运行rope_demo.cpp。若要编译运行rope_qwen_demo_0.cpp、rope_qwen_demo_1.cpp、rope_qwen_demo_2.cpp等其他示例只需将脚本中编译命令的源文件从rope_demo.cpp替换为对应的.cpp文件名即可。运行成功后终端会输出Rope demo success!字样。四、四个 Demo 场景与输入输出规格四个 demo 文件分别对应不同的应用场景读者可根据自身模型规模选择参考。下表汇总了它们的核心差异参数均取cosFormat 0Demo 文件数据精度Q/K 规模旋转系数适用形态rope_demo.cppfloat164×16 / 4×16rotaryCoeff 4最小功能验证示例rope_qwen_demo_0.cppbf161024×640 / 1024×128rotaryCoeff 2长序列 Qwen 场景prefillrope_qwen_demo_1.cppbf161×640 / 1×128rotaryCoeff 2单 token Qwen 场景decoderope_qwen_demo_2.cppbf165×640 / 5×128rotaryCoeff 2短批量 Qwen 场景适用范围提示rope_qwen_demo_0/1/2.cpp三个 Qwen 示例仅适用于 Atlas A2/A3 训练产品、Atlas 800I A2 推理产品以及 Atlas A3 推理产品。rope_demo.cpp为通用最小示例默认构建脚本即可直接编译运行。4.1 rope_demo.cppfloat16 最小示例参数设置cosFormat 0rotaryCoeff 4。输入张量Tensor数据类型数据格式Shapequeryfloat16nd[4, 16]keyfloat16nd[4, 16]cosfloat16nd[4, 8]sinfloat16nd[4, 8]seqlenuint32nd[1]输出张量Tensor数据类型数据格式ShaperopeQfloat16nd[4, 16]ropeKfloat16nd[4, 16]对应源码中定义的常量rope_demo.cppBATCH_SIZE 1、NTOKENS 4、HIDDENSIZEQ 16、HIDDENSIZEK 16、HEAD_SIZE 16。注意这里 cos/sin 的列数 8 HEAD_SIZE / 2这是因为对半旋转时每两个通道共用一个旋转角。4.2 rope_qwen_demo_0.cpp1024 token 批量场景参数设置cosFormat 0rotaryCoeff 2。输入张量Tensor数据类型数据格式Shapequerybf16nd[1024, 640]keybf16nd[1024, 128]cosbf16nd[1024, 128]sinbf16nd[1024, 128]seqlenuint32nd[1]输出张量Tensor数据类型数据格式ShaperopeQbf16nd[1024, 640]ropeKbf16nd[1024, 128]该场景 Q 隐藏层 640、K 隐藏层 128对应 Qwen 系列 40 个头 × 16 维 head dimK 侧 8 个头的典型配置。4.3 rope_qwen_demo_1.cpp单 token 解码场景参数设置cosFormat 0rotaryCoeff 2。输入张量Tensor数据类型数据格式Shapequerybf16nd[1, 640]keybf16nd[1, 128]cosbf16nd[1, 128]sinbf16nd[1, 128]seqlenuint32nd[1]输出张量Tensor数据类型数据格式ShaperopeQbf16nd[1, 640]ropeKbf16nd[1, 128]4.4 rope_qwen_demo_2.cpp5 token 短批量场景参数设置cosFormat 0rotaryCoeff 2。输入张量Tensor数据类型数据格式Shapequerybf16nd[5, 640]keybf16nd[5, 128]cosbf16nd[5, 128]sinbf16nd[5, 128]seqlenuint32nd[1]输出张量Tensor数据类型数据格式ShaperopeQbf16nd[5, 640]ropeKbf16nd[5, 128]五、Demo 代码核心流程逐段拆解四个 demo 的代码结构完全一致仅常量与参数不同。下面以 rope_demo.cpp 为例剖析完整调用链路。5.1 算子创建PrepareOperationatb::Status PrepareOperation(atb::Operation **ropeOp) { atb::infer::RopeParam opParam; opParam.rotaryCoeff 4; opParam.cosFormat 0; return atb::CreateOperation(opParam, ropeOp); }通过atb::CreateOperation(opParam, ropeOp)模板特化接口创建算子实例。从源码 rope_operation.cpp 可以看到该特化实现会执行预留参数校验OP_PARAM_RSV_CHECK并构造RopeOperation对象在 ASCEND_950 平台上还会额外校验cosFormat必须为 0并加载对应的 aclnn 执行函数。5.2 输入张量准备PrepareInTensorstd::vectorfloat qData(NTOKENS * HIDDENSIZEQ, 1.0); atb::Tensor tensorQ; CHECK_STATUS(CreateTensorFromVector(contextPtr, stream, qData, ACL_FLOAT16, aclFormat::ACL_FORMAT_ND, {NTOKENS, HIDDENSIZEQ}, tensorQ)); // ... key、cos、sin 同理 std::vectorint32_t seqLenHost(BATCH_SIZE, 4); atb::Tensor tensorSeqLen; CHECK_STATUS(CreateTensor(ACL_INT32, aclFormat::ACL_FORMAT_ND, {BATCH_SIZE}, tensorSeqLen)); tensorSeqLen.hostData seqLenHost.data(); inTensors {tensorQ, tensorK, tensorCos, tensorSin, tensorSeqLen};这里使用了仓库提供的公共工具头文件 example/op_demo/demo_util.hCreateTensorFromVector把宿主侧std::vectorfloat数据拷入设备内存并创建atb::Tensor。对于 float16/bf16 目标类型内部会先按 float 创建中间 tensor再调用 Elewise 的ELEWISE_CAST算子完成数据类型转换对应CastOp函数CreateTensor仅按数据类型、格式与 shape 分配设备内存并填充atb::Tensor描述desc.dtype、desc.format、desc.shape并调用atb::Utils::GetTensorSize计算dataSizeCHECK_STATUS宏统一错误码检查非 0 时打印文件与行号并根据错误码区间10 万99 万为 ACL 错误码提示查阅对应错误码文档。注意 seqlen 张量在这里通过hostData直接挂在宿主内存上值恒为 4它描述每个 batch 内实际参与旋转的 token 数。5.3 执行与资源释放Setup / Execute / 清理atb::Tensor ropeQ; CHECK_STATUS(CreateTensor(ACL_FLOAT16, aclFormat::ACL_FORMAT_ND, {NTOKENS, HIDDENSIZEQ}, ropeQ)); atb::Tensor ropeK; CHECK_STATUS(CreateTensor(ACL_FLOAT16, aclFormat::ACL_FORMAT_ND, {NTOKENS, HIDDENSIZEK}, ropeK)); variantPack.outTensors {ropeQ, ropeK}; uint64_t workspaceSize 0; CHECK_STATUS(ropeOp-Setup(variantPack, workspaceSize, context)); uint8_t *workspacePtr nullptr; if (workspaceSize 0) { CHECK_STATUS(aclrtMalloc((void **)(workspacePtr), workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST)); } ropeOp-Execute(variantPack, workspacePtr, workspaceSize, context); CHECK_STATUS(aclrtSynchronizeStream(stream));Setup阶段会执行输入/输出张量的形状与参数校验详见下一节并返回算子所需的 workspace 大小Execute阶段传入已分配的 workspace 在指定 stream 上异步执行aclrtSynchronizeStream保证 device 侧任务完成之后依次释放输入张量、workspace、销毁算子与 context、aclFinalize收尾。主函数开头还有一段标准的 ACL/ATB 初始化序列读者在编写自己的调用程序时可直接复用CHECK_STATUS(aclInit(nullptr)); int32_t deviceId 0; CHECK_STATUS(aclrtSetDevice(deviceId)); atb::Context *context nullptr; CHECK_STATUS(atb::CreateContext(context)); void *stream nullptr; CHECK_STATUS(aclrtCreateStream(stream)); context-SetExecuteStream(stream);六、参数与张量校验源码级约束demo 只是合法输入的示例实际开发中若想修改 shape 或参数务必遵守算子实现中的校验规则。以下约束均可在 src/ops/ops_infer/rope/rope_operation.cpp 中找到实现1. rotaryCoeff 取值约束ParamCheckrotaryCoeff只支持 2、4 或cos.dims[1]headDim同时要求cos.dims[1]与sin.dims[1]都能被rotaryCoeff整除否则返回ERROR_INVALID_PARAM。2. cosFormat 取值约束cosFormat仅支持 0 或 1其余取值报ERROR_INVALID_PARAM在 ASCEND_950 平台上还额外限制只能为 0见 CreateOperation 实现。3. 维度约束DimCheckquery / key 支持 2 维或 4 维 shapecos / sin 必须为 2 维seqlen 必须为 1 维query 与 key 的 token 数必须一致4 维时取dims[0] * dims[1]cos 与 sin 的 shape 必须完全相同非 ASCEND_950 平台下cos/sin 的列数head_size 或 head_size/2必须在 [16, 4096] 区间内。4. 隐藏层约束HiddenSizeCheckhiddenSizeQ必须是hiddenSizeK的整数倍对应hiddenSizeQ headDim * headNum的约束hiddenSizeK必须是 headDimcos 列数的整数倍。5. shape 推导InferShapeImpl输出 ropeQ、ropeK 的 shape 直接继承输入 query、key 的 shape无需额外计算见 InferShapeImpl。6. Runner 分派CreateRunner算子执行时ASCEND_950 平台走RopeAclnnRunner其余平台走 RunnerPool 中的RopeOpsRunner见 CreateRunner两者共用同一套宿主侧 shape 校验逻辑。七、数据生成与精度验证Python 测试参考README 明确提示示例中生成的数据并不代表真实输出demo 用全 1 填充输入仅用于验证流程可跑通。如需验证算子的实际计算结果请参考仓库根目录下的 Python 用例目录tests/apitest/opstest/python/operations/rope/该目录下的测试用例如 test_rope_operation.py、test_rope_operation2.py覆盖了多种rotaryCoeff/cosFormat组合rotaryCoeff 4的基础用例test_rope_operation.pyrotaryCoeff 64、rotaryCoeff 2, cosFormat 0、按 headDim 推导 rotaryCoeff 的用例test_rope_operation2.pyrotaryCoeff 2, cosFormat 1的训练场景用例test_rope_operation3.py。从测试代码可以看出rotaryCoeff与headDim的关系为rotaryStride headDim // rotaryCoeff、rotaryTimesPerHead rotaryCoeff / 2这与前文rotaryCoeff支持headDim / 2的约束相互印证。编写自己的 C 测试时可参考这些 Python 用例中的输入构造与参考实现进行端到端对比。八、常见问题与使用建议环境变量未加载编译时找不到atb/atb_infer.h头文件或链接时找不到-latb请确认已按第二节顺序 source 了 CANN 与 NNAL或源码编译产物的set_env.sh。ABI 不匹配若编译通过但运行崩溃或行为异常优先检查_GLIBCXX_USE_CXX11_ABI是否与库一致推荐直接使用仓库提供的 build.sh 自动探测。shape 不满足校验修改 Q/K 规模时务必保证hiddenSizeQ是hiddenSizeK的整数倍、hiddenSizeK是 cos 列数的整数倍且 cos/sin shape 完全一致否则Setup会返回ERROR_INVALID_PARAM/ERROR_INVALID_TENSOR_DIM等错误码。平台适配Qwen 三个 demo 仅适用于 Atlas A2/A3 训练产品、Atlas 800I A2 推理产品及 Atlas A3 推理产品rope_demo.cpp作为最小示例适用面更广。ASCEND_950 平台下cosFormat仅支持 0。数据真实性demo 输出为流程验证数据真实精度对比请使用tests/apitest/opstest/python/operations/rope/下的 Python 用例。通过本文的示例代码、参数表格与源码校验逻辑你可以在此基础上快速定制自己的 RopeOperation 调用程序并将其接入更大规模的 Transformer 推理或训练流水线。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考