flashinfer.cake_rmsnorm_train.cake_rmsnorm

flashinfer.cake_rmsnorm_train.cake_rmsnorm(x: Tensor, w: Tensor, eps: float = 1e-06, residual: Tensor | None = None, deterministic: bool = True)

Autograd-enabled fused BF16 RMSNorm (optionally with a fused residual add).

Parameters:

  • x (torch.Tensor) – BF16 [T, H] input on a supported CUDA device.

  • w (torch.Tensor) – BF16 [H] RMSNorm weight.

  • eps (float) – Variance epsilon evaluated in FP32 by the kernel.

  • residual (torch.Tensor, optional) – BF16 [T, H] update added to x before normalization. When provided, the normalized input h_new is returned with the output.

  • deterministic (bool) – Whether to use the deterministic fixed-order weight-gradient reduction. Only True is supported; False is rejected during backward.

Returns:

BF16 y with shape [T, H], or (y, h_new) when residual is given. Backward launches the fused dx and deterministic dw kernels of cake_rmsnorm_train_backward().

Return type:

torch.Tensor or tuple[torch.Tensor, torch.Tensor]