算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载导读MhcPost 是 CANN ops-transformer 算子库中 mHCMulti-Head Collaboration多头协同架构的关键后处理算子它将 mHC 架构上一层的输出h_out做 Post Mapping、上一层的输入x做 Res Mapping 后执行残差连接一步融合生成下一层输入x_{l1}。本文基于 mhc/mhc_post/README.md 及配套的 aclnnMhcPost 接口文档、PyTorch API 文档结合仓库源码与测试用例完整讲解其计算原理、算子规格、两种调用方式aclnn C 接口与 PyTorch API以及底层 tiling/vector 实现帮助开发者在 Ascend NPU 上正确配置与调用该算子。MhcPost 在 mHC 架构中的定位mHCMulti-Head Collaboration是面向 Transformer 大模型的多头协同计算架构。在 mHC 中每一层 Transformer 块的前后都会插入映射计算用于在多头注意力/MLP 子层与层间残差流之间做信息变换。MhcPost 正是承担这一层后处理职责的算子它位于某一层Atten/MLP 子层之后将该层输出与上一层输入重新组织成下一层的输入从而将后处理投影 残差映射 残差相加三类计算融合为单次算子调用避免多次独立算子调用带来的额外搬运与调度开销。从 PyTorch API 文档 的功能描述可以确认其定位实现 MHC Post 组件的前向计算用于 Transformer 模型中多层残差连接的后处理阶段。该算子将残差矩阵变换与输出状态投影融合为单次计算避免多次独立算子调用带来的额外开销。计算原理与公式完整计算h_res 提供时设上一层输入为 $x_l$上一层输出Atten/MLP 层输出为 $h_{l}^{out}$mHC 的残差映射矩阵为 $H_{l}^{res}$sinkhorn 变换后的双随机矩阵后处理映射矩阵为 $H_{t}^{post}$则下一层输入为$$ x_{l1} (H_{l}^{res})^{T} \times x_l h_{l}^{out} \otimes H_{t}^{post} $$其中 $\otimes$ 表示逐元素乘法与广播两部分语义如下对应 torchapi_mhc_post.md 中的逐 head 展开式Res Mapping对输入 $x_l$ 做残差矩阵转置乘法对输出的第 $i$ 个 head行有$$ x_{l1}[i] \sum_{j0}^{n-1} H_{l}^{res}[j, i] \cdot x_{l}[j] $$即 $H_{l}^{res}$ 以转置方式与 $x_l$ 做矩阵乘法$H_{l}^{res}[j, i]$ 为标量对 $x_l$ 的第 $j$ 行做标量乘后累加到第 $i$ 行输出。Post Mapping对上一层输出 $h_{l}^{out}$ 做后处理投影对第 $i$ 个 head$$ x_{l1}[i] \mathrel{} H_{t}^{post}[i] \cdot h_{l}^{out} $$即 $H_{t}^{post}[i]$ 为标量对 $h_{l}^{out}$ 整行做标量乘后加到第 $i$ 行输出。综合逐 head 完整形式为$$ x_{l1}[i, :] H_{t}^{post}[i] \cdot h_{l}^{out}[:] \sum_{j0}^{n-1} H_{l}^{res}[j, i] \cdot x_{l}[j, :] $$其中 $x_l$ 对应参数x$H_{l}^{res}$ 对应h_res$h_{l}^{out}$ 对应h_out$H_{t}^{post}$ 对应h_post$x_{l1}$ 对应输出y/out。退化路径h_res 缺省时当h_res传入nullptr/None时跳过 Res Mapping公式退化为直接残差连接$$ x_{l1} x_l h_{l}^{out} \otimes H_{t}^{post} $$该退化路径仅 Ascend 950PR/Ascend 950DT 支持见下文产品支持与约束。数据布局约定算子支持两种输入维度格式详见 torchapi_mhc_post.md 的维度说明BSND4 维(B, S, n, D)B 为 Batch批量大小S 为 Seq-Length序列长度n 为 head 数D 为每个 head 的隐藏维度headdimTND3 维(T, n, D)T 为所有 Batch 序列长度的累加和$T B \times S$。产品支持情况与约束说明产品支持矩阵依据 README.md 与 aclnnMhcPost.md支持情况如下产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×约束说明h_res 可选性Ascend 950PR/Ascend 950DTh_res支持传入nullptr此时退化为直接残差连接Atlas A2/A3 训练与推理系列产品h_res为必传参数不支持传入nullptr传入会报错。规格约束规格项规格规格说明n4固定为 4d范围 1 到 100000128 的倍数确定性计算aclnnMhcPost默认确定性实现见 aclnnMhcPost.md 约束说明。数据类型约束来自 torchapi_mhc_post.mdx与h_out数据类型必须相同输出y数据类型与x保持一致h_res、h_post为 FLOAT32。Shape 一致性以 BSND 格式为例h_res的(B, S)需与x一致、后两维为(n, n)h_out的(B, S)需与x一致、D 维与x的 D 维一致h_post的(B, S)与x一致、n 维与x的 n 维一致。所有输入 Tensor 各维度值必须为正数大于 0。图模式限制h_res传入None仅支持单算子模式调用图模式torch.compile下h_res必须传入否则会在 GE 编译阶段报错。参数说明算子输入/输出依据 README.md 的参数表算子输入输出定义如下参数名输入/输出描述数据类型数据格式x输入待计算的张量表示网络中 mHC 层的输入数据FLOAT16、BFLOAT16NDh_res输入可选mHC 的 h_res 变换矩阵是做完 sinkhorn 变换后的双随机矩阵。缺省时退化为直接残差连接仅 Ascend 950 支持FLOAT32NDh_out输入Atten/MLP 层的输出FLOAT16、BFLOAT16NDh_post输入mHC 的 h_post 变换矩阵FLOAT32NDout输出网络中 mHC 层的输出数据作为下一层的输入FLOAT16、BFLOAT16ND从算子注册源码 mhc/mhc_post/op_host/mhc_post_def.cpp 可以看到x/h_out为 REQUIRED 且数据类型为DT_FLOAT16/DT_BF16h_res为 OPTIONAL 且为DT_FLOATh_post为 REQUIRED 且为DT_FLOAT输出y与x同类型同时配置了ascend910b、ascend910_93、ascend950、ascend350四类 AICore 配置其中 950/350 走mhc_post_apt实现见ExtendCfgInfo(opFile.value, mhc_post_apt)910b/910_93 走mhc_post实现。调用方式一aclnn C 接口两段式接口原型MhcPost 的 aclnn 调用采用 CANN 标准两段式接口先调用aclnnMhcPostGetWorkspaceSize获取 workspace 大小与执行器再调用aclnnMhcPost执行计算。aclnnStatus aclnnMhcPostGetWorkspaceSize( const aclTensor *x, const aclTensor *hRes, // 可选Ascend 950 上可传 nullptr const aclTensor *hOut, const aclTensor *hPost, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMhcPost( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)GetWorkspaceSize 参数与返回码aclnnMhcPostGetWorkspaceSize中各张量参数详见 aclnnMhcPost.md的 shape 与约束如下参数输入/输出数据类型维度(shape)非连续 Tensorx输入FLOAT16、BFLOAT16[B,S,N,D]、[T,N,D]√hRes输入可选FLOAT32[B,S,N,N]、[T,N,N]√hOut输入与 x 相同[B,S,D]、[T,D]√hPost输入FLOAT32[B,S,N]、[T,N]√out输出与 x 相同[B,S,N,D]、[T,N,D]-workspaceSize输出---executor输出---第一段接口完成入参校验常见返回码如下完整返回码语义参见 aclnn 返回码说明返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001x、hOut、hPost、out 存在空指针ACLNN_ERR_PARAM_INVALID161002x、hRes非空时、hOut、hPost、out 的数据类型不在支持范围内或 shape 维度不在支持范围内或数据类型/shape 不匹配ACLNN_ERR_INNER_NULLPTR561103n 不等于 4d 不在 [1, 100000] 范围内d 不能被 128 整除需要特别说明的是hRes nullptr的 SoC 校验在源码 aclnn_mhc_post.cpp 的aclnnMhcPostGetWorkspaceSize中显式实现当hRes为nullptr且当前 NPU 架构不是DAV_3510即 Ascend 950时直接返回ACLNN_ERR_RUNTIME_ERROR。这是h_res 缺省仅 Ascend 950 支持这一约束在接口层的落地保证。完整调用示例仓库提供了可直接参考的样例 mhc/mhc_post/examples/test_aclnn_mhc_post.cppREADME 中链接为examples/test_aclnn_mhc_post.cpp完整的调用流程与 aclnnMhcPost.md 中的示例代码一致核心步骤设备与流初始化aclInit→aclrtSetDevice(deviceId)→aclrtCreateStream构造输入输出 Tensor以 BSND 格式为例典型 shape 为x {1, 1024, 4, 5120}BSND、hRes {1, 1024, 4, 4}、hOut {1, 1024, 5120}BSD、hPost {1, 1024, 4}BSn、out {1, 1024, 4, 5120}。通过aclrtMalloc申请 Device 内存、aclrtMemcpy将 Host 数据拷入再用aclCreateTensor以ACL_FORMAT_ND创建aclTensorx/hOut 用ACL_FLOAT16hRes/hPost 用ACL_FLOAT调用第一段接口获取 workspace 大小与执行器若workspaceSize 0则aclrtMalloc申请 workspace 内存调用第二段接口aclnnMhcPost(workspaceAddr, workspaceSize, executor, stream)执行计算同步并取回结果aclrtSynchronizeStream后aclrtMemcpyDEVICE_TO_HOST拷贝输出释放资源aclDestroyTensor、aclrtFree、aclrtDestroyStream、aclrtResetDevice、aclFinalize。从接口实现源码 aclnn_mhc_post.cpp 可以确认两段式接口内部的完整执行链第一段接口依次完成格式校验CheckFormat拒绝私有格式、数据类型校验CheckDtypex/hOut 限定 FP16/BF16 且同型hRes/hPost 限定 FP32、Shape 校验CheckShape区分 3D TND 与 4D BSND 两套约束以及空 Tensor 短路处理随后通过l0op::Contiguous将非连续输入转连续README 中标注非连续 Tensor √即由此支持再调用底层l0op::MhcPost构建算子执行器最后用l0op::ViewCopy将结果写回输出并返回 workspace 大小。第二段接口则通过CommonOpExecutorRun提交 AICore 任务。调用方式二PyTorch API函数原型cann_ops_transformer.mhc_post(x, h_res, h_out, h_post) - Tensor其中h_res为可选输入可传入None仅 Ascend 950PR/Ascend 950DT 的单算子模式支持传入None图模式不支持。参数与返回值依据 torchapi_mhc_post.md参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选当前层的输入 token 特征对应公式中的 $x_l$bfloat16、float16(B, S, n, D)、(T, n, D)h_resTensor可选残差连接矩阵对应 $H_{l}^{res}$传入 None 时跳过 Res Mappingfloat32(B, S, n, n)、(T, n, n)h_outTensor必选上一层的输出状态对应 $H_{l}^{out}$bfloat16、float16(B, S, D)、(T, D)h_postTensor必选后处理权重矩阵对应 $H_{t}^{post}$float32(B, S, n)、(T, n)返回值yMHC Post 计算输出对应公式中的 $x_{l1}$数据类型与x一致shape 与x一致(B, S, n, D) 或 (T, n, D)。单算子模式调用示例import torch import torch_npu from cann_ops_transformer.ops import mhc_post B 2 S 8 n 4 D 128 x torch.randn(B, S, n, D, dtypetorch.bfloat16).npu() h_res torch.randn(B, S, n, n, dtypetorch.float32).npu() h_out torch.randn(B, S, D, dtypetorch.bfloat16).npu() h_post torch.randn(B, S, n, dtypetorch.float32).npu() y mhc_post(x, h_res, h_out, h_post) print(foutput shape: {y.shape}) # h_res 缺省时仅 Ascend 950PR/Ascend 950DT 支持且仅支持单算子模式 y mhc_post(x, None, h_out, h_post) print(foutput shape: {y.shape})图模式调用示例import torch import torch_npu import torchair from cann_ops_transformer.ops import mhc_post torch_npu.npu.set_device(0) B 2 S 8 n 4 D 128 class MhcPostModel(torch.nn.Module): def forward(self, x, h_res, h_out, h_post): return mhc_post(x, h_res, h_out, h_post) model MhcPostModel().npu() npu_backend torchair.get_npu_backend() model torch.compile(model, backendnpu_backend, dynamicFalse) x torch.randn(B, S, n, D, dtypetorch.bfloat16, devicenpu) h_res torch.randn(B, S, n, n, dtypetorch.float32, devicenpu) h_out torch.randn(B, S, D, dtypetorch.bfloat16, devicenpu) h_post torch.randn(B, S, n, dtypetorch.float32, devicenpu) y model(x, h_res, h_out, h_post)Torch 扩展实现要点从 torch_extension/mhc_post.py 源码可以看到 PyTorch 侧的封装机制通过OpBuilder动态编译csrc/mhc/mhc_post.cpp并注册自定义算子mhc_post(Tensor x, Tensor? hRes, Tensor hOut, Tensor hPost) - Tensor注意 schema 中hRes声明为可空Tensor?当任一输入需要梯度时自动走torch.autograd.FunctionMhcPostFunction其backward会调用仓库中配套的mhc_post_backwardmhc/mhc_post_backward计算四个输入的梯度h_res is None时梯度返回None反向实现中还将grad_output先contiguous()物化以规避 h_res 为 None 时 0-stride 的扩展梯度如sum().backward()产生触发aclnnMhcPostBackward异常的问题。底层实现Tiling 与 Vector 计算算子注册与 Shape 推导算子定义位于 mhc_post_def.cpp通过OpDef注册输入输出、数据类型与各 SoC 的 AICore 配置Shape/数据类型推导位于 mhc_post_infershape.cpp输出y的 shape 与x逐维一致数据类型与x相同x的维度必须为 3TND或 4BSND否则推导失败。Tiling 机制mhc_post_tiling_base.cpp 通过TilingRegistryArch按架构分发 tiling 实现arch22对应 Atlas A2 系列ascend910bmhc_post_tiling_base_arch22.cpparch35对应 Ascend 950/A3 系列ascend950/ascend350分为三套实现 mhc_post_tiling_base_arch35.cpp常规路径、mhc_post_tiling_nohres_arch35.cpph_res 缺省退化路径与 mhc_post_tiling_regbase_arch35.cppregbase 路径。Tiling 的核心任务是把(BS, n, d)维度的计算按d 方向切块dInner/dTail和多核均分normalCoreProcessNum/tailCoreProcessNum展开从 kernel 代码中的bsIdx globalItemIdx / dOuter、dIdx globalItemIdx % dOuter可看出每个计算 item 是一个(bs, d 块)组合。Vector Kernel 实现要点主 kernel 入口 mhc_post.cpp 通过模板参数usePermanentX区分两种路径计算实现集中在 arch22/mhc_post_arch22.h关键优化点包括Double Buffer 流水输入队列hOutTileQueue_、xTileQueue_与输出队列outputTileQueue_深度均为 2注释明确说明Double Buffer 提升 Memory Bound 算子性能FP32 中间计算hOut/x先经Cast转为 F32hOutF32Buf_、xF32Buf_Post Mapping 用Muls(outF32, hOutF32, hPost[i], dNum)实现标量乘Res Mapping 用Axpy(outF32, xF32, hRes[j*ni], dNum)实现转置矩阵乘累加下标j * n i正是 $H^{res}$ 转置访问的体现最后Cast(CAST_RINT)回合回 FP16/BF16 输出usePermanentX 优化当USE_PERMANENT_X 1时一次性把 n 行x全部搬入DataCopyExtParams多行 stride 拷贝内层循环只做Axpy累加避免反复搬 x 数据。h_res 缺省路径则由 arch35/mhc_post_nohres.h 实现MhcPostNoHRes类公式x_{l1} x_l h_{l}^{out} * H_{t}^{post}同样采用 Double Buffer 队列hOutTileQueue_/xTileQueue_/hPostTileQueue_与DoMulAndAdd完成乘加融合。单元测试仓库为算子提供了较完整的 UT 覆盖mhc/mhc_post/tests/utShape 推导测试test_mhc_post_infershape.cpp各架构 Tiling 测试arch22/test_mhc_post_tiling.cpp、arch35/test_mhc_post_tiling.cppaclnn 接口测试op_api/test_aclnn_mhc_post.cpp。小结与选型建议MhcPost 将 mHC 架构层间后处理所需的 Post Mapping、Res Mapping 与残差连接融合为单算子既可用于 aclnn 两段式 C 编程调用样例也可通过cann_ops_transformer.mhc_post在 PyTorch 单算子/图模式下直接调用。开发时需重点把握以下几点平台差异h_res缺省的退化路径仅 Ascend 950PR/Ascend 950DT 支持接口层由DAV_3510架构校验保证Atlas A2/A3 产品必须显式传入h_res规格合规n 固定为 4d 必须在 [1, 100000] 且为 128 的倍数否则接口层返回 561103 错误类型与 Shapex/h_out必须同为 FP16 或 BF16h_res/h_post为 FP32BSND 与 TND 两套 shape 需按上文约束严格对齐输出与x同型同 shape模式差异h_resNone仅支持单算子模式图模式torch.compile下必须传入h_res否则 GE 编译阶段报错。对需要反向传播的场景可直接使用 PyTorch 扩展层自动接入 mhc_post_backward 完成梯度计算无需手工编写反向逻辑。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer 算子 aclnnMhcPost 使用指南mHC 架构 Post Mapping 与残差连接的 NPU 融合实现CANN ops transformer 算子 aclnnMhcPost 使用指南mHC 架构 Post Mapping 与残差连接的 NPU 融合实现 导读算子库人工智能深度学习AscendCANN ops-transformer 量化算子深度解析GroupedMatmulSwigluQuantV2 融合算子原理与 aclnn 调用实战CANN ops transformer 量化算子深度解析GroupedMatmulSwigluQuantV2 融合算子原理与 aclnn 调用实战 导读 本算子库人工智能深度学习AscendCANN ops-transformer 的 AllGatherMatmul 通算融合算子原理、参数与 aclnn 调用实战CANN ops transformer 的 AllGatherMatmul 通算融合算子原理、参数与 aclnn 调用实战 导读 AllGatherMatm算子库人工智能深度学习Ascend上一篇Fiddler中文版性能优化如何分析网站加载速度和瓶颈下一篇性能之巅SyncTrayzor与主流Syncthing管理工具深度对比测试创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考