diff --git a/packages/runntime/bench/kernels/matmul.ab.ts b/packages/runntime/bench/kernels/matmul.ab.ts index 79e4d46..e1c42c0 100644 --- a/packages/runntime/bench/kernels/matmul.ab.ts +++ b/packages/runntime/bench/kernels/matmul.ab.ts @@ -49,6 +49,25 @@ function quant(name: string, t: number, e: number, k: number, o: number, bits: 8 }; } +/** [m, k] × a whole quantW weight [k, n]. */ +function quantWhole( + name: string, + m: number, + k: number, + n: number, + bits: 8 | 4, + groupSize: number, +): AbCase { + return { + name, + build: (version) => { + const a = randMat(version, m, k); + const { w, scales } = quantWeight(version, 1, k, n, bits, groupSize); + return () => version.lib.matmul(a, w, { scales, bits, groupSize }); + }, + }; +} + function gather(name: string, t: number, e: number, k: number, n: number): AbCase { return { name, @@ -158,6 +177,11 @@ abSuite('matmul', [ quant('quantW int4, privacy-filter expert [32,640]×[640,640]', 32, 16, 640, 640, 4), quant('quantW int8, decode [1,640]×[640,640]', 1, 16, 640, 640, 8), + // Quantized weights, whole: the tiled route from 8 rows, plain below. + quantWhole('quantW int4 whole, q|k projection [256,1024]×[1024,2304]', 256, 1024, 2304, 4, 32), + quantWhole('quantW int8 whole, per channel [256,2048]×[2048,1024]', 256, 2048, 1024, 8, 2048), + quantWhole('quantW int4 whole, decode [1,1024]×[1024,2304]', 1, 1024, 2304, 4, 32), + // matmulGather: one expert per row. gather('matmulGather f32, privacy-filter [64,640], 16 experts', 64, 16, 640, 640), diff --git a/packages/runntime/src/core/gpu/dispatch/matmul.ts b/packages/runntime/src/core/gpu/dispatch/matmul.ts index 844e600..083a791 100644 --- a/packages/runntime/src/core/gpu/dispatch/matmul.ts +++ b/packages/runntime/src/core/gpu/dispatch/matmul.ts @@ -40,6 +40,11 @@ import { createMatmulQuantWPipeline, matmulQuantWHandle, } from '../../kernels/matmul/matmulQuantW.ts'; +import { + createMatmulQuantTiledPipeline, + matmulQuantTiledEligible, + matmulQuantTiledHandle, +} from '../../kernels/matmul/matmulQuantTiled.ts'; import { createMatmulTiledVec4Pipeline, matmulTiledVec4Eligible, @@ -144,34 +149,83 @@ export const matmulSpec = defineSpec({ export const matmulQuantWSpec = defineSpec({ attrs: (p) => { - const [baseRow, scaleBase, bits, groupSize] = p.attrs as [number, number, number, number]; - return { baseRow, scaleBase, bits, groupSize }; + const [baseRow, scaleBase, bits, groupSize, hasBias, hasAdd, act] = p.attrs as [ + number, + number, + number, + number, + number, + number, + number, + ]; + return { baseRow, scaleBase, bits, groupSize, hasBias, hasAdd, act }; + }, + cfg: (node, attrs) => { + const [m, k] = node.pending!.inputs[0]!.shape.dims as [number, number]; + const n = node.shape.dims![1]!; + const rows = node.pending!.inputs[1]!.shape.dims![0]!; + const epilogue = attrs.hasBias + attrs.hasAdd + attrs.act > 0; + // A whole weight goes to the tiled kernel from 8 rows, and always for an + // f16 input or an epilogue, which the plain kernel lacks. Expert views + // stay on the plain kernel. + const tiled = + attrs.baseRow === 0 && + attrs.scaleBase === 0 && + rows === k && + matmulQuantTiledEligible(k, n, attrs.groupSize) && + (node.shape.dtype === 'f16' || epilogue || m >= MATMUL_TILED_MIN_M); + return { + k, + n, + bits: attrs.bits, + groupSize: attrs.groupSize, + hasBias: attrs.hasBias, + hasAdd: attrs.hasAdd, + act: attrs.act, + tiled: tiled ? 1 : 0, + }; + }, + pipeline: (root, dtype, cfg) => + cfg.tiled + ? createMatmulQuantTiledPipeline(root, cfg, elemFor(dtype)) + : createMatmulQuantWPipeline(root, cfg), + encode: ({ node, attrs, cfg, pipeline, inputs, out, ctx }) => { + const m = node.pending!.inputs[0]!.shape.dims![0]!; + const dtype = node.shape.dtype; + if (cfg.tiled) { + const addendAt = attrs.hasBias === 1 ? 4 : 3; + return [ + matmulQuantTiledHandle( + ctx.root, + pipeline, + m, + cfg.n, + { + a: narrowFloat(inputs[0]!, dtype), + w: narrow(inputs[1]!, 'u32'), + scales: narrow(inputs[2]!, 'f32'), + bias: attrs.hasBias === 1 ? narrowFloat(inputs[3]!, dtype) : undefined, + addend: attrs.hasAdd === 1 ? narrowFloat(inputs[addendAt]!, dtype) : undefined, + out: narrowFloat(out, dtype), + }, + elemFor(dtype), + ), + ]; + } + return [ + matmulQuantWHandle( + ctx.root, + pipeline, + { m, n: cfg.n, baseRow: attrs.baseRow, scaleBase: attrs.scaleBase }, + { + a: narrow(inputs[0]!, 'f32'), + w: narrow(inputs[1]!, 'u32'), + scales: narrow(inputs[2]!, 'f32'), + out: narrow(out, 'f32'), + }, + ), + ]; }, - cfg: (node, attrs) => ({ - k: node.pending!.inputs[0]!.shape.dims![1]!, - n: node.shape.dims![1]!, - bits: attrs.bits, - groupSize: attrs.groupSize, - }), - pipeline: (root, _dtype, cfg) => createMatmulQuantWPipeline(root, cfg), - encode: ({ node, attrs, cfg, pipeline, inputs, out, ctx }) => [ - matmulQuantWHandle( - ctx.root, - pipeline, - { - m: node.pending!.inputs[0]!.shape.dims![0]!, - n: cfg.n, - baseRow: attrs.baseRow, - scaleBase: attrs.scaleBase, - }, - { - a: narrow(inputs[0]!, 'f32'), - w: narrow(inputs[1]!, 'u32'), - scales: narrow(inputs[2]!, 'f32'), - out: narrow(out, 'f32'), - }, - ), - ], }); export const moeUpSpec = defineSpec({ diff --git a/packages/runntime/src/core/graph/ops/matmul.ts b/packages/runntime/src/core/graph/ops/matmul.ts index 169c639..67079a9 100644 --- a/packages/runntime/src/core/graph/ops/matmul.ts +++ b/packages/runntime/src/core/graph/ops/matmul.ts @@ -4,11 +4,16 @@ import { pending, Value } from '../value.ts'; import { cat } from './shape.ts'; import { MOE_SLOTS } from '../../kernels/matmul/moeQuantW.ts'; +import { + MATMUL_QUANT_TILED_BK, + matmulQuantTiledEligible, +} from '../../kernels/matmul/matmulQuantTiled.ts'; import { ACT_CODE, FLOAT_DTYPES, want2d, type SlotAct } from './shared.ts'; -/** Options for matmul. `bias` and `addend` are fused epilogue adds, not - * available for quantW. `baseRow` views into a taller weight without copying, - * so MoE experts stay views into one fused tensor. */ +/** Options for matmul. `bias` and `addend` are fused epilogue adds; on a + * quantW weight they need the whole weight, no baseRow or scaleBase view. + * `baseRow` views into a taller weight without copying, so MoE experts stay + * views into one fused tensor. */ export interface MatmulOpts { bias?: Value; addend?: Value; @@ -21,19 +26,14 @@ export interface MatmulOpts { groupSize?: number; } -/** Activation codes shared by the matmul and conv kernels. */ /** 2D matrix multiply, a [M,K] × b [K,N] to [M,N]. One name for every weight * storage; b's dtype picks the kernel. * - * Float b supports the fused bias and addend epilogue. quantW b is - * dequantized with `scales`, needs bits and groupSize, and has no epilogue. */ + * b takes the fused bias, activation and addend epilogue. quantW b is + * dequantized with `scales`, needs bits and groupSize, and takes the + * epilogue only as a whole weight. */ export function matmul(a: Value, b: Value, opts: MatmulOpts = {}): Value { if (b.shape.dtype === 'quantW') { - if (opts.bias || opts.addend || opts.activation) { - throw new Error( - 'matmul: bias/addend/activation epilogue is not supported for quantW weights', - ); - } const { scales, bits, groupSize } = opts; if (!scales || !bits || !groupSize) { throw new Error('matmul: quantW weights need scales, bits and groupSize'); @@ -43,6 +43,9 @@ export function matmul(a: Value, b: Value, opts: MatmulOpts = {}): Value { scaleBase: opts.scaleBase ?? 0, bits, groupSize, + bias: opts.bias, + addend: opts.addend, + activation: opts.activation, }); } if (opts.scales || opts.bits || opts.groupSize) { @@ -93,15 +96,24 @@ export function matmul(a: Value, b: Value, opts: MatmulOpts = {}): Value { ]); } -/** matmul's quantW path. bits and groupSize travel as attrs, so one pipeline - * serves every quant format. */ +/** matmul's quantW path. The result takes a's dtype. An f16 a, a bias, an + * addend or an activation need the tiled kernel: the whole weight (baseRow + * and scaleBase 0) and a shape it supports. */ function matmulQuantW( a: Value, w: Value, scales: Value, - opts: { baseRow: number; scaleBase: number; bits: 8 | 4; groupSize: number }, + opts: { + baseRow: number; + scaleBase: number; + bits: 8 | 4; + groupSize: number; + bias?: Value; + addend?: Value; + activation?: SlotAct; + }, ): Value { - const [m, k] = want2d('matmul', a, ['f32']); + const [m, k, dtype] = want2d('matmul', a, FLOAT_DTYPES); const wd = w.shape.dims; if (!wd || wd.length !== 2) { throw new Error('matmul: quantW weight must be 2D'); @@ -135,12 +147,43 @@ function matmulQuantW( } if (!Number.isInteger(scaleBase) || scaleBase < 0) throw new Error(`matmul: bad scaleBase ${scaleBase}`); + const { bias, addend, activation } = opts; + const whole = baseRow === 0 && scaleBase === 0 && r === k; + if (whole && scales.shape.elems !== (k / groupSize) * n) { + throw new Error(`matmul: scales must be f32 [K/groupSize, N] = [${k / groupSize}, ${n}]`); + } + if (dtype === 'f16' || bias || addend || activation) { + if (!whole) { + throw new Error( + 'matmul: an f16 input or an epilogue on a quantW weight needs the whole weight', + ); + } + if (!matmulQuantTiledEligible(k, n, groupSize)) { + throw new Error( + `matmul: an f16 input or an epilogue on a quantW weight needs K a multiple of ${MATMUL_QUANT_TILED_BK} and N and groupSize multiples of 4, got K ${k}, N ${n}, groupSize ${groupSize}`, + ); + } + } + if (bias && (bias.shape.dtype !== dtype || bias.shape.elems !== n)) { + throw new Error(`matmul: bias must be ${dtype} [N] = ${n} elems`); + } + if (addend && (addend.shape.dtype !== dtype || addend.shape.elems !== m * n)) { + throw new Error(`matmul: addend must be ${dtype} [M,N] = ${m * n} elems`); + } return pending( - { elems: m * n, dtype: 'f32', dims: [m, n] }, + { elems: m * n, dtype, dims: [m, n] }, 'matmulQuantW', - [a, w, scales], + [a, w, scales, ...(bias ? [bias] : []), ...(addend ? [addend] : [])], undefined, - [baseRow, scaleBase, bits, groupSize], + [ + baseRow, + scaleBase, + bits, + groupSize, + bias ? 1 : 0, + addend ? 1 : 0, + activation ? ACT_CODE[activation] : 0, + ], ); } diff --git a/packages/runntime/src/core/kernels/matmul/matmulQuantTiled.ts b/packages/runntime/src/core/kernels/matmul/matmulQuantTiled.ts new file mode 100644 index 0000000..d4b5409 --- /dev/null +++ b/packages/runntime/src/core/kernels/matmul/matmulQuantTiled.ts @@ -0,0 +1,223 @@ +/** Matmul with a quantW weight: out = act(a·W + bias) + addend, in 32×32 + * output tiles. W holds int8 or int4 values in the blocked-unit layout, times + * f32 group scales. a, bias, addend and out are f16 or f32; sums are f32. + * + * Each step unpacks a slice of W into workgroup memory as floats, so the + * inner loop is a plain float matmul with matmulTiledVec4's tiling. */ + +import tgpu, { d, std } from 'typegpu'; +import type { TgpuRoot } from 'typegpu'; +import { + type F32Buffer, + type FloatBuffer, + type KernelHandle, + makeHandle, + type U32Buffer, +} from '../common.ts'; +import { cachedBindGroup, cachedUniform } from '../../gpu/dispatchCache.ts'; +import { activationFor, actSlot } from '../activations.ts'; +import { type Elem, F32_ELEM } from '../elem.ts'; +import { columnWord, deqRow } from '../quantCommon.ts'; + +const Dims = d.struct({ m: d.u32 }); + +function makeLayout(elem: Elem) { + return tgpu.bindGroupLayout({ + a: { storage: d.arrayOf(elem.vec4), access: 'readonly' }, + w: { storage: d.arrayOf(d.u32), access: 'readonly' }, + scales: { storage: d.arrayOf(d.vec4f), access: 'readonly' }, // [K/group, N/4] + bias: { storage: d.arrayOf(elem.vec4), access: 'readonly' }, // [N/4]; dead when hasBias=0 + addend: { storage: d.arrayOf(elem.vec4), access: 'readonly' }, // [M,N/4]; dead when hasAdd=0 + out: { storage: d.arrayOf(elem.vec4), access: 'mutable' }, + dims: { uniform: Dims }, + }); +} + +const layouts = new Map>(); +function layoutFor(elem: Elem): ReturnType { + let l = layouts.get(elem.key); + if (!l) { + l = makeLayout(elem); + layouts.set(elem.key, l); + } + return l; +} + +// Each workgroup computes a BM×BN tile of out, each thread TM rows × TN +// groups of 4 columns. +const BM = 32; +const BN = 32; +const TM = 4; +const TN = 1; + +/** K per step of the tile loop: K must be a multiple of it. */ +export const MATMUL_QUANT_TILED_BK = 8; +const BK = MATMUL_QUANT_TILED_BK; + +/** Whether a [M,K]·[K,N] quantW matmul can use this kernel. */ +export function matmulQuantTiledEligible(k: number, n: number, groupSize: number): boolean { + return k % BK === 0 && n % 4 === 0 && groupSize % 4 === 0 && k % groupSize === 0; +} + +export interface MatmulQuantTiledCfg { + k: number; + n: number; + bits: number; + groupSize: number; + hasBias: number; + hasAdd: number; + act?: number; +} + +function makeMatmulQuantTiledKernel(cfg: MatmulQuantTiledCfg, elem: Elem) { + const { k, n, bits, groupSize, hasBias, hasAdd } = cfg; + const layout = layoutFor(elem); + const storeVec4 = elem.vec4; + const bk4 = BK / 4; + const bn4 = BN / 4; + // Threads per row of the tile, each owning TN units spaced colThreads apart. + const colThreads = bn4 / TN; + const lanes = colThreads * (BM / TM); + const k4 = k / 4; + const n4 = n / 4; + const aVecs = Math.ceil((BM * bk4) / lanes); + const bUnits = Math.ceil((BK * bn4) / lanes); + const tiles = k / BK; + // Words per packed 4×4 block of W. + const tileWords = bits / 2; + const cw1 = columnWord(1, bits); + const cw2 = columnWord(2, bits); + const cw3 = columnWord(3, bits); + const tileA = tgpu.workgroupVar(d.arrayOf(elem.vec4, BM * bk4)); + const tileB = tgpu.workgroupVar(d.arrayOf(elem.vec4, BK * bn4)); + return tgpu.computeFn({ + in: { lid: d.builtin.localInvocationIndex, wid: d.builtin.workgroupId }, + workgroupSize: [lanes], + })(({ lid, wid }) => { + 'use gpu'; + const m = layout.$.dims.m; + const tx = lid % colThreads; + const ty = std.intdiv(lid, colThreads); + const mBase = wid.y * BM; + const nBase4 = wid.x * bn4; + + const accs = d.arrayOf(d.vec4f, TM * TN)(); + const bv = d.arrayOf(d.vec4f, 4 * TN)(); + for (const t of std.range(tiles)) { + for (const i of tgpu.unroll(std.range(aVecs))) { + const e = lid + i * lanes; + if (e < BM * bk4) { + const am = mBase + std.intdiv(e, bk4); + let v = storeVec4(); + if (am < m) { + v = storeVec4(layout.$.a[am * k4 + t * bk4 + (e % bk4)]!); + } + tileA.$[e] = storeVec4(v); + } + } + for (const i of tgpu.unroll(std.range(bUnits))) { + const e = lid + i * lanes; + if (e < BK * bn4) { + const bk = t * BK + std.intdiv(e, bn4); + const bc4 = nBase4 + (e % bn4); + let wv = storeVec4(); + if (bc4 < n4) { + const wb = (bc4 * k4 + std.intdiv(bk, 4)) * tileWords; + const q = deqRow( + layout.$.w[wb]!, + layout.$.w[wb + cw1]!, + layout.$.w[wb + cw2]!, + layout.$.w[wb + cw3]!, + bk % 4, + bits, + ); + wv = storeVec4(q * layout.$.scales[std.intdiv(bk, groupSize) * n4 + bc4]!); + } + tileB.$[e] = storeVec4(wv); + } + } + std.workgroupBarrier(); + + for (const kv of tgpu.unroll(std.range(bk4))) { + for (const j of tgpu.unroll(std.range(TN))) { + for (const r of tgpu.unroll(std.range(4))) { + bv[j * 4 + r] = d.vec4f(tileB.$[(kv * 4 + r) * bn4 + tx + j * colThreads]!); + } + } + for (const i of tgpu.unroll(std.range(TM))) { + const av = d.vec4f(tileA.$[(ty * TM + i) * bk4 + kv]!); + for (const j of tgpu.unroll(std.range(TN))) { + accs[i * TN + j] = + accs[i * TN + j]! + + av.x * bv[j * 4]! + + av.y * bv[j * 4 + 1]! + + av.z * bv[j * 4 + 2]! + + av.w * bv[j * 4 + 3]!; + } + } + } + std.workgroupBarrier(); + } + + for (const j of tgpu.unroll(std.range(TN))) { + const c4 = nBase4 + tx + j * colThreads; + if (c4 < n4) { + let bias = d.vec4f(); + if (hasBias > 0) { + bias = d.vec4f(layout.$.bias[c4]!); + } + for (const i of tgpu.unroll(std.range(TM))) { + const row = mBase + ty * TM + i; + if (row < m) { + let v = actSlot.$(accs[i * TN + j]! + bias); + if (hasAdd > 0) { + v += d.vec4f(layout.$.addend[row * n4 + c4]!); + } + layout.$.out[row * n4 + c4] = storeVec4(v); + } + } + } + } + }); +} + +/** One pipeline per dtype, K, N, bits, group size, epilogue and activation. */ +export function createMatmulQuantTiledPipeline( + root: TgpuRoot, + cfg: MatmulQuantTiledCfg, + elem: Elem = F32_ELEM, +) { + return root + .with(actSlot, activationFor(cfg.act)) + .createComputePipeline({ compute: makeMatmulQuantTiledKernel(cfg, elem) }); +} + +/** One workgroup per 32×32 output tile. */ +export function matmulQuantTiledHandle( + root: TgpuRoot, + pipeline: ReturnType, + m: number, + n: number, + buffers: { + a: FloatBuffer; + w: U32Buffer; + scales: F32Buffer; + out: FloatBuffer; + bias?: FloatBuffer; + addend?: FloatBuffer; + }, + elem: Elem = F32_ELEM, +): KernelHandle { + const dims = cachedUniform(root, Dims, { m }); + // Absent epilogue operands alias `a`; those pipelines never read them. + const bindGroup = cachedBindGroup(root, layoutFor(elem), { + ...buffers, + bias: buffers.bias ?? buffers.a, + addend: buffers.addend ?? buffers.a, + dims, + }); + return makeHandle(pipeline, 'matmulQuantTiled', bindGroup, [ + Math.ceil(n / BN), + Math.ceil(m / BM), + ]); +} diff --git a/packages/runntime/src/core/kernels/quantCommon.ts b/packages/runntime/src/core/kernels/quantCommon.ts index 1d88462..739716c 100644 --- a/packages/runntime/src/core/kernels/quantCommon.ts +++ b/packages/runntime/src/core/kernels/quantCommon.ts @@ -1,7 +1,8 @@ import { d, std } from 'typegpu'; /** The physical layout contract for quantW weight buffers, shared by the - * packer and both quant matmuls. Change it in one place or not at all. + * packer and the matmulQuantW, moeQuantW and matmulQuantTiled kernels. + * Change it in one place or not at all. * * The logical [E·K, N] weight is stored as 16-value tiles, one per (expert, * 4-column unit, 4-row chunk). Within a tile values are column-major and @@ -12,9 +13,11 @@ import { d, std } from 'typegpu'; * The point: one thread owning a 4-column unit streams its weight run * K-contiguously. Scales and biases are unaffected. */ +/** The four values of one column of a tile, from the word that holds it and + * the column's bit offset in that word. */ export const deqVec4 = (word: number, off: number, bits: number) => { 'use gpu'; - const w = std.bitcastU32toI32(word); + const w = std.bitcast(d.u32, d.i32)(word); return d.vec4f( d.f32(std.extractBits(w, off, bits)), d.f32(std.extractBits(w, off + bits, bits)), @@ -22,3 +25,19 @@ export const deqVec4 = (word: number, off: number, bits: number) => { d.f32(std.extractBits(w, off + 3 * bits, bits)), ); }; + +/** The word of a tile that holds column c, from its first word. */ +export const columnWord = (c: number, bits: number) => Math.floor((c * 4 * bits) / 32); + +/** Row r of a tile, one value per column, from the words that hold columns + * 0 to 3. */ +export const deqRow = (c0: number, c1: number, c2: number, c3: number, r: number, bits: number) => { + 'use gpu'; + const off = r * bits; + return d.vec4f( + d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(c0), off, bits)), + d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(c1), (4 * bits + off) % 32, bits)), + d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(c2), (8 * bits + off) % 32, bits)), + d.f32(std.extractBits(std.bitcast(d.u32, d.i32)(c3), (12 * bits + off) % 32, bits)), + ); +}; diff --git a/packages/runntime/src/core/nn/module.ts b/packages/runntime/src/core/nn/module.ts index 9393003..42e262f 100644 --- a/packages/runntime/src/core/nn/module.ts +++ b/packages/runntime/src/core/nn/module.ts @@ -190,11 +190,11 @@ export abstract class Module { } /** Stores every float parameter and buffer as f16; reducing kernels still - * accumulate in f32. + * accumulate in f32. quantW weights and their pinned f32 scales keep their + * storage. * * Call it before loadStateDict — afterwards the dtype is baked into every - * pipeline that read the buffer, so it throws. Also throws on quantW - * weights, whose kernels need f32. */ + * pipeline that read the buffer, so it throws. */ half(root?: TgpuRoot): this { // Graph-only contexts have no device and never compile a shader; let the // first real dispatch fail there instead. @@ -205,13 +205,6 @@ export abstract class Module { "half(): this device has no shader-f16 — check supportsF16() first, and request it with tgpu.init({ device: { optionalFeatures: ['shader-f16'] } })", ); } - for (const [name, p] of this.namedParameters()) { - if (p.shape.dtype === 'quantW') { - throw new Error( - `half(): '${name}' is a quantized (quantW) weight — quantized models don't support f16 activations, leave this model at f32`, - ); - } - } retypeFloats(this, 'f32', 'f16'); return this; } diff --git a/packages/runntime/test/core/graph/ops/matmul.test.ts b/packages/runntime/test/core/graph/ops/matmul.test.ts index e1e682f..b81ec91 100644 --- a/packages/runntime/test/core/graph/ops/matmul.test.ts +++ b/packages/runntime/test/core/graph/ops/matmul.test.ts @@ -108,14 +108,28 @@ describe('matmul', () => { describe('quantized weights', () => { const a = input([2, 8]); - const q = { scales: input([16]), bits: 8 as const, groupSize: 4 }; + // K 8 in groups of 4: two scale rows of 16 columns. + const q = { scales: input([2, 16]), bits: 8 as const, groupSize: 4 }; it.each([8, 4] as const)('int%i gives f32 [M, N]', (bits) => { const c = matmul(a, quantWeight(8, 16, bits), { ...q, bits }); expect(c.shape).toEqual(meta([2, 16])); expect(c.pending?.op).toBe('matmulQuantW'); - // attrs: baseRow, scaleBase, bits, groupSize - expect(c.pending?.attrs).toEqual([0, 0, bits, 4]); + // attrs: baseRow, scaleBase, bits, groupSize, hasBias, hasAdd, act + expect(c.pending?.attrs).toEqual([0, 0, bits, 4, 0, 0, 0]); + }); + + it('gives f16 [M, N] for f16 activations', () => { + const c = matmul(input([2, 8], 'f16'), quantWeight(8, 16, 8), q); + expect(c.shape).toEqual(meta([2, 16], 'f16')); + }); + + it('takes a bias, an addend and an activation on the whole weight', () => { + const bias = input([16]); + const addend = input([2, 16]); + const c = matmul(a, quantWeight(8, 16, 8), { ...q, bias, addend, activation: 'silu' }); + expect(c.pending?.inputs.slice(3)).toEqual([bias, addend]); + expect(c.pending?.attrs?.slice(4)).toEqual([1, 1, ACT_CODE.silu]); }); it('reads one expert of a fused weight with a K-aligned baseRow', () => { @@ -124,10 +138,16 @@ describe('matmul', () => { it.each([ { - name: 'an epilogue', + name: 'an epilogue on one expert', + w: quantWeight(24, 16, 8), + opts: { ...q, baseRow: 16, bias: input([16]) }, + error: /an epilogue on a quantW weight needs the whole weight/, + }, + { + name: 'a bias of the wrong width', w: quantWeight(8, 16, 8), - opts: { ...q, bias: input([16]) }, - error: /not supported for quantW/, + opts: { ...q, bias: input([8]) }, + error: /bias must be f32 \[N\] = 16 elems/, }, { name: 'missing scales', @@ -174,11 +194,37 @@ describe('matmul', () => { expect(() => matmul(a, w, q)).toThrow(/elems 16 != rows\*cols\/4 = 32/); }); - it('needs f32 activations', () => { - expect(() => matmul(input([2, 8], 'f16'), quantWeight(8, 16, 8), q)).toThrow( - /expected a 2D f32 tensor/, + it('needs the whole weight for f16 activations', () => { + expect(() => + matmul(input([2, 8], 'f16'), quantWeight(24, 16, 8), { ...q, baseRow: 16 }), + ).toThrow(/an f16 input or an epilogue on a quantW weight needs the whole weight/); + }); + + it('rejects scales that do not match the whole weight', () => { + expect(() => matmul(a, quantWeight(8, 16, 8), { ...q, scales: input([16]) })).toThrow( + /scales must be f32 \[K\/groupSize, N\] = \[2, 16\]/, ); }); + + describe('with K not a multiple of the tiled K step', () => { + // K 12 in groups of 4: three scale rows of 16 columns. + const q12 = { scales: input([3, 16]), bits: 8 as const, groupSize: 4 }; + const w = quantWeight(12, 16, 8); + + it('runs f32 with no epilogue', () => { + expect(matmul(input([2, 12]), w, q12).shape).toEqual(meta([2, 16])); + }); + + it.each([ + { name: 'an f16 input', a: input([2, 12], 'f16'), opts: q12 }, + { name: 'a bias', a: input([2, 12]), opts: { ...q12, bias: input([16]) } }, + { name: 'an activation', a: input([2, 12]), opts: { ...q12, activation: 'gelu' as const } }, + ])('rejects $name', ({ a, opts }) => { + expect(() => matmul(a, w, opts)).toThrow( + /needs K a multiple of 8 and N and groupSize multiples of 4, got K 12, N 16, groupSize 4/, + ); + }); + }); }); }); diff --git a/packages/runntime/test/core/kernels/matmul/matmulQuantTiled.gpu.test.ts b/packages/runntime/test/core/kernels/matmul/matmulQuantTiled.gpu.test.ts new file mode 100644 index 0000000..a19238e --- /dev/null +++ b/packages/runntime/test/core/kernels/matmul/matmulQuantTiled.gpu.test.ts @@ -0,0 +1,93 @@ +import { describe, it } from 'vitest'; +import { toArray } from '../../../../src/core/gpu/eval.ts'; +import { matmul, type SlotAct } from '../../../../src/core/graph/ops.ts'; +import { matrix } from '../../../../src/core/graph/value.ts'; +import { expectClose, gpu, randomF32, requireFeature, roundF16, upload } from '../../../gpu.ts'; +import { ACTIVATION_REF } from './checkMatmul.ts'; +import { randomQuantWeight } from './quantW.ts'; + +const { root, ex } = await gpu(); + +interface Case { + m: number; + k: number; + n: number; + bits: 4 | 8; + groupSize: number; + bias?: boolean; + addend?: boolean; + activation?: SlotAct; +} + +/** act(a · dequant(W) + bias) + addend on the whole quantW weight, through the + * tiled route: an f16 input, an epilogue, or at least 8 rows. */ +async function checkQuantTiled(c: Case, dtype: 'f32' | 'f16') { + const { m, k, n, bits, groupSize } = c; + const round = dtype === 'f16' ? roundF16 : (v: Float32Array) => v; + const meta = (rows: number, cols: number) => ({ ...matrix(rows, cols), dtype }); + const weight = randomQuantWeight(root, { experts: 1, k, n, bits, groupSize }, 2); + const a = round(randomF32(m * k, 1)); + const bias = c.bias ? round(randomF32(n, 3)) : undefined; + const addend = c.addend ? round(randomF32(m * n, 4)) : undefined; + const got = await toArray( + matmul(upload(root, a, meta(m, k)), weight.w, { + scales: upload(root, weight.scales, [k / groupSize, n]), + bits, + groupSize, + bias: bias && upload(root, bias, meta(1, n)), + addend: addend && upload(root, addend, meta(m, n)), + activation: c.activation, + }), + ex, + ); + const want = new Float64Array(m * n); + for (let i = 0; i < m; i++) { + for (let col = 0; col < n; col++) { + let sum = bias ? bias[col]! : 0; + for (let j = 0; j < k; j++) sum += a[i * k + j]! * weight.dequant[j * n + col]!; + if (c.activation) sum = ACTIVATION_REF[c.activation](sum); + if (addend) sum += addend[i * n + col]!; + want[i * n + col] = sum; + } + } + // The f16 route stores dequantized weights in f16 before the f32 sums. + const tol = + dtype === 'f16' ? { atol: 2e-3 * Math.sqrt(k), rtol: 5e-3 } : { atol: 1e-6 * k, rtol: 1e-5 }; + expectClose(got, want, tol); +} + +describe('matmulQuantTiled kernel', () => { + // tiled: a whole weight from 8 rows, K % 8 == 0 + const SHAPES = [ + { name: 'one tile', m: 8, k: 8, n: 8, groupSize: 4 }, + { name: 'partial row and column tiles', m: 37, k: 48, n: 40, groupSize: 8 }, + { name: 'one group per column', m: 9, k: 64, n: 16, groupSize: 64 }, + { name: 'a fused q|k projection', m: 97, k: 1024, n: 2304, groupSize: 32 }, + { name: 'a per-channel down projection', m: 40, k: 2048, n: 1024, groupSize: 2048 }, + ]; + it.for( + (['f32', 'f16'] as const).flatMap((dtype) => + ([4, 8] as const).flatMap((bits) => SHAPES.map((s) => ({ dtype, bits, ...s }))), + ), + )('$dtype int$bits $name: [$m, $k] × [$k, $n], groupSize $groupSize', async (c, ctx) => { + if (c.dtype === 'f16') requireFeature(ctx, root, 'shader-f16'); + await checkQuantTiled(c, c.dtype); + }); + + // tiled: an f16 input or an epilogue on a whole weight, whatever m + it.for( + (['f32', 'f16'] as const).flatMap((dtype) => + ([4, 8] as const).flatMap((bits) => + [ + { name: 'bias', bias: true }, + { name: 'addend', addend: true }, + { name: 'silu', activation: 'silu' as const }, + { name: 'bias, gelu and addend', bias: true, activation: 'gelu' as const, addend: true }, + ].map((e) => ({ dtype, bits, ...e })), + ), + ), + )('$dtype int$bits with $name on 3 rows', async ({ dtype, ...e }, ctx) => { + if (dtype === 'f16') requireFeature(ctx, root, 'shader-f16'); + await checkQuantTiled({ m: 3, k: 64, n: 24, groupSize: 16, ...e }, dtype); + }); +}); diff --git a/packages/runntime/test/core/kernels/matmul/matmulQuantW.gpu.test.ts b/packages/runntime/test/core/kernels/matmul/matmulQuantW.gpu.test.ts index 57d6322..c007bad 100644 --- a/packages/runntime/test/core/kernels/matmul/matmulQuantW.gpu.test.ts +++ b/packages/runntime/test/core/kernels/matmul/matmulQuantW.gpu.test.ts @@ -45,25 +45,39 @@ async function checkQuantMatmul({ m, k, n, bits, groupSize, experts = 1, expert } describe('matmulQuantW kernel', () => { - // One route: a thread per (row, 4-column unit) streams its tiles. bits and - // groupSize are pipeline config; K and N are multiples of 4, and of 8 for - // int4. + // Routes are picked in gpu/dispatch/matmul.ts. This kernel takes an f32 + // input with no epilogue on one expert, below 8 rows, or when K is not a + // multiple of 8. Other calls on a whole weight take matmulQuantTiled. K and + // groupSize are multiples of 4, N of 4 for int8 and of 8 for int4. it.for( ([8, 4] as const).flatMap((bits) => [ + // plain: m < 8 { name: 'one tile', m: 1, k: 4, n: 8, groupSize: 4 }, { name: 'one group per column', m: 3, k: 12, n: 8, groupSize: 12 }, { name: 'a group every 4 rows', m: 3, k: 12, n: 16, groupSize: 4 }, { name: 'groups of 8', m: 5, k: 64, n: 24, groupSize: 8 }, - { name: 'rows past a workgroup of units', m: 70, k: 32, n: 8, groupSize: 16 }, { name: 'a wide row', m: 2, k: 640, n: 640, groupSize: 32 }, - // Past GRID_Y_STRIDE units the dispatch goes 2-D. + { name: 'one row short of the tiled route', m: 7, k: 32, n: 8, groupSize: 16 }, + // tiled: a whole weight, m >= 8 and K % 8 == 0 + { name: 'the first m on the tiled route', m: 8, k: 32, n: 8, groupSize: 16 }, + // plain: one expert of a fused weight, whatever m + { + name: 'rows past a workgroup of units on one expert', + m: 70, + k: 32, + n: 8, + groupSize: 16, + experts: 2, + expert: 1, + }, + // plain: K % 8 != 0. Past GRID_Y_STRIDE units the dispatch goes 2-D. { name: 'units past GRID_Y_STRIDE', m: GRID_Y_STRIDE / 2 + 5, k: 4, n: 8, groupSize: 4 }, ].map((c) => ({ bits, ...c })), ), )('int$bits $name: [$m, $k] × [$k, $n], groupSize $groupSize', (c) => checkQuantMatmul(c)); - // baseRow and scaleBase read one expert of a fused [E·K, N] weight. + // plain: baseRow and scaleBase read one expert of a fused [E·K, N] weight. it.for(([8, 4] as const).flatMap((bits) => [0, 1, 2].map((expert) => ({ bits, expert }))))( 'int$bits expert $expert of 3', ({ bits, expert }) => @@ -78,7 +92,11 @@ describe('matmulQuantW kernel', () => { const fused = { experts: 2, k: 8, n: 8, bits: 8 as const, groupSize: 4 }; it.each([ - { name: 'a bias', opts: { bias: true }, err: /epilogue is not supported for quantW/ }, + { + name: 'a bias on one expert', + opts: { bias: true }, + err: /an epilogue on a quantW weight needs the whole weight/, + }, { name: 'no scales', opts: { scales: false }, @@ -132,13 +150,13 @@ describe('matmulQuantW kernel', () => { ); }); - it('refuses an f16 input', (ctx) => { + it('refuses an f16 input on one expert', (ctx) => { requireFeature(ctx, root, 'shader-f16'); const weight = randomQuantWeight(root, fused, 2); const a = upload(root, new Float32Array(16), { ...matrix(2, 8), dtype: 'f16' }); const scales = upload(root, weight.scales, [4, 8]); expect(() => matmul(a, weight.w, { scales, bits: 8, groupSize: 4 })).toThrow( - /expected a 2D f32 tensor/, + /an f16 input or an epilogue on a quantW weight needs the whole weight/, ); }); }); diff --git a/packages/runntime/test/core/nn/module.test.ts b/packages/runntime/test/core/nn/module.test.ts index 521da69..a5ea77a 100644 --- a/packages/runntime/test/core/nn/module.test.ts +++ b/packages/runntime/test/core/nn/module.test.ts @@ -219,11 +219,18 @@ describe('half() and float()', () => { expect(m.half().table.shape.dtype).toBe('f16'); }); - it('refuses a model with quantized weights', () => { + it('keeps quantized weights and pinned scales as they are', () => { class Quant extends Module { w = new Parameter({ elems: 4, dtype: 'quantW', dims: [4, 4] }); + scales = new Parameter(meta([4]), undefined, true); + bias = param(); } - expect(() => new Quant().half()).toThrow(/'w' is a quantized \(quantW\) weight/); + const m = new Quant().half(); + expect([m.w.shape.dtype, m.scales.shape.dtype, m.bias.shape.dtype]).toEqual([ + 'quantW', + 'f32', + 'f16', + ]); }); it('refuses a device without shader-f16', () => {