flashinfer.cake_rmsnorm_silu_gate_train.cake_rmsnorm_silu_gate_train_backward

flashinfer.cake_rmsnorm_silu_gate_train.cake_rmsnorm_silu_gate_train_backward(grad_out: Tensor, x: Tensor, z: Tensor, weight: Tensor, eps: float = 1e-06, *, grad_x: Tensor | None = None, grad_z: Tensor | None = None, grad_weight: Tensor | None = None, accumulate_grad_weight: bool = False, compute_grad_weight: bool = True, workspace: Tensor | None = None) → tuple[Tensor, Tensor, Tensor | None]

SiLU-gated RMSNorm backward: dx, dz and the weight gradient.

With n = x * rstd, s = sigmoid(z) and a = dy * silu(z):

dz = BF16(dy * (n * weight) * s * (1 + z * (1 - s)))
dx = BF16(rstd * a * weight - x * rstd^3 * mean(a * weight * x))
dw = sum over every token and head of a * n

in FP32, the statistics and the gate recomputed from x and z. dw is accumulated in FP32 per CTA chunk and in FP64 over the chunks, in an order fixed by the row count, and rounded once.

Parameters:

  • grad_out (torch.Tensor) – BF16 tensors of one [..., 128] shape (any strides, see the module).

  • x (torch.Tensor) – BF16 tensors of one [..., 128] shape (any strides, see the module).

  • z (torch.Tensor) – BF16 tensors of one [..., 128] shape (any strides, see the module).

  • weight (torch.Tensor) – BF16 [128].

  • eps (float) – The forward’s epsilon.

  • grad_x (torch.Tensor, optional) – Preallocated BF16 outputs of x’s shape.

  • grad_z (torch.Tensor, optional) – Preallocated BF16 outputs of x’s shape.

  • grad_weight (torch.Tensor, optional) – A BF16 or FP32 [128] tensor that receives the weight gradient (rounded once into its dtype). Without it a new BF16 tensor is returned.

  • accumulate_grad_weight (bool) – Add the gradient to grad_weight’s current value before the one rounding (needs grad_weight).

  • compute_grad_weight (bool) – False skips the weight-gradient reduction and returns None for it.

  • workspace (torch.Tensor, optional) – FP32 tensor of at least cake_rmsnorm_silu_gate_train_backward_workspace_bytes() bytes, contiguous; allocated on the current stream when omitted.

Return type:

(grad_x, grad_z, grad_weight).