flashinfer.norm.gemma_fused_add_rmsnorm

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

Gemma-style fused add root mean square normalization.

Step 1: residual[i] += input[i]

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

Parameters:

• input (torch.Tensor) – Input tensor, shape (batch_size, hidden_size).

• residual (torch.Tensor) – Residual tensor, shape (batch_size, hidden_size).

• weight (torch.Tensor) – Weight tensor, shape (hidden_size,).

• eps (float) – Epsilon for numerical stability.