企业官网建设流程全解析

CANN ops-transformer 算子解析：BiasGateGelu —— Transformer FFN 层 GeGLU 融合算子的设计与实现

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

导读

BiasGateGelu 是 CANN ops-transformer 仓库experimental/moe/biasgategelu目录下提供的一个 NPU 自定义融合算子，用于在 Transformer 模型的 FFN（Feed-Forward Network，前馈网络）层中一次性完成Gate/Value 切分、双路加偏置、GELU 激活与逐元素乘的 GeGLU（GELU Gated Linear Unit）融合计算。本文以该算子官方 README 为主体，结合仓库内 Kernel 源码、构建脚本与测试用例，完整讲解其功能语义、参数约束、Python 调用方式、AscendC 内核实现原理、编译集成方式与 CPU/NPU 结果一致性验证方法，帮助读者掌握在 CANN 生态中开发、集成与验证此类融合算子的完整链路。

一、算子定位与背景：GeGLU 为什么需要融合

在 Transformer 大模型的 FFN 层中，激活函数已经从早期的 ReLU 演化为多种门控线性单元（GLU）变体，其中GeGLU（GELU Gated Linear Unit）是当前主流大模型（包括各类 MoE 模型）中最常用的形式之一。

GeGLU 的典型计算方式为：将隐藏层输出按列均分为Gate和Value两部分，对 Gate 施加 GELU 激活后再与 Value 逐元素相乘。若再叠加各自的偏置项，则完整数学表达式为：

$$\text{out} = \text{GELU}(gate + gateBias) \odot (value + valueBias)$$

如果按 PyTorch 原生算子逐个执行，至少需要 Split、Add、GELU、Mul 多次 Kernel 启动与中间张量的显存读写；而在 NPU 上将其融合为单个算子，可以显著减少 AI Core 启动开销与 Global Memory 访问次数。BiasGateGelu 正是这一融合思路的实现——它位于仓库的experimental/moe目录下，与该目录中的 BiasSigmoid、gategelu_quant、moe_ffn、fusedexpert 等算子共同构成 MoE（Mixture of Experts）场景下的实验性算子集合。

二、产品支持情况

从构建配置看，该算子按 Atlas A2 训练系列中的Ascend910B系列进行编译适配（详见下文构建集成章节），其 CMake 编译参数中指定了--cce-soc-version=Ascend910B1与--cce-soc-core-type=VecCore。

三、功能说明

BiasGateGelu 在 Transformer 模型的 FFN 层中实现GeGLU（GELU Gated Linear Unit）的融合计算，其核心流程为：

1. 均分：将输入张量in（shape 为(gbH, gbW)）按列均分为 Gate 与 Value 两部分，即前gbW/2列为 Gate，后gbW/2列为 Value；
2. 加偏置：偏置向量bias的 shape 为(gbW,)，其前半部分为 Gate 偏置，后半部分为 Value 偏置，分别与 Gate、Value 逐元素相加；
3. 激活：对 Gate 部分施加 GELU 激活函数；
4. 相乘：将激活后的 Gate 与加偏置后的 Value 逐元素相乘，得到 shape 为(gbH, gbW/2)的输出out。

对应计算流程可表示为：

$$\text{out} = \text{GELU}(in[:, :gbW/2] + bias[:gbW/2]) \ \odot \ (in[:, gbW/2:] + bias[gbW/2:])$$

从仓库源码 experimental/moe/biasgategelu/tests/biasgategelu.py 中的 CPU 参考实现可以更直观地确认这一语义：

def bias_gate_gelu_cpu(in_tensor, bias): """CPU侧实现 BiasGateGelu 算子""" gbH, gbW = in_tensor.shape half_dim = gbW // 2 gate = in_tensor[:, :half_dim] value = in_tensor[:, half_dim:] gate_bias = bias[:half_dim] value_bias = bias[half_dim:] gate = torch.nn.functional.gelu(gate + gate_bias, approximate='tanh') value = value + value_bias out = gate * value return out

值得注意的是，测试中的 CPU 参考实现使用了approximate='tanh'（tanh 近似 GELU），而 Kernel 源码中调用的是Gelu<half, true, false>(...)（第一个模板参数true同样指示 tanh 近似模式），两者语义保持一致——即本算子采用的 GELU 是tanh 近似版本，而非精确 erf 版本。

四、参数说明

参数语义的补充说明（依据 BIAS_GATE_GELU.cpp 中的实际实现）：

• blockDim：参与计算的 AI Core 数量。Kernel 启动时会将行维按 AI Core 数进行任务切分（blockNum_ = GetBlockNum()），如果gbH < blockDim，启动函数会自动把blockDim收敛为gbH（见 BIAS_GATE_GELU.cpp），避免无效核的空转；
• gbW 必须是偶数：只有列数为偶数才能被均分为等宽的 Gate 与 Value 两半；同时gbW代表2 × hidden_size，因此输出的列数恒为gbW/2 = hidden_size；
• bias 长度与列对齐：bias 长度为gbW，与in的列数一一对应，前gbW/2个元素作用于 Gate 列，后gbW/2个元素作用于 Value 列。

五、约束说明

• 输入张量的列数gbW必须为偶数，以便均分为 Gate 和 Value 两部分；
• 输入、偏置、输出的数据类型均为 float16（half）；
• 输入in与bias的 shape 与维度语义必须与gbH、gbW保持一致；
• 相关张量必须位于 NPU 设备（PrivateUse1 设备类型）上。

除 README 中声明的约束外，Host 侧入口函数bias_gate_gelu_npu还通过TORCH_CHECK在运行时执行了一系列校验（见 BIAS_GATE_GELU.cpp）：

TORCH_CHECK(in.device().type() == torch::kPrivateUse1, "input must be on NPU"); TORCH_CHECK(bias.device().type() == torch::kPrivateUse1, "bias must be on NPU"); TORCH_CHECK(out.device().type() == torch::kPrivateUse1, "out must be on NPU"); TORCH_CHECK(gbH > 0, "gbH must be positive"); TORCH_CHECK(gbW > 0, "gbW must be positive"); TORCH_CHECK(gbW % 2 == 0, "gbW must be even for gate/value split"); TORCH_CHECK(block_dim > 0, "block_dim must be positive, got ", block_dim); TORCH_CHECK(in.sizes() == torch::IntArrayRef({gbH, gbW}), "in tensor shape mismatch");

这意味着：三个张量必须显式搬移到 NPU；gbH、gbW、block_dim必须为正数；gbW必须为偶数；且in的实际 shape 必须严格等于(gbH, gbW)。

六、调用说明

该算子通过自定义算子扩展库ascend_ops以 PyTorch 自定义算子（Custom Op）的形式暴露，调用方式如下：

torch.ops.ascend_ops.bias_gate_gelu( block_dim, in_tensor, bias, out, gbH, gbW)

其中out需要预先分配好 shape 为(gbH, gbW // 2)的 NPU 张量。一个完整的最小调用示例（结合仓库测试脚本）为：

import torch import torch_npu import ascend_ops gbH, gbW = 8, 256 block_dim = 40 dtype = torch.float16 in_tensor = torch.randn(gbH, gbW, dtype=dtype) bias = torch.randn(gbW, dtype=dtype) in_tensor_npu = in_tensor.npu() bias_npu = bias.npu() out_npu = torch.zeros(gbH, gbW // 2, dtype=dtype).npu() torch.ops.ascend_ops.bias_gate_gelu(block_dim, in_tensor_npu, bias_npu, out_npu, gbH, gbW) result = out_npu.cpu()

调用前需要确保环境中已安装并加载torch_npu（PyTorch NPU 适配层）以及包含本算子的ascend_ops扩展包（由本仓库构建产出）。

从算子注册机制看，该算子通过TORCH_LIBRARY_IMPL宏注册到ascend_ops库的PrivateUse1（即 NPU）后端实现（见 BIAS_GATE_GELU.cpp）：

TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m) { m.impl("bias_gate_gelu", BiasGateGelu::bias_gate_gelu_npu); }

七、源码实现解析：从 Host 启动到 AI Core 内核

该算子的完整实现位于单文件 experimental/moe/biasgategelu/BIAS_GATE_GELU.cpp 中，整体分为三层：Host 侧入口（参数校验与 Kernel 启动）、AscendC 内核类（BIAS_GATE_GELU）与设备侧核函数（biasGateGelu_kernel）。

7.1 内核入口与启动逻辑

设备侧核函数为：

extern "C" __global__ __aicore__ void biasGateGelu_kernel(const int64_t gbH, const int64_t gbW, GM_ADDR in, GM_ADDR bias, GM_ADDR out) { BIAS_GATE_GELU op; op.Init(gbH, gbW, in, bias, out); op.Process(); }

Host 侧启动函数biasGateGelu_lanuch先对blockDim做保护性收敛（if (gbH < blockDim) { blockDim = gbH; }），再通过<<<blockDim, nullptr, stream>>>在当前 NPU 流上启动内核。而bias_gate_gelu_npu则负责获取当前 NPU 流（c10_npu::getCurrentNPUStream()）并将torch::Tensor转换为裸指针后传入启动函数。

7.2 Init：UB 空间预算与分块策略

内核初始化阶段（BIAS_GATE_GELU.cpp）完成三件关键工作：

1. 任务分块（Block 级）：bkH_ = 1、bkW_ = gbW_ / 2，即每个 Block 每次处理一行数据中的 Gate/Value 半宽；行维任务量bkLoop_ = ceil(gbH / blockNum_)，按 AI Core 编号轮转切分（i * blockNum_ + blockIdx_），保证负载均衡。

2. UB 空间预算（Tile 级）：算子使用UB_MAX_BYTES = 184 * 1024（约 184KB 的 Unified Buffer 预算）与BUFFER_NUM = 2的双缓冲模式。初始化时依据双缓冲下三类缓冲区（Gate 输入、Value 输入、输出各一份，外加拼接的 Bias 缓冲区）估算单次 Tile 最大宽度：

constexpr int64_t UB_MAX_BYTES = 184 * 1024; constexpr int64_t BUFFER_NUM = 2; // ... int64_t temp = BUFFER_NUM * bkH_ * sizeof(half) * 3; // inOne / inTwo / out 三份 temp += BUFFER_NUM * 2 * sizeof(half); // bias 缓冲（Gate+Value 两段） tlMaxW_ = UB_MAX_BYTES / temp; tlMaxW_ = tlMaxW_ / 64 * 64; // 向下对齐到 64

3. 宽度方向的 Tile 循环：若bkW_超过tlMaxW_，则按ceil(bkW_ / tlMaxW_)等分宽度；最终tlW_向上对齐到 64 元素，尾块宽度tlTailW_单独记录，由tlLoop_控制宽度方向的外层循环。

7.3 Process 主流程：双缓冲流水

Process()（BIAS_GATE_GELU.cpp）采用典型的三段式流水结构：

• 每轮宽度 Tile 先通过DataCopyPad将当前列的Gate 偏置段（biasGm_[j*tlW_]）与Value 偏置段（biasGm_[bkW_+j*tlW_]）拷入同一个本地张量bias_local的前后两段，入队后由biasLm_持有；
• 内层遍历行块，对每个有效行调用CopyIn → Compute → CopyOut；
• 内层结束后释放偏置张量，进入下一宽度 Tile。

针对尾块（tlTailW_ > 0且为最后一轮），代码会将DataCopyParams的blockLen收缩为实际字节数，同时用对齐后的宽度参与向量计算，避免越界与性能损失。

7.4 CopyIn / Compute / CopyOut：数据搬运与向量计算

CopyIn（BIAS_GATE_GELU.cpp）按输入布局一次性取回 Gate 与 Value 两段：输入第i行、第j个宽度 Tile 的首地址为(bkl * blockNum_ + blockIdx_) * (bkW_ * 2) + (tll * tlW_)，其中inGm_[offset]段是 Gate、inGm_[offset + bkW_]段是 Value。

Compute（BIAS_GATE_GELU.cpp）完整对应 GeGLU 数学表达式，仅用三条向量指令完成：

Add(in_one_local, in_one_local, biasLm_, real_tlAlignW); // gate + gateBias Gelu<half, true, false>(in_one_local, in_one_local, real_tlAlignW); // GELU(gate + gateBias) Add(in_two_local, in_two_local, biasLm_[tlW_], real_tlAlignW); // value + valueBias Mul(out_local, in_one_local, in_two_local, real_tlAlignW); // GELU(...) * (value + valueBias)

注意这里所有计算都是in-place 与就地复用的，Gate 与 Value 的本地缓冲区在计算后即被释放，最大限度降低 UB 占用。

CopyOut（BIAS_GATE_GELU.cpp）将结果写回输出地址(bkl * blockNum_ + blockIdx_) * bkW_ + (tll * tlW_)——由于输出列宽恰为bkW_ = gbW/2，输出地址与输入地址形成紧凑的对应关系。

整体来看，该 Kernel 通过“AI Core 行维并行 + UB 宽度分块 + 双缓冲流水 + 就地计算”的组合，实现了单次 Kernel 启动完成偏置、激活、相乘的全流程融合，这正是融合算子在 NPU 上减少访存与启动开销的核心手段。

八、构建集成方式

该算子以独立的 CMake 子工程形式组织在experimental/moe/biasgategelu/目录下，其 CMakeLists.txt 的关键配置为：

message(STATUS "BUILD_TORCH_OPS ON in biasgategelu") # BIAS_GATE_GELU operation sources file(GLOB BIAS_GATE_GELU_NPU_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp") set(BIAS_GATE_GELU_SOURCES ${BIAS_GATE_GELU_NPU_SOURCES}) # Mark .cpp files with special properties set_source_files_properties( ${BIAS_GATE_GELU_NPU_SOURCES} PROPERTIES LANGUAGE CXX COMPILE_FLAGS "--cce-soc-version=Ascend910B1 --cce-soc-core-type=VecCore --cce-auto-sync -xcce" ) # Create object library add_library(bias_gate_gelu_objects OBJECT ${BIAS_GATE_GELU_SOURCES}) target_compile_options(bias_gate_gelu_objects PRIVATE ${COMMON_COMPILE_OPTIONS}) target_include_directories(bias_gate_gelu_objects PRIVATE ${COMMON_INCLUDE_DIRS})

从中可以看出：

• 源码以.cpp后缀编写，但通过-xcce与--cce-*系列编译选项以 CCE（AscendC）编译器进行编译，面向Ascend910B1芯片与VecCore（向量核）执行；
• --cce-auto-sync启用自动同步，简化了异步数据搬运的同步管理；
• 算子以OBJECT 库形式产出，由上层 experimental/moe/CMakeLists.txt 通过遍历子目录统一add_subdirectory纳入构建，最终随ascend_ops扩展包集成到 PyTorch 生态中。

因此，构建该算子属于本仓库 torch 算子扩展的整体构建流程的一部分：在具备 CANN 工具链与 Ascend910B 系列环境的前提下，按仓库顶层构建说明编译即可将torch.ops.ascend_ops.bias_gate_gelu注册进 Python 侧。

九、测试验证：CPU 参考实现与 NPU 结果比对

仓库为该算子提供了可直接运行的验证脚本 experimental/moe/biasgategelu/tests/biasgategelu.py，其验证策略是标准的“CPU 参考实现 vs NPU 算子输出”一致性测试：

1. 构造输入：gbH=8, gbW=256, block_dim=40，输入与偏置均为torch.randn生成的 float16 随机张量；
2. CPU 参考：调用bias_gate_gelu_cpu，按gate/value 切分 → 加偏置 → tanh 近似 GELU → 相乘得到期望输出；
3. NPU 执行：将张量.npu()搬移到 NPU，预分配out_npu，调用torch.ops.ascend_ops.bias_gate_gelu(...)；
4. 结果比对：将 NPU 输出拷回 CPU，与参考实现做torch.allclose(z_cpu, z_npu_cpu, rtol=1e-2, atol=1e-2)对比，并打印输出形状、前两行示例及最大差异。

脚本输出示例结构：

=== BiasGateGelu 算子测试 === [CPU计算] 输出形状: torch.Size([8, 128]) ... [NPU计算] 输出形状: torch.Size([8, 128]) ... [结果对比] ✓ CPU与NPU结果一致！

由于 float16 精度有限，比对容差设置为rtol=1e-2, atol=1e-2。该测试同时印证了：输出列宽为gbW // 2 = 128、block_dim=40正好匹配 Ascend910B 的 40 个 AI Core，以及算子对外暴露的调用签名为(block_dim, in, bias, out, gbH, gbW)的六参数形式。

十、应用场景与小结

BiasGateGelu 面向的是 Transformer / MoE 大模型 FFN 层中高频出现的 GeGLU 计算模式。在 MoE 场景下，每个 expert 的 FFN 都可能执行一次 Gate/Value 门控计算，将其融合为单算子可以在批量推理与训练中显著减少算子调度与显存搬移开销。它作为experimental/moe实验性算子集合的一员（与 BiasSigmoid、gategelu_quant、moe_ffn 等同目录算子协同），体现了 CANN ops-transformer 在 MoE 推理链路上的融合优化思路。

总结本文核心要点：

• 功能语义：输入按列均分 Gate/Value，双路加偏置后执行GELU(gate+bias) * (value+bias)，输出列宽为输入一半；
• 调用方式：torch.ops.ascend_ops.bias_gate_gelu(block_dim, in, bias, out, gbH, gbW)，要求 float16、ND 布局、NPU 设备；
• 实现要点：AscendC 单文件 Kernel，采用 AI Core 行维并行、UB 宽度分块（184KB 预算、双缓冲）、DataCopyPad搬移、Add/Gelu/Mul三条向量指令就地完成融合计算；
• 验证手段：CPU 参考实现 + NPU 执行 +allclose容差比对，可直接运行experimental/moe/biasgategelu/tests/biasgategelu.py复现；
• 运行前提：Atlas A2 训练系列（Ascend910B）环境，安装torch_npu与由本仓库构建的ascend_ops扩展包。

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer