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, backend: Literal['auto', 'cake', 'cute'] = 'auto') → 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. The Cake package ships a plain and a programmatic-dependent-launch build of each program and runs the one the flag selects.

  • backend (str) – "auto" (default) runs the Cake kernels (flashinfer.cake_muse_spark_norm_backend) where the package ships a program for the device, the call’s inputs (weight present or not, enable_pdl) and its row count (the package’s routes may leave row-count gaps where the CuTe DSL kernel measured ahead) and the tensors sit in 32-byte aligned storage (the package’s programs read and write 256-bit vectors; storage aligned to 16 bytes only runs the CuTe DSL kernel), else the CuTe DSL kernel; "cake" and "cute" select one and raise ValueError when it is not available ("cake" at a row count in a gap, or on storage aligned to 16 bytes only, raises too).

Returns:

The fused result.

Return type:

torch.Tensor