LayerNorm 是现代神经网络的标配——Transformer 的每一层都有它。公式简单：μ = mean(x), σ² = var(x), y = (x-μ) / √(σ²+ε) * γ + β。但 NPU 上的实现有三个陷阱：FP16 精度下 mean/variance 计算不稳定、Warp reduce 的并行归约需要跨 lane 同步、反向传播的梯度计算涉及 5 个中间变量。

ops-math 的 LayerNorm 算子用 Welford 在线算法 + Warp reduce + 前向-反向融合三层优化解决这些问题。

标准 LayerNorm 的数值不稳定

FP16 精度下，先算 mean 再算 variance 会累积误差：

标准方法（数值不稳定）
mean = Σx_i / N
variance = Σ(x_i - mean)² / N

问题 1：x_i - mean 可能下溢（FP16 最小值 ~6e-5）
问题 2：Σ(x_i - mean)² 可能溢出（FP16 最大值 65504）
问题 3：variance 接近 0 时，1/√variance 精度损失严重

Welford 在线算法（数值稳定）：

Welford 在线算法
初始化：mean = 0, M2 = 0

for i = 1 to N:
    delta = x_i - mean
    mean += delta / i
    delta2 = x_i - mean  # 注意：这里的 mean 已经更新
    M2 += delta * delta2

variance = M2 / N

优势：
1. 不需要先计算 mean 再算 variance（单次扫描）
2. delta 和 delta2 不会同时很大（数值稳定性）
3. 适合 online 计算（不需要存储全部 x_i）

ops-math LayerNorm 的实现

// ops-math/kernels/layer_norm.cpp

__aicore__ void LayerNormKernel(
    GlobalTensor<float16>& input,   // [batch, seq, hidden]
    GlobalTensor<float16>& gamma,   // [hidden]
    GlobalTensor<float16>& beta,    // [hidden]
    GlobalTensor<float16>& output,  // [batch, seq, hidden]
    float epsilon,
    int batch, int seq_len, int hidden
) {
    // 每个 block 处理一个 (batch, seq) 位置的 hidden 维向量
    for (int b = blockIdx.x; b < batch * seq_len; b += gridDim.x) {
        int batch_id = b / seq_len;
        int seq_id = b % seq_len;

        // ===== 阶段 1：Welford 在线算法算 mean + variance =====
        // 用 Welford 算法（单次扫描，数值稳定）
        float mean = 0.0f;
        float M2 = 0.0f;

        // 分块加载 hidden 维（每次 256 个元素，L1 缓存友好）
        for (int h_start = 0; h_start < hidden; h_start += 256) {
            int h_end = min(h_start + 256, hidden);
            int num_elements = h_end - h_start;

            // 加载一块到 L1
            LocalTensor<float16> input_block(256);
            DataCopy(input_block, input[b * hidden + h_start], num_elements);

            // Welford 更新（256 个 lane 并行）
            for (int i = 0; i < num_elements; i++) {
                int lane_id = (h_start + i) % 256;
                if (lane_id == __lane_id__) {
                    float x = float(input_block[i]);

                    float delta = x - mean;
                    mean += delta / float(h_start + i + 1);
                    float delta2 = x - mean;
                    M2 += delta * delta2;
                }
            }

            // Warp reduce（跨 lane 归约 mean 和 M2）
            mean = WarpReduceMean(mean, num_elements);
            M2 = WarpReduceSum(M2, num_elements);
        }

        float variance = M2 / float(hidden);
        float inv_std = rsqrtf(variance + epsilon);

        // ===== 阶段 2：归一化 + 仿射变换 =====
        for (int h_start = 0; h_start < hidden; h_start += 256) {
            int h_end = min(h_start + 256, hidden);
            int num_elements = h_end - h_start;

            LocalTensor<float16> input_block(256);
            LocalTensor<float16> gamma_block(256);
            LocalTensor<float16> beta_block(256);

            DataCopy(input_block, input[b * hidden + h_start], num_elements);
            DataCopy(gamma_block, gamma[h_start], num_elements);
            DataCopy(beta_block, beta[h_start], num_elements);

            // 归一化：y = (x - mean) * inv_std * gamma + beta
            LocalTensor<float16> output_block(256);
            for (int i = 0; i < num_elements; i++) {
                int lane_id = (h_start + i) % 256;
                if (lane_id == __lane_id__) {
                    float x = float(input_block[i]);
                    float g = float(gamma_block[i]);
                    float be = float(beta_block[i]);

                    float normalized = (x - mean) * inv_std;
                    output_block[i] = float16(normalized * g + be);
                }
            }

            DataCopy(output[b * hidden + h_start], output_block, num_elements);
        }
    }
}

Welford 算法的关键：单次扫描同时算 mean 和 variance，不需要存储全部输入——节省 HBM 带宽。

Warp Reduce 的优化

LayerNorm 需要对 hidden 维做归约（求和、求均值）。256 个 lane 各算一部分，需要合并结果。

// ops-math/kernels/warp_reduce.cpp

// Warp 内归约（butterfly 模式，5 次 shuffle）
__aicore__ float WarpReduceSum(float val) {
    // NPU 的 __lane_shuffle_xor 是硬件原语
    // 延迟 < 4 cycles（直接走 Cross-Lane 交换网络）

    float peer;
    peer = __lane_shuffle_xor(val, 16);  val += peer;  // 步长 16
    peer = __lane_shuffle_xor(val, 8);   val += peer;  // 步长 8
    peer = __lane_shuffle_xor(val, 4);   val += peer;  // 步长 4
    peer = __lane_shuffle_xor(val, 2);   val += peer;  // 步长 2
    peer = __lane_shuffle_xor(val, 1);   val += peer;  // 步长 1

    return val;  // 所有 lane 的 val 都相同（归约结果）
}

__aicore__ float WarpReduceMean(float val, int count) {
    float sum = WarpReduceSum(val);
    return sum / float(count);  // 归约后除以 count
}

Butterfly 归约的并行度：

• 第 1 步：128 对 lane 并行交换（步长 16）
• 第 2 步：64 对 lane 并行交换（步长 8）
• …
• 第 5 步：1 对 lane 交换（步长 1）

总延迟：5 × 4 cycles = 20 cycles。对比从 HBM 逐个累加快 ~30×。

前向-反向融合 Kernel

训练时，LayerNorm 的前向和反向需要分别启动——但反向需要前向往的 mean/variance/inv_std 中间结果。标准实现把中间结果写回 HBM，反向时再读——读写开销大。

ops-math 的融合 kernel：

// ops-math/kernels/layer_norm_bprop_fused.cpp

__aicore__ void LayerNormBpropFused(
    GlobalTensor<float16>& grad_output,  // [batch, seq, hidden]
    GlobalTensor<float16>& input,        // [batch, seq, hidden]（前向输入）
    GlobalTensor<float16>& gamma,        // [hidden]
    GlobalTensor<float16>& grad_input,    // [batch, seq, hidden] 输出
    GlobalTensor<float16>& grad_gamma,   // [hidden] 输出
    GlobalTensor<float16>& grad_beta,    // [hidden] 输出
    float epsilon,
    int batch, int seq_len, int hidden
) {
    for (int b = blockIdx.x; b < batch * seq_len; b += gridDim.x) {
        // ===== 重新计算前向的 mean/inv_std（不读 HBM）=====
        float mean = 0.0f;
        float M2 = 0.0f;

        // 复用前向的 Welford 算法（不写 HBM）
        for (int h = 0; h < hidden; h += 256) {
            // ... Welford 更新（同前向）...
        }

        float inv_std = rsqrtf(M2 / float(hidden) + epsilon);

        // ===== 反向计算 =====
        // 公式（推导略）：
        // dL/dx = (1/N) * (N*dL/dy - ΣdL/dy - (x-mean)*ΣdL/dy*(x-mean)/Σ(x-mean)²) * gamma * inv_std
        // dL/dgamma = Σ(dL/dy * (x-mean) * inv_std)
        // dL/dbeta = Σ(dL/dy)

        float sum_dy = 0.0f;
        float sum_dy_x = 0.0f;
        float sum_dy_x2 = 0.0f;

        for (int h = 0; h < hidden; h += 256) {
            // 加载数据
            LocalTensor<float16> dy_block(256);
            LocalTensor<float16> x_block(256);
            LocalTensor<float16> gamma_block(256);

            DataCopy(dy_block, grad_output[b * hidden + h], 256);
            DataCopy(x_block, input[b * hidden + h], 256);
            DataCopy(gamma_block, gamma[h], 256);

            // 计算中间变量（融在一个 kernel 内）
            for (int i = 0; i < 256; i++) {
                float dy = float(dy_block[i]);
                float x = float(x_block[i]);
                float g = float(gamma_block[i]);

                sum_dy += dy;
                sum_dy_x += dy * (x - mean) * inv_std;
                sum_dy_x2 += dy * (x - mean) * inv_std * (x - mean) * inv_std;

                // 累计梯度（写回 grad_gamma/grad_beta）
                AtomicAdd(grad_gamma[h + i], dy * (x - mean) * inv_std);
                AtomicAdd(grad_beta[h + i], dy);
            }
        }

        // Warp reduce 汇总
        sum_dy = WarpReduceSum(sum_dy);
        sum_dy_x = WarpReduceSum(sum_dy_x);
        sum_dy_x2 = WarpReduceSum(sum_dy_x2);

        // 计算 dL/dx
        for (int h = 0; h < hidden; h += 256) {
            // ... 用 sum_dy/sum_dy_x/sum_dy_x2 算 grad_input ...
        }
    }
}

融合的收益：

• 标准流程：前向写 mean/inv_std 到 HBM（hidden × 4 bytes）→ 反向读（hidden × 4 bytes）→ 总 HBM 读写 = 2 × hidden × 4 bytes
• 融合流程：前向不写 HBM，反向重新计算 mean/inv_std（复用 L1 中的 input）→ HBM 读写 = 0

踩坑一：Welford 算法的 FP16 中间溢出

Welford 的 delta * delta2 在 FP16 下可能溢出（x_i = 65500, mean = 0 → delta = 65500, delta2 = 65500 - 65500/2 = 32750 → delta * delta2 = 2.1e9 → 溢出）。

修复：强制转 FP32 做中间计算

// 错误：FP16 中间结果
float16 delta = x - mean;
float16 delta2 = x - mean2;  // mean2 是更新后的
M2 += float16(delta * delta2);  // 可能溢出

// 正确：FP32 中间结果
float delta = float(x) - mean;
float mean2 = mean + delta / float(count);
float delta2 = float(x) - mean2;
M2 += delta * delta2;  // FP32，不会溢出

踩坑二：Warp Reduce 的 inactive lane 处理

hidden 不被 256 整除时，最后一个 Warp 的部分 lane 是 inactive 的（没有有效数据）。Warp reduce 把它们也归约进去了——结果错误。

修复：inactive lane 贡献 0

// 错误：
float sum = 0.0f;
for (int i = 0; i < hidden; i += 256) {
    float x = input[i + __lane_id__];  // inactive lane 读了垃圾值
    sum += WarpReduceSum(x);
}

// 正确：
float sum = 0.0f;
for (int i = 0; i < hidden; i += 256) {
    int valid_lanes = min(256, hidden - i);
    float x = (__lane_id__ < valid_lanes) ? float(input[i + __lane_id__]) : 0.0f;
    sum += WarpReduceSum(x);  // inactive lane 贡献 0
}

踩坑三：LayerNorm 的 epsilon 选择

inv_std = 1 / sqrt(variance + epsilon)。epsilon 太小 → variance 接近 0 时 inv_std 溢出。epsilon 太大 → 归一化效果变差（梯度消失）。

经验值：

• FP32：epsilon = 1e-5（PyTorch 默认）
• FP16：epsilon = 1e-3（防止 sqrt(variance + 1e-5) 下溢）

// ops-math 自动根据数据类型选择 epsilon
if (dtype == FP16) {
    epsilon = 1e-3f;  // FP16 用更大的 epsilon
} else {
    epsilon = 1e-5f;  // FP32 用标准 epsilon
}

---

LayerNorm 看起来简单——减均值、除标准差、乘 gamma、加 beta。但 NPU 上的优化是三层叠加：Welford 在线算法保证数值稳定、Warp reduce 并行归约隐藏延迟、前向-反向融合消除 HBM 读写。Transformer 的每一层都依赖 LayerNorm——这个算子的性能直接决定模型的训练/推理速度。