RmsNormGated 算子

概述

RmsNormGated 是一个基于 Triton 实现的融合门控 RMS 归一化算子，专为门控线性注意力等序列模型设计。该算子将 RMSNorm 与门控激活函数融合到单个 kernel 中，计算公式为 y = RMSNorm(x) * activation(g)，其中 activation 支持 silu（即 swish）和 sigmoid。算子利用昇腾 NPU 的并行计算能力实现高性能的归一化操作，同时支持残差连接、prenorm 模式及可选偏置。

函数签名

函数式接口

def rms_norm_gated(
    x: torch.Tensor,
    g: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    activation: str = "swish",
    residual: torch.Tensor | None = None,
    prenorm: bool = False,
    residual_in_fp32: bool = False,
    eps: float = 1e-6,
) -> torch.Tensor:
    """执行融合门控 RMS 归一化前向计算

    Args:
        x: 输入张量，形状 [B, T, D] 或 [T, D]
        g: 门控张量，形状与 x 相同
        weight: 可学习缩放权重，形状 [D]
        bias: 偏置，形状 [D]（可选，当前实现固定为 None）
        activation: 门控激活函数类型，支持 'swish'/'silu'/'sigmoid'
        residual: 残差输入，形状与 x 相同（可选）
        prenorm: 是否启用 prenorm 模式，启用时返回 (y, residual_out)
        residual_in_fp32: 是否将残差以 fp32 精度存储
        eps: 防止除零的小常数

    Returns:
        y: 归一化输出，形状与 x 相同
        residual_out: prenorm=True 时额外返回残差输出，形状与 x 相同
    """

参数说明

前向输入

参数	类型	形状	描述	是否必须
x	torch.Tensor	[B, S, D] 或 [T, D]	输入张量，B 为 batch 大小，S 为序列长度，T 为BS长度，D 为特征维度	是
g	torch.Tensor	同 x	门控张量，与 x 形状完全一致	是
weight	torch.Tensor	[D]	可学习缩放权重（gamma）	是
bias	torch.Tensor	[D]	偏置（当前实现固定为 None）	否
activation	str	-	门控激活函数，支持 'swish'/'silu'/'sigmoid'	否，默认 "swish"
residual	torch.Tensor	同 x	残差输入，在归一化前与 x 相加	否
prenorm	bool	-	是否启用 prenorm 模式	否，默认 False
residual_in_fp32	bool	-	是否将残差以 fp32 精度存储	否，默认 False
eps	float	-	防止除零的小常数	否，默认 1e-6

前向输出

参数	形状	描述
y	同 x	归一化输出，y = RMSNorm(x) * activation(g)
residual_out	同 x	prenorm=True 时返回残差输出

反向输入

参数	形状	描述
dy	同 x	上层传入的输出梯度
x	同 x	前向保存的原始输入
g	同 x	前向保存的门控输入
weight	[D]	前向保存的缩放权重
bias	[D]	前向保存的偏置（可选）
mean	[T]	前向保存的均值（LayerNorm 模式）
rstd	[T]	前向保存的 1/std
dresidual	同 x	残差梯度（prenorm 模式）

反向输出

参数	形状	描述
dx	同 x	输入 x 的梯度
dg	同 x	门控 g 的梯度
dw	[D]	缩放权重的梯度
db	[D]	偏置的梯度（可选）
dres_in	同 x	残差输入的梯度（可选）

支持的数据类型

输入/输出数据类型

数据类型	PyTorch dtype	说明
32位浮点	torch.float32	最高精度，推荐用于精度敏感场景
16位浮点	torch.float16	推荐用于训练和推理，性能与精度平衡
BFloat16	torch.bfloat16	推荐用于大模型训练，动态范围优于 fp16

内部计算精度

• 归一化计算（mean、variance、rstd）统一在 float32 精度下进行，确保数值稳定性
• weight/bias 的梯度（dw/db）以 float32 累加后转换回原始 dtype
• 输入数据在 kernel 内部自动提升为 float32 进行计算，输出时转换回输入 dtype

类型约束

约束项	说明
x 与 g	必须相同 dtype
weight	可与 x 不同 dtype，kernel 内部自动转换
residual	可与 x 不同 dtype，通过 residual_in_fp32=True 强制 fp32 存储
输出 y	与 x 相同 dtype（除非指定 out_dtype）

实现原理

数学定义

RMSNormGated 的前向计算公式为：

y = (x / sqrt(mean(x^2) + eps)) * weight * activation(g)

其中：

• RMSNorm 部分：x_norm = x * rsqrt(mean(x^2) + eps)，不进行均值中心化
• 门控部分：activation(g) 对门控张量施加激活函数
  • silu/swish：g * sigmoid(g)
  • sigmoid：sigmoid(g)
• 最终输出：y = x_norm * weight * activation(g)

前向计算

算子根据特征维度 D 的大小选择不同的 kernel 路径：

1. D ≤ 512（kernel0 路径）：采用分块策略，将行维度按 [BT, D] 的 tile 切分，每个 program 处理 BS 行，通过 grid-stride loop 遍历所有 tile
2. D > 512（kernel1 路径）：采用分块策略，将行维度按 [1, BD] 的 tile 切分，每个 program 处理 BS 行，通过 grid-stride loop 遍历所有 tile

前向 kernel 的核心步骤：

1. 加载输入：从 HBM 加载 x 和 g 的数据块到 SRAM
2. 残差加法：若存在 residual，与 x 相加
3. 归一化：计算方差 var = mean(x^2)，然后 x_norm = x * rsqrt(var + eps)
4. 仿射变换：乘以 weight（可选加 bias）
5. 门控激活：将归一化结果乘以 activation(g)
6. 写回输出：将结果写回 HBM

反向计算

反向 kernel 同样根据 D 的大小选择 kernel0/kernel1 路径，核心步骤：

1. 重计算前向：从保存的 x、mean、rstd 重新计算 x_hat 和 y
2. 门控梯度：
   • silu：dg = dy * y * (sigmoid(g) + g * sigmoid(g) * (1 - sigmoid(g)))，dy = dy * g * sigmoid(g)
   • sigmoid：dg = dy * y * sigmoid(g) * (1 - sigmoid(g))，dy = dy * sigmoid(g)
3. 权重梯度：dw = sum(dy * x_hat)，db = sum(dy)
4. 输入梯度：
   • RMSNorm：dx = (wdy - x_hat * mean(x_hat * wdy)) * rstd
   • LayerNorm：dx = (wdy - (x_hat * c1 + c2)) * rstd
5. 多 program 梯度归约：每个 program 计算局部 dw/db，最后在 HBM 上求和

使用示例

基础用法（函数式接口）

import torch
from mindspeed_ops.api.triton.rmsnormgated import rms_norm_gated

B, T, D = 4, 2048, 128
x = torch.randn(B, T, D, device='npu', dtype=torch.bfloat16)
g = torch.randn(B, T, D, device='npu', dtype=torch.bfloat16)
weight = torch.randn(D, device='npu', dtype=torch.bfloat16)

y = rms_norm_gated(x, g, weight, bias=None, activation='silu')
assert y.shape == (B, T, D)

带残差连接

residual = torch.randn(B, T, D, device='npu', dtype=torch.bfloat16)
y = model(x, g, residual=residual)

Prenorm 模式

y, residual_out = model(x, g, prenorm=True)

性能对比

方法	性能
PyTorch 实现	基准
Triton 实现	4x

注：性能数据基于昇腾 NPU 环境测试，具体数值可能因硬件配置、输入形状不同而有所差异。

常见训练 Shape

以下为模型训练时的典型输入形状：

B	T	D	x / g / y 形状
32	2048	128	[32, 2048, 128]
16	2048	256	[16, 2048, 256]
8	2048	512	[8, 2048, 512]
4	2048	1024	[4, 2048, 1024]
2	4096	1200	[2, 4096, 1200]
1	2097152	128	[1, 2097152, 128]

注意事项

1. x 与 g 必须具有完全相同的形状
2. weight 的长度必须等于 x 的最后一维 D
3. 激活函数名称 swish 和 silu 等价，均表示 x * sigmoid(x)
4. 当前实现中 bias 固定为 None，如需偏置需扩展实现
5. 反向传播时 kernel 会重计算前向输出 y，这是用计算换显存的设计，在训练显存敏感场景下属于预期行为
6. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销可能超过计算收益
7. 该算子不支持 arch35 平台