Triton LayerNorm Forward Kernel

hardPROReductions & NormalizationTriton Programming

Problem Statement

Implement the forward pass of LayerNorm over the last dimension of a (M, N) tensor: subtract the row mean, divide by the row standard deviation, then scale by weight and shift by bias.

Background

Per row: mean = sum(x)/N, var = sum((x-mean)^2)/N, rstd = 1/sqrt(var + eps), out = (x-mean)rstdweight + bias. Mask invalid columns to 0 before the variance sum.

Your Task

Implement layernorm_kernel and run(M=64, N=256, eps=1e-5) comparing to torch.nn.functional.layer_norm.

How it is tested

Your solution must define a top-level function run(...) that allocates inputs on the GPU, launches your Triton kernel, and returns a boolean from torch.allclose(triton_out, torch_reference, ...). The grader prints run(...); the expected output is True.

Example

Input

M = 64, N = 256, eps = 1e-5

Output

True

Reasoning

The input values are used to launch the layernorm_kernel function on a GPU, which implements the forward pass of LayerNorm over the last dimension of a (M, N) tensor.
For each row, the mean is calculated as $mean = \sum(x)/N$ , and the var is calculated as $var = \sum((x-mean)^2)/N$ , masking invalid columns to 0 before the variance sum.
The row standard deviation rstd is then calculated as $rstd = 1/\sqrt{var + eps}$ , and the output out is computed as $out = (x-mean)*rstd*weight + bias$ .
The output of the layernorm_kernel function is compared to the output of torch.nn.functional.layer_norm using torch.allclose, which returns True if the two outputs are close enough, indicating that the implementation is correct.

Constraints

Compute mean and variance per row using tl.sum
rstd = 1 / tl.sqrt(var + eps)
out = (x - mean) * rstd * weight + bias

Learn this in

Advanced Kernels & OptimizationGPU Basics with Triton

More from Triton Programming

Triton Single-Query Attention Kernelhard Triton Row Cumulative Sumhard Triton Stable LogSumExp Kernelhard Triton Seeded Fused Dropouthard Triton Autotuned Vector Addmedium

Concept primerPREVIEW

Background Knowledge

Introduction to Layer Normalization

Layer normalization is a technique used in deep learning to normalize the inputs of each layer. It is similar to batch normalization, but instead of normalizing over the batch dimension, it normalizes over the feature dimension. This helps to reduce the effect of internal covariate shift, which can speed up training and improve the stability of the model.

Mathematical Formulation

The mathematical formulation of layer normalization involves calculating the mean and variance of the input data over the feature dimension. The mean is calculated as $mean = \frac{1}{N} \sum_{i=1}^{N} x_i$ , where $x_i$ is the $i^{th}$ feature of the input data. The variance is calculated as $var = \frac{1}{N} \sum_{i=1}^{N} (x_i - mean)^2$ . The normalized data is then calculated as $out = \frac{x - mean}{\sqrt{var + \epsilon}} \cdot weight + bias$ , where $\epsilon$ is a small value added for numerical stability.

Importance of Epsilon

The $\epsilon$ value is added to the variance to prevent division by zero. This is especially important when the variance is zero, which can happen when the input data is constant. The $\epsilon$ value is typically set to a small value, such as $1e-5$ , to ensure that the division is stable.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Calculate the mean of the input data over the feature dimension
Calculate the variance of the input data over the feature dimension
Calculate the normalized data using the mean and variance
Scale and shift the normalized data using the weight and bias

This approach can be parallelized using a GPU, which can significantly speed up the computation.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

GPU · T4

Test results