算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载导读BlockAttnResUpdate是 CANN ops-transformerTransformer 类大模型算子库中用于历史残差注意力Attention Residuals两段式计算阶段二的 NPU 算子它把增量delta原地累加进局部注意力结果partial_block计算更新后结果与pseudo_query的归一化相关分数RMSNorm score再与阶段一算子block_attn_res_prepare产出的在线 Softmax 状态numerator、logit_max、exp_sum融合输出注意力结果h。本文基于 attention/block_attn_res_update/README.md 展开结合本仓库的 Kernel、Host Tiling、算子定义与测试代码完整讲解其数学原理、参数约束、aclnn / PyTorch 两种调用方式与源码级实现细节。读完本文你将掌握该算子在 Ascend 950PR/950DT 上的正确使用姿势、两阶段在线 Softmax 融合的数值逻辑以及 Kernel 端 UB 双缓冲与向量寄存器的优化思路。算子定位两段式注意力残差计算的第二阶段在 block 注意力残差Block Attn Residuals这类大模型推理/训练场景中注意力结果往往需要先准备、再更新两个步骤阶段一block_attn_res_prepare基于历史信息生成在线 Softmax 状态——numerator历史加权和对应 O、logit_max历史最大 logit对应 M、exp_sum历史指数和分母对应 L详见同仓库的 block_attn_res_prepare 模块。阶段二BlockAttnResUpdate将新到来的增量delta累加到已累计的局部残差partial_block上并在此基础上完成与新查询pseudo_query的分数计算最后把当前结果与历史在线 Softmax 状态做数值稳定的融合产出当前层的注意力输出h同时把更新后的残差原地写回partial_block。从算子注册代码 block_attn_res_update_def.cpp 可以看到partial_block同时被声明为输入与输出同名引用关系Kernel 通过该引用直接原地读写这正是更新语义的实现基础。计算原理与数学公式设 $T$ 为 token 数$D$ 为隐藏维度$t \in [0,T)$$j \in [0,D)$。算子整体分为四个步骤1. 更新局部注意力结果对每个 token $t$将增量delta以 FP32 精度累加到partial_block$$ p_t partial_block_t \operatorname{float}(delta_t) $$注意partial_block是 FLOATFP32、delta是 BFLOAT16Kernel 端通过Cast BF16→FP32后以 FP32 完成加法见 block_attn_res_update_full_d.h 中BARU_CAST_BF16_TO_FP32的 CastTrait 定义保证累加精度。2. 计算 RMS 归一化相关分数先计算更新结果 $p_t$ 的 RMS$$ rms_t \sqrt{\frac{1}{D}\sum_{j0}^{D-1}p_{t,j}^{2} eps} $$再计算 $p_t$ 与伪查询pseudo_query的归一化相关分数$$ score_t \frac{\sum_{j0}^{D-1}p_{t,j} \cdot pseudo_query_j}{rms_t} $$eps为 RMS 计算中的数值稳定项必须是有限正数默认 $10^{-6}$。从源码看Host Tiling 会预先计算 $invD 1/D$ 并写入 BlockAttnResUpdateTilingDataKernel 端使用Muls(squareSum, invD)以乘法替代除法再Sqrt求根、用DivPRECISION_0ULP_FTZ_TRUE模式得到 score。3. 与历史在线 Softmax 状态融合对每个 token将当前score_t与历史logit_max_t做数值稳定的在线 Softmax 合并$$ max_t \max(logit_max_t, score_t) $$$$ alpha_t \exp(logit_max_t-max_t), \qquad beta_t \exp(score_t-max_t) $$$$ denominator_t exp_sum_t \cdot alpha_t beta_t $$Kernel 端依次使用Max、ExpSub$\exp(a-b)$ 原子指令、MulDstAdd完成上述链路的矢量计算随后对alpha、beta做归一化乘以 $1/denominator$使每个 D 循环只需MulMulDstAdd即可完成加权融合见 Phase 2 三个特化路径的注释。4. 生成输出$$ partial_block_out_t p_t $$$$ h_t \operatorname{bfloat16}\left( numerator_t \cdot \frac{alpha_t}{denominator_t} p_t \cdot \frac{beta_t}{denominator_t}\right) $$其中 $h_t$ 为 $D$ 维向量输出数据类型为 BFLOAT16Kernel 端使用Cast FP32→BF16、CAST_RINT舍入模式。更新的partial_block通过partial_block_out输出与输入同名引用、复用存储。历史状态的语义约束对于 token $t$ 的非空历史集合 $\mathcal{H}t$设历史 logit 为 $s{t,i}$、对应 $D$ 维 value 为 $v_{t,i}$则输入的在线 Softmax 状态必须满足$$ logit_max_t \max_{i \in \mathcal{H}t}s{t,i} $$$$ exp_sum_t \sum_{i \in \mathcal{H}t}\exp(s{t,i}-logit_max_t) 0 $$$$ numerator_t \sum_{i \in \mathcal{H}t}\exp(s{t,i}-logit_max_t)v_{t,i} $$本算子只读取这份历史状态并将当前 $p_t$ 与其融合不更新numerator、logit_max或exp_sum——阶段二只产出当层输出状态的推进由外层准备阶段/更新阶段交替负责。参数说明下表来自 attention/block_attn_res_update/README.md 并补充接口文档细节参数名输入/输出/属性描述数据类型数据格式partial_block输入待更新的局部注意力结果公式中的 partial_blockshape 为 $(T, D)$与同名输出复用存储FLOATNDdelta输入局部注意力结果的增量公式中的 deltashape 为 $(T, D)$BFLOAT16NDpseudo_query输入用于计算归一化相关分数的伪 Query 向量公式中的 pseudo_queryshape 为 $(D)$FLOATNDnumerator输入非空历史在线 Softmax 状态对应的分子公式中的 numeratorshape 为 $(T, D)$本算子只读FLOATNDlogit_max输入非空历史在线 Softmax 状态中每个 token 的有限最大值公式中的 logit_maxshape 为 $(T)$本算子只读FLOATNDexp_sum输入非空历史在线 Softmax 状态中每个 token 的有限正指数和公式中的 exp_sumshape 为 $(T)$本算子只读FLOATNDpartial_block输出更新后的局部注意力结果公式中的 partial_block_outshape 为 $(T, D)$与同名输入构成引用关系FLOATNDh输出融合更新后的局部结果与已有在线 Softmax 状态得到的注意力结果公式中的 hshape 为 $(T, D)$BFLOAT16NDeps可选属性RMS 计算中的稳定项必须为有限值且大于 0默认值为 $10^{-6}$FLOAT-对应的算子定义可在 block_attn_res_update_def.cpp 中核实所有 Tensor 均声明AutoContiguous()eps以Attr(eps).AttrType(OPTIONAL).Float(1e-6F)注册且不作为 Kernel Tensor 参数而是被序列化进BlockAttnResUpdateTilingDatafloat eps字段随 Tiling 数据下发。约束说明使用该算子前必须满足以下约束shape 要求$T$ 和 $D$ 必须是已知正整数$D$ 取值范围为 $[1, 8192]$接口文档层面则放宽为 $T \ge 0$、$0 \le D \le 8192$支持空 Tensor。partial_block、delta、numerator及输出partial_block、h的 shape 必须相同均为 $(T, D)$pseudo_query、logit_max、exp_sum的 shape 必须分别为 $(D)$、$(T)$、$(T)$。格式要求所有输入的原始格式和存储格式均必须为 ND所有输入、输出的 StorageShape 必须与 OriginShape 完全一致所有 Tensor 必须连续非连续 Tensor 会在第一段接口校验时报ACLNN_ERR_PARAM_INVALID。历史状态要求numerator、logit_max、exp_sum是位于 GMGlobal Memory中的运行时只读状态Host Tiling 不读取或校验其元素值。调用者必须保证每个 token 的历史状态非空logit_max和numerator中的元素为有限值exp_sum为有限值且严格大于 0并且三者来自同一份历史状态。当前不支持以logit_max-inf、exp_sum0和零numerator表示的空历史状态。寻址范围$T \times D$ 不能超出有符号 64 位 Kernel GM 元素偏移的表示范围按核切分后的 $T$ 大小不能超出uint32_t的表示范围。UB 容量一行完整的 $D$ 维数据必须能够以双缓冲方式放入统一缓冲区Unified BufferUB否则 Tiling 失败。原地更新语义partial_block为原地更新参数。框架调用时由引用关系保证输入、输出复用存储Kernel 直接通过partial_block输入地址读写不访问对应的输出 ABI 地址见 Kernel 入口中对partial_block_ref参数的(void)partial_block_ref;显式忽略处理。workspace当前 Host Tiling 申请的 workspace 大小为 0Kernel 不访问 workspace 地址。数值行为RMS 计算使用 FP32 的sqrt和除法不进行牛顿迭代也不对输入 Tensor 中的零值、极值、NaN 或 Inf 做额外处理。eps必须是有限正数否则第一段接口返回ACLNN_ERR_PARAM_INVALID。产品支持方面本算子当前仅支持Ascend 950PR/Ascend 950DTAtlas A3、A2、200I/500 A2、推理系列、训练系列产品均不支持产品不在支持范围时返回ACLNN_ERR_RUNTIME_ERROR错误码 361001。调用方式一aclnn 接口C与 CANN 惯例一致该算子采用两段式接口。必须先调用aclnnBlockAttnResUpdateGetWorkspaceSize获取 workspace 大小和执行器再调用aclnnBlockAttnResUpdate执行计算。函数原型如下aclnnStatus aclnnBlockAttnResUpdateGetWorkspaceSize( aclTensor *partialBlockRef, const aclTensor *delta, const aclTensor *pseudoQuery, const aclTensor *numerator, const aclTensor *logitMax, const aclTensor *expSum, double eps, aclTensor *h, uint64_t *workspaceSize, aclOpExecutor **executor); aclnnStatus aclnnBlockAttnResUpdate( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);第一段接口完成入参校验主要错误码返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001partialBlockRef、delta、pseudoQuery、numerator、logitMax、expSum、h、workspaceSize或executor为空ACLNN_ERR_PARAM_INVALID161002dtype/format/dimension/shape 关系不满足约束$T0$ 或 $D$ 不在 $[0,8192]$非空场景 $T \times D$ 超出int64_t任一 Tensor 非连续eps不是有限正数ACLNN_ERR_RUNTIME_ERROR361001产品型号不在支持范围内完整的可编译调用示例见 test_aclnn_block_attn_res_update.cpp该文件也收录于接口文档 aclnnBlockAttnResUpdate.md。其调用骨架为aclInit/aclrtSetDevice/aclrtCreateStream初始化环境用aclrtMalloc分配各 Tensor 的 Device 内存aclrtMemcpy载入 Host 数据aclCreateTensor构造aclTensorstrides 按行主序计算保证连续调用第一段接口aclnnBlockAttnResUpdateGetWorkspaceSize(partialBlockRef, delta, pseudoQuery, numerator, logitMax, expSum, eps, h, workspaceSize, executor)按返回的workspaceSize若大于 0分配 workspace调用第二段接口aclnnBlockAttnResUpdate(workspaceAddr, workspaceSize, executor, stream)aclrtSynchronizeStream同步后把原地更新的partialBlockRef与输出h拷回 Host 验证。样例中partialBlockRef初始化为 0.25F、delta为 BF16(0.125F)执行后打印partialBlockRef[0]原地更新结果与h[0]。接口文档同时说明所有 Tensor 输入都支持空 Tensor但必须继续满足彼此的 shape 关系不能单独任意置空该接口默认为确定性实现。调用方式二PyTorch API在安装了torch_npu与cann_ops_transformer的 Python 环境中可通过 Python 封装直接调用实现见 torch_extension/block_attn_res_update.py 及 torchapi_block_attn_res_update.md 接口文档cann_ops_transformer.block_attn_res_update( partial_block, delta, pseudo_query, numerator, logit_max, exp_sum, *, eps1e-6, ) - Tensor参数与数据类型组合必须严格匹配| partial_block | delta | pseudo_query | numerator | logit_max | exp_sum | h | | ------------- | ----- | ------------ | --------- | --------- | ------- | - | |torch.float32|torch.bfloat16|torch.float32|torch.float32|torch.float32|torch.float32|torch.bfloat16|h为返回值 Tensor数据类型和 shape 与delta一致且连续。eps省略时使用默认值1e-6Torch C bridge 在调用 ACLNN 前会将其显式转换为 float32。单算子模式调用示例import torch import torch_npu import cann_ops_transformer T 32 D 7168 partial_block torch.rand((T, D), dtypetorch.float32).npu() delta torch.rand((T, D), dtypetorch.bfloat16).npu() pseudo_query torch.rand((D,), dtypetorch.float32).npu() numerator torch.rand((T, D), dtypetorch.float32).npu() logit_max torch.rand((T,), dtypetorch.float32).npu() exp_sum torch.rand((T,), dtypetorch.float32).npu() h cann_ops_transformer.block_attn_res_update( partial_block, delta, pseudo_query, numerator, logit_max, exp_sum, )aclgraph 模式调用示例通过torch.compile与npugraph_ex后端编译适合将算子编入静态图获得更好性能import torch import torch_npu import cann_ops_transformer class OneOp(torch.nn.Module): def forward(self, partial_block, delta, pseudo_query, numerator, logit_max, exp_sum): return cann_ops_transformer.block_attn_res_update( partial_block, delta, pseudo_query, numerator, logit_max, exp_sum, eps1e-6, ) compiled_op torch.compile( OneOp(), backendnpugraph_ex, fullgraphTrue, dynamicFalse, options{static_kernel_compile: True}, ) T 8 D 7168 h compiled_op( torch.rand((T, D), dtypetorch.float32).npu(), torch.rand((T, D), dtypetorch.bfloat16).npu(), torch.rand((D,), dtypetorch.float32).npu(), torch.rand((T, D), dtypetorch.float32).npu(), torch.rand((T,), dtypetorch.float32).npu(), torch.rand((T,), dtypetorch.float32).npu(), )该接口支持推理场景、单算子模式与 aclgraph 模式默认支持确定性计算。源码级实现解析Kernel 入口与两阶段调度Kernel 入口 block_attn_res_update.cpp 是一个模板化的__global__ __aicore__函数SINGLE_TILE模板参数区分单 Tile/多 Tile 两种执行形态模板参数声明见 block_attn_res_update_tiling_key.h。Kernel 通过GET_TILING_DATA_WITH_STRUCT解析 Host 下发的BlockAttnResUpdateTilingData随后实例化BlockAttnResUpdateFullDSINGLE_TILE完成全部计算。值得注意的是入口对partial_block_ref与workspace参数显式(void)忽略——再次印证原地更新 零 workspace的约束描述。核心计算类 block_attn_res_update_full_d.h 将计算划分为两个阶段Phase 1对 UB 中每个 tile加载partial与deltaBF16→FP32 后相加并立即写回partial原地更新同时累积平方和与点积得到rms与score并将每个 T 行的 score 存入 stats 平面。Phase 2从 stats 平面读取logit_max、exp_sum与 Phase 1 产出的score执行Max/ExpSub/MulDstAdd链计算alpha、beta与归一化分母再对partial与numerator做Mul MulDstAdd加权融合最后 FP32→BF16 输出到h。按 D 宽度的三条特化路径Kernel 根据 $D$ 与向量寄存器宽度一个 256 字节向量寄存器可容纳 64 个 FP32 元素即BARU_VREG_FP32_ELEMENTS 64的关系选择计算路径见Process()中的分支One-VREG$D \le 64$query 只占一个向量寄存器在 T 循环外加载一次并做一次 tail 清零循环体内只维护一个累加链Two-VREG$64 D \le 128$两个相邻向量寄存器组成一对两条独立寄存器链暴露给 RVEC 双发射dual issue第二个向量寄存器也只在循环外做一次 tail 清零Generic$D 128$按2 × VL成对处理完整向量奇数完整向量与 tail 配对并区分纯奇数完整与纯 tail两种互斥的单 VL 剩余分支始终维持两条独立累加链以提高指令级并行。UB 布局与多 Tile 双缓冲InitUbLayout()定义的 UB 布局为[query][buffer 0][buffer 1]每个 buffer 依次包含partial、delta/h、numerator与 3 个 stats 平面logit_max、exp_sum、score按 T 主序排列见BARU_LOGIT_MAX_PLANE_INDEX0、BARU_EXP_SUM_PLANE_INDEX1、BARU_SCORE_PLANE_INDEX2。stats 平面的行间步长statsTStride由 Host 对齐到 32 字节边界便于整行加载/广播。多 Tile 模式下两个互不重叠的 UB 缓冲交替使用bufferId ^ 1U并通过MTE3_MTE2、MTE2_V、V_MTE3事件令牌保证Phase 1 的拷贝入与计算、partial的拷贝出、Phase 2 的独立拷贝入写不同的 UB 区域以及最终h的拷贝出能够互相重叠流水。注释明确说明Phase 1 与 Phase 2 在 V 上按程序顺序执行因此delta在其 UB 区域被h复用前必然已被完全消费。Host 端Tiling 与 InferShapeHost 端 Tiling 入口 通过TilingRegistryArch::DoTilingImpl派发到架构相关实现最终填充 BlockAttnResUpdateTilingData字段含义dSize隐藏维度 DtPerCore/lastTPerCore除最后一个核外每核处理的 T 行数 / 最后一个核处理的剩余行数tileT单个 UB tile 最多处理的 T 行数statsTStrideFP32 stats 平面间的元素步长32 字节对齐epsRMS 稳定项由属性序列化而来invDHost 预计算的 $1/D$供 RMS 归一化使用usedCoreNum实际参与计算的核数Kernel 端按blockIdx * tPerCore切分 T 维最后一个核处理lastTPerCore行blockIdx usedCoreNum的核直接返回。InferShape 实现见 block_attn_res_update_infershape.cpp输出partial_block_ref的 shape 直接复制输入partial_block的 shape输出h的 shape 同样继承partial_block数据类型分别为 FLOAT引用输入与 BF16。这与 README 中输出与输入 shape 一致的约束一一对应。测试覆盖仓库为该算子提供了完整的单测与 ST 用例可结合源码与测试进一步验证行为UTop_apitest_aclnn_block_attn_res_update.cpp 与对应的 CSV 用例表 test_aclnn_block_attn_res_update.csv覆盖 aclnn 接口的入参校验与数值正确性UTop_hosttest_block_attn_res_update_infershape.cpp 与 test_block_attn_res_update_tiling.cpp分别校验 shape/dtype 推断与 Tiling 参数UTop_kerneltest_block_attn_res_update.cpp配合 CPU debug stub 校验 Kernel 计算路径STttk_aclnn_block_attn_res_update_st.csv、ttk_e2e_block_attn_res_update_st.csv、ttk_kernel_block_attn_res_update_st.csv 三张用例表覆盖 aclnn 层、端到端与 Kernel 层三个层面的昇腾环境验证。快速启动与构建在已配置好 CANN 环境变量的前提下于仓库根目录执行以下命令即可单独构建本算子README 原始说明# 在仓库根目录执行假设已准备好环境变量 bash build.sh --soc${soc_version} --opsblock_attn_res_update其中${soc_version}需替换为目标昇腾 SoC 型号如ascend950系列且必须落在上文产品支持范围内的产品型号。构建产物可用于运行tests/下的 UT/ST 用例或在 examples/arch35/test_aclnn_block_attn_res_update.cpp 基础上编写自己的调用程序编译与运行步骤可参考仓库 docs/QUICKSTART.md 及接口文档中的编译运行样例指引。小结BlockAttnResUpdate是 block 注意力残差两段式计算中的更新与输出环节以 FP32 精度原地累加增量、以 RMSNorm 相关分数刻画当前残差与查询的匹配程度再通过数值稳定的在线 Softmax 公式把当前结果与历史状态numerator/logit_max/exp_sum融合成当层输出h。从本文可见该算子在数学上精确对应一整套在线 Softmax 增量融合协议在工程实现上则通过按 D 宽度的向量化特化、UB 双缓冲与事件流水、Host 侧invD/statsTStride预计算等手法把这一更新-打分-融合流程压进 Ascend 950 的 AIV 流水线。理解上述约束与调用范式后即可将该算子正确接入自定义的注意力残差计算图中。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐cann-samples BlockAttnResAscend 950 上 Prepare Update 两阶段残差注意力算子的性能优化实战cann samples BlockAttnResAscend 950 上 Prepare Update 两阶段残差注意力算子的性能优化实战 本文围绕 b示例工程人工智能CANNCANN ops-transformer 中的 RainFusionAttention 算子块级稀疏注意力原理、ACLNN 两段式接口与实战调用CANN ops transformer 中的 RainFusionAttention 算子块级稀疏注意力原理、ACLNN 两段式接口与实战调用 本篇技术指南算子库人工智能深度学习AscendCANN ops-transformer 算子解析FlashAttentionScoreGrad 反向注意力算子原理与 aclnn 调用实战CANN ops transformer 算子解析FlashAttentionScoreGrad 反向注意力算子原理与 aclnn 调用实战 导读 Flash算子库人工智能深度学习Ascend上一篇ballcat事件驱动架构基于Spring事件的解耦设计下一篇MSBuild 反模式 AP-07分析器与构建工具包必须使用 PrivateAssetsall——阻止私有依赖泄漏给库消费者创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考