flashinfer.norm.gated_residual_rmsnorm

flashinfer.norm.gated_residual_rmsnorm(branch: Tensor, residual: Tensor, weight: Tensor | None, alpha: Tensor, beta: Tensor, eps: float = 1e-06, gain_center: float = 0.0, out: Tensor | None = None, *, enable_pdl: bool = True) → Tensor

Gated residual RMSNorm for hidden size 8192: blend a normalized branch into the stream.

The operation is alpha * residual + beta * RMSNorm(branch). When weight is supplied, weight + gain_center is applied to branch before its RMS reduction. When it is None, the branch is unweighted.

This optimized implementation currently supports contiguous BF16 tensors with hidden size 8192 on SM100, SM103, and SM107 GPUs.

Parameters:

  • branch (torch.Tensor) – Two-dimensional BF16 and FP32 tensors, respectively, with shape (rows, 8192).

  • residual (torch.Tensor) – Two-dimensional BF16 and FP32 tensors, respectively, with shape (rows, 8192).

  • weight (Optional[torch.Tensor]) – Optional BF16 pre-normalization weight with shape (8192,).

  • alpha (torch.Tensor) – FP32 gate tensors with shape (8192,).

  • beta (torch.Tensor) – FP32 gate tensors with shape (8192,).

  • eps (float) – Epsilon added to the RMS denominator.

  • gain_center (float) – Constant added to weight in FP32.

  • out (Optional[torch.Tensor]) – Optional FP32 output tensor with the same shape and layout as residual.

  • enable_pdl (bool) – Whether to enable programmatic dependent launch. Defaults to True; pass False to disable it.

Returns:

The fused result.

Return type:

torch.Tensor