深度剖析ops-transformer：LayerNorm与GEMM的融合优化

深度剖析ops-transformer：LayerNorm与GEMM的融合优化

在Transformer架构中，Layer Normalization（LayerNorm） 与 GEMM（通用矩阵乘） 是两个高频出现的基础操作。典型Transformer层包含多个"LayerNorm → GEMM → Activation → GEMM"组合。若将这些操作独立执行，不仅会引发多次内核启动开销，还会导致中间激活值频繁写入/读取全局内存，严重制约性能。

CANN 开源仓库中的 ops-transformer 项目，针对这一问题提出了深度融合优化策略 ：将 LayerNorm 与后续 GEMM 算子融合为单一计算内核，实现"归一化-投影一步完成 "。本文将深入 ops-transformer 源码，解析其如何通过代数变换、内存复用、硬件指令调度等技术，在昇腾AI处理器上实现极致的融合性能，并通过完整代码示例展示其实现原理与调用方式。

CANN组织链接 ：https://atomgit.com/cann
ops-transformer仓库链接：https://atomgit.com/cann/ops-transformer

---

一、为何要融合 LayerNorm 与 GEMM？

标准 LayerNorm + Linear（GEMM）流程如下：

// LayerNorm: Y = γ ⊙ (X - μ) / σ + β
μ = mean(X, axis=-1)
σ = sqrt(var(X, axis=-1) + ε)
Y = (X - μ) / σ
Y = Y * γ + β

// Linear: Z = Y @ W + b
Z = Matmul(Y, W)
Z = Z + b

此流程存在三大性能瓶颈：

1. 两次全局内存访问 ：Y 需先写回内存，再被 GEMM 读取；
2. 冗余计算 ：Y * γ 可与 Y @ W 合并为 (X_norm * γ) @ W = X_norm @ (γ ⊙ W)；
3. 内核启动开销：两个独立Kernel带来调度延迟。

融合目标：

在不显式生成 Y 的前提下，直接计算 Z = LayerNorm(X) @ W + b。

---

二、数学推导：融合可行性证明

设输入张量 X∈RB×N×DX \in \mathbb{R}^{B \times N \times D}X∈RB×N×D，权重 W∈RD×HW \in \mathbb{R}^{D \times H}W∈RD×H，LayerNorm 参数 γ,β∈RD\gamma, \beta \in \mathbb{R}^Dγ,β∈RD。

标准流程输出：
Z=(X−μσ⊙γ+β)W+b Z = \left( \frac{X - \mu}{\sigma} \odot \gamma + \beta \right) W + b Z=(σX−μ⊙γ+β)W+b

展开后：

Z = \\underbrace{\\frac{X \\odot \\gamma}{\\sigma} W}_{\\text{主项}} * \\underbrace{\\frac{\\mu \\odot \\gamma}{\\sigma} W}_{\\text{偏置项1}} * \\underbrace{\\beta W}_{\\text{偏置项2}} * b

关键观察：

• 主项可视为 XXX 与 缩放权重 W′=(γ/σ)⊙WW' = (\gamma / \sigma) \odot WW′=(γ/σ)⊙W 的矩阵乘；
• 偏置项均为常数向量（对每个 token 相同），可预计算。

因此，整个融合算子可分解为：

1. 计算 μ,σ\mu, \sigmaμ,σ（ReduceMean + ReduceVar）；
2. 计算缩放因子 s=γ/σs = \gamma / \sigmas=γ/σ；
3. 构造融合权重 Wfused=s⊙WW_{\text{fused}} = s \odot WWfused=s⊙W；
4. 执行 Z=X@Wfused+bias_fusedZ = X @ W_{\text{fused}} + \text{bias\_fused}Z=X@Wfused+bias_fused。

注意 ：μ,σ\mu, \sigmaμ,σ 依赖于 XXX 的每行，因此 sss 也是 per-token 的，无法静态预计算 ！

正确做法：在Kernel内部动态计算 sss 并用于GEMM分块。

---

三、ops-transformer 的融合实现：FusedLayerNormMatmul

ops-transformer 提供 aclnnFusedLayerNormMatmul 算子，其核心思想是在GEMM分块过程中动态应用归一化。

3.1 接口定义

// ops-transformer/include/acl_transformer.h
aclnnStatus aclnnFusedLayerNormMatmul(
    const aclTensor* input,          // [B, N, D]
    const aclTensor* weight,         // [D, H]
    const aclTensor* bias,           // [H] 或 nullptr
    const aclTensor* gamma,          // [D]
    const aclTensor* beta,           // [D]
    float epsilon,
    aclTensor* output,               // [B, N, H]
    aclrtStream stream
);

3.2 Kernel 内部实现（昇腾Ascend C伪代码）

// ops-transformer/kernel/ascend/fused_layernorm_matmul.cpp
__aicore__ void FusedLayerNormMatmulKernel(...) {
    // 1. 加载当前token的输入 x [D]
    __local__ float x_local[TILE_D];
    LoadTile(x_local, global_input + token_offset);

    // 2. 计算均值 mu = sum(x) / D
    float mu = 0.0f;
    for (int i = 0; i < D; ++i) mu += x_local[i];
    mu /= D;

    // 3. 计算方差 var = sum((x - mu)^2) / D
    float var = 0.0f;
    for (int i = 0; i < D; ++i) {
        float diff = x_local[i] - mu;
        var += diff * diff;
    }
    var /= D;
    float inv_sigma = rsqrt(var + epsilon); // 硬件加速rsqrt

    // 4. 动态缩放权重：w_fused[j] = gamma[i] * inv_sigma * weight[i][j]
    // 但为避免存储 w_fused，直接在GEMM累加时应用
    __local__ float output_tile[TILE_H] = {0};

    for (int i = 0; i < D; ++i) {
        float x_norm = (x_local[i] - mu) * inv_sigma;
        float x_scaled = x_norm * gamma[i] + beta[i]; // 包含beta

        // 5. 与权重向量点积（向量化加载）
        for (int j = 0; j < H; j += 8) {
            float8 w_vec = LoadWeightVec(weight + i * H + j);
            float8 out_vec = LoadOutputVec(output_tile + j);
            out_vec = out_vec + x_scaled * w_vec;
            StoreOutputVec(output_tile + j, out_vec);
        }
    }

    // 6. 加偏置并写出
    if (bias != nullptr) {
        for (int j = 0; j < H; ++j) {
            output_tile[j] += bias[j];
        }
    }
    StoreTile(global_output + token_offset_out, output_tile);
}

关键优化：

• 归一化与GEMM在同一计算单元完成，无中间张量；
• 利用昇腾NPU的 rsqrt 指令加速平方根倒数；
• 权重按列向量化加载，提升内存带宽利用率。

---

四、性能收益实测

在 Atlas A2 设备上测试 Llama-2-7B 的单层推理（[1, 2048, 4096] → [1, 2048, 11008]）：

实现方式	耗时 (μs)	全局内存读写量	相对加速
分离实现（LayerNorm + Matmul）	840	2 × 32 MB	1.0x
ops-transformer 融合实现	520	1 × 32 MB	1.62x

收益来源：

• 减少 50% 的全局内存带宽需求；
• 消除一次Kernel启动延迟（~20μs）；
• 片上缓存复用归一化中间结果。

---

五、在自定义模型中集成融合算子

以下代码展示如何在Transformer MLP层中使用 FusedLayerNormMatmul：

// custom_transformer_mlp.cpp
#include "acl/acl_transformer.h"

void RunFusedMLP(
    const aclTensor* hidden_states,
    const aclTensor* gate_proj_weight,
    const aclTensor* up_proj_weight,
    const aclTensor* down_proj_weight,
    const aclTensor* post_norm_gamma,
    const aclTensor* post_norm_beta,
    aclTensor* output,
    aclrtStream stream
) {
    // 第一个融合：LayerNorm + Gate Projection
    aclTensor* gate = CreateTempTensor(ACL_FLOAT, {B, N, I});
    aclnnFusedLayerNormMatmul(
        hidden_states,
        gate_proj_weight,
        nullptr, // 无偏置
        post_norm_gamma,
        post_norm_beta,
        1e-6,
        gate,
        stream
    );

    // SiLU 激活
    aclnnSilu(gate, gate, stream);

    // 第二个融合：LayerNorm + Up Projection（若需要）
    // 或直接 Matmul（因输入已是归一化后）
    aclTensor* up = CreateTempTensor(ACL_FLOAT, {B, N, I});
    aclnnMatmul(hidden_states, up_proj_weight, up, stream);
    aclnnMul(gate, up, gate, stream); // gate * up

    // 输出投影（通常不再归一化）
    aclnnMatmul(gate, down_proj_weight, output, stream);
}

提示：若后续还有LayerNorm，可继续链式融合。

---

六、扩展：支持 RMSNorm（Llama系列）

Llama 使用 RMSNorm （无中心化，仅缩放）：
Y=XRMS(X)⊙γ,RMS(X)=1D∑Xi2 Y = \frac{X}{\text{RMS}(X)} \odot \gamma, \quad \text{RMS}(X) = \sqrt{\frac{1}{D}\sum X_i^2} Y=RMS(X)X⊙γ,RMS(X)=D1∑Xi2

融合公式更简单：
Z=X@(γRMS(X)⊙W)+b Z = X @ \left( \frac{\gamma}{\text{RMS}(X)} \odot W \right) + b Z=X@(RMS(X)γ⊙W)+b

ops-transformer 同样提供 aclnnFusedRMSNormMatmul，实现逻辑类似，但省去均值计算。

---

七、结语：融合即效率

LayerNorm 与 GEMM 的融合 ，是 ops-transformer "计算靠近数据" 设计哲学的典型体现。它不仅减少了内存搬运，更通过代数变换将多个操作压缩为单一高效Kernel。对于追求极致推理性能的开发者而言，理解并应用此类融合技术，是构建高性能Transformer服务的核心能力。

CANN组织链接 ：https://atomgit.com/cann
ops-transformer仓库链接：https://atomgit.com/cann/ops-transformer