ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

CANN ops-nn `aclnnFusedMatmulGelu` 融合算子接口详解:MatMul + Bias + GELU 的 NPU 融合计算与两阶段 aclnn 调用

CANN ops-nn `aclnnFusedMatmulGelu` 融合算子接口详解:MatMul + Bias + GELU 的 NPU 融合计算与两阶段 aclnn 调用 CANN ops-nnaclnnFusedMatmulGelu融合算子接口详解MatMul Bias GELU 的 NPU 融合计算与两阶段 aclnn 调用【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnaclnnFusedMatmulGelu是 CANN ops-nn 仓库中FusedMatmulGelu融合算子的对外 aclnn 接口用于在 NPU 上一次性完成全连接线性变换MatMul/Linear、可选偏置加法与 GELU 激活函数的融合计算从而减少 kernel 启动开销与中间结果搬运。本文以 experimental/matmul/fused_matmul_gelu/docs/aclnnFusedMatmulGelu.md 为骨架结合仓库中 op_api、op_host、op_kernel 的源码实现与单元测试完整讲解接口原型、参数语义、约束条件、两阶段调用流程以及从参数校验到 Tiling 分片、再到 AICore Kernel 执行的底层原理帮助开发者在 Atlas A2 系列产品上正确、高效地使用该接口。算子功能与融合动机aclnnFusedMatmulGelu接口用于调用自定义FusedMatmulGelu融合算子实现MatMul/Linear 矩阵乘、可选偏置加法Bias Add和 GELU 激活函数的融合计算。相比MatMul → Add → GELU三段独立算子串联融合算子将整个计算链路在一个 kernel 内完成避免了中间结果反复写回/读取 GMGlobal Memory同时复用 Cube 核AIC的矩阵计算能力与 Vector 核AIV的激活计算能力显著降低整体执行延迟。计算公式如下y GELU(x * weight^T bias)当bias为空时计算公式为y GELU(x * weight^T)其中x输入张量shape 为[..., K]weight权重张量shape 为[N, K]计算时按weight^T即[K, N]参与矩阵乘bias可选偏置张量shape 为[N]为空时跳过偏置加法y输出张量shape 为[..., N]。从功能语义上看该算子等价于对一个[..., K]的输入执行一次[N, K]权重的线性投影后套用 GELU 激活是 Transformer/MLP 类网络结构中 Linear-GELU 子结构的典型落点。支持的产品型号根据关联文档当前接口支持的产品型号如下产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√作为补充算子目录下的 README.md 给出了更细粒度的产品支持情况除 Atlas A2 训练/推理系列产品外还支持 Ascend 950PR/Ascend 950DT、Atlas A3 训练/推理系列产品而 Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持。这与算子注册配置见下文 算子定义与注册中仅为ascend910b、ascend910_93、ascend950三个平台添加 AICore 配置的实现是一致的——从源码结构看平台适配由fused_matmul_gelu_def.cpp中的AddFusedMatmulGeluAicoreConfig显式声明。实际部署时请以目标环境配套的 CANN 版本所支持的硬件为准。接口原型该接口采用 aclnn 标准的两阶段调用模型第一阶段计算 workspace 大小并构建执行器executor第二阶段携带 workspace 与 stream 真正执行算子。aclnnStatus aclnnFusedMatmulGeluGetWorkspaceSize( const aclTensor* x, const aclTensor* weight, const aclTensor* bias, int64_t approximate, aclTensor* y, uint64_t* workspaceSize, aclOpExecutor** executor); aclnnStatus aclnnFusedMatmulGelu( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream);接口声明位于 op_api/aclnn_fused_matmul_gelu.h头文件中以extern C导出保证 C/C 均可链接实现位于 op_api/aclnn_fused_matmul_gelu.cpp。参数说明参数名输入/输出说明x输入输入张量shape 为[..., K]数据类型支持 FLOAT16、BFLOAT16数据格式为 ND。weight输入权重张量shape 为[N, K]数据类型需要与x保持一致数据格式为 ND。bias输入可选偏置张量shape 为[N]数据类型需要与x保持一致数据格式为 ND可为空传nullptr。approximate输入GELU 计算模式属性当前仅支持取值为 1tanh 近似模式。y输出输出张量shape 为[..., N]数据类型需要与x保持一致数据格式为 ND。workspaceSize输出返回执行该算子所需的 workspace 大小字节。executor输出返回算子执行器。workspace输入算子执行所需的 workspace 地址由调用方按workspaceSize申请。stream输入ACL runtime stream指定算子提交到的执行流。各参数的源码级语义结合 op_api/aclnn_fused_matmul_gelu.cpp 的实现可以进一步确认各参数在底层的精确语义x的秩约束源码中定义MIN_X_DIM_NUM 2即x至少是 2 维[..., K]同时x的非最后一维会被连乘展开为 MatMul 的 M 维见 op_host/fused_matmul_gelu_tiling.cpp 中mSize_的累乘逻辑因此x支持任意 ≥2 维的高维 batch 形态。weight的秩约束WEIGHT_DIM_NUM 2必须恰好为 2 维[N, K]。bias的秩约束BIAS_DIM_NUM 1必须为 1 维[N]且N与weight的第 0 维一致。y的 shape 推导由InferYShape函数推导——y的 shape 为x去掉最后一维后拼接weight的第 0 维N即[..., N]op_api 层的l0op::FusedMatmulGelu见 op_api/fused_matmul_gelu.cpp会按照同样的规则为输出张量分配内存。约束说明使用aclnnFusedMatmulGelu时需满足以下约束x、weight、bias和y的数据类型需要保持一致。当前支持 FLOAT16、BFLOAT16 数据类型。当前支持 ND 数据格式。x的最后一维大小K需要与weight的最后一维大小K保持一致。bias为可选输入当bias不为空时其 shape 需要为[N]。approximate当前仅支持取值为 1tanh 近似模式。约束在源码中的落地上述约束并非文档层面的建议而是被 op_api 层与 tiling 层双重强制校验的硬性条件op_api 层参数校验aclnn_fused_matmul_gelu.cpp 中CheckParam及其子函数CheckNotNullL85-L95x、weight、y、workspaceSize、executor为空时返回ACLNN_ERR_PARAM_NULLPTRbias允许为空CheckDtypeValidL97-L111通过FLOAT_DTYPE_SUPPORT_LIST {DT_FLOAT16, DT_BF16}校验数据类型并强制x/weight/bias/y四者类型一致CheckShapeValidL113-L169校验x维度 ≥2、weight恰好 2 维、K 对齐、N/K 0、bias为[N]、y与推导 shape 完全一致CheckAttrValidL171-L178approximate必须等于APPROXIMATE_TANH 1否则返回ACLNN_ERR_PARAM_INVALID。tiling 层二次校验fused_matmul_gelu_tiling.cppCheckAndParseInputShape额外限制 K 不能超过MAX_K_SIZE 65534CheckAndParseDtype、CheckAndParseAttr在运行时再次确认数据类型与approximate取值合法。因此一旦传入不满足约束的输入如approximate2、K 不匹配、biasshape 错误GetWorkspaceSize阶段就会直接报错而不会进入后续执行阶段。对应行为可参见 tests/ut/op_api/test_aclnn_fused_matmul_gelu.cpp 中的invalid_approximate、invalid_weight_k_mismatch、invalid_bias_shape、invalid_output_shape、invalid_dtype_mismatch、input_nullptr_invalid等测试用例。调用说明与完整调用流程关联文档给出的调用样例为 examples/test_aclnn_fused_matmul_gelu.cpp其入口仅包含对头文件aclnn_fused_matmul_gelu.h的引用。结合 aclnn 标准调用范式与仓库内的 UT 用例一个完整的调用流程包含以下步骤1. 包含头文件#include aclnn_fused_matmul_gelu.h2. 准备输入/输出张量使用aclCreateTensor创建x[..., K]、weight[N, K]、y[..., N]三个张量描述并按需创建bias[N]// 以 2 维输入为例x[M, K]weight[N, K]bias[N]y[M, N] aclTensor* x aclCreateTensor(...); // shape [M, K], 类型 FLOAT16/BFLOAT16, ND aclTensor* weight aclCreateTensor(...); // shape [N, K], 类型与 x 一致, ND aclTensor* bias aclCreateTensor(...); // shape [N], 可选不使用可传 nullptr aclTensor* y aclCreateTensor(...); // shape [M, N], 类型与 x 一致, NDbias为空即不参与计算调用时将bias置为nullptr即可。UT 用例fp16_without_bias_successtest_aclnn_fused_matmul_gelu.cpp#L59-L62验证了无 bias 场景。3. 第一阶段查询 workspace 并构建执行器uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; aclnnStatus ret aclnnFusedMatmulGeluGetWorkspaceSize( x, weight, bias, approximate /*1*/, y, workspaceSize, executor);该阶段会完成参数合法性校验 → 将x/weight/bias通过l0op::Contiguous归一化为连续内存视图 → 调用l0op::FusedMatmulGelu构建算子图节点 → 通过l0op::ViewCopy将融合结果写回用户输出张量y最终从 executor 中查询得到 workspace 大小见 aclnn_fused_matmul_gelu.cpp#L193-L230。4. 申请 workspace 内存void* workspace nullptr; if (workspaceSize 0) { // 通过 aclrtMalloc 在设备侧申请 workspaceSize 字节的内存 aclrtMalloc(workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); }5. 第二阶段执行算子ret aclnnFusedMatmulGelu(workspace, workspaceSize, executor, stream);该阶段通过CommonOpExecutorRun将构建好的执行器提交到指定stream上异步执行见 aclnn_fused_matmul_gelu.cpp#L232-L236。6. 资源释放与同步执行完成后使用aclrtFree(workspace)释放 workspace使用aclDestroyTensor释放张量描述使用aclrtSynchronizeStream(stream)同步流确保算子执行完成后再读取y数据。可验证的调用组合示例仓库 UT 中的FusedMatmulGeluCommonTesttests/ut/op_api/test_aclnn_fused_matmul_gelu.cpp#L22-L47展示了不同合法组合用例x shapeweight shapebias shapey shape期望结果fp16_with_bias_success[2, 4][3, 4][3][2, 3]ACL_SUCCESSfp16_without_bias_success[2, 4][3, 4]空[2, 3]ACL_SUCCESSbf16_with_bias_success[2, 4][3, 4][3][2, 3]ACL_SUCCESS其中x[M, K]、weight[N, K]、bias[N]、y[M, N]即最典型的 2 维线性层场景高维输入如[B, S, K]会将前面各维连乘后作为 M行为与 2 维场景一致。源码实现纵深剖析算子定义与注册FusedMatmulGelu算子通过 op_host/fused_matmul_gelu_def.cpp 完成算子原型注册输入x、weight为REQUIREDbias为OPTIONAL数据类型均限定为{ge::DT_BF16, ge::DT_FLOAT16}格式限定为FORMAT_ND输出y为REQUIRED约束同上属性approximate为可选属性默认值Int(1)tanh 模式AICore 配置通过MakeFusedMatmulGeluAicoreConfig设置动态编译静态标志、动态 rank/shape 支持、PrecisionReduceFlag(true)并注册到ascend910b、ascend910_93、ascend950三个平台L55-L61。此外op_host/config/ascend910b/fused_matmul_gelu_binary.json 定义了按数据类型编译的二进制内核清单FusedMatmulGelu_BF16_BF16_BF16_BF16_TANH与FusedMatmulGelu_FP16_FP16_FP16_FP16_TANH两个 bin 分别覆盖 BF16 与 FP16 场景bias为 optionalgen_placeholder模式approximate固定为 1。也就是说GELU 的 tanh 近似模式在编译期即被固化进二进制内核运行时仅按 dtype 选择对应 bin。从参数到算子图op_api 层构建链路aclnnFusedMatmulGeluGetWorkspaceSize的完整构建链路如下aclnn_fused_matmul_gelu.cpp#L193-L230创建唯一 executorCREATE_EXECUTOR执行CheckParam完成前述全部参数校验对x、weight以及非空的bias分别调用l0op::Contiguous将非连续视图归一化为连续内存布局——这是 ND 格式下算子高效执行的前提调用l0op::FusedMatmulGelu(xContiguous, weightContiguous, biasContiguous, approximate, executor)构建融合算子节点op_api/fused_matmul_gelu.cpp#L43-L56该函数内部按[..., N]规则分配输出张量内存并通过ADD_TO_LAUNCHER_LIST_AICORE将x、weight、bias、y、approximate组装进 AICore 启动列表通过l0op::ViewCopy将融合算子产生的中间输出视图拷贝到调用方传入的y张量保证输出落在用户指定的存储上从 executor 获取 workspace 大小GetWorkspaceSize将 executor 所有权移交ReleaseTo给调用方。第二阶段aclnnFusedMatmulGelu则直接调用CommonOpExecutorRun完成异步提交。Tiling 层workspace 计算与 shape 感知分片Tiling 逻辑位于 op_host/fused_matmul_gelu_tiling.cpp核心流程RunKernelTilingL382-L455依次完成shape/dtype/attr 解析与校验确定mSize_x非最后一维连乘、nSize_weight第 0 维、kSize_x最后一维且 ≤ 65534shape 感知的 baseN 调优L398-L413针对小 N 场景自动调整分片粒度——N 1024且M 1时baseN 128N 1024且M 1时baseN 64默认baseM 128、baseN 256、baseK 64以在小 shape 下获得更充分的 N 向并行度核数裁剪当总的 Cube 分块数mBlockNum * nBlockNum小于 AIC 核数时按实际任务量收缩 AIC/AIV 核数并让 AIV 核数为 AIC 的两倍MIX_AIC_1_2模式的基础Vector 侧切分SetVectorTiling基于 UB 容量计算每轮循环元素数elemsPerVecLoop_上限 81928 元素对齐并将totalElement M * N均匀切分到各 AIV 核workspace 计算L428-L444MatMul 中间结果写回 workspace 所需大小为totalElement * dtypeBytes按 512 字节对齐再叠加SYS_WORKSPACE_BYTES 16MB的系统 workspace即workspaces[0] 16MB matmulWorkspaceSize——这就是第一阶段workspaceSize的来源TilingKey 设置tilingKey_ approximate当前恒为 1tanh 模式kernel 侧据此选择 GELU 实现分支tiling 数据落盘与调度将FusedMatmulGeluTilingData定义见 op_host/fused_matmul_gelu_tiling.h包含 m/k/n、各核任务数、hasBias、approximate、mmTiling等字段写入原始 tiling buffer并设置 block dim 与BATCH_MODE调度模式。其中MatMul 部分通过matmul_tiling::MatmulApiTiling完成 Cube 侧 tilingA 矩阵x按 GM/ND 布局B 矩阵weight标记转置对应x weight^TC 矩阵中间结果直接写到用户 workspace 供后续 Vector 核使用FP16 场景下偏置由 MatMul 的 bias 通道吸收mmTiling.SetBias(hasBias_ dtype FP16)BF16 场景偏置则放到 AIV 尾处理阶段完成L284-L307。Kernel 层Cube Vector 异构协同AICore kernel 入口位于 op_kernel/fused_matmul_gelu.cpp声明为extern C __global__ __aicore__签名接收x、weight、bias、y、workspace、tiling六个 GM 地址参数KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2)表明该 kernel 以1 个 AIC 搭配 2 个 AIV的混合调度模式运行Cube 核执行 MatMul 分块计算将中间结果写入 workspaceVector 核随后读取中间结果完成偏置加法BF16 场景与 GELU 激活GetUserWorkspace(workspace)从 tiling 阶段预留的系统 workspace 之后取得用户 workspace 起始地址与 tiling 侧16MB matmulWorkspaceSize的布局一一对应运行时通过宏INIT_AND_PROCESS依据TILING_KEY_IS(1)实例化 tanh 近似模式的算子对象并执行InitProcess。GELU 的 tanh 近似公式approximate1对应标准定义0.5 * x * (1 tanh(sqrt(2/π) * (x 0.044715 * x^3)))融合后整体计算在 Vector 核上以向量化方式完成。测试与验证该算子在仓库中配套了三个层级的单元测试见 experimental/matmul/fused_matmul_gelu/tests/utop_api 层tests/ut/op_api/test_aclnn_fused_matmul_gelu.cpp覆盖 FP16/BF16、有无 bias 的成功路径以及approximate非法、K 不匹配、biasshape 错误、输出 shape 错误、dtype 不一致、空指针等异常路径验证接口的参数校验行为op_host 层tests/ut/op_host/test_fused_matmul_gelu_tiling.cpp验证 tiling 数据生成与 shape 解析逻辑op_kernel 层tests/ut/op_kernel/test_fused_matmul_gelu.cpp验证 AICore kernel 的计算正确性。开发者可通过仓库的测试框架如 CMake 组织的 op_api/op_host/op_kernel UT 目标在对应硬件环境上运行这些用例作为接入aclnnFusedMatmulGelu前后的回归验证手段。总结aclnnFusedMatmulGelu是 CANN ops-nn 提供的、面向 Atlas A2及按 README 所述的 Ascend 950、Atlas A3平台的 MatMulBiasGELU 融合算子接口其核心价值在于功能上一次调用完成y GELU(x * weight^T bias)bias可选x支持任意 ≥2 维的高维输入接口上遵循标准的两阶段 aclnn 范式GetWorkspaceSize → 申请 workspace → 执行参数约束在 op_api 与 tiling 两层被硬性校验实现上Cube 核负责 MatMul、Vector 核负责偏置与 GELU 激活的MIX_AIC_1_2异构流水中间结果通过 workspace 在核间传递approximate与 dtype 在编译期/运行期分别通过 binary json 与 TilingKey 精确选择内核分支。在接入该接口时只需确保四张量类型一致、ND 格式、K 维对齐、biasshape 为[N]、approximate1即可按两阶段流程稳定调用详细参数校验与异常行为可对照 op_api/aclnn_fused_matmul_gelu.cpp 与 tests/ut/op_api/test_aclnn_fused_matmul_gelu.cpp 进一步确认。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表