flashinfer.cake_rmsnorm_silu_gate_train.cake_rmsnorm_silu_gate_train_forward

flashinfer.cake_rmsnorm_silu_gate_train.cake_rmsnorm_silu_gate_train_forward(x: Tensor, z: Tensor, weight: Tensor, eps: float = 1e-06, *, out: Tensor | None = None) → Tensor

SiLU-gated RMSNorm forward: y = BF16(((x * rstd) * weight) * silu(z)).

rstd = 1 / sqrt(mean(x^2) + eps) over the last dimension (128), in FP32; silu(z) = z / (1 + exp(-z)) in FP32; one rounding to BF16 at the end.

Parameters:

  • x (torch.Tensor) – BF16 [..., 128] input ([T, H, 128] or [T * H, 128]).

  • z (torch.Tensor) – BF16 gate input of the same shape as x (its own strides).

  • weight (torch.Tensor) – BF16 [128] weight shared by every token and head.

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

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

Returns:

BF16 y of x’s shape (out when given).

Return type:

torch.Tensor