Quantization

Symmetric int8 quantization by absolute maximum, one scale a row: what a kernel does to an f32 activation before an integer contraction such as dot_t over i8 or the quantized built-ins.

kernel int8_linear(X: tensor<f32>[M, 256], Wq: tensor<i8>[N, 256], Ws: tensor<f32>[1, N], Y: tensor<f32>[M, N]) {
  let m0 = ctx.block.x * 16
  let n0 = ctx.block.y * 16
  var xq: tile<i8>[16, 256]
  var xs: tile<f32>[16, 1]
  quantize_i8(xq, xs, X[m0 :+ 16, :])
  var acc: tile<i32>[16, 16] = dot_t(xq, Wq[n0 :+ 16, :])
  Y[m0 :+ 16, n0 :+ 16] = f32(acc) * xs * Ws[0 :+ 1, n0 :+ 16]
}

The activation’s rows are quantized into shared memory, the product is exact in i32, and the epilogue scales each element by its row’s scale and its column’s.

quantize_i8

fn quantize_i8(q, s, x)statement

Each row of x ([R, C]) as i8 in q ([R, C]) and its scale in s ([R, 1], f32), so x is about q * s: the row's largest magnitude maps to 127. q and s may be tiles or tensor slices.

Source, std/quant.ph:8

fn quantize_i8(q, s, x) {
  var m = rowmax(abs(x))
  q = i8(i32(round(x * (127.0 / (m + 1.0e-30)))))
  s = m / 127.0
}

dequantize_i8

fn dequantize_i8(x, q, s)statement

x = q * s, the inverse of quantize_i8.

Source, std/quant.ph:15

fn dequantize_i8(x, q, s) {
  x = f32(q) * s
}