ops-nn 的激活函数：ReLU、GELU 与 SiLU 在昇腾上的实现

神经网络去掉激活函数就是一个线性变换的堆叠。多层线性变换等价于一层线性变换。激活函数在每层之间引入非线性是网络能逼近任意函数的关键。CANN 的 ops-nn 仓库管理着所有神经网络相关的算子——卷积、归一化、池化、激活函数。其中激活函数看似简单逐元素操作但在 Transformer 时代 GELU 和 SiLU 替代了 ReLU计算量和数据搬运模式都变了。为什么神经网络需要激活函数一个没有激活函数的双层网络y W2 (W1 x)。矩阵乘法是线性的两层展开就是y (W2 W1) x——等价于单层网络。激活函数在每层之间打断线性关系。ReLU 把负数全部清零GELU 根据输入值做概率加权保留。有了激活函数网络才能真正学习到数据中的复杂模式。在 Transformer 中激活函数出现在 FFN 层FFN 子层: X → GEMM_W1 → GELU → GEMM_W2 → Output ↑ ↑ ↑ 线性投影 非线性激活 线性投影GELU 夹在两个 GEMM 之间。它的计算量和数据量都不大但它的位置决定了 GEMM1 的中间结果必须完整写出到 DDR 才能喂给 GELUGELU 的输出又必须写回 DDR 才能喂给 GEMM2。这个写出再读入的模式是激活函数在推理中对性能的真正影响——不是激活计算慢而是中间 Tensor 的搬运。GELU 为什么适合 TransformerReLU 的问题在于它对负数完全抑制。Transformer 训练时某些神经元会学出始终为负的输出ReLU 让这些神经元永远输出 0——神经元死亡。GELU 的曲线在负数区域不平滑也不完全归零保留了负值的部分信息GELU(x) x * Φ(x) 其中 Φ 是标准正态分布的 CDF SiLU(x) x * σ(x) 其中 σ 是 Sigmoid 函数GELU 和 SiLU 在深层 Transformer 中的收敛速度和最终精度都优于 ReLU现在已经成为 FFN 激活函数的标准选择。LLaMA 用 SiLUSwiGLU 的门控函数BERT 用 GELU。昇腾NPU如何执行激活算子在昇腾上激活算子在 Vector Unit 上执行。Vector Unit 是达芬奇架构中专做逐元素运算的计算单元——跟 Cube Unit矩阵乘并行工作。Vector Unit 执行激活函数的方式是一个 SIMD 流水线。输入 Tensor 从 DDRGM搬到片上 L1 BufferVector Unit 在 L1 上逐元素计算激活函数结果写回 DDR。// 昇腾 Vector Unit 上的 ReLU Kernel简化__vector__voidrelu_kernel(LocalTensorfloatoutput,constLocalTensorfloatinput,intlen){// Vec 指令一次处理 128 个 float 元素for(inti0;ilen;i128){// 从 L1 读取 128 个元素float16 vec[128]input[i:i128];// 每个元素做 max(x, 0)float16 result[128]max(vec,0);// 写回 L1output[i:i128]result;}}ReLU 就是一次 max 指令。GELU 需要更多的 Vector 指令——指数运算、误差函数近似、乘法。ops-nn 把 GELU 实现为一个融合的 Vector Kernel在 L1 上完成全部计算。ops-nn 如何优化激活融合ops-nn 的激活优化不是提升激活函数本身的计算速度——Vector Unit 算 GELU 已经够快了。优化点在于不要让激活函数独立成一个 Kernel Launch。ops-nn 提供的融合接口让 GE 在编译时将激活函数合并到前面或后面的算子中// ops-nn 的激活融合接口// 把 GELU 融合到 GEMM1 的 Epilogue 中GemmParam gemm_param;gemm_param.activationACTIVATION_GELU;// GEMM 算完后立即激活// Runtime 执行时GEMM → GELU → 写回 DDR一笔搬运aclrtLaunchOp(gemm_kernel,input,output,stream,gemm_param);融合后GEMM1 的 Cube Unit 每算出一个小 Tile 的结果Vector Unit 立即对这个 Tile 做 GELU结果直接作为 GEMM2 的输入——中间 Tensor 不写 DDR。实测数据LLaMA-7B 的 FFN 层配置搬运量延迟GEMM1 GELU GEMM2 独立执行3 次 GM 读写0.45msGEMM1GELU 融合 GEMM2 独立2 次 GM 读写0.38msGEMM1GELUGEMM2 三算子融合1 次 GM 读写0.31ms三算子完全融合后搬运量从 3 次降到 1 次延迟降低 31%。激活函数对推理速度的影响激活函数本身的浮点计算量很少。以 GELU 为例n 个元素的激活计算约3n次 FLOPs指数近似的乘法。对比 GEMM 的2×M×N×KFLOPs激活函数的计算量可以忽略不计。但激活函数因为打断算子流水线而成为性能瓶颈独立 Kernel 的 Launch 开销。每次 Kernel 调用有 5-15μs 的 Runtime 调度开销。融合后这些开销被消除。中间 Tensor 搬运。不融合时激活函数的输入输出各写一次 DDR。融合后中间 Tensor 在 L1 上流转。执行单元切换。独立的激活 Kernel 意味着 Cube Unit 在等激活算完才能进行下一个 GEMM。融合后 Cube 和 Vector 可以流水线工作。ReLU 的 Vector Unit 实现ReLU 在 Vector Unit 上的实现就是一条max指令——对输入向量的每个元素做max(x, 0)。Ascend 910 的 Vector Unit 单次处理 128 个 float16 元素所以计算 4096 个元素的 ReLU 只需要 32 条 Vector 指令约 0.1μs。实际性能瓶颈不在计算——计算的 0.1μs 几乎可以忽略——而在 DataCopy把 Tensor 从 DDR 搬到 L1 需要约 3-5μs。这就是激活算子在推理管线中的真实代价不是激活计算慢是数据搬运的时间远大于计算时间。SiLU 与 SwiGLULLaMA 系列使用的 SwiGLU 激活函数是 SiLU 的变体SwiGLU(x) SiLU(x_gate) * x_upscale。它把 FFN 的一个 GEMM 拆成两个——gate 和 upscale——然后用 SiLU 对 gate 做激活再跟 upscale 逐元素相乘。SwiGLU 的计算量比 ReLU 大一次 SiLU 激活 一次逐元素乘但因为引入了门控机制模型在相同参数量下的精度更好。ops-nn 的 SiLU 实现跟 GELU 类似——在 Vector Unit 上用 Sigmoid 近似 乘法完成一次 Kernel Launch 搞定。参考仓库ops-nn 神经网络算子库LayerNorm 归一化算子