flashinfer.cake_rmsnorm_train.cake_rmsnorm_train_backward

flashinfer.cake_rmsnorm_train.cake_rmsnorm_train_backward(g: Tensor, x: Tensor, w: Tensor, rstd: Tensor, *, deterministic: bool = True, g_residual: Tensor | None = None, workspace: Tensor | None = None, dx: Tensor | None = None, dw: Tensor | None = None) → tuple[Tensor, Tensor]

Fused BF16 RMSNorm backward: dx and the FP32 weight gradient dw.

Parameters:

  • g (torch.Tensor) – BF16 [T, H] gradient of y of any layout. Views with a contiguous feature dimension, 16-byte aligned storage and a row stride that is a multiple of 16 bytes (expanded or sliced autograd gradients included) are read in place; other layouts are materialized once.

  • x (torch.Tensor) – The BF16 [T, H] tensor that was normalized (x of the plain forward, h_new of the residual forward); same layout rules as g.

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

  • rstd (torch.Tensor) – FP32 [T] reciprocal RMS returned by the forward.

  • deterministic (bool) – Only the deterministic fixed-order dw reduction is exported; False is rejected.

  • g_residual (torch.Tensor, optional) – BF16 [T, H] gradient flowing directly through h_new (residual forward), any layout. It is added once so that the returned dx is the gradient of both residual-add inputs.

  • workspace (torch.Tensor, optional) – Caller-owned uint8 buffer from cake_rmsnorm_train_backward_workspace(); allocated when omitted.

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

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

Returns:

  • (dx, dw) with dw in FP32; the caller decides whether to cast it.

  • For a fixed token count and architecture dw is bitwise reproducible.