Rows and columns

Operations on a rank-2 tile [R, C], each row on its own, as a softmax over the keys of each query is. They are built from the row reductions rowmax, rowsum and rowargmax, which fold a tile to a [R, 1] column, and elementwise statements that broadcast that column back over the rows.

@autotune(R in [8])
kernel softmax_rows(S: tensor<f32>[M, 128], P: tensor<f32>[M, 128], n: i32) {
  let r = ctx.block.x * R
  var s = S[r :+ R, :]
  mask_cols(s, n, -1.0e30)
  softmax(s)
  P[r :+ R, :] = s
}

Only the first n of the 128 columns hold scores. A maximum and an exponential would count the rest, so they are set to -1.0e30 first and softmax gives them no weight. A tile that runs past a tensor’s end needs the same: its slice reads as 0 there, which a sum ignores but a softmax does not.

softmax(s) is four statements spliced in: a row maximum, exp(s - m) over the tile with the [R, 1] maximum broadcast along its row, a row sum, and a scale by its reciprocal, which is computed once a row. A row reduction folds each row with warp shuffles, and a barrier separates each step from the next, since each reads what the one before stored.

softmax: 4 statements, each one pass over the tile or its rows.

shared memory
weight row

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

Reductions

`rowmean`

fn rowmean(t)expression

The mean of each row, [R, C] -> [R, 1]. t may be an elementwise expression, as in rowmean(x * x).

Source, std/reduce.ph:11

fn rowmean(t) = rowsum(t) / f32(extent(t, 1))

`rowmin`

fn rowmin(t)expression

The minimum of each row, [R, C] -> [R, 1].

Source, std/reduce.ph:14

fn rowmin(t) = -rowmax(-t)

`rowargmin`

fn rowargmin(t)expression

The column of each row's minimum, [R, C] -> [R, 1] (i32), the smaller column on a tie.

Source, std/reduce.ph:18

fn rowargmin(t) = rowargmax(-t)

`rowvar`

fn rowvar(t)expression

The variance of each row, [R, C] -> [R, 1]: the mean squared distance from the row's mean, the biased variance a layer norm takes.

Source, std/reduce.ph:22

fn rowvar(t) = rowsum(square(t - rowmean(t))) / f32(extent(t, 1))

`logsumexp`

fn logsumexp(out, t)statement

out ([R, 1]) is log sum_j e^t[i, j] of each row, computed around the row maximum.

Source, std/reduce.ph:50

fn logsumexp(out, t) {
  var m = rowmax(t)
  out = m + log(rowsum(exp(t - m)))
}

logsumexp is the normalizer of a log-softmax, and with the target logit subtracted a cross-entropy, which cross_entropy is:

kernel nll(Logits: tensor<f32>[M, 512], Target: tensor<f32>[M, 1], Loss: tensor<f32>[M, 1]) {
  let r = ctx.block.x * 8
  var x = Logits[r :+ 8, :]
  var lse: tile<f32>[8, 1]
  logsumexp(lse, x)
  Loss[r :+ 8, :] = lse - Target[r :+ 8, :]
}

Softmax

`softmax`

fn softmax(t)statement

Each row's softmax, in place: e^(t - max) / sum e^(t - max). The row maximum is taken out first, so no row overflows.

Source, std/reduce.ph:33

fn softmax(t) {
  var m = rowmax(t)
  t = exp(t - m)
  var r = 1.0 / rowsum(t)
  t = t * r
}

`log_softmax`

fn log_softmax(t)statement

Each row's log-softmax, in place: t - logsumexp(t).

Source, std/reduce.ph:41

fn log_softmax(t) {
  var m = rowmax(t)
  t = t - m
  var lse = log(rowsum(exp(t)))
  t = t - lse
}

Masks

`mask_cols`

fn mask_cols(t, n, fill)statement

Sets every column at or past n to fill, in place: the mask a ragged last tile needs before a reduction that a zero would change. A slice past the end of a tensor reads as 0, so a softmax over one is masked with mask_cols(s, n, -1.0e30).

Source, std/reduce.ph:59

fn mask_cols(t, n, fill) {
  t = where(iota(1) < n, t, fill)
}

`causal_mask`

fn causal_mask(t, q0, k0)statement

The causal mask of attention scores, in place: element [i, j] is query q0 + i against key k0 + j, and a key after its query is set to -1.0e30, so a softmax gives it no weight.

Source, std/reduce.ph:66

fn causal_mask(t, q0, k0) {
  t = where(k0 + iota(1) <= q0 + iota(0), t, -1.0e30)
}

Causal attention uses it on each block of scores.

Down the columns

A column reduction folds the rows of a tile into a [1, C] row, as a batch norm folds a batch of samples. The language reduces rows only, so each of these transposes the tile through shared memory first and reduces the rows of that: one more copy of the tile than a row reduction.

kernel column_stats(X: tensor<f32>[32, 128], Mean: tensor<f32>[1, 128], Top: tensor<f32>[1, 128]) {
  var x = X[0 :+ 32, :]
  Mean[0 :+ 1, :] = colmean(x)
  Top[0 :+ 1, :] = colmax(x)
}

`colsum`, `colmax`, `colmin`, `colmean`

fn colsum(t)fn colmax(t)fn colmin(t)fn colmean(t)expression

Reductions down the columns, [R, C] -> [1, C], as a batch norm folds a batch. Each transposes t through shared memory and reduces its rows.

Source, std/reduce.ph:26

fn colsum(t) = transpose(rowsum(transpose(t)))
fn colmax(t) = transpose(rowmax(transpose(t)))
fn colmin(t) = -colmax(-t)
fn colmean(t) = colsum(t) / f32(extent(t, 0))

Selection

`topk`

fn topk(values, indices, t)statement

The K largest values of each row of t into values ([R, K]) and their columns into indices ([R, K], i32), largest first, the smaller column on a tie. t is consumed: each pick is overwritten with -1.0e30. values and indices must be names.

Source, std/reduce.ph:74

fn topk(values, indices, t) {
  for j in range(0, extent(indices, 1)) {
    var i = rowargmax(t)
    values[:, j :+ 1] = rowmax(t)
    indices[:, j :+ 1] = i
    t = where(iota(1) == i, -1.0e30, t)
  }
}

K rounds of rowargmax, each masking its pick, as a sampler’s top-k needs:

kernel top4(Logits: tensor<f32>[M, 256], Values: tensor<f32>[M, 4], Indices: tensor<i32>[M, 4]) {
  let r = ctx.block.x * 8
  var t = Logits[r :+ 8, :]
  var v: tile<f32>[8, 4]
  var i: tile<i32>[8, 4]
  topk(v, i, t)
  Values[r :+ 8, :] = v
  Indices[r :+ 8, :] = i
}

Rows and columns

Reductions

rowmean

rowmin

rowargmin

rowvar

logsumexp