Normalization Operators

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

LayerNorm

tileops.norm.LayerNormFwdOp

Layer Normalization operator.

Computes layer normalization over the trailing normalized_shape axes:

\[ y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} \cdot w + b \]

Follows torch.nn.functional.layer_norm: an absent weight scales by one and an absent bias shifts by zero.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

__init__

__init__(
    normalized_shape,
    eps=1e-05,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• normalized_shape (Sequence[int]) –

  Trailing-axis shape over which the statistics run.

• eps (float, default: 1e-05 ) –

  Epsilon for numerical stability.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    weight=None,
    bias=None,
)

Apply layer normalization over the trailing normalized_shape.

Parameters:

• x (Tensor) –

  Input tensor whose trailing shape equals normalized_shape.

• weight (Optional[Tensor], default: None ) –

  Affine scale of shape normalized_shape, or None.

• bias (Optional[Tensor], default: None ) –

  Affine shift of shape normalized_shape, or None.

Returns:

• Tensor –

  Normalized tensor of the same shape as x.

tileops.norm.FusedAddLayerNormFwdOp

Fused residual addition and Layer Normalization operator.

Computes the residual sum followed by layer normalization in a single fused kernel:

\[ \begin{aligned} r &= x + \mathrm{residual} \\ y &= \frac{r - \mathrm{E}[r]}{\sqrt{\mathrm{Var}[r] + \epsilon}} \cdot w + b \end{aligned} \]

Returns dual outputs (y, residual_out) so downstream residual connections can reuse the pre-norm sum without recomputation.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

Note

Supports arbitrary leading dimensions (3-D+) via flatten/unflatten. Handles non-contiguous inputs and non-power-of-two hidden dims by padding to 256-element alignment.

__init__

__init__(
    eps=1e-05,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• eps (float, default: 1e-05 ) –

  Epsilon for numerical stability (manifest params.eps).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    residual,
    weight,
    bias,
)

Apply fused residual addition and normalization.

Parameters:

• x (Tensor) –

  Input tensor of shape (*leading, N).

• residual (Tensor) –

  Residual tensor of the same shape as x.

• weight (Tensor) –

  Affine scale of shape \([N]\).

• bias (Tensor) –

  Affine shift of shape \([N]\).

Returns:

• Tensor –

  (y, residual_out), where residual_out is x + residual, both of the

• Tensor –

  same shape as x.

Raises:

• ValueError –

  Dtypes or shapes disagree. Raised by the generated signature checks.

RMSNorm

tileops.norm.RMSNormFwdOp

Standalone Root Mean Square (RMS) Normalization operator.

Follows torch.nn.functional.rms_norm. Computes::

y = x * rsqrt(mean(x ** 2, trailing_axes) + eps) * weight

where the reduction runs over the trailing len(normalized_shape) axes. An absent weight scales by one; eps=None is the machine epsilon of the float32 accumulation, as in torch.

Example

op = RMSNormFwdOp(normalized_shape=(4096,), eps=1e-6)
x = torch.randn(1024, 4096, dtype=torch.float16, device="cuda")
w = torch.randn(4096, dtype=torch.float16, device="cuda")
y = op(x, w)  # shape: (1024, 4096)

__init__

__init__(
    normalized_shape,
    eps=None,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• normalized_shape (Sequence[int]) –

  Trailing-axis shape over which the reduction runs.

• eps (Optional[float], default: None ) –

  Epsilon for numerical stability; None takes the machine epsilon of float32, the dtype the reduction accumulates in, as torch does.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    weight=None,
)

Apply RMS normalization over the trailing normalized_shape.

Parameters:

• x (Tensor) –

  Input tensor whose trailing shape equals normalized_shape.

• weight (Optional[Tensor], default: None ) –

  Affine scale of shape normalized_shape, or None.

Returns:

• Tensor –

  Normalized tensor of the same shape as x.

tileops.norm.FusedAddRMSNormFwdOp

Fused residual addition and RMS Normalization operator.

Computes the residual sum followed by RMS normalization in a single fused kernel:

\[ \begin{aligned} r &= x + \mathrm{residual} \\ y &= \frac{r}{\sqrt{\mathrm{mean}(r^2) + \epsilon}} \cdot w \end{aligned} \]

Returns dual outputs (y, residual_out) so downstream residual connections can reuse the pre-norm sum without recomputation.

Supported dtypes

torch.float16, torch.bfloat16.

Note

Supports arbitrary leading dimensions (3-D+) via flatten/unflatten. Handles non-contiguous inputs and non-power-of-two hidden dims by padding to 256-element alignment.

__init__

__init__(
    eps=1e-06,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• eps (float, default: 1e-06 ) –

  Epsilon for numerical stability (manifest params.eps).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    residual,
    weight,
)

Apply fused residual addition and normalization.

Parameters:

• x (Tensor) –

  Input tensor of shape (*leading, N).

• residual (Tensor) –

  Residual tensor of the same shape as x.

• weight (Tensor) –

  Affine scale of shape \([N]\).

Returns:

• Tensor –

  (y, residual_out), where residual_out is x + residual, both of the

• Tensor –

  same shape as x.

Raises:

• ValueError –

  Dtypes or shapes disagree. Raised by the generated signature checks.

Adaptive LayerNorm

tileops.norm.AdaLayerNormFwdOp

Adaptive Layer Normalization (AdaLN) operator.

Applies layer normalization with per-token adaptive scale and shift:

\[ y = s \cdot \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} + d \]

where s (scale) and d (shift) are per-token tensors of shape (M, N), pre-computed by the caller from a conditioning signal. Linear projection from the conditioning input to scale/shift is the caller's responsibility.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

Note

Supports arbitrary leading dimensions (3-D+) via flatten/unflatten. Handles non-contiguous inputs and non-power-of-two hidden dims.

__init__

__init__(
    eps=1e-05,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• eps (float, default: 1e-05 ) –

  Epsilon for numerical stability (manifest params.eps).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    scale,
    shift,
)

Apply adaptive layer normalization.

Parameters:

• x (Tensor) –

  Tensor of shape (*leading, N).

• scale (Tensor) –

  Tensor of shape (*leading, N).

• shift (Tensor) –

  Tensor of shape (*leading, N).

Returns:

• Tensor –

  Tensor of the same shape as x.

Raises:

• ValueError –

  Dtypes or shapes disagree. Raised by the generated signature checks.

tileops.norm.AdaLayerNormZeroFwdOp

Adaptive Layer Normalization-Zero (AdaLN-Zero) operator.

Applies layer normalization with per-token adaptive scale, shift, and gating:

\[ y = g \cdot \left( s \cdot \frac{x - \mathrm{E}[x]} {\sqrt{\mathrm{Var}[x] + \epsilon}} + d \right) \]

where s (scale), d (shift), and g (gate) are per-token tensors of shape \([M \times N]\), pre-computed by the caller from a conditioning signal. Linear projection from the conditioning input to scale/shift/gate is the caller's responsibility.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

Note

Supports arbitrary leading dimensions (3-D+) via flatten/unflatten. Handles non-contiguous inputs and non-power-of-two hidden dims.

__init__

__init__(
    eps=1e-05,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• eps (float, default: 1e-05 ) –

  Epsilon for numerical stability (manifest params.eps).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    scale,
    shift,
    gate,
)

Apply adaptive layer normalization with zero-init gating.

Parameters:

• x (Tensor) –

  Tensor of shape (*leading, N).

• scale (Tensor) –

  Tensor of shape (*leading, N).

• shift (Tensor) –

  Tensor of shape (*leading, N).

• gate (Tensor) –

  Tensor of shape (*leading, N).

Returns:

• Tensor –

  Tensor of the same shape as x.

Raises:

• ValueError –

  Dtypes or shapes disagree. Raised by the generated signature checks.

BatchNorm

tileops.norm.BatchNormFwdOp

Batch Normalization forward operator (training and inference).

Computes batch normalization over the channel dimension:

\[ y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} \cdot \gamma + \beta \]

where the mean and variance are computed per channel over (N, *spatial) elements in training, and read from the running statistics in inference.

Follows torch.nn.functional.batch_norm: forward takes (input, running_mean, running_var, weight, bias) in PyTorch's positional order and returns only the normalized output. The running statistics are optional in training, where they are updated in place when passed, and required in inference. An absent weight scales by one and an absent bias shifts by zero.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

__init__

__init__(
    training=False,
    momentum=0.1,
    eps=1e-05,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• training (bool, default: False ) –

  Whether the batch statistics come from this call's input, which is also what decides whether passed running statistics are written.

• momentum (float, default: 0.1 ) –

  Running-stat update momentum (used in training mode).

• eps (float, default: 1e-05 ) –

  Epsilon for numerical stability.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    running_mean=None,
    running_var=None,
    weight=None,
    bias=None,
)

Run batch normalization forward pass.

The training mode is bound at ctor time. Construct a separate op instance to switch between training and inference.

Parameters:

• x (Tensor) –

  Input tensor of shape (N, C, *spatial).

• running_mean (Optional[Tensor], default: None ) –

  Running mean of shape \([C]\), torch.float32; updated in place in training, required in inference.

• running_var (Optional[Tensor], default: None ) –

  Running variance, the same; passed exactly when running_mean is.

• weight (Optional[Tensor], default: None ) –

  Affine scale (gamma) of shape \([C]\), torch.float32, or None.

• bias (Optional[Tensor], default: None ) –

  Affine shift (beta) of shape \([C]\), torch.float32, or None.

Returns:

• Tensor –

  Normalized output tensor with the same shape as x.

tileops.norm.BatchNormBwdOp

Batch Normalization backward operator.

Computes the gradients of the training forward with respect to its input, scale and shift, from the per-channel mean and rstd that forward computed.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

__init__

__init__(
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    grad_out,
    x,
    weight,
    mean,
    rstd,
)

Run batch normalization backward pass.

Parameters:

• grad_out (Tensor) –

  Upstream gradient of shape (N, C, *spatial).

• x (Tensor) –

  Original input tensor of the same shape.

• weight (Tensor) –

  Affine scale (gamma) of shape \([C]\), torch.float32.

• mean (Tensor) –

  Per-channel batch mean from the forward pass, \([C]\), torch.float32.

• rstd (Tensor) –

  Per-channel reciprocal std from the forward pass, \([C]\), torch.float32.

Returns:

• Tensor –

  Tuple of (grad_x, grad_weight, grad_bias) where grad_x

• Tensor –

  has the same shape as x, grad_weight has shape \([C]\),

• Tensor –

  and grad_bias has shape \([C]\).

GroupNorm and InstanceNorm

tileops.norm.GroupNormFwdOp

Group Normalization forward operator.

Computes group normalization over (C/num_groups, *spatial) slices:

\[ y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} \cdot w + b \]

where the mean and variance are computed per group over (C/num_groups, *spatial) elements. weight and bias are independent, as in torch: an absent one scales by one or shifts by zero.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

__init__

__init__(
    num_groups,
    eps=1e-05,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• num_groups (int) –

  Number of groups; it divides the channel count.

• eps (float, default: 1e-05 ) –

  Epsilon for numerical stability.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    weight=None,
    bias=None,
)

Apply group normalization.

Parameters:

• x (Tensor) –

  Input tensor of shape (N, C, *spatial).

• weight (Optional[Tensor], default: None ) –

  Affine scale of shape \([C]\), or None.

• bias (Optional[Tensor], default: None ) –

  Affine shift of shape \([C]\), or None.

Returns:

• Tensor –

  Normalized tensor of the same shape as x.

tileops.norm.InstanceNormFwdOp

Instance Normalization forward operator.

Computes instance normalization over spatial dimensions for each (batch, channel) independently:

\[ y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} \cdot w + b \]

where the mean and variance are computed over *spatial for each sample-channel pair, or read from the running statistics when use_input_stats=False. An absent weight scales by one and an absent bias shifts by zero. With use_input_stats=True passed running statistics take the batch mean of the instance means and unbiased variances with momentum. As in torch, the running statistics are read and updated in the input's dtype.

Supported dtypes

torch.float32, torch.float16, torch.bfloat16.

__init__

__init__(
    use_input_stats=True,
    momentum=0.1,
    eps=1e-05,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• use_input_stats (bool, default: True ) –

  Whether each instance is normalized by its own statistics (the default) or by the running statistics.

• momentum (float, default: 0.1 ) –

  Weight of this call's statistics in the running-statistics update.

• eps (float, default: 1e-05 ) –

  Epsilon for numerical stability.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x,
    running_mean=None,
    running_var=None,
    weight=None,
    bias=None,
)

Apply instance normalization.

Parameters:

• x (Tensor) –

  Input tensor of shape (N, C, *spatial).

• running_mean (Optional[Tensor], default: None ) –

  Per-channel running mean of shape \([C]\), torch.float32. Required when use_input_stats=False; updated in place otherwise.

• running_var (Optional[Tensor], default: None ) –

  Per-channel running variance, passed exactly when running_mean is.

• weight (Optional[Tensor], default: None ) –

  Affine scale of shape \([C]\) in x's dtype, or None.

• bias (Optional[Tensor], default: None ) –

  Affine shift of shape \([C]\) in x's dtype, or None.

Returns:

• Tensor –

  Normalized tensor of the same shape as x.