flashinfer.cake_rmsnorm_det_train.cake_rmsnorm_det

flashinfer.cake_rmsnorm_det_train.cake_rmsnorm_det(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, or [..., H] under the layout rule of cake_rmsnorm_det_train_forward() (its leading dimensions collapse into rows).

  • w (torch.Tensor) – BF16 or FP32 [H] RMSNorm weight (FP32 on the routes that ship the FP32-weight programs, see weight_dtypes()); the weight gradient is returned in w.dtype.

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

  • residual (torch.Tensor, optional) – BF16 update of x’s shape 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 of x’s shape, or (y, h_new) when residual is given. Backward launches the fused dx and deterministic dw kernels of cake_rmsnorm_det_train_backward().

Return type:

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