CANN ops-transformer 中 aclnnMhcPreV2 算子接口全解析：MHC 架构 H 投影矩阵与 hIn 计算实战指南

CANN ops-transformer 中 aclnnMhcPreV2 算子接口全解析MHC 架构 H 投影矩阵与 hIn 计算实战指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读aclnnMhcPreV2 是 CANN ops-transformer 算子库中 mhc/mhc_pre 模块面向 Ascend NPU 提供的 MhcPre 算子两段式 aclnn 接口。该算子负责在 Mamba-Hybrid-ConvMHC混合架构中一次性计算出 hidden 层的 H^res、H^post 投影矩阵以及送入 Attention/MLP 层的输入矩阵 h^in是 MHC 层前向计算的关键前置算子。与第一代 aclnnMhcPre 相比V2 版本新增了opImplMode参数允许开发者在 Cube 单元的 FP32 与 HF32 两种矩阵乘精度模式之间显式切换从而在精度与性能之间取得平衡。读完本文你将掌握 aclnnMhcPreV2 的数学语义、全部入参/出参约束、两段式调用流程、错误码含义、规格限制以及一份可直接编译运行的完整 C 调用示例。一、算子定位MHC 架构中的前置投影计算MhcPre 算子位于mhc/mhc_pre目录是 MHC一种将 Mamba 状态空间模型与卷积/Attention 组件混合的架构网络中 hidden 层的前置计算单元。其核心职责是基于输入 x 与参数矩阵 phi计算得到 hidden 层的 H^res、H^post 投影矩阵以及作为 Attention 或 MLP 层输入的 h^in 矩阵为后续的 sinkhorn 变换、注意力计算等环节提供中间结果。从算子注册文件 mhc_pre_def.cpp 可以看到该算子的输入输出定义输入xBF16/FP16、phiFP32、alphaFP32、biasFP32、可选gammaFP32输出hinBF16/FP16、h_postFP32、h_resFP32以及可选的inv_rms、h_mix、h_pre属性out_flag默认 0、norm_eps默认 1e-6f、hc_eps默认 1e-6f、op_impl_mode默认 0。其中默认属性值与文档中 normEps、hcEps 的建议值 1e-6 完全一致。该算子在 Graph 层同时注册了ascend910bA2 系列与ascend950Ascend 950PR/950DT两套 AICore 配置但在 aclnn 接口文档中产品支持矩阵以接口文档为准。二、V2 新增能力opImplMode 计算模式开关aclnnMhcPreV2 相较于 aclnnMhcPre 的核心差异是新增opImplMode参数用于指定 MhcPre 算子在 Cube矩阵乘单元中的计算模式opImplMode含义说明0FP32 模式建议值Cube 中以 FP32 精度完成矩阵乘计算精度更高1HF32 模式Cube 中以 HF32 模式计算可降低循环开销、提升性能该参数在接口实现 aclnn_mhc_pre.cpp 中被定义为两个常量MHC_PRE_USE_FP32 0、MHC_PRE_USE_HF32 1并在参数校验阶段对取值合法性进行检查opImplMode不为 0 或 1 时报参数非法错误。其作用最终传递到 Kernel 侧在 mhc_pre_cube_compute.h 中当implMode MHC_PRE_IMPL_MODE_HF32时通过AscendC::SetHF32Mode(1)与AscendC::SetHF32TransMode(1)开启 HF32 模式且限制仅在 AIC 侧写入避免 AIV 侧非法操作矩阵乘结束后再通过AscendC::SetHF32Mode(0)关闭同一文件中的注释还揭示了两种模式的底层差异FP32 模式保持 K32 的累加顺序以保证精度HF32 模式则消费完整的 baseK 以减少循环开销。从源码结构看HF32 模式正是为追求更高吞吐而设计的。三、数学语义完整的计算公式MhcPre 的计算过程可拆解为 RmsNorm、投影、逐元素变换、Sinkhorn 前处理等多个环节。文档给出的完整公式如下其中 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 归一化对 xFlat 求均方后加 normEps 再开方取倒数得到 invRms即 1/r。若提供了 gamma则先对 xFlat 做逐元素缩放得到 xGamma投影xGamma 与 phi 的转置做矩阵乘得到 hMix再与 invRms 逐元素相乘得到 w切分根据 alpha.shape 决定切分方式。alpha.shape[3] 时把 w 切分为 [n, n, n²] 三段分别对应 pPre、pPost、pResalpha.shape[2] 时切分为 [n, n] 两段pRes 恒为 0非线性变换pPre 经 sigmoid 加偏置缩放后叠加 hcEps 得到 hPrepPost 经同样变换后乘 2 得到 hPostalpha.shape[2] 时同样叠加 hcEpspRes 经缩放加偏置得到 hResalpha.shape[2] 时为 0即无残差路径汇聚hIn 由 hPre 与原始 x 在 n 维上加权求和得到即 $hIn_{d} \sum_{i0}^{n-1} hPre_{i} x_{i,d}$。其中 hPre 对应的 sigmoid 路径在 mhc_pre_cube_compute.h 等 Kernel 实现中有对应的向量计算Vector环节与 Cube 矩阵乘环节共同构成完整的数据流。四、函数原型与两段式调用模型所有 aclnn 算子都遵循两段式接口规范先调用 GetWorkspaceSize 接口获取 workspace 大小并创建包含算子计算流程的执行器再调用执行接口真正下发计算。aclnnMhcPreV2 亦不例外aclnnStatus aclnnMhcPreV2GetWorkspaceSize( const aclTensor *x, const aclTensor *phi, const aclTensor *alpha, const aclTensor *bias, const aclTensor *gammaOptional, double normEps, double hcEps, int64_t opImplMode, aclTensor *hIn, aclTensor *hPost, aclTensor *hRes, aclTensor *invRmsOptional, aclTensor *hMixOptional, aclTensor *hPreOptional, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMhcPreV2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)该接口声明位于 aclnn_mhc_pre_v2.h实现位于 aclnn_mhc_pre.cpp。从接口实现看参数被封装为MhcParamsBase结构体默认opImplMode MHC_PRE_USE_FP32通过MhcBuilder建造器模式组装随后经过一连串校验函数非空校验、空 Tensor 校验、维度校验、shape 校验、可选输出组校验等后才进入 workspace 计算与算子下发阶段。五、aclnnMhcPreV2GetWorkspaceSize 参数详解第一段接口的参数可分为四类必选输入、可选输入、输出、以及 workspace/executor 返回参数。5.1 必选输入参数名描述使用说明数据类型数据格式维度(shape)非连续Tensorx待计算数据网络中 MHC 层的输入数据对应公式中的 x不能为空 TensorBFLOAT16 或 FLOAT16ND[B,S,n,D] 或 [T,n,D]√phiMHC 的参数矩阵对应公式中的 phi不能为空 TensorFLOAT32ND[n²2n, nD] 或 [2n, nD]√alphaMHC 的缩放参数对应公式中的 alpha不能为空 TensorFLOAT32-[3] 或 [2]-biasMHC 的 bias 参数对应公式中的 bias不能为空 TensorFLOAT32-[n²2n] 或 [2n]-其中 x 支持 [B,S,n,D] 与 [T,n,D] 两种布局前者对应带 Batch 与序列长度的训练/推理输入后者对应 Token 维度摊平的变长序列输入。x、phi 均支持非连续 Tensor调用侧会先做 contiguous 转换alpha 与 bias 为一维向量。5.2 可选输入参数名描述使用说明数据类型gammaOptionalRmsNorm 计算的缩放因子对应公式中的 gamma推理场景可传空指针FLOAT32normEpsRmsNorm 的防除零参数对应公式中的 normEps建议值 1e-6DOUBLEhcEpsh_pre 的 sigmoid 后 eps 参数alpha.shape[2] 时同时作用于 h_post建议值 1e-6DOUBLEopImplMode指定 MhcPre 算子的计算模式0建议值Cube 中 FP32 计算1Cube 中 HF32 计算INT64normEps与hcEps在算子注册中均有默认值 1e-6f与文档建议值一致gammaOptional在推理场景可省略公式中 xGamma 退化为 xFlat。5.3 输出参数名描述使用说明数据类型维度(shape)hIn输出的 h_in作为 Attention/MLP 层的输入不能为空 TensorBFLOAT16 或 FLOAT16[B,S,D] 或 [T,D]hPost输出的 MHC 的 h_post 变换矩阵不能为空 TensorFLOAT32[B,S,n] 或 [T,n]hRes输出的 h_res 变换矩阵未做 sinkhorn 变换alpha[3] 时不能为空alpha[2] 时可传空指针FLOAT32[B,S,n,n] 或 [T,n,n]invRmsOptionalRmsNorm 计算得到的 1/r与 hMixOptional、hPreOptional 互存FLOAT32[B,S] 或 [T]hMixOptionalxGamma 与 phi 矩阵乘的结果同上FLOAT32[B,S,n²2n] 或 [T,n²2n] 或 [B,S,2n] 或 [T,2n]hPreOptional做完 sigmoid 计算之后的 h_pre 矩阵同上FLOAT32[B,S,n] 或 [T,n]workspaceSize返回需要在 Device 侧申请的 workspace 大小由第一段接口输出--executor返回 op 执行器包含算子计算流程由第一段接口输出--注意输出维度的“降维”关系x 为 4 维 [B,S,n,D] 时hIn 为 [B,S,D]、hPost 为 [B,S,n]x 为 3 维 [T,n,D] 时hIn 为 [T,D]、hPost 为 [T,n]。这一映射逻辑在接口实现的CheckOutputShape4D/CheckOutputShape3D等校验函数中均有完整校验。5.4 返回值第一段接口第一段接口完成入参校验返回aclnnStatus状态码具体取值参见 aclnn 返回码。校验失败场景如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001x、phi、alpha、bias、hIn、hPost、hRes 存在空指针ACLNN_ERR_PARAM_INVALID161002x、phi、alpha、bias、gammaOptional、hIn、hPost、hRes、invRmsOptional、hMixOptional、hPreOptional 的数据类型不在支持范围内ACLNN_ERR_PARAM_INVALID161002x、phi、alpha、bias、gammaOptional、hIn、hPost、hRes、invRmsOptional、hMixOptional、hPreOptional 的 shape 维度不在支持范围内ACLNN_ERR_PARAM_INVALID161002opImplMode 不为 0 或 1ACLNN_ERR_RUNTIME_ERROR361001API 内部调用 npu runtime 的接口异常从实现源码看这些校验对应 aclnn_mhc_pre.cpp 中的CheckNotNull非空校验、CheckEmptyTensor空 Tensor 校验、CheckInputDims/CheckOutputDims维度校验、CheckParamShapesphi/bias/gamma shape 校验以及可选输出组校验等函数。六、aclnnMhcPreV2 执行接口参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMhcPreV2GetWorkspaceSize 获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream第二段接口返回aclnnStatus状态码具体参见 aclnn 返回码。七、约束说明7.1 确定性计算aclnnMhcPreV2 默认确定性实现即相同输入在同一硬件条件下计算结果可复现。相关背景可参考 确定性计算。7.2 Batch 一致性默认非 Batch 一致性实现不支持通过 aclrtSetSysParamOpt 开启 Batch 一致性。背景知识参见 Batch 一致性。7.3 规格约束n 目前支持 4、6、8D 支持 1~16384 范围以内且需满足 D 为 16 对齐当 alpha[3] 时支持 hRes 输出必须满足输入 phi[n²2n, nD]、bias[n²2n]输出 hMixOptional[B,S,n²2n] 或 [T,n²2n]当 alpha[2] 时 hRes 输出为 0必须满足输入 phi[2n, nD]、bias[2n]输出 hMixOptional[B,S,2n] 或 [T,2n]可选输出 invRmsOptional、hMixOptional、hPreOptional 为互存关系需同时输出或全部不输出不支持仅返回其中部分输出条件invRmsOptional!nullptr hMixOptional!nullptr hPreOptional!nullptr。上述“互存关系”在实现中由CheckOptionalOutputGroup函数强制执行见 aclnn_mhc_pre.cpp若仅提供部分可选输出会报参数非法错误。需要说明的是mhc_pre/README.md 中针对 Atlas A2/A3 系列产品还列出了另一组规格n 仅支持 4D 支持 100000 范围内且 128 对齐op_impl_mode 仅支持 0而 aclnnMhcPreV2 接口文档的产品支持矩阵显示该 V2 接口当前仅支持 Ascend 950PR/Ascend 950DT实际规格以当前接口文档为准。八、产品支持情况aclnnMhcPreV2 接口当前的产品支持矩阵如下产品支持情况Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品不支持Atlas A2 训练系列产品/Atlas A2 推理系列产品不支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持九、完整调用示例以下示例代码来源于仓库 examples/test_aclnn_mhc_pre_v2.cpp与接口文档中的示例一致展示了从环境初始化、Tensor 创建、两段式调用到结果打印、资源释放的完整流程。具体编译与执行过程请参考 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_mhc_pre_v2.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) // 计算Tensor形状对应的总元素数 int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t size 1; for (int64_t dim : shape) { size * dim; } return size; } // 将Device侧Tensor数据拷贝到Host侧并打印float类型 void PrintTensorDataFloat(const std::vectorint64_t shape, void *device_addr) { int64_t size GetShapeSize(shape); std::vectorfloat host_data(size, 0.0f); aclError ret aclrtMemcpy(host_data.data(), size * sizeof(float), device_addr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Memcpy device to host failed, error: %d\n, ret); return); LOG_PRINT(Tensor data (first 10 elements): ); for (int64_t i 0; i std::min((int64_t)10, size); i) { LOG_PRINT(%f , host_data[i]); } LOG_PRINT(\n); } // 将Device侧Tensor数据拷贝到Host侧并打印float16类型 void PrintTensorDataFloat16(const std::vectorint64_t shape, void *device_addr) { int64_t size GetShapeSize(shape); std::vectoraclFloat16 host_fp16(size); std::vectorfloat host_data(size, 0.0f); aclError ret aclrtMemcpy(host_fp16.data(), size * sizeof(aclFloat16), device_addr, size * sizeof(aclFloat16), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Memcpy device to host failed, error: %d\n, ret); return); for (int64_t i 0; i size; i) { host_data[i] aclFloat16ToFloat(host_fp16[i]); } LOG_PRINT(Tensor data (first 10 elements): ); for (int64_t i 0; i std::min((int64_t)10, size); i) { LOG_PRINT(%f , host_data[i]); } LOG_PRINT(\n); } // 初始化AscendCL环境Device/Context/Stream int InitAcl(int32_t device_id, aclrtContext context, aclrtStream stream) { aclError ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed, error: %d\n, ret); return -1); ret aclrtSetDevice(device_id); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed, error: %d\n, ret); return -1); ret aclrtCreateContext(context, device_id); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateContext failed, error: %d\n, ret); return -1); ret aclrtSetCurrentContext(context); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetCurrentContext failed, error: %d\n, ret); return -1); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed, error: %d\n, ret); return -1); return 0; } // 创建FLOAT32类型Device侧aclTensor含数据拷贝 int CreateAclTensorFloat32(const std::vectorfloat host_data, const std::vectorint64_t shape, void *device_addr, aclTensor *tensor) { int64_t size GetShapeSize(shape) * sizeof(float); aclError ret aclrtMalloc(device_addr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed, error: %d\n, ret); return -1); ret aclrtMemcpy(device_addr, size, host_data.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed, error: %d\n, ret); return -1); std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; --i) { strides[i] strides[i 1] * shape[i 1]; } tensor aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT, strides.data(), 0, ACL_FORMAT_ND, shape.data(), shape.size(), device_addr); CHECK_RET(tensor ! nullptr, LOG_PRINT(aclCreateTensor failed\n); return -1); return 0; } // 创建FLOAT16类型Device侧aclTensor含数据拷贝 int CreateAclTensorFloat16(const std::vectorfloat host_data, const std::vectorint64_t shape, void *device_addr, aclTensor *tensor) { int64_t size GetShapeSize(shape); std::vectoraclFloat16 host_data_fp16(size); for (int64_t i 0; i size; i) { host_data_fp16[i] aclFloat16(host_data[i]); } int64_t byte_size size * sizeof(aclFloat16); aclError ret aclrtMalloc(device_addr, byte_size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed, error: %d\n, ret); return -1); ret aclrtMemcpy(device_addr, byte_size, host_data_fp16.data(), byte_size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed, error: %d\n, ret); return -1); std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; --i) { strides[i] strides[i 1] * shape[i 1]; } tensor aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT16, strides.data(), 0, ACL_FORMAT_ND, shape.data(), shape.size(), device_addr); CHECK_RET(tensor ! nullptr, LOG_PRINT(aclCreateTensor failed\n); return -1); return 0; } // 创建FLOAT16类型输出aclTensor仅申请内存 int CreateAclTensorFloat16Output(const std::vectorint64_t shape, void *device_addr, aclTensor *tensor) { int64_t size GetShapeSize(shape); int64_t byte_size size * sizeof(aclFloat16); aclError ret aclrtMalloc(device_addr, byte_size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed, error: %d\n, ret); return -1); tensor aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT16, nullptr, 0, ACL_FORMAT_ND, shape.data(), shape.size(), device_addr); CHECK_RET(tensor ! nullptr, LOG_PRINT(aclCreateTensor failed\n); return -1); return 0; } // 创建FLOAT32类型输出aclTensor仅申请内存 int CreateAclTensorFloat32Output(const std::vectorint64_t shape, void *device_addr, aclTensor *tensor) { int64_t size GetShapeSize(shape); int64_t byte_size size * sizeof(float); aclError ret aclrtMalloc(device_addr, byte_size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed, error: %d\n, ret); return -1); tensor aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT, nullptr, 0, ACL_FORMAT_ND, shape.data(), shape.size(), device_addr); CHECK_RET(tensor ! nullptr, LOG_PRINT(aclCreateTensor failed\n); return -1); return 0; } struct Tensors { void *x_addr nullptr, *phi_addr nullptr, *alpha_addr nullptr, *bias_addr nullptr, *gamma_addr nullptr; void *hin_addr nullptr, *h_post_addr nullptr, *h_res_addr nullptr, *inv_rms_addr nullptr; void *h_mix_addr nullptr, *h_pre_addr nullptr; aclTensor *x nullptr, *phi nullptr, *alpha nullptr, *bias nullptr, *gamma nullptr; aclTensor *hin nullptr, *h_post nullptr, *h_res nullptr, *inv_rms nullptr; aclTensor *h_mix nullptr, *h_pre nullptr; }; int CreateInputTensors(const std::vectorint64_t x_shape, const std::vectorint64_t phi_shape, const std::vectorint64_t alpha_shape, const std::vectorint64_t bias_shape, const std::vectorint64_t gamma_shape, Tensors tensors) { std::vectorfloat x_host_data(GetShapeSize(x_shape), 1.0f); std::vectorfloat phi_host_data(GetShapeSize(phi_shape), 1.0f); std::vectorfloat alpha_host_data(3, 1.0f); std::vectorfloat bias_host_data(GetShapeSize(bias_shape), 1.0f); std::vectorfloat gamma_host_data(GetShapeSize(gamma_shape), 1.0f); int ret CreateAclTensorFloat16(x_host_data, x_shape, tensors.x_addr, tensors.x); CHECK_RET(ret 0, LOG_PRINT(Create x_tensor failed\n); return -1); ret CreateAclTensorFloat32(phi_host_data, phi_shape, tensors.phi_addr, tensors.phi); CHECK_RET(ret 0, LOG_PRINT(Create phi_tensor failed\n); return -1); ret CreateAclTensorFloat32(alpha_host_data, alpha_shape, tensors.alpha_addr, tensors.alpha); CHECK_RET(ret 0, LOG_PRINT(Create alpha_tensor failed\n); return -1); ret CreateAclTensorFloat32(bias_host_data, bias_shape, tensors.bias_addr, tensors.bias); CHECK_RET(ret 0, LOG_PRINT(Create bias_tensor failed\n); return -1); ret CreateAclTensorFloat32(gamma_host_data, gamma_shape, tensors.gamma_addr, tensors.gamma); CHECK_RET(ret 0, LOG_PRINT(Create gamma_tensor failed\n); return -1); return 0; } int CreateOutputTensors(const std::vectorint64_t hin_shape, const std::vectorint64_t h_post_shape, const std::vectorint64_t h_res_shape, const std::vectorint64_t inv_rms_shape, const std::vectorint64_t h_mix_shape, const std::vectorint64_t h_pre_shape, Tensors tensors) { int ret CreateAclTensorFloat16Output(hin_shape, tensors.hin_addr, tensors.hin); CHECK_RET(ret 0, LOG_PRINT(Create hin_tensor failed\n); return -1); ret CreateAclTensorFloat32Output(h_post_shape, tensors.h_post_addr, tensors.h_post); CHECK_RET(ret 0, LOG_PRINT(Create h_post_tensor failed\n); return -1); ret CreateAclTensorFloat32Output(h_res_shape, tensors.h_res_addr, tensors.h_res); CHECK_RET(ret 0, LOG_PRINT(Create h_res_tensor failed\n); return -1); ret CreateAclTensorFloat32Output(inv_rms_shape, tensors.inv_rms_addr, tensors.inv_rms); CHECK_RET(ret 0, LOG_PRINT(Create inv_rms_tensor failed\n); return -1); ret CreateAclTensorFloat32Output(h_mix_shape, tensors.h_mix_addr, tensors.h_mix); CHECK_RET(ret 0, LOG_PRINT(Create h_mix_tensor failed\n); return -1); ret CreateAclTensorFloat32Output(h_pre_shape, tensors.h_pre_addr, tensors.h_pre); CHECK_RET(ret 0, LOG_PRINT(Create h_pre_tensor failed\n); return -1); return 0; } void DestroyTensors(Tensors tensors) { aclDestroyTensor(tensors.x); aclDestroyTensor(tensors.phi); aclDestroyTensor(tensors.alpha); aclDestroyTensor(tensors.bias); aclDestroyTensor(tensors.gamma); aclDestroyTensor(tensors.hin); aclDestroyTensor(tensors.h_post); aclDestroyTensor(tensors.h_res); aclDestroyTensor(tensors.inv_rms); aclDestroyTensor(tensors.h_mix); aclDestroyTensor(tensors.h_pre); } void FreeDeviceMemory(Tensors tensors) { aclrtFree(tensors.x_addr); aclrtFree(tensors.phi_addr); aclrtFree(tensors.alpha_addr); aclrtFree(tensors.bias_addr); aclrtFree(tensors.gamma_addr); aclrtFree(tensors.hin_addr); aclrtFree(tensors.h_post_addr); aclrtFree(tensors.h_res_addr); aclrtFree(tensors.inv_rms_addr); aclrtFree(tensors.h_mix_addr); aclrtFree(tensors.h_pre_addr); } int main() { int32_t device_id 0; aclrtContext context nullptr; aclrtStream stream nullptr; Tensors tensors; 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); CHECK_RET(ret 0, LOG_PRINT(InitAcl failed, error: %d\n, ret); return -1); ret CreateInputTensors(x_shape, phi_shape, alpha_shape, bias_shape, gamma_shape, tensors); CHECK_RET(ret 0, return -1); ret CreateOutputTensors(hin_shape, h_post_shape, h_res_shape, inv_rms_shape, h_mix_shape, h_pre_shape, tensors); CHECK_RET(ret 0, return -1); uint64_t workspace_size 0; aclOpExecutor *executor nullptr; int64_t opImplMode 1; aclnnStatus aclnn_ret aclnnMhcPreV2GetWorkspaceSize( tensors.x, tensors.phi, tensors.alpha, tensors.bias, tensors.gamma, 1e-6, 1e-6, opImplMode, tensors.hin, tensors.h_post, tensors.h_res, tensors.inv_rms, tensors.h_mix, tensors.h_pre, workspace_size, executor); CHECK_RET(aclnn_ret ACL_SUCCESS, LOG_PRINT(aclnnMhcPreV2GetWorkspaceSize failed, error: %d\n, aclnn_ret); return -1); void *workspace_addr nullptr; if (workspace_size 0) { ret aclrtMalloc(workspace_addr, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnrtMalloc workspace failed, error: %d\n, ret); return -1); } aclnn_ret aclnnMhcPreV2(workspace_addr, workspace_size, executor, stream); CHECK_RET(aclnn_ret ACL_SUCCESS, LOG_PRINT(aclnnMhcPreV2 failed, error: %d\n, aclnn_ret); return -1); CHECK_RET(aclrtSynchronizeStream(stream) ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed\n); return -1); LOG_PRINT(MhcPre compute success!\nOutput tensor data: \n); 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(); LOG_PRINT(All resources released successfully!\n); return 0; }9.1 示例关键点解读shape 设计示例采用 3 维 x 布局 [B*S, n, D]即 TB*S 的摊平形式n4、D2560满足文档“D 为 16 对齐”的约束alpha[3] 对应带残差路径hRes 非 0的完整三投影模式可选输出全量开启示例同时创建了 invRms、hMix、hPre 三个可选输出符合“互存关系”约束opImplMode 选择示例将opImplMode设为 1HF32 模式演示了 V2 新增能力的用法如追求更高精度可改为 0两段式调用先aclnnMhcPreV2GetWorkspaceSize获取 workspace 大小与 executor再按需aclrtMalloc申请 workspace随后aclnnMhcPreV2异步下发计算最后aclrtSynchronizeStream同步等待结果资源释放依次销毁 TensoraclDestroyTensor、释放 Device 内存aclrtFree、销毁 Stream/Context、aclrtResetDevice与aclFinalize收尾避免资源泄漏。十、源码级验证从 aclnn 到 Kernel 的调用链在 mhc/mhc_pre/op_host/op_api/aclnn_mhc_pre.cpp 中第一段接口的实现采用建造器模式组织参数校验MhcParamsBase保存全部入参默认opImplMode MHC_PRE_USE_FP32CheckNotNull校验 x、phi、alpha、bias、hIn、hPost 非空并额外校验 alpha[3] 时 hRes 非空对应文档“hRes 当 alpha[3] 时不能为空”的约束CheckEmptyTensor禁止空 TensorCheckInputDims/CheckOutputDims校验维度x 为 3 或 4 维、phi 2 维、alpha/bias 1 维、gammaOptional 2 维CheckParamShapes校验 phi[n²2n, nD]、bias[n²2n]、gamma[n,D] 等精确 shapeCheckOptionalOutputGroup强制三个可选输出互存CheckOutputShape4D/CheckOutputShape3D校验 hIn/hPost/hRes/可选输出的精确 shape。Tiling 侧由 mhc_pre_tiling.cpp 等文件负责针对 Ascend 950 架构实现了多种切分策略如 mhc_pre_m_k_split.h、mhc_pre_split_bs.h、mhc_pre_split_nd.hKernel 侧则分为 Cube 矩阵乘环节mhc_pre_cube_compute.h与 Vector 计算环节mhc_pre_vector_compute.h共同完成公式中的 RmsNorm、矩阵乘与 sigmoid 等运算。仓库同时提供了配套的单测如 tests/ut/op_host/op_api/test_aclnn_mhc_pre.cpp与 infershape/tiling 单测可作为行为对照与回归验证依据。十一、总结与实践建议aclnnMhcPreV2 为 MHC 架构在 Ascend NPU 上的前置投影计算提供了完整的 aclnn 接入能力功能完整一次调用同时产出 hIn、hPost、hRes 三个核心结果并可选导出 invRms、hMix、hPre 中间量便于上层二次计算或调试模式可选opImplMode在 FP32高精度与 HF32高性能之间提供显式权衡实际使用时可结合模型精度需求与性能指标进行选择约束清晰n∈{4,6,8}、D 16 对齐、可选输出互存等约束均在接口层强制校验调用前按文档核对 shape 即可规避 161002 类错误推理友好gammaOptional 在推理场景可传空指针支持 [T,n,D] 摊平布局适配变长序列推理。当前接口仅支持 Ascend 950PR/Ascend 950DT且默认确定性实现、不支持 Batch 一致性。若需在 A2/A3 等产品上使用 MhcPre可参考 aclnnMhcPre 接口及 mhc/mhc_pre/README.md 中的产品规格说明。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考