feat(core): tiled matmul for quantW weights with f16 activations and the epilogue

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),
 

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({

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,
+    ],
   );
 }