Normalization

Normalizations over the rows of a rank-2 tile [R, C], in place, as a transformer normalizes each token’s hidden state. Each row’s statistic is a row reduction, a [R, 1] column, and the normalization one elementwise sweep that broadcasts it.

rms_norm: 2 statements, each one pass over the tile or its rows.

• shared memory
• weight row

softmaxrms_normlayer_norm

Schematic. The statements are the library's source. A barrier separates each from the next, since each reads what the one before stored.

@autotune(R in [4])
kernel rms_norm_rows(X: tensor<f32>[M, 1024], G: tensor<f32>[1, 1024], Y: tensor<f32>[M, 1024], eps: f32) {
  let r = ctx.block.x * R
  var x = X[r :+ R, :]
  rms_norm(x, G[0 :+ 1, :], eps)
  Y[r :+ R, :] = x
}

The weight is a [1, C] row, here a slice of the tensor G read where the sweep needs it. A scalar works too: rms_norm(x, 1.0, eps) has no weight.

The sums are taken in the tile’s own type. To normalize an f16 activation, widen it into an f32 tile first; the store narrows it back:

@autotune(R in [4])
kernel layer_norm_rows(X: tensor<f16>[M, 768], G: tensor<f32>[1, 768], B: tensor<f32>[1, 768],
                       Y: tensor<f16>[M, 768]) {
  let r = ctx.block.x * R
  var x: tile<f32>[R, 768] = f32(X[r :+ R, :])
  layer_norm(x, G[0 :+ 1, :], B[0 :+ 1, :], 0.00001)
  Y[r :+ R, :] = x
}

rms_norm

fn rms_norm(x, w, eps)statement

Root-mean-square norm: x / sqrt(mean(x^2) + eps) * w, LLaMA's.

Source, std/norm.ph:10

fn rms_norm(x, w, eps) {
  var r = rsqrt(rowmean(x * x) + eps)
  x = x * r * w
}

layer_norm

fn layer_norm(x, w, b, eps)statement

Layer norm: (x - mean(x)) / sqrt(var(x) + eps) * w + b, the biased variance, as torch.nn.functional.layer_norm.

Source, std/norm.ph:17

fn layer_norm(x, w, b, eps) {
  var mu = rowmean(x)
  x = x - mu
  var r = rsqrt(rowmean(x * x) + eps)
  x = x * r * w + b
}

l2_normalize

fn l2_normalize(x, eps)statement

Each row scaled to unit length: x / max(||x||, eps), as torch.nn.functional.normalize.

Source, std/norm.ph:41

fn l2_normalize(x, eps) {
  var r = 1.0 / max(sqrt(rowsum(x * x)), eps)
  x = x * r
}

Batch norm

A batch norm normalizes each column over the batch, the rows. At inference it applies running statistics, an elementwise expression that fuses into the store; in training a block whose rows are the whole batch computes them first:

kernel batch_norm_train(X: tensor<f32>[32, 64], G: tensor<f32>[1, 64], B: tensor<f32>[1, 64],
                        Y: tensor<f32>[32, 64], Stats: tensor<f32>[2, 64]) {
  var x = X[0 :+ 32, :]
  var mean: tile<f32>[1, 64]
  var variance: tile<f32>[1, 64]
  batch_norm_stats(mean, variance, x)
  Y[0 :+ 32, :] = batch_norm(x, mean, variance, G[0 :+ 1, :], B[0 :+ 1, :], 0.00001)
  Stats[0 :+ 1, :] = mean
  Stats[1 :+ 1, :] = variance
}

batch_norm

fn batch_norm(x, mean, variance, w, b, eps)expression

Batch norm as inference applies it: (x - mean) / sqrt(variance + eps) * w + b, with mean, variance, w and b the running statistics and the affine parameters, each a [1, C] row or a scalar. An expression function, so it fuses into the statement that stores it.

Source, std/norm.ph:28

fn batch_norm(x, mean, variance, w, b, eps) = (x - mean) * rsqrt(variance + eps) * w + b

batch_norm_stats

fn batch_norm_stats(mean, variance, x)statement

The batch statistics of the rows of x, a batch of R samples: each column's mean into mean and its biased variance into variance, both [1, C] names. The statistics of a batch spread over several blocks are the merge of each block's.

Source, std/norm.ph:34

fn batch_norm_stats(mean, variance, x) {
  mean = colmean(x)
  variance = colmean(square(x - mean))
}