flashinfer.norm.gated_residual_next_rmsnorm

flashinfer.norm.gated_residual_next_rmsnorm(branch: Tensor, residual: Tensor, next_weight: Tensor, alpha: Tensor | float, beta: Tensor, post_eps: float = 1e-06, next_eps: float = 1e-06, next_gain_center: float = 0.0, residual_out: Tensor | None = None, ffn_input: Tensor | None = None, *, post_weight: Tensor | None = None, post_gain_center: float = 0.0, enable_pdl: bool = True) → Tuple[Tensor, Tensor]

Gated residual RMSNorm plus the next RMSNorm in one launch, for hidden size 8192.

This computes residual_out = alpha * residual + beta * RMSNorm(branch) in FP32. It then computes ffn_input = RMSNorm(residual_out) * (next_weight + next_gain_center) and rounds only that output to BF16. When post_weight is supplied, branch is multiplied by post_weight + post_gain_center in FP32 before its RMSNorm. The optimized implementation supports contiguous hidden-size-8192 tensors on SM100, SM103, and SM107 GPUs.

Parameters:

  • branch (torch.Tensor) – Two-dimensional BF16 branch tensor.

  • residual (torch.Tensor) – FP32 residual stream with the same shape as branch.

  • next_weight (torch.Tensor) – BF16 weight for the second RMSNorm, with shape (8192,).

  • alpha (Union[torch.Tensor, float]) – FP32 stream gate with shape (8192,), or a scalar coefficient. Use 1.0 for an ungated stream. Scalar values specialize the kernel without allocating or loading an alpha tensor.

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

  • post_eps (float) – Epsilons for the branch and residual-stream RMSNorms.

  • next_eps (float) – Epsilons for the branch and residual-stream RMSNorms.

  • next_gain_center (float) – Constant added to next_weight in FP32.

  • residual_out (Optional[torch.Tensor]) – Optional FP32 destination with the same shape as residual.

  • ffn_input (Optional[torch.Tensor]) – Optional BF16 destination with the same shape as branch.

  • post_weight (Optional[torch.Tensor]) – Optional BF16 input scale for the first RMSNorm, with shape (8192,).

  • post_gain_center (float) – Constant added to post_weight in FP32; unused when it is None.

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

Returns:

The FP32 residual stream and BF16 normalized FFN input.

Return type:

Tuple[torch.Tensor, torch.Tensor]