Attention

Flash attention streams the keys and values through a CTA one block at a time and never holds the whole score matrix. What makes that possible is the online softmax: for each row it keeps the largest score seen so far m, the sum of the exponentials so far l, and the output so far acc, and when a block brings a new maximum it rescales what it has by exp(m - mnew) before adding the block’s share. The library has that step as online_softmax and the whole pass over the keys as attention.

After the last block, acc / l is the softmax-weighted sum of every value.

• global memory
• shared memory

Schematic. Each step names what it computes: the scores are the caller's, the rest is online_softmax, up to two of its statements to a line. Under @tensorcore the scores and acc live in registers instead of shared memory.

@tensorcore
@launch(128)
@autotune(D in [64], BR in [64], BC in [32])
@aligned(Nq = BR)
kernel flash_attention(Q: tensor<f16>[Nq, D], K: tensor<f16>[Nk, D], V: tensor<f16>[Nk, D],
                       O: tensor<f16>[Nq, D], scale: f32) {
  let row = ctx.block.x * BR
  let q = Q[row :+ BR, :]
  var acc: tile<f32>[BR, D] = 0.0
  attention(acc, q, K, V, scale, BR, BC)
  O[row :+ BR, :] = acc
}

This compiles to the same MLIR as Flash attention written out by hand, as in the Flash attention example with a masked last block, on every target: under @tensorcore the scores and acc stay in the tensor cores’ registers through the softmax, as the written-out kernel’s do. That needs acc declared with its 0.0, which is why attention takes it zeroed rather than zeroing it itself.

online_softmax

fn online_softmax(acc, m, l, s, v)statement

One block of keys' step of the online softmax, Flash attention's inner update. s ([BR, BC]) holds this block's scaled scores, masked where needed, and is overwritten with their unnormalized probabilities. The running state is the row maximum m and the denominator l, both [BR, 1], and the unnormalized output acc ([BR, D]), which gains the block's values v ([BC, D]) weighted by s. Start with m at -1.0e30, l and acc at 0, and divide acc by l after the last block.

Source, std/attention.ph:12

fn online_softmax(acc, m, l, s, v) {
  var mnew = rowmax(s)
  mnew = tmax(m, mnew)
  s = exp(s - mnew)
  var corr = exp(m - mnew)
  l = l * corr
  l += rowsum(s)
  acc = acc * corr
  acc += dot(s, v)
  m = mnew
}

attention

fn attention(acc, q, K, V, scale, BR, BC)statement

Attention of the queries q ([BR, D]) over every key of K and value of V (both [Nk, D] tensors), BC keys a step: acc ([BR, D], f32) becomes softmax(q K^T * scale) V. acc must start at 0, as a declaration var acc: tile<f32>[BR, D] = 0.0 leaves it, which keeps it in the tensor cores' registers under @tensorcore. A last block of keys past Nk is masked.

Source, std/attention.ph:30

fn attention(acc, q, K, V, scale, BR, BC) {
  var m: tile<f32>[BR, 1] = -1.0e30
  var l: tile<f32>[BR, 1] = 0.0
  for kt in range(0, extent(K, 0), BC) {
    let k = K[kt :+ BC, :]
    let v = V[kt :+ BC, :]
    var s: tile<f32>[BR, BC] = dot_t(q, k)
    s = s * scale
    mask_cols(s, extent(K, 0) - kt, -1.0e30)
    online_softmax(acc, m, l, s, v)
  }
  acc = acc / l
}

attention takes BR although q has BR rows. Inside a function a tile type is sized only by the function’s parameters, and the score tile, tile<f32>[BR, BC], has to be declared with a type to be f32 over f16 inputs.

Causal attention

A decoder’s queries attend only to the keys at or before them. online_softmax is the same step; the kernel stops the loop at its last query and masks the scores of each block with causal_mask:

@autotune(D in [64], BR in [32], BC in [32])
@aligned(N = BR)
kernel causal_attention(Q: tensor<f32>[N, D], K: tensor<f32>[N, D], V: tensor<f32>[N, D],
                        O: tensor<f32>[N, D], scale: f32) {
  let row = ctx.block.x * BR
  let q = Q[row :+ BR, :]
  var acc: tile<f32>[BR, D] = 0.0
  var m: tile<f32>[BR, 1] = -1.0e30
  var l: tile<f32>[BR, 1] = 0.0
  // Keys past the block's last query never count, so the loop stops there.
  for kt in range(0, row + BR, BC) {
    let k = K[kt :+ BC, :]
    let v = V[kt :+ BC, :]
    var s: tile<f32>[BR, BC] = dot_t(q, k)
    s = s * scale
    causal_mask(s, row, kt)
    online_softmax(acc, m, l, s, v)
  }
  acc = acc / l
  O[row :+ BR, :] = acc
}

Rotary position embeddings

rope

fn rope(x, cos, sin)statement

Rotary position embedding, in place, in the rotate-half form of GPT-NeoX and Hugging Face's LLaMA: each row of x ([R, D]) is split into halves a and b, and becomes [a cos - b sin, b cos + a sin]. cos and sin are [R, D / 2], the angles of each row's position. GPT-J's form, which rotates adjacent pairs, is rope_pairs.

Source, std/attention.ph:49

fn rope(x, cos, sin) {
  var a = x[:, 0 :+ extent(x, 1) / 2]
  var b = x[:, extent(x, 1) / 2 :+ extent(x, 1) / 2]
  x[:, 0 :+ extent(x, 1) / 2] = a * cos - b * sin
  x[:, extent(x, 1) / 2 :+ extent(x, 1) / 2] = b * cos + a * sin
}

kernel rotary(X: tensor<f32>[T, 128], Cos: tensor<f32>[T, 64], Sin: tensor<f32>[T, 64], Y: tensor<f32>[T, 128]) {
  let t = ctx.block.x * 8
  var x = X[t :+ 8, :]
  rope(x, Cos[t :+ 8, :], Sin[t :+ 8, :])
  Y[t :+ 8, :] = x
}

Both halves are copied out of x before either is written back, so the second half’s rotation reads the first half as it was.

rope_pairs

fn rope_pairs(x0, x1, cos, sin)statement

Rotary position embedding of the pairs (x0, x1), in place: x0 becomes x0 cos - x1 sin and x1 becomes x1 cos + x0 sin, all four the same shape. GPT-J's interleaved pairs are two names when the kernel declares the heads as [T, D / 2, 2]: x0 is X[t :+ R, :, 0] and x1 is X[t :+ R, :, 1], with cos and sin declared [T, D / 2, 1] and sliced the same way.

Source, std/attention.ph:62

fn rope_pairs(x0, x1, cos, sin) {
  var a = x0
  var b = x1
  x0 = a * cos - b * sin
  x1 = b * cos + a * sin
}

GPT-J, and llama.cpp by default, rotate adjacent pairs instead: elements 2i and 2i + 1. Declared as [T, D / 2, 2], a head’s pairs are two slices of the last axis, and rope_pairs rotates them in the tensor itself:

kernel rotary_pairs(X: tensor<f32>[T, 64, 2], Cos: tensor<f32>[T, 64, 1], Sin: tensor<f32>[T, 64, 1]) {
  let t = ctx.block.x * 8
  rope_pairs(X[t :+ 8, :, 0], X[t :+ 8, :, 1], Cos[t :+ 8, :, 0], Sin[t :+ 8, :, 0])
}