flashinfer.norm.gemma_fused_add_rmsnorm

flashinfer.norm.gemma_fused_add_rmsnorm(input: Tensor, residual: Tensor, weight: Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None

Gemma 风格的融合加根均方归一化。

步骤 1: residual[i] += input[i]

步骤 2: input[i] = (residual[i] / RMS(residual)) * (weight + 1)

参数:

• input (torch.Tensor) – 输入张量，形状为 (batch_size, hidden_size)。

• residual (torch.Tensor) – 残差张量，形状为 (batch_size, hidden_size)。

• weight (torch.Tensor) – 权重张量，形状为 (hidden_size,)。

• eps (float) – 用于数值稳定的 epsilon。

• enable_pdl (bool) – 是否启用 程序依赖启动