算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载导读block_attention_residuals是 CANN ops-transformer 算子库mhcMulti-Head-Chunking 类注意力残差方向提供的高性能融合算子它把partial_block与block_res按 block 维拼接后一次完成RMS 归一化 → 投影打分 → Softmax 加权融合直接产出hidden_states用于大模型注意力残差分支的轻量融合计算。本文以 torchapi_block_attention_residuals.md 为骨架完整展开其计算公式、Python 接口签名、参数约束与调用示例并结合本仓库的 torch 扩展封装、算子定义、动态 tiling 与 kernel 分派源码讲清该算子在 NPU 上的实现原理与使用要点。读完本文你将能够独立完成该算子的环境准备、参数组织、单算子调用与结果校验并理解其 RESIDENT / RELOAD / HSLICE 三种 kernel 变体的选择逻辑。算子定位与产品支持情况block_attention_residuals属于mhc/block_attention_residuals算子目录仓库内同时提供两种调用形态PyTorch 接口cann_ops_transformer.block_attention_residuals(...)即本文讲解对象aclnn 底层接口aclnnBlockAttentionResiduals两段式调用详见 aclnnBlockAttentionResiduals.md。产品支持情况以文档为准语义为前三个平台支持后三个不支持产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持需要说明的是Ascend 950PR/Ascend 950DT、Atlas A3 系列、Atlas A2 系列三个平台上的计算语义一致可以放心在支持的产品间迁移代码。功能说明与数学原理接口功能将partial_block与block_res按 block 维拼接后完成RMS 归一化、norm_weight ⊙ proj_weight投影打分、Softmax 加权融合输出hidden_states。Torch 接口当前仅返回hidden_states这一个输出反向所需的inv_norm、probs由底层在需要时内部保存见后文“反向支持”小节。计算公式整体计算可以写作$$ \operatorname{block_attention_residuals} \left(partial_block, block_res, proj_weight, norm_weight; valid_block_num, norm_eps\right) \longrightarrow hidden_states. $$记 token 数为 $T$block 数为 $N$hidden size 为 $H$拼接后逻辑行数 $B N 1$。行序与 golden 一致前 $N$ 行来自block_res最后一行来自partial_block$$ v_{t,i,h} \begin{cases} block_res_{t,i,h}, 0 \le i N \ partial_block_{t,h}, i N \end{cases} $$拼接后按行做 RMS 归一化对每行 $H$ 个元素求均方根$$ variance_{t,i} \frac{1}{H}\sum_{h0}^{H-1} v_{t,i,h}^{2} $$$$ inv_rms_{t,i} \frac{1}{\sqrt{variance_{t,i} norm_eps}} $$$$ k_{t,i,h} v_{t,i,h} \cdot inv_rms_{t,i} $$将norm_weight与proj_weight逐元素相乘得到打分权重再对归一化后的 $k$ 做投影求和得到每个逻辑行的标量得分$$ score_weight_{h} norm_weight_{h} \cdot proj_weight_{0,h} $$$$ s_{t,i} \sum_{h0}^{H-1} k_{t,i,h} \cdot score_weight_{h} $$对 $N1$ 行得分做 Softmax$$ probs_{t,i} \frac{e^{s_{t,i}}}{\sum_{j0}^{N} e^{s_{t,j}}} $$最后用概率对原始$v$注意不是 RMS 归一化后的 $k$做加权融合得到输出$$ hidden_states_{t,h} \sum_{i0}^{N} probs_{t,i} \cdot v_{t,i,h} $$其中各符号含义$partial_block \in \mathbb{R}^{T \times H}$输入partial_block作为拼接后的第 $N$ 行 value$block_res \in \mathbb{R}^{T \times N \times H}$输入block_res作为拼接后的前 $N$ 行 value$proj_weight$输入proj_weightshape 为 $[1, H]$$proj_weight_{0,h}$ 表示其第 $h$ 个元素$norm_weight \in \mathbb{R}^{H}$输入norm_weightRMS 缩放权重$valid_block_num$输入valid_block_num可选预留属性默认-1当前仅支持传入-1不参与计算$norm_eps$输入norm_epsRMS 归一化因子中的数值稳定项$t \in [0, T)$、$i \in [0, N]$、$h \in [0, H)$token、逻辑行、hidden dimension 索引$v_{t,i,h}$拼接后的 value$inv_rms_{t,i}$逐行 RMS 归一化系数$score_weight_{h}$norm_weight与proj_weight逐元素相乘结果$s_{t,i}$投影得分$probs_{t,i}$Softmax 概率$hidden_states_{t,h}$加权融合输出。从上述公式可以看到该算子本质上是“attention 的轻量替代”不做 Q/K 点积而是用一组可学习的norm_weight/proj_weight把每行的 value 压成一个得分再用 Softmax 做融合特别适合注意力残差分支中“从多个 block 候选里软选一个综合结果”的场景。函数原型cann_ops_transformer.block_attention_residuals( partial_block, block_res, proj_weight, norm_weight, valid_block_num-1, norm_eps1.0e-6, ) - Tensor该签名与仓库 torch 扩展中注册的 schema 完全一致torch_extension/block_attention_residuals.py 中schema()方法定义了同样的默认值valid_block_num-1、norm_eps1e-6常量定义在文件头DEFAULT_NORM_EPS 1.0e-6。参数说明参数名参数类型可选/必选描述数据类型维度(shape)partial_blockTensor必选拼接后的第 $N$ 行 value对应公式中的 $partial_block$。数据格式为 ND。float16、bfloat16、float32[T, H]block_resTensor必选前 $N$ 行 value对应公式中的 $block_res$。数据格式为 NDdtype 必须与partial_block一致。同 partial_block[T, N, H]proj_weightTensor必选投影权重与norm_weight共同构成 $score_weight$。数据格式为 NDdtype 必须与partial_block一致。同 partial_block[1, H]norm_weightTensor必选RMS 缩放权重。数据格式为 NDdtype 必须与partial_block一致。同 partial_block[H]valid_block_numint可选预留属性默认值为-1当前仅支持传入-1不参与计算。int-norm_epsfloat可选RMS 归一化的数值稳定项必须为有限正数默认值为1.0e-6。float-补充实现细节源码佐证在 torch_extension/block_attention_residuals.py 的_check_inputs中可以看到proj_weight实际同时接受 1 维[H]或 2 维[1, H]两种 shapePROJ_WEIGHT_RANK 1、PROJ_WEIGHT_RANK_2D 2且首维为 1接口更宽松四个输入 Tensor 必须位于同一 device且block_res、proj_weight、norm_weight的 dtype 必须与partial_block一致支持的数据类型枚举为SUPPORTED_DTYPES (torch.float16, torch.bfloat16, torch.float32)。返回值说明返回hidden_states。参数名参数类型可选/必选描述数据类型维度(shape)hidden_statesTensor必选加权融合结果对应公式中的 $hidden_states$。数据格式为 ND。同 partial_block[T, H]在 torch_extension/csrc/block_attention_residuals.cpp 的 pybind 绑定中函数实际返回三元组(hiddenStates, invNorm, probs)当needBackwardfalse时invNorm、probs为空 Tensorshape[0]Python 层只把hiddenStates暴露给用户invNorm/probs仅在反向需要时在内部流转。约束说明该接口支持单算子模式调用暂不支持 TorchAir 图模式调用。主路径 dtype 必须一致支持 float16 / bfloat16 / float32。shape 要求T 0H 01 N 100。T 0或H 0时返回对应 shape 的空输出。block_res.shape[0]、block_res.shape[2]必须分别等于partial_block的 $T$、$H$proj_weight最后一维与norm_weight长度必须等于 $H$。valid_block_num为可选预留属性默认值为-1当前仅支持传入-1不参与计算。norm_eps必须为有限正数。partial_block或block_res含NaN、Inf或平方、乘法及累加发生溢出时结果可能包含NaN或Inf。几点来自源码的补充说明关于H 0文档明确T 0或H 0时返回空输出不过当前 Python 封装实现里_check_inputs对hidden_size的检查是hidden_size MIN_HIDDEN_SIZE即H 1因此实际调用时空 H 场景会被校验拦截使用时应以H 1为准。关于valid_block_num文档限定仅支持-1底层 C 绑定torch_extension/csrc/block_attention_residuals.cpp与 README 中放宽为“-1默认使用 blockRes 的 N或等于 N”两者皆可但语义上等价——传入-1时内部就用N。关于N的上限 100tiling 侧的MAX_NUM_BLOCKS 100与 shape 推导侧的MAX_NUM_BLOCKS 100保持一致见 op_host/block_attention_residuals_tiling.cpp 与 op_host/block_attention_residuals_infershape.cpp。输入 Tensor 允许非连续接口内部会转 Contiguous 后计算输出不支持非连续 Tensor——该约束在 aclnn 文档中有明确说明。确定性与 Batch 一致性确定性说明默认支持确定性计算同输入同输出结果可复现。Batch 一致性说明默认 Batch 一致性实现即 batch 维度上的处理策略一致不会因 batch 大小变化引入结果差异。调用示例单算子模式以下示例来自原文档可直接复制运行需已安装torch_npu与cann_ops_transformer并在一张受支持的 NPU 上执行import torch import torch_npu import cann_ops_transformer torch_npu.npu.set_device(0) partial_block torch.tensor( [[1.0, 2.0, 3.0, 4.0]], dtypetorch.bfloat16, devicenpu, ) block_res torch.tensor( [ [ [1.0, 2.0, 3.0, 4.0], [2.0, 0.0, -1.0, 1.0], ] ], dtypetorch.bfloat16, devicenpu, ) proj_weight torch.tensor([[0.5, 0.25, 0.25, 0.0]], dtypetorch.bfloat16, devicenpu) norm_weight torch.ones(4, dtypetorch.bfloat16, devicenpu) # valid_block_num / norm_eps 为可选参数默认值分别为 -1、1e-6仅返回 hidden_states。 hidden_states cann_ops_transformer.block_attention_residuals( partial_block, block_res, proj_weight, norm_weight, ) torch_npu.npu.synchronize() print(hidden_states:, hidden_states)运行要点示例中T1、N2、H4拼接后逻辑行数为B N 1 3前两行来自block_res最后一行来自partial_blocknorm_weight全为 1、proj_weight [0.5, 0.25, 0.25, 0.0]打分权重即proj_weight本身可选参数可省略valid_block_num默认-1、norm_eps默认1e-6调用后建议torch_npu.npu.synchronize()确保异步流上的算子执行完成再读取结果。源码级实现剖析1. Python 前端入参校验与自动微分挂钩cann_ops_transformer.block_attention_residuals的完整调用链在 torch_extension/block_attention_residuals.py先执行_check_inputs完成 rank、dtype、device、shape 范围、norm_eps有限正数等全部校验若处于梯度开启状态且任一输入requires_gradTrue则走_BlockAttentionResidualsFunction.apply(...)torch.autograd.Function前向时以need_backwardTrue请求底层同时产出inv_norm、probs并save_for_backward反向时调用cann_ops_transformer.block_attention_residuals_backward对应仓库中的 mhc/block_attention_residuals_grad 目录否则以need_backwardFalse直接执行前向只返回hidden_states。也就是说文档中“Torch 接口仅返回hidden_states”并不意味着算子不支持反向——反向所需的中间量由底层在需要时自动产出Python 侧对用户透明。2. 算子定义输入输出与属性默认值算子注册在 op_host/block_attention_residuals_def.cpp四个必选输入partial_block、block_res、proj_weight、norm_weight均支持DT_BF16 / DT_FLOAT16 / DT_FLOAT与 ND 格式输出hidden_states必选与inv_norm、probs可选仅DT_FLOAT三个属性默认值valid_block_num-1、norm_eps1e-6、need_backwardfalse平台配置ascend910b、ascend910_93Atlas A2/A3 对应与ascend950950 平台使用block_attention_residuals_apt文件。3. Shape 推导与 dtype 推导op_host/block_attention_residuals_infershape.cpp 中校验partial_block为 2 维、block_res为 3 维、proj_weight为 1 或 2 维[H]或[1,H]、norm_weight为 1 维校验T 0、H 0、1 N 100以及block_res的 T/H 与partial_block对齐输出hidden_states的 shape 设为[T, H]当need_backwardtrue时inv_norm、probs的 shape 推导为[T, N1]dtype 固定为DT_FLOAT。4. 动态 TilingRESIDENT / RELOAD / HSLICE 三分支host 侧 tilingop_host/block_attention_residuals_tiling.cpp会按运行时的BN1与H在三种 kernel 变体中选择TK-RESIDENT当B CalcMaxResidentRows(H)按 UB 容量 192KB 保守估算、扣除权重与元数据开销后能容纳全部 block 行时所有 block 行常驻 UB一次加载复用适合小 N 场景TK-RELOAD无法常驻但整体 UB 预算够用时逐行/分块从 GM 重新加载TK-HSLICEH 过大时对 hidden 维做切片CalcHSliceChunk计算满足 UB 预算的最大 chunk按 chunk 循环计算。此外 tiling 中把 token 均分到多个 AIV coretokensPerCore ceil(T / usedCoreNum)并在blockCount numBlocks 1、invHiddenSize 1/H等字段写入 tiling 数据供 kernel 使用。5. Kernel 入口与分派kernel 入口在 op_kernel/block_attention_residuals.cppblock_attention_residuals内核函数读取 tiling 数据后按 tiling key 分派到RunReload/RunResident/RunHSlice三个实现分别对应 arch22/block_attention_residuals_reload.h、resident.h、hslice.h以及 arch35 平台的对应实现输入 dtype 由编译期宏DTYPE_PARTIAL_BLOCK注入未注入时默认bfloat16_t框架为每个 dtype 生成独立二进制。底层 aclnn 接口与反向支持除 Python 接口外仓库还提供两段式 aclnn 接口aclnnBlockAttentionResiduals.mdaclnnBlockAttentionResidualsGetWorkspaceSize(...)完成入参校验返回 workspace 大小与 op 执行器aclnnBlockAttentionResiduals(workspace, workspaceSize, executor, stream)执行计算。aclnn 层暴露了 Python 接口隐藏的needBackward参数与invNorm/probs两个可选输出needBackwardtrue时二者 shape 为[T, N1]dtype 为 FLOAT32。完整可编译示例见 examples/test_aclnn_block_attention_residuals.cpp该示例使用T2、N4、H64的 BF16 数据走完“Init → 建 Tensor → 两段式调用 → 同步 → 回拷结果”全流程。Python 封装正是通过 pybind 直调该 aclnn 接口实现见 torch_extension/csrc/block_attention_residuals.cpp 的ACLNN_CMD(aclnnBlockAttentionResiduals, ...)。需要反向时前向保存的inv_norm、probs会传给 mhc/block_attention_residuals_grad 目录下的反向算子由 Python 封装在backward()中通过cann_ops_transformer.block_attention_residuals_backward调用从而支持对partial_block、block_res、proj_weight、norm_weight的梯度回传。使用建议与注意事项平台选择在 Ascend 950PR/DT、Atlas A3、Atlas A2 系列上使用Atlas 200I/500 A2、Atlas 推理系列、Atlas 训练系列不支持。shape 规划N建议控制在 1100 之间T可以为 0返回空输出H以 1为准。dtype 一致性四个输入与输出必须同 dtypefp16/bf16/fp32 三选一inv_norm/probs内部固定为 fp32。数值稳定norm_eps需为有限正数默认1e-6输入含 NaN/Inf 或计算溢出会导致输出出现 NaN/Inf接入前建议对partial_block/block_res做数值检查。调用模式当前仅支持单算子模式暂不支持 TorchAir 图模式如需在反向传播中使用可直接对输出调用.backward()前向中间量由算子内部自动保存。参考资料本文主文档torchapi_block_attention_residuals.md算子 READMEmhc/block_attention_residuals/README.mdaclnn 接口文档aclnnBlockAttentionResiduals.mdPython 封装与自动微分torch_extension/block_attention_residuals.pypybind 绑定torch_extension/csrc/block_attention_residuals.cpp算子定义op_host/block_attention_residuals_def.cppshape/dtype 推导op_host/block_attention_residuals_infershape.cpp动态 tilingop_host/block_attention_residuals_tiling.cppKernel 入口分派op_kernel/block_attention_residuals.cppaclnn 调用示例examples/test_aclnn_block_attention_residuals.cpp反向算子目录mhc/block_attention_residuals_grad赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer BlockAttentionResiduals 算子全解析分块残差的 RMS 归一化与 Softmax 加权融合CANN ops transformer BlockAttentionResiduals 算子全解析分块残差的 RMS 归一化与 Softmax 加权融合 B算子库人工智能深度学习AscendCANN ops-transformer BlockAttentionResiduals 算子深度解析aclnn 两段式接口、RMS 加权融合原理与调用实践CANN ops transformer BlockAttentionResiduals 算子深度解析aclnn 两段式接口、RMS 加权融合原理与调用实践算子库人工智能深度学习AscendCANN ops-transformer 算子详解MoeGatingTopKSoftmaxV2 门控 TopKSoftmax 融合算子CANN ops transformer 算子详解MoeGatingTopKSoftmaxV2 门控 TopKSoftmax 融合算子 MoeGatingT算子库人工智能深度学习Ascend上一篇GitHub PR Tree 用户指南掌握文件树浏览、已查看文件跟踪与暗模式设置下一篇3分钟解锁Sentry搜索黑科技从崩溃日志到性能瓶颈的精准定位创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考