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(...).

Layer norm¶

tileops.ops.norm.layer_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 \]

Mirrors torch.nn.functional.layer_norm. normalized_shape is the only entry point (the manifest spec).

Supported dtypes

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

init ¶

__init__(
    normalized_shape,
    eps=DEFAULT_EPS,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

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

Parameters:

normalized_shape (Sequence[int]) –

Trailing-axis shape tuple over which the reduction runs (manifest params.normalized_shape).
eps (Optional[float], default: DEFAULT_EPS ) –

Epsilon for numerical stability (manifest params.eps). None uses the PyTorch default 1e-5.
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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

forward ¶

forward(
    x,
    weight,
    bias,
)

Apply layer normalization.

Parameters:

x (Tensor) –

Input tensor with trailing shape equal to normalized_shape.
weight (Tensor) –

Affine scale of shape normalized_shape.
bias (Tensor) –

Affine shift of shape normalized_shape.

Returns:

Tensor –

Normalized tensor of the same shape as x.

Raises:

ValueError –

Dtypes or devices disagree, or shapes are incompatible with the configured normalized_shape. Raised from inside the operator, by _eager_forward.

tileops.ops.norm.fused_add_layer_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,
    kernel_map=None,
    tune=False
)

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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

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 from inside the operator, by _eager_forward.

RMS norm¶

tileops.ops.norm.rms_norm.RMSNormFwdOp ¶

Standalone Root Mean Square (RMS) Norm operator.

Mirrors 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; normalized_shape is the only entry point (the manifest spec).

Example

op = RMSNormFwdOp(normalized_shape=(4096,))
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=DEFAULT_EPS,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

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

Parameters:

normalized_shape (Sequence[int]) –

Trailing-axis shape tuple over which the reduction runs (manifest params.normalized_shape).
eps (Optional[float], default: DEFAULT_EPS ) –

Epsilon for numerical stability (manifest params.eps). None selects the same default the signature carries. Normalized here, so a backend is handed the number rather than None.
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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

Whether to autotune (default False).

forward ¶

forward(
    x,
    weight,
)

Apply RMS normalization over the trailing normalized_shape.

Parameters:

x (Tensor) –

Input tensor whose trailing shape equals normalized_shape.
weight (Tensor) –

Affine scale of shape normalized_shape.

Returns:

Tensor –

Normalized tensor of the same shape as x.

Raises:

ValueError –

Dtypes or devices disagree, or shapes are incompatible with the configured normalized_shape. Raised from inside the operator, by _eager_forward.

tileops.ops.norm.fused_add_rms_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,
    kernel_map=None,
    tune=False
)

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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

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 from inside the operator, by _eager_forward.

Adaptive layer norm¶

tileops.ops.norm.ada_layer_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,
    kernel_map=None,
    tune=False
)

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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

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 from inside the operator, by _eager_forward.

tileops.ops.norm.ada_layer_norm_zero.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,
    kernel_map=None,
    tune=False
)

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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

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 from inside the operator, by _eager_forward.

Batch norm¶

tileops.ops.norm.batch_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.

Mirrors torch.nn.functional.batch_norm: forward accepts (input, running_mean, running_var, weight, bias) in PyTorch's positional order and returns only the normalized output. Internal mean/rstd computed in training mode stay private; callers needing them for the backward pass recompute on the original input.

Supported dtypes

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

init ¶

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

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 the running statistics are written (manifest params.training).
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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

forward ¶

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

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) on CUDA.
running_mean (Tensor) –

Running mean of shape \([C]\) on the same CUDA device as x, with dtype torch.float32. Updated in-place during training.
running_var (Tensor) –

Running variance of shape \([C]\) on the same CUDA device as x, with dtype torch.float32. Updated in-place during training.
weight (Tensor) –

Affine scale (gamma) of shape \([C]\) on the same CUDA device as x.
bias (Tensor) –

Affine shift (beta) of shape \([C]\) on the same CUDA device as x.

Returns:

Tensor –

Normalized output tensor with the same shape as x.

tileops.ops.norm.batch_norm.BatchNormBwdOp ¶

Batch Normalization backward operator.

Computes gradients with respect to input, scale, and shift for batch normalization.

Supported dtypes

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

init ¶

__init__(
    *,
    target=None,
    kernel_map=None,
    tune=False
)

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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

forward ¶

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

Run batch normalization backward pass.

All inputs must reside on the same CUDA device.

Parameters:

grad_out (Tensor) –

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

Original input tensor of shape (N, C, *spatial).
weight (Tensor) –

Affine scale (gamma) of shape \([C]\) on the same CUDA device as x. Internally cast to torch.float32 for the backward kernel.
mean (Tensor) –

Per-channel batch mean from the forward pass, shape (C,). Expected as torch.float32.
rstd (Tensor) –

Per-channel reciprocal std from the forward pass, shape \([C]\). Expected as 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]\).

Group and instance norm¶

tileops.ops.norm.group_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.

Supported dtypes

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

Note

Supports arbitrary spatial dimensions (1-D, 2-D, 3-D+). Handles non-contiguous inputs via explicit contiguous() call. The per-channel affine is applied inside the kernel, so the op does no post-kernel arithmetic.

weight and bias are one switch: pass both for the affine form, pass neither for torch.nn.GroupNorm(affine=False). Passing one alone is an error — the manifest states the same in shape_rules.

init ¶

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

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

Parameters:

num_groups (int) –

Number of groups (manifest params.num_groups). Must divide C evenly.
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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

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. weight and bias are one switch: give both or neither.

Returns:

Tensor –

Normalized tensor of the same shape as x.

Raises:

ValueError –

Only one of weight / bias is given, dtypes disagree, or a shape is incompatible with x. Raised from inside the operator, by _eager_forward.

tileops.ops.norm.instance_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, and the trailing affine applies only when weight and bias are passed. Equivalent to Group Normalization with num_groups = C.

Supported dtypes

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

Note

Supports arbitrary spatial dimensions (1-D, 2-D, 3-D+). The affine call delegates to GroupNormKernel with one group per channel, which applies the per-channel affine itself; without affine it delegates to InstanceNormNoAffineKernel.

init ¶

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

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

Parameters:

use_input_stats (bool, default: True ) –

Mirrors torch.nn.functional.instance_norm. When True (the default), per-instance statistics are computed from the input. False normalizes by the passed running stats, and is implemented for the affine-free call only.
momentum (float, default: 0.1 ) –

Mirrors torch.nn.functional.instance_norm. Stored on the op instance for API parity with PyTorch but unused: neither path updates the running stats.
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.
kernel_map (Optional[Dict[str, Kernel]], default: None ) –

Optional kernel override dictionary.
tune (bool, default: False ) –

If True, autotune tile configurations.

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]\), dtype torch.float32, on x's device. Required when use_input_stats=False.
running_var (Optional[Tensor], default: None ) –

Per-channel running variance, same constraints.
weight (Optional[Tensor], default: None ) –

Affine scale of shape \([C]\), x's dtype. Must be passed together with bias.
bias (Optional[Tensor], default: None ) –

Affine shift, same constraints as weight.

Returns:

Tensor –

Normalized tensor of the same shape as x.

Raises:

ValueError –

A dtype mismatches, a shape is incompatible, one half of a pair is passed, or use_input_stats=False without running stats.
NotImplementedError –

use_input_stats=False combined with the affine tensors. Both raised from inside the operator, by _eager_forward.

Normalization Operators¶

Layer norm¶

tileops.ops.norm.layer_norm.LayerNormFwdOp ¶

__init__ ¶

forward ¶

tileops.ops.norm.fused_add_layer_norm.FusedAddLayerNormFwdOp ¶

__init__ ¶

forward ¶

RMS norm¶

tileops.ops.norm.rms_norm.RMSNormFwdOp ¶

__init__ ¶

forward ¶

tileops.ops.norm.fused_add_rms_norm.FusedAddRMSNormFwdOp ¶

__init__ ¶

forward ¶

Adaptive layer norm¶

tileops.ops.norm.ada_layer_norm.AdaLayerNormFwdOp ¶

__init__ ¶

forward ¶

tileops.ops.norm.ada_layer_norm_zero.AdaLayerNormZeroFwdOp ¶

__init__ ¶

forward ¶

Batch norm¶

tileops.ops.norm.batch_norm.BatchNormFwdOp ¶

__init__ ¶

forward ¶

tileops.ops.norm.batch_norm.BatchNormBwdOp ¶

__init__ ¶

forward ¶

Group and instance norm¶

tileops.ops.norm.group_norm.GroupNormFwdOp ¶

__init__ ¶

forward ¶

tileops.ops.norm.instance_norm.InstanceNormFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶