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 and backend ("auto": the Cake package or the CuTe DSL kernel) can run the call on their device.
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 and backend ("auto": the Cake package or the CuTe DSL kernel) can run the call on their device.
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 a BF16, head-dimension-64 call can run on device (the current one when None) through backend ("auto": the Cake package or the CuTe DSL kernel) in the launch form enable_pdl selects.

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_det(x, w[, eps, residual, ...])	Autograd-enabled fused BF16 RMSNorm (optionally with a fused residual add).
cake_rmsnorm_det_train_forward(x, w, eps, *)	Fused BF16 RMSNorm forward with FP32 statistics.
cake_rmsnorm_det_train_backward(g, x, w, rstd, *)	Fused BF16 RMSNorm backward: dx and the FP32 weight gradient dw.
cake_rmsnorm_det_train_backward_workspace_bytes(...)	Bytes of the caller-owned backward workspace of width hidden on device for up to rows tokens.
cake_rmsnorm_det_train_backward_workspace(...)	Allocate a zero-initialized backward workspace (see ..._workspace_bytes).
program_windows(arch, hidden, kind, residual)	The row-count windows of one route and the program each launches.
weight_dtypes(arch, hidden, kind, residual)	The weight dtypes (tokens of WEIGHT_DTYPES) the route ships on arch.
CakeRMSNormDetFunction(*args, **kwargs)	Autograd wrapper saving only (x_or_h_new, w, rstd) for the backward.

Training RMSNorm, opt-in non-deterministic weight gradient (Cake)

The names the training RMSNorm had before the deterministic package took its det names. Every entry point takes deterministic (default True):

• deterministic=True runs the deterministic package above and returns its bits.

• deterministic=False runs this package’s backward in the form its registry lists for each token-count window (the backward’s docstring lists the windows). In the windows that ship the single-kernel form, one kernel per call computes dx row by row and adds the dw sums of its CTAs into FP64 sums in the workspace, in the order the CTAs finish: dw is not guaranteed bitwise reproducible; dx is, and it is within one BF16 ulp of max(|dx|, rms(dx)) of the deterministic backward’s dx (not always bitwise equal to it). In the windows the backward’s notes name, a twosum single kernel keeps each thread’s dw column sums exact (error-free FP32 products, TwoSum compensation) before the FP64 additions, so its dw error does not grow with the token count. The other windows run the deterministic package’s backward pair (main kernel and fixed-order reduce), rendered into this package. The backward workspace (cake_rmsnorm_train_backward_workspace_bytes(..., deterministic=False)) has one layout with fixed offsets: the pair’s counters, the single kernel’s counters and FP64 column sums and, for the pair windows, the partial rows. Zeroed once, it serves any sequence of deterministic=False calls it is large enough for with no host write: the kernels leave the counters at rest and the sums at zero. It serves calls of one mode.

• Under torch.use_deterministic_algorithms(True), deterministic=False raises RuntimeError. With warn_only=True it warns; the forward runs this package’s (deterministic) forward programs and the backward runs the deterministic package with its own workspace (a caller’s workspace is not used), or raises one RuntimeError when that package cannot serve the call ([..., H] inputs, an FP32 weight, a width or residual route it does not export). The workspace helpers with deterministic=False warn and return the deterministic package’s workspace in that mode.

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 for up to rows tokens.
cake_rmsnorm_train_backward_workspace(rows, ...)	Allocate a zero-initialized backward workspace (see ..._workspace_bytes, including its warn_only fallback: a warning and the deterministic package's workspace, which that fallback's backward does not use).
CakeRMSNormFunction(*args, **kwargs)	Autograd wrapper saving only (x_or_h_new, w, rstd) for the backward.
is_cake_rmsnorm_train_supported(device[, ...])	Whether device (and optionally hidden) has an exported route.

Names before the det rename and their deterministic names

Name (deterministic=True by default)	Deterministic name (always deterministic)
flashinfer.cake_rmsnorm_train	flashinfer.cake_rmsnorm_det_train
cake_rmsnorm	cake_rmsnorm_det
cake_rmsnorm_train_forward	cake_rmsnorm_det_train_forward
cake_rmsnorm_train_backward	cake_rmsnorm_det_train_backward
cake_rmsnorm_train_backward_workspace_bytes	cake_rmsnorm_det_train_backward_workspace_bytes
cake_rmsnorm_train_backward_workspace	cake_rmsnorm_det_train_backward_workspace
CakeRMSNormFunction	CakeRMSNormDetFunction
is_cake_rmsnorm_train_supported	is_cake_rmsnorm_det_train_supported

Training SiLU-gated RMSNorm (Cake)

RMSNorm(x) * weight * SiLU(z) with a separate gate input z, forward / backward for training on SM100, SM103 and SM107: rows of 128 ([T, H, 128] or [T * H, 128]) normalized per head with FP32 statistics and gate, BF16 inputs, outputs and gradients, a weight shared by every token and head whose gradient is summed in a fixed order (repeated calls return the same bits), and an autograd wrapper that saves only its inputs.

cake_rmsnorm_silu_gate_det(x, z, weight[, eps])	Differentiable RMSNorm(x) * weight * SiLU(z) over the last dimension (128): CakeRMSNormSiLUGateDetFunction applied to the arguments.
cake_rmsnorm_silu_gate_det_train_forward(x, ...)	SiLU-gated RMSNorm forward: y = BF16(((x * rstd) * weight) * silu(z)).
cake_rmsnorm_silu_gate_det_train_backward(...)	SiLU-gated RMSNorm backward: dx, dz and the weight gradient.
cake_rmsnorm_silu_gate_det_train_backward_workspace_bytes(...)	Bytes of the backward's FP32 weight-gradient partial buffer for rows rows (T * H) on device: chunks * 128 * 4.
cake_rmsnorm_silu_gate_det_train_backward_workspace(...)	A caller-owned backward workspace for rows rows: an uninitialized FP32 tensor of cake_rmsnorm_silu_gate_det_train_backward_workspace_bytes() bytes (the backward writes every element it reads).
CakeRMSNormSiLUGateDetFunction(*args, **kwargs)	Autograd of cake_rmsnorm_silu_gate_det(): saves x, z and weight (no FP32 intermediate) and recomputes the rest in the backward, so torch.utils.checkpoint(..., use_reentrant=False) sees only the inputs as saved state.

Training SiLU-gated RMSNorm, opt-in non-deterministic weight gradient (Cake)

The names the training SiLU-gated RMSNorm had before the deterministic package took its det names. Every entry point takes the keyword deterministic (default True; a value that is not a bool raises TypeError), and CakeRMSNormSiLUGateFunction takes it as its last forward input:

• deterministic=True runs the deterministic package above and returns its bits.

• deterministic=False runs the standalone non-deterministic package (flashinfer.cake_rmsnorm_silu_gate_train_backend, whose compute and workspace entry points take deterministic=False only): the same forward kernel (up to 12 rows per SM on SM100 and SM103 and up to 59 on SM107, a forward with 16 threads to a row, and from 976 rows per SM on SM100 and SM103 and from 800 on SM107, a forward whose CTAs each walk several row blocks: the same bits) and a one-kernel backward: dx and dz are the deterministic backward’s bits, and the CTAs take row blocks (each CTA its own block when the grid has a CTA for every block; past one wave of CTAs they claim the blocks in chunks from the workspace’s claim counter) and add their FP64 weight-gradient column sums into FP64 accumulators of the workspace in the order they finish. The claimed blocks, and so the FP32 partial sums, and that order depend on scheduling: the weight gradient can differ between calls, always within the deterministic backward’s error bound (for FP32 normal gates, terms and sums; with subnormal values both backwards can leave it). Up to 256 rows the backward is one thread-block cluster of 16 CTAs that adds the CTAs’ partial sums in a fixed order on chip: the same bits on every call, within the same bound, and no workspace. The workspace (cake_rmsnorm_silu_gate_train_backward_workspace(..., deterministic=False), a few KiB whatever the token count) is zero after every call: zero it once and reuse it, one per concurrently running stream. A buffer that anything else wrote into is zeroed with cake_rmsnorm_silu_gate_train_reset_workspace(..., deterministic=False).

• Under torch.use_deterministic_algorithms(True), deterministic=False raises RuntimeError from the compute, workspace and autograd entry points; with warn_only=True they warn and run the deterministic kernels. The backward of a deterministic=False autograd node that runs after the mode was enabled raises PyTorch’s alert (warn_only=True: a warning, then the non-deterministic backward). The support query and the workspace reset do not depend on that setting.

cake_rmsnorm_silu_gate(x, z, weight[, eps, ...])	Differentiable RMSNorm(x) * weight * SiLU(z) over the last dimension (128).
cake_rmsnorm_silu_gate_train_forward(x, z, ...)	SiLU-gated RMSNorm forward: y = BF16(((x * rstd) * weight) * silu(z)).
cake_rmsnorm_silu_gate_train_backward(...[, ...])	SiLU-gated RMSNorm backward: dx, dz and the weight gradient.
cake_rmsnorm_silu_gate_train_backward_workspace_bytes(...)	Bytes of the caller-owned backward workspace for rows rows (T * H) on device.
cake_rmsnorm_silu_gate_train_backward_workspace(...)	A caller-owned backward workspace for rows rows (T * H).
cake_rmsnorm_silu_gate_train_reset_workspace(...)	Return a caller-owned backward workspace to its first-use state.
CakeRMSNormSiLUGateFunction(*args, **kwargs)	Autograd of RMSNorm(x) * weight * SiLU(z) with the numerical mode as an input: saves x, z and weight (no FP32 intermediate) and recomputes the rest in the backward, so torch.utils.checkpoint(..., use_reentrant=False) sees only the inputs as saved state.
is_cake_rmsnorm_silu_gate_train_supported(...)	Whether device is a CUDA device of an exported architecture.