Module layernorm 

LayerNorm Triton kernels.

Layout: input x is [M, N] row-major where M = product of batch / outer dimensions and N = product of the normalized dimensions. Each CTA handles one row (one sample), reading all N elements in BLOCK_N-wide tiles.

Forward: y[m, n] = (x[m, n] − mean_m) / sqrt(var_m + eps) * γ[n] + β[n]

Training launches a single forward kernel that also writes out the saved mean and rstd buffers for the backward pass.

Structs

LayerNormBackward

Backward pass for LayerNorm.

LayerNormForward

Forward pass that also saves per-row mean and rstd for the backward pass.

LayerNormForwardInference

Forward pass using pre-computed running statistics (inference only).

LayerNormForwardInferenceRuntimeOp

RuntimeOp for LayerNorm inference.

Functions

layer_norm_backward

Backward pass for LayerNorm.

layer_norm_forward

Forward pass that also saves per-row mean and rstd for the backward pass.

layer_norm_forward_inference

Forward pass using pre-computed running statistics (inference only).