CANN ops-transformer MhcPre 算子解析MHC 架构 hidden 层投影与 h_in 生成原理及 aclnn 调用实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读MhcPreMHC Pre-Process是 CANN ops-transformer 算子库中用于 MHCMulti-head Hyper Connection架构前置计算的 NPU 算子它基于 RmsNorm、参数矩阵乘法与分段激活一次性产出 hidden 层的残差投影矩阵 H^res、H^post 以及 Attention/MLP 层的输入矩阵 h^in。本文以 mhc/mhc_pre/README.md 为主体结合 aclnnMhcPre 与 aclnnMhcPreV2 两套 aclnn 接口文档及仓库源码完整讲解算子数学模型、全部输入输出参数、各产品规格约束并给出可运行的两段式接口调用示例帮助读者在 NPU 上正确配置与调用 MhcPre。MhcPre 在 MHC 网络中的位置与算子功能MhcPre 是 MHC多头超连接架构中 hidden 层的前置计算算子。其核心功能是基于输入 x 与参数矩阵 phi、缩放参数 alpha、偏置 bias及可选的 RmsNorm 缩放因子 gamma通过一次融合计算得到H^res 残差投影矩阵对应输出 h_res未做 sinkhorn 变换H^post 变换矩阵对应输出 h_posth^in 输入矩阵作为后续 Attention 层或 MLP 层的输入对应输出 h_in。从算子定义源码 mhc/mhc_pre/op_host/mhc_pre_def.cpp 可以看到MhcPre 输入为xBF16/FP16、phiFP32、alphaFP32、biasFP32以及可选输入gammaFP32输出为必选输出hin、h_post、h_res与可选输出inv_rms、h_mix、h_pre并通过out_flag、norm_eps、hc_eps、op_impl_mode四个可选属性控制输出开关与计算精度。整个算子数据均采用 ND 格式。数学模型与计算流程MhcPre 的计算公式mhc/mhc_pre/README.md 与两篇接口文档一致如下。其中xFlat表示将 x 的最后两维 n 和 D 视作长度为 nD 的向量gammaFlat表示将 gamma 视作长度为 nD 的向量表示矩阵乘法⊙表示逐元素乘法$$ \begin{aligned} invRms \left(mean(xFlat^{2}) normEps\right)^{-\frac{1}{2}}\ xGamma \begin{cases} xFlat \odot gammaFlat, gamma \ne null \ xFlat, gamma null \end{cases}\ hMix xGamma phi^{T}\ w hMix \odot invRms\ (pPre, pPost, pRes) \begin{cases} split(w, (n, n, n^{2})), alpha.shape(3) \ (split(w, (n, n)), 0), alpha.shape(2) \end{cases}\ hPre \sigma(pPre \odot alpha0 bias0) hcEps\ hPost \begin{cases} 2\sigma(pPost \odot alpha1 bias1), alpha.shape(3) \ 2\sigma(pPost \odot alpha1 bias1) hcEps, alpha.shape(2) \end{cases}\ hRes \begin{cases} pRes \odot alpha2 bias2, alpha.shape(3) \ 0, alpha.shape(2) \end{cases}\ hIn_{d} \sum_{i0}^{n-1} hPre_{i} x_{i,d} \end{aligned} $$可以将计算过程拆解为四个阶段理解RmsNorm 归一化先对展平后的 x 求均方mean(xFlat²)加上防除零项normEps再取-1/2次幂得到invRms即 1/r若提供 gamma则先做逐元素缩放xFlat ⊙ gammaFlat。矩阵乘投影将归一化前的xGamma与phi的转置做矩阵乘得到hMix再逐元素乘上invRms得到w。phi的行维度由 alpha 的形状决定alpha.shape(3)时 phi 为(n²2n, nD)或(n!2n, nD)alpha.shape(2)时 phi 为(2n, nD)。分段激活将w按(n, n, n²)alpha 为 3 维或(n, n)alpha 为 2 维切分为 pPre、pPost、pRes 三部分分别施加带 alpha 缩放与 bias 偏置的 sigmoid 激活hPre σ(pPre·alpha0 bias0) hcEpshPost 2σ(pPost·alpha1 bias1)alpha(2) 时额外加 hcEpshRes pRes·alpha2 bias2alpha(2) 时 hRes 恒为 0。输入矩阵合成hIn_d Σ_{i0}^{n-1} hPre_i · x_{i,d}即用激活后的 hPre 对原始 x 沿 n 维做加权求和得到可直接送入 Attention/MLP 层的输入矩阵。从 infershape 源码 mhc/mhc_pre/op_host/mhc_pre_infershape.cpp 可以看出实现细节alpha形状为[2]时hasResifalse无残差分支phi第一维必须为2nalpha形状为[3]时允许phi第一维为n!2n阶乘分支或n²2n平方分支其中n!由源码中的Factorial(n)计算得到。参数说明下表完整列出 MhcPre 算子的全部输入、输出与属性参数数据来源于 mhc/mhc_pre/README.md 参数说明表参数名输入/输出/属性描述数据类型数据格式x输入待计算数据表示网络中 mHC 层的输入数据对应公式中的 xBFLOAT16, FLOAT16NDphi输入mHC 的参数矩阵对应公式中的 phiFLOAT32NDalpha输入mHC 的缩放参数对应公式中的 alphaFLOAT32-bias输入mHC 的 bias 参数对应公式中的 biasFLOAT32-gamma可选输入表示进行 RmsNorm 计算的缩放因子对应公式中的 gammaFLOAT32NDout_flag可选输入表示是否输出公式中的 invRms、hMix 和 hPre默认为 0 表示不输出为 1 表示全部输出INT64-norm_eps可选输入RmsNorm 的防除零参数对应公式中的 normEpsFLOAT-hc_eps可选输入h_pre 的 sigmoid 后的 eps 参数以及当 alpha.shape(2) 时 h_post 的 sigmoid 后的 eps 参数对应公式中的 hcEpsFLOAT-op_impl_mode可选输入指定 MhcPre 算子的计算模式0 表示 Cube 使用 FP32 模式计算1 表示 Cube 使用 HF32 模式计算默认值为 0INT64-h_in输出输出的 h_in 作为 Attention/MLP 层的输入对应公式中的 hInBFLOAT16, FLOAT16NDh_post输出输出的 mHC 的 h_post 变换矩阵对应公式中的 hPostFLOAT32NDh_res输出输出的 mHC 的 h_res 变换矩阵未做 sinkhorn 变换对应公式中的 hResFLOAT32NDinv_rms可选输出RmsNorm 计算得到的 1/r对应公式中的 invRmsFLOAT32NDh_mix可选输出xGamma 与 phi 矩阵乘的结果对应公式中的 hMixFLOAT32NDh_pre可选输出做完 sigmoid 计算之后的 h_pre 矩阵对应公式中的 hPreFLOAT32ND参数默认值与内部实现在算子定义 mhc/mhc_pre/op_host/mhc_pre_def.cpp 中三个可选属性的默认值被显式声明out_flag默认0不输出可选输出norm_eps默认1e-6fhc_eps默认1e-6fop_impl_mode默认0FP32 模式。在 tiling 侧mhc/mhc_pre/op_host/op_tiling/arch35/mhc_pre_tiling.cpp同样定义了IMPL_MODE_FP32 0与IMPL_MODE_HF32 1两个常量并在 tiling 数据结构mhc/mhc_pre/op_host/op_tiling/arch35/mhc_pre_tiling.h中通过implMode字段下发给 kernel。tiling 阶段支持三种切分模式SPLIT_BS、SPLIT_ND、SPLIT_M_K并由 tiling keymhc/mhc_pre/op_kernel/arch35/mhc_pre_tiling_key.h与HAS_RESI是否含残差分支组合出多套模板实例kernel 侧根据 tiling key 选择对应的 AIC/AIV 混编执行路径。产品支持情况与规格约束MhcPre 在不同硬件产品上的支持情况如下mhc/mhc_pre/README.md产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×对应的规格约束分为两档Ascend 950PR/Ascend 950DTn 目前支持 4、6、8D 支持 1~16384 范围以内需满足 D 为 16 对齐。Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品参数op_impl_mode仅支持配置为 0n 目前支持 4D 支持 100000 范围以内需满足 D 为 128 对齐。从 aclnnMhcPre 的约束说明中还可以补充以下使用要点确定性计算aclnnMhcPre及 aclnnMhcPreV2默认确定性实现Batch 一致性默认非 Batch 一致性实现不支持通过aclrtSetSysParamOpt开启 Batch 一致性alpha(2) 时的输出约定hRes 输出为 0此时要求输入 phi 为(2n, nD)、bias 为(2n)hMixOptional 为(B, S, 2n)或(T, 2n)alpha(3) 时的输出约定支持 hRes 输出要求 phi 为(n²2n, nD)或 A3/A2 下额外支持(n!2n, nD)、bias 对应为(n²2n)或(n!2n)、hMixOptional 对应为(B, S, n²2n)或(B, S, n!2n)可选输出互存关系可选输出 invRmsOptional、hMixOptional、hPreOptional 为互存关系需同时输出或全部不输出不支持仅返回其中部分判定条件为invRmsOptional ! nullptr hMixOptional ! nullptr hPreOptional ! nullptr。aclnn 两段式接口调用说明MhcPre 算子通过 aclnn 两段式接口参见 docs/zh/context/two_phase_api.md调用先调用aclnnMhcPreGetWorkspaceSize获取计算所需 workspace 大小与执行器再调用aclnnMhcPre执行计算。仓库提供两种接口对应两个可编译的示例程序调用方式调用样例说明aclnn 调用test_aclnn_mhc_pre通过 aclnnMhcPre 接口方式调用 MhcPre 算子aclnn 调用test_aclnn_mhc_pre_v2通过 aclnnMhcPreV2 接口方式调用 MhcPre 算子并通过 opImplMode 选择 Cube 的 FP32 或 HF32 计算模式接口函数原型aclnnMhcPre的第一段接口aclnnStatus aclnnMhcPreGetWorkspaceSize( const aclTensor *x, const aclTensor *phi, const aclTensor *alpha, const aclTensor *bias, const aclTensor *gammaOptional, double normEps, double hcEps, aclTensor *hIn, aclTensor *hPost, aclTensor *hRes, aclTensor *invRmsOptional, aclTensor *hMixOptional, aclTensor *hPreOptional, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口aclnnStatus aclnnMhcPre( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnMhcPreV2与上述原型唯一区别在于多出一个int64_t opImplMode参数位于hcEps与hIn之间用于指定计算模式opImplMode 0建议值在 Cube 中使用 FP32 模式计算opImplMode 1在 Cube 中使用 HF32 模式计算。在 aclnn 层源码 mhc/mhc_pre/op_host/op_api/aclnn_mhc_pre.cpp 中对应定义了MHC_PRE_USE_FP32 0与MHC_PRE_USE_HF32 1参数校验函数CheckOpImplMode会对取值合法性进行检查非 0/1 时返回ACLNN_ERR_PARAM_INVALID错误码 161002。aclnnMhcPre内部实质上是opImplMode固定为MHC_PRE_USE_FP32的特例。需要特别注意的是aclnnMhcPreV2 仅支持 Ascend 950PR/Ascend 950DT见 aclnnMhcPreV2 的产品支持情况而 Atlas A2/A3 系列只能使用 aclnnMhcPre 且op_impl_mode固定为 0。各张量 shape 与使用说明以 aclnnMhcPre 的参数表为据各参数的 shape 与使用约束如下括号内为支持的维度数参数输入/输出数据类型shape 约束非连续 Tensorx输入BFLOAT16、FLOAT16(B, S, n, D)或(T, n, D)3-4√phi输入FLOAT32(n!2n, nD)、(n²2n, nD)或(2n, nD)2√alpha输入FLOAT32(3)或(2)1×bias输入FLOAT32(n!2n)、(n²2n)或(2n)1√gammaOptional可选输入FLOAT32(n, D)2推理场景可传空指针√normEps可选输入DOUBLE-建议值 1e-6-hcEps可选输入DOUBLE-建议值 1e-6-hIn输出BFLOAT16、FLOAT16与 x 一致(B, S, D)或(T, D)2-3×hPost输出FLOAT32(B, S, n)或(T, n)2-3×hRes输出FLOAT32(B, S, n, n)、(T, n, n)、(B, S, n!)或(T, n!)2-4alpha(2) 时可传空指针×invRmsOptional可选输出FLOAT32(B, S)或(T)1-2×hMixOptional可选输出FLOAT32(B, S, n!2n)、(T, n!2n)、(B, S, n²2n)、(T, n²2n)、(B, S, 2n)或(T, 2n)2-3×hPreOptional可选输出FLOAT32(B, S, n)或(T, n)2-3×其中 Ascend 950PR/950DT 上 phi 的 shape 仅支持(n²2n, nD)或(2n, nD)不支持 n! 分支对应 bias 为(n²2n)或(2n)hRes 为(B, S, n, n)或(T, n, n)。输入 x 不支持空 Tensorphi、alpha、bias、hIn、hPost、hRes 均为必填项传空指针时第一段接口返回ACLNN_ERR_PARAM_NULLPTR错误码 161001数据类型或 shape 维度不在支持范围内时返回ACLNN_ERR_PARAM_INVALID错误码 161002API 内部调用 npu runtime 接口异常时返回ACLNN_ERR_RUNTIME_ERROR错误码 361001。返回码的完整说明可参考 docs/zh/context/aclnn_return_code.md。完整调用示例aclnnMhcPre以下代码节选自仓库示例 mhc/mhc_pre/examples/test_aclnn_mhc_pre.cpp 的主流程展示了从初始化 AscendCL、创建输入输出 Tensor、两段式调用到释放资源的完整链路完整代码请直接查看该文件编译与运行方式参考 docs/zh/context/compile_and_run_sample.md#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_mhc_pre.h // ...GetShapeSize / PrintTensorDataFloat / PrintTensorDataFloat16 / // InitAcl / CreateAclTensorFloat32 / CreateAclTensorFloat16 / // CreateAclTensorFloat16Output / CreateAclTensorFloat32Output / // CreateInputTensors / CreateOutputTensors / DestroyTensors / FreeDeviceMemory 等辅助函数见原文件 int main() { int32_t device_id 0; aclrtContext context nullptr; aclrtStream stream nullptr; Tensors tensors; // 示例规格B1, S2048, n4, D2560使用 alpha.shape(3) 的残差分支 int B 1, S 2048, n 4, D 2560; std::vectorint64_t x_shape {B * S, n, D}, phi_shape {n * n 2 * n, n * D}, alpha_shape {3}, bias_shape {n * n 2 * n}, gamma_shape {n, D}; std::vectorint64_t hin_shape {B * S, D}, h_post_shape {B * S, n}, h_res_shape {B * S, n, n}, inv_rms_shape {B * S}, h_mix_shape {B * S, n * n 2 * n}, h_pre_shape {B * S, n}; int ret InitAcl(device_id, context, stream); ret CreateInputTensors(x_shape, phi_shape, alpha_shape, bias_shape, gamma_shape, tensors); ret CreateOutputTensors(hin_shape, h_post_shape, h_res_shape, inv_rms_shape, h_mix_shape, h_pre_shape, tensors); // 第一段接口获取 workspace 大小与执行器 uint64_t workspace_size 0; aclOpExecutor *executor nullptr; aclnnStatus aclnn_ret aclnnMhcPreGetWorkspaceSize( tensors.x, tensors.phi, tensors.alpha, tensors.bias, tensors.gamma, 1e-6, 1e-6, tensors.hin, tensors.h_post, tensors.h_res, tensors.inv_rms, tensors.h_mix, tensors.h_pre, workspace_size, executor); // 按需申请 workspace void *workspace_addr nullptr; if (workspace_size 0) { ret aclrtMalloc(workspace_addr, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST); } // 第二段接口执行计算 aclnn_ret aclnnMhcPre(workspace_addr, workspace_size, executor, stream); CHECK_RET(aclrtSynchronizeStream(stream) ACL_SUCCESS, ...); // 打印输出结果hIn 为 FP16其余为 FP32 PrintTensorDataFloat16(hin_shape, tensors.hin_addr); PrintTensorDataFloat(h_post_shape, tensors.h_post_addr); PrintTensorDataFloat(h_res_shape, tensors.h_res_addr); PrintTensorDataFloat(inv_rms_shape, tensors.inv_rms_addr); PrintTensorDataFloat(h_mix_shape, tensors.h_mix_addr); PrintTensorDataFloat(h_pre_shape, tensors.h_pre_addr); // 释放资源 DestroyTensors(tensors); FreeDeviceMemory(tensors); if (workspace_size 0) aclrtFree(workspace_addr); aclrtDestroyStream(stream); aclrtDestroyContext(context); aclrtResetDevice(device_id); aclFinalize(); return 0; }如需切换到 HF32 计算模式将头文件改为aclnnop/aclnn_mhc_pre_v2.h并在第一段接口调用中加入int64_t opImplMode 1;参数置于hcEps之后参见 mhc/mhc_pre/examples/test_aclnn_mhc_pre_v2.cpp。调用示例参数核对示例中B1, S2048, n4, D2560完全满足约束n4 在所有支持产品范围内D2560 满足 16 对齐与 128 对齐要求2560 16×160 128×20phi_shape{n²2n, nD}{24, 10240}、bias_shape{24}、h_mix_shape{B*S, 24}符合 alpha(3) 残差分支的 shape 约定输出 h_res 为(B*S, n, n)与 950 平台的输出约束一致。示例中同时传入 gamma 与全部三个可选输出对应out_flag1的完整输出场景。源码实现要点速览除上文已引用的算子定义、infershape、tiling 与 aclnn 接口源码外以下文件可帮助进一步理解 MhcPre 的底层实现kernel 基础实现mhc/mhc_pre/op_kernel/mhc_pre_base.h 定义了 kernel 侧的基础常量如BLOCK_SIZE32、REPEAT_SIZE256、双缓冲等与逐元素乘/减/除等向量原语封装其中Factorial(n)用于阶乘分支架构相关 kernelmhc/mhc_pre/op_kernel/arch35/mhc_pre_common.h、mhc_pre_cube_compute.h、mhc_pre_vector_compute.h、mhc_pre_m_k_split.h、mhc_pre_split_bs.h、mhc_pre_split_nd.h 分别对应 Cube 矩阵乘、Vector 向量计算以及三种切分模式的实现tiling 主流程mhc/mhc_pre/op_host/op_tiling/arch35/mhc_pre_tiling.cpp 负责平台资源读取、shape 校验、切分策略选择SPLIT_BS/SPLIT_ND/SPLIT_M_K与 TilingData 填充单元测试mhc/mhc_pre/tests/ut/op_host/op_api/test_aclnn_mhc_pre.cpp、test_mhc_pre_infershape.cpp 与 test_mhc_pre_tiling.cpp 分别覆盖 aclnn 接口、shape 推导与 tiling 逻辑。从源码结构看MhcPre 采用了 AICCube 矩阵乘与 AIVVector 向量计算混编的架构tiling key 中ASCENDC_TPL_MIX_AIC_1_2即一个 AIC 搭配两个 AIV 的模板组合将 RmsNorm、矩阵乘、分段 sigmoid 激活与 h_in 合成融合为单算子执行避免中间结果多次读写 Global Memory。总结MhcPre 是 MHC 架构 transformer 网络中位于 hidden 层的关键前置算子将 RmsNorm、参数投影、分段激活与输入矩阵合成融合为一次 NPU 计算一次产出 H^res、H^post 与 h^in。使用时需重点确认三点其一根据alpha形状(2) 或 (3)选择匹配的phi/bias/hRes/hMixshape其二根据目标产品核对 n 与 D 的规格约束950 平台支持 n4/6/8、D 16 对齐A2/A3 平台 n4、D 128 对齐且仅支持 FP32 模式其三可选输出 inv_rms、h_mix、h_pre 必须同时输出或全部不输出。需要 HF32 计算模式时仅 Ascend 950PR/950DT 支持通过 aclnnMhcPreV2 的opImplMode1开启。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考