flashinfer.norm

Kernels for normalization layers.

rmsnorm(input, weight[, eps, out, enable_pdl])	Root mean square normalization.
rmsnorm_quant(out, input, weight, scale[, ...])	Root mean square normalization + fp8 quantization.
fused_add_rmsnorm(input, residual, weight[, ...])	Fused add root mean square normalization.
fused_add_rmsnorm_quant(out, input, ...[, ...])	Fused add root mean square normalization + fp8 quantization.
fused_add_rmsnorm_fp8_block_quant(out, ...)	Fused add-residual + RMSNorm + 1x128 fp8 block quantization, in one pass.
gemma_rmsnorm(input, weight[, eps, out, ...])	Gemma-style root mean square normalization.
gemma_fused_add_rmsnorm(input, residual, weight)	Gemma-style fused add root mean square normalization.
layernorm(input, gemma, beta[, eps])	Layer normalization.
layernorm_quant(out, input, gemma, beta, scale)	Layer normalization + fp8 quantization.
fused_rmsnorm_silu(input, weight[, eps, ...])	Fused RMSNorm + SiLU activation.
fused_qk_rmsnorm_rope(qkv, q_weight, ...[, ...])	Fused QK RMSNorm + 3D RoPE + V copy for video generation DIT self-attention.
fused_dit_residual_layernorm_scale_shift(...)	Fused residual + LayerNorm + scale/shift for DIT self-attention.
fused_dit_gate_residual_layernorm_scale_shift(...)	Fused gate + residual + LayerNorm + scale/shift for DIT self-attention.
fused_dit_gate_residual_layernorm_gamma_beta(...)	Fused gate + residual + LayerNorm(gamma, beta) for DIT self-attention.

Fixed-Width Fused Norms

Fused gated-residual and per-head gated RMSNorm kernels for SM100, SM103 and SM107. They support only the widths in their docstrings (hidden size 8192, head dim 64); check the matching supports_* helper and fall back otherwise.

gated_residual_rmsnorm(branch, residual, ...)	Gated residual RMSNorm for hidden size 8192: blend a normalized branch into the stream.
supports_gated_residual_rmsnorm(branch, ...)	Return whether tensors satisfy the SM100/SM103/SM107 fused-kernel contract.
gated_residual_next_rmsnorm(branch, ...[, ...])	Gated residual RMSNorm plus the next RMSNorm in one launch, for hidden size 8192.
supports_gated_residual_next_rmsnorm(branch, ...)	Return whether inputs satisfy the SM100/SM103/SM107 two-norm kernel contract.
rmsnorm_sigmoid_gate(input, gate[, eps, ...])	Per-head RMSNorm and sigmoid gate, in place, for head dim 64.
supports_rmsnorm_sigmoid_gate(dtype, head_dim)	Whether the SM100/SM103/SM107 fused kernel supports this dtype and head size.

Training RMSNorm (Cake)

Fused BF16 RMSNorm forward / backward for training on SM100, SM103 and SM107: FP32 statistics and accumulation, one FP32 reciprocal RMS saved per row, fused dx + deterministic FP32 dw backward, optional residual-add fusion and an autograd wrapper.

cake_rmsnorm(x, w[, eps, residual, ...])	Autograd-enabled fused BF16 RMSNorm (optionally with a fused residual add).
cake_rmsnorm_train_forward(x, w, eps, *[, ...])	Fused BF16 RMSNorm forward with FP32 statistics.
cake_rmsnorm_train_backward(g, x, w, rstd, *)	Fused BF16 RMSNorm backward: dx and the FP32 weight gradient dw.
cake_rmsnorm_train_backward_workspace_bytes(...)	Bytes of the caller-owned backward workspace of width hidden on device.
cake_rmsnorm_train_backward_workspace(rows, ...)	Allocate a zero-initialized backward workspace (see ..._workspace_bytes).
CakeRMSNormFunction(*args, **kwargs)	Autograd wrapper saving only (x_or_h_new, w, rstd) for the backward.