flashinfer.cake_rmsnorm_train.cake_rmsnorm_train_forward

flashinfer.cake_rmsnorm_train.cake_rmsnorm_train_forward(x: Tensor, w: Tensor, eps: float, *, residual: Tensor | None = None, out: Tensor | None = None, rstd: Tensor | None = None, residual_out: Tensor | None = None) → tuple[Tensor, ...]

Fused BF16 RMSNorm forward with FP32 statistics.

Parameters:

  • x (torch.Tensor) – BF16 [T, H] input of any layout. Row-strided views with a contiguous feature dimension, 32-byte aligned storage and a row stride that is a multiple of 32 bytes are read in place; other layouts (odd column offset, odd row stride, transposed storage) are materialized once. With residual this is the residual stream h and the normalized value is h_new = BF16(h + residual).

  • w (torch.Tensor) – BF16 [H] weight (copied when not contiguous).

  • eps (float) – Variance epsilon (FP32 inside the kernel).

  • residual (torch.Tensor, optional) – BF16 [T, H] update added to x before normalization (same layout rules as x). Enables the fused residual-add route.

  • out (torch.Tensor, optional) – Preallocated outputs: BF16 [T, H] y, FP32 [T] reciprocal RMS and (residual only) BF16 [T, H] h_new. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

  • rstd (torch.Tensor, optional) – Preallocated outputs: BF16 [T, H] y, FP32 [T] reciprocal RMS and (residual only) BF16 [T, H] h_new. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

  • residual_out (torch.Tensor, optional) – Preallocated outputs: BF16 [T, H] y, FP32 [T] reciprocal RMS and (residual only) BF16 [T, H] h_new. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

Returns:

  • (y, rstd) or (y, rstd, h_new) when residual is given.

  • rstd is the FP32 per-row 1 / sqrt(mean(h_new^2) + eps) the backward

  • needs; no other [T, H] temporary is produced.