Ascend C 实战：开发高性能自定义 RMSNorm 算子替代 LayerNorm 加速 LLaMA 模型

关键洞察：

• 使用 rsqrtf() 替代 sqrt() + 除法
• 所有中间结果保留在 Local Memory 或寄存器
• FP32 累加 避免 FP16 下溢

三、第一步：定义算子原型

3.1 JSON 原型文件

文件：rmsnorm_custom.json

{"op":"RMSNormCustom","input_desc":[{"name":"x","type":"float16","format":"ND"},{"name":"gamma","type":"float16","format":"ND"}],"output_desc":[{"name":"y","type":"float16","format":"ND"}],"attr":[{"name":"eps","type":"float","default":1e-6}]}

📝 说明：gamma 形状为 [D]，广播到输入最后一维eps 默认为 1e-6（LLaMA 官方配置）

四、第二步：生成工程模板

msopgen gen \ -i rmsnorm_custom.json \ -c ai_core-Ascend910B \ -lan cpp \ -out ./RMSNormCustom

五、第三步：编写核函数（NPU 侧）

5.1 完整核函数代码

文件：kernel/rmsnorm_custom_kernel.cpp

#include "common.h"
extern "C" __global__ __aicore__ void RMSNormKernel(
    __gm__ half* x,// 输入 [total_size]
    __gm__ half* gamma,// 缩放参数 [D]
    __gm__ half* y,// 输出 [total_size]
    uint32_t total_size,// 总元素数
    uint32_t D,// 归一化维度大小（如 hidden_size）
    uint32_t outer_size,// 外层维度积（如 B * seq_len）
    float eps
){
    uint32_t block_idx = GetBlockIdx();
    uint32_t block_num = GetBlockNum();
    uint32_t samples_per_block = (outer_size + block_num - 1) / block_num;
    uint32_t start_sample = block_idx * samples_per_block;
    uint32_t end_sample = min(start_sample + samples_per_block, outer_size);
    const int TILE_SIZE = 256;
    __local__ half x_tile[TILE_SIZE];
    __local__ half gamma_tile[TILE_SIZE];
    __local__ half y_tile[TILE_SIZE];
    for(uint32_t sample = start_sample; sample < end_sample; sample++){
        // === 第一阶段：计算平方和 sum(x^2) ===
        float sum_sq = 0.0f;
        for(uint32_t i = 0; i < D; i += TILE_SIZE){
            int copy_len = min(TILE_SIZE, static_cast<int>(D - i));
            dma_copy(x_tile, x + sample * D + i, copy_len * sizeof(half));
            for(int j = 0; j < copy_len; j++){
                float val = static_cast<float>(x_tile[j]);
                sum_sq += val * val;// FP32 累加，避免下溢
            }
        }
        // 计算 1 / sqrt(mean_sq + eps)
        float mean_sq = sum_sq / D;
        float inv_rms = rsqrtf(mean_sq + eps);// 关键：硬件加速倒数平方根
        // === 第二阶段：归一化 + 缩放 ===
        for(uint32_t i = 0; i < D; i += TILE_SIZE){
            int copy_len = min(TILE_SIZE, static_cast<int>(D - i));
            dma_copy(x_tile, x + sample * D + i, copy_len * sizeof(half));
            dma_copy(gamma_tile, gamma + i, copy_len * sizeof(half));
            for(int j = 0; j < copy_len; j++){
                float x_f32 = static_cast<float>(x_tile[j]);
                float g_f32 = static_cast<float>(gamma_tile[j]);
                // y = (x * inv_rms) * gamma
                float normalized = x_f32 * inv_rms;
                y_tile[j] = static_cast<half>(normalized * g_f32);
            }
            dma_copy(y + sample * D + i, y_tile, copy_len * sizeof(half));
        }
    }
}

5.2 关键优化点

1. 单次平方和累加：避免存储 (x^2)
2. rsqrtf() 硬件指令：比 sqrt() + 除法快 3 倍
3. FP32 中间累加：保证数值稳定性（尤其对小值）
4. 零中间全局存储：所有临时数据在 Local Memory

六、第四步：向量化生产级优化

上述标量循环仅用于教学。实际部署必须向量化：

6.1 向量化版本（关键片段）

// 在第二阶段循环内
for(int j = 0; j < copy_len; j += 8){
    __vector__ half x_vec, gamma_vec;
    vector_load(x_vec, x_tile + j);
    vector_load(gamma_vec, gamma_tile + j);
    // 转为 float 向量（展开）
    float x_f32[8], g_f32[8];
    for(int k = 0; k < 8; k++){
        x_f32[k] = static_cast<float>(x_vec[k]);
        g_f32[k] = static_cast<float>(gamma_vec[k]);
    }
    // 向量化计算：y = x * inv_rms * gamma
    half y_vec[8];
    for(int k = 0; k < 8; k++){
        y_vec[k] = static_cast<half>(x_f32[k] * inv_rms * g_f32[k]);
    }
    vector_store(y_tile + j, y_vec);
}

✅ 效果：充分利用 Vector Core 的 8-way FP16 并行能力。

七、第五步：Tiling 与 Host 封装

7.1 Tiling 策略

文件：tiling/rmsnorm_custom_tiling.h

void ComputeTiling(const std::vector<TensorDesc>& inputs,
                   const std::map<std::string, std::any>& attrs,
                   std::vector<Tiling>& tilings){
    auto shape = inputs[0].GetShape();
    uint64_t D = shape.GetDim(shape.GetDimNum()-1); // 最后一维
    uint64_t outer_size = shape.Size() / D;
    uint32_t block_num = min(32U, static_cast<uint32_t>(outer_size));
    tilings[0].Set("block_num", block_num);
    tilings[0].Set("D", static_cast<uint32_t>(D));
    tilings[0].Set("outer_size", static_cast<uint32_t>(outer_size));
    tilings[0].Set("total_size", static_cast<uint32_t>(shape.Size()));
    tilings[0].Set("eps", std::any_cast<float>(attrs.at("eps")));
}

7.2 Host 封装

文件：host/rmsnorm_custom.cpp

class RMSNormCustomOp : public OpKernel{
public:
    Status Compute(const OpKernelContext* context) override {
        const Tensor* x = context->Input(0);
        const Tensor* gamma = context->Input(1);
        Tensor* y = context->Output(0);
        auto tiling = GetTilingData();
        uint32_t block_num = tiling.Get<uint32_t>("block_num");
        uint32_t D = tiling.Get<uint32_t>("D");
        uint32_t outer_size = tiling.Get<uint32_t>("outer_size");
        uint32_t total_size = tiling.Get<uint32_t>("total_size");
        float eps = tiling.Get<float>("eps");
        void* args[] = {
            const_cast<half*>(x->data<half>()),
            const_cast<half*>(gamma->data<half>()),
            y->data<half>(),
            &total_size, &D, &outer_size, &eps
        };
        aclrtLaunchKernel("RMSNormKernel", dim3(block_num), dim3(1), args, 0, nullptr);
        return Status::OK();
    }
};

八、第六步：编译与集成

cd RMSNormCustom
bash build.sh
cp librmsnorm_custom.so $ASCEND_HOME/python/site-packages/torch_npu/libs/

九、第七步：PyTorch 集成与验证

9.1 Python 调用示例

import torch
import torch_npu

torch.ops.load_library("librmsnorm_custom.so")

# LLaMA-7B 配置
B, L, H = 1, 512, 4096
x = torch.randn(B, L, H, dtype=torch.float16).npu()
gamma = torch.ones(H, dtype=torch.float16).npu()

# 自定义 RMSNorm
y_custom = torch.ops.custom.rmsnorm_custom(x, gamma, eps=1e-6)

# 对标 HuggingFace 实现
def rms_norm_ref(x, gamma, eps=1e-6):
    variance = x.pow(2).mean(-1, keepdim=True)
    x = x * torch.rsqrt(variance + eps)
    return x * gamma

y_ref = rms_norm_ref(x, gamma)

# 验证
max_diff = torch.max(torch.abs(y_custom - y_ref)).item()
print(f"Max difference: {max_diff:.6f}") # 应 < 1e-3

9.2 性能对比（LLaMA-7B 单层）

✅ 延迟降低 68%，显存减少 33%，完美适配 LLaMA 推理

十、高级技巧：支持无 gamma 版本

部分模型（如早期 LLaMA）使用 无缩放 RMSNorm（即 γ = 1）。我们可通过属性控制：

// 修改 JSON 原型
"attr":[{"name":"eps","type":"float","default":1e-6},{"name":"has_gamma","type":"bool","default":true}]

核函数中增加分支：

if(has_gamma){
    // 读取 gamma 并相乘
}else{
    // 直接输出 x * inv_rms
}

⚠️ 注意：避免运行时分支影响性能，建议编译两个 Kernel。

十一、总结与展望

通过本文，你已掌握：

1. RMSNorm 数学原理与 LLaMA 适配性
2. Ascend C 两阶段融合设计
3. rsqrtf 硬件指令高效使用
4. 动态 Shape 与多 Batch 支持

下一步建议：实现 RMSNorm + Linear 融合算子，探索 INT8 量化 RMSNorm，贡献至 Qwen / LLaMA 昇腾适配项目。

参考资料

1. LLaMA 论文（arXiv:2302.13971）
2. RMSNorm 原始论文（arXiv:1910.07467）
3. HuggingFace Transformers RMSNorm 实现