From 865b4df4c3aa26f734c235979121d6cf942deb0f Mon Sep 17 00:00:00 2001 From: chmjkb Date: Tue, 27 Oct 2026 11:35:04 +0100 Subject: [PATCH 1/4] refactor(zoo): drop privacy-filter's unused CPU-routed MoE path --- .../runntime/src/zoo/privacy-filter/moeMlp.ts | 85 +----------------- .../src/zoo/privacy-filter/routing.ts | 88 ------------------- 2 files changed, 4 insertions(+), 169 deletions(-) delete mode 100644 packages/runntime/src/zoo/privacy-filter/routing.ts diff --git a/packages/runntime/src/zoo/privacy-filter/moeMlp.ts b/packages/runntime/src/zoo/privacy-filter/moeMlp.ts index c4b4e45..6dfc962 100644 --- a/packages/runntime/src/zoo/privacy-filter/moeMlp.ts +++ b/packages/runntime/src/zoo/privacy-filter/moeMlp.ts @@ -1,8 +1,6 @@ -/** The privacy filter's MoE MLP block, with topk routing on the GPU. - * - * A forward runs in two eval phases: `selection` computes the routing, and the - * caller reads it back and groups it to shape `forward`, which builds the - * expert-compute DAG. `forwardPacked` does both in one DAG instead. +/** The privacy filter's MoE MLP block. Topk routing stays on the GPU: the + * expert matmuls read the selected ids straight from topk's output, so a + * forward is one DAG with no readback. * * The residual is folded in here, so composing callers must not re-add it. * K is assumed equal to H. */ @@ -10,11 +8,7 @@ import { astype, add, - cat, - gatherRows, - matmul, matmulGather, - matrix, moeDown, moeUp, mul, @@ -25,7 +19,6 @@ import { type Value, type ValueMeta, } from '../../core/index.ts'; -import type { Grouping } from './routing.ts'; const RMS_EPS = 1e-5; // legacy rmsnorm kernel epsilon (not the eager 1e-6 default) @@ -36,7 +29,7 @@ export interface MoeDims { groupSize?: number; } -export class EagerMoeMlp extends nn.Module<[Value, Grouping, Value[]], Value> { +export class EagerMoeMlp extends nn.Module { readonly norm: InstanceType; readonly gate: InstanceType; readonly gluWeight: InstanceType; @@ -79,33 +72,6 @@ export class EagerMoeMlp extends nn.Module<[Value, Grouping, Value[]], Value> { } } - routerLogits(x: Value): Value { - return this.gate.forward(this.norm.forward(x)); - } - - selection(x: Value): Value { - return topk(this.routerLogits(x), 4); - } - - private expertMatmul(x: Value, w: Value, scales: Value | undefined, expert: number): Value { - const h = this.dims.hidden; - const f = this.dims.weightFormat; - if (f === 'f32') return matmul(x, w, { baseRow: expert * h }); - if (f === 'quantInt8' || f === 'quantInt4') { - const bits = f === 'quantInt8' ? 8 : 4; - const groupSize = this.dims.groupSize ?? h; - const kg = h / groupSize; - return matmul(x, w, { - scales: scales!, - baseRow: expert * h, - scaleBase: expert * kg * h, - bits, - groupSize, - }); - } - return matmul(x, slice(w, 0, expert * h, (expert + 1) * h)); - } - private get actDtype(): 'f32' | 'f16' { return this.norm.scale.shape.dtype === 'f16' ? 'f16' : 'f32'; } @@ -154,47 +120,4 @@ export class EagerMoeMlp extends nn.Module<[Value, Grouping, Value[]], Value> { groupSize, }); } - - override forward(x: Value, grouping: Grouping, gateColumns: Value[]): Value { - if (!grouping) - throw new Error( - 'EagerMoeMlp.forward needs (x, grouping, gateColumns) — evaluate selection(x) and run parseTopk/groupByExpert first (see routing.ts)', - ); - if (grouping.groups.length === 0) throw new Error('EagerMoeMlp.forward: empty grouping'); - if (gateColumns.length !== grouping.groups.length) - throw new Error( - `EagerMoeMlp.forward: ${gateColumns.length} gate columns for ${grouping.groups.length} expert groups`, - ); - const xn = this.norm.forward(x); - const contribs: Value[] = []; - grouping.groups.forEach((g, gi) => { - const rows = gatherRows(xn, g.tokens); // [n_g, H] - const e = g.expert; - const glu = add( - this.expertMatmul(rows, this.gluWeight.value, this.gluScales?.value, e), - slice(this.gluBias.value, 0, e, e + 1), - ); - const lin = add( - this.expertMatmul(rows, this.linWeight.value, this.linScales?.value, e), - slice(this.linBias.value, 0, e, e + 1), - ); - const act = swiglu(glu, lin); // [n_g, H] - const o = add( - this.expertMatmul(act, this.outWeight.value, this.outScales?.value, e), - slice(this.outBias.value, 0, e, e + 1), - ); - contribs.push(mul(o, gateColumns[gi]!)); // [n_g,1] col-broadcast - }); - const pile = cat(contribs, 0); // [4T, H] in group order - let y: Value | undefined; - for (let s = 0; s < 4; s++) { - const slot = gatherRows(pile, grouping.slotPos[s]!); // [T, H] - y = y ? add(y, slot) : slot; - } - return add(x, y!); // residual folded in - } -} - -export function gateColumnShape(n: number) { - return matrix(n, 1); } diff --git a/packages/runntime/src/zoo/privacy-filter/routing.ts b/packages/runntime/src/zoo/privacy-filter/routing.ts deleted file mode 100644 index 81ee594..0000000 --- a/packages/runntime/src/zoo/privacy-filter/routing.ts +++ /dev/null @@ -1,88 +0,0 @@ -/** CPU-side helpers around the GPU `topk` primitive. Selection MATH happens - * on the GPU (core `topk`: strict-`>` top-4 like torch.topk + softmax over - * the selected logits); its small packed [T,2k] output is read back and - * parsed here, because the expert-compute DAG's structure (which tokens - * gather to which expert) is built in JS. `routeTop4` is the CPU reference - * mirror of the kernel, kept for tests. */ - -export interface Selection { - idx: number[][]; - w: number[][]; -} - -export function parseTopk(packed: Float32Array, tokens: number, k: number): Selection { - if (packed.length !== tokens * 2 * k) { - throw new Error(`parseTopk: packed length ${packed.length} != tokens·2k = ${tokens * 2 * k}`); - } - const idx: number[][] = []; - const w: number[][] = []; - for (let t = 0; t < tokens; t++) { - const base = t * 2 * k; - idx.push(Array.from(packed.subarray(base, base + k))); - w.push(Array.from(packed.subarray(base + k, base + 2 * k))); - } - return { idx, w }; -} - -export function routeTop4(logits: Float32Array, tokens: number, experts: number): Selection { - const idx: number[][] = []; - const w: number[][] = []; - for (let t = 0; t < tokens; t++) { - const vals = [-Infinity, -Infinity, -Infinity, -Infinity]; - const ids = [0, 0, 0, 0]; - for (let e = 0; e < experts; e++) { - const v = logits[t * experts + e]!; - if (v > vals[3]!) { - vals[3] = v; - ids[3] = e; - for (let s = 3; s > 0 && vals[s]! > vals[s - 1]!; s--) { - [vals[s], vals[s - 1]] = [vals[s - 1]!, vals[s]!]; - [ids[s], ids[s - 1]] = [ids[s - 1]!, ids[s]!]; - } - } - } - const m = vals[0]!; - const exps = vals.map((v) => Math.exp(v - m)); - const z = exps.reduce((s, v) => s + v, 0); - idx.push(ids); - w.push(exps.map((v) => v / z)); - } - return { idx, w }; -} - -export interface ExpertGroup { - expert: number; - tokens: number[]; - weights: number[]; -} - -export interface Grouping { - groups: ExpertGroup[]; - slotPos: number[][]; -} - -export function groupByExpert(sel: Selection, experts: number): Grouping { - const tokens = sel.idx.length; - const groups: ExpertGroup[] = []; - const slotPos: number[][] = Array.from({ length: 4 }, () => new Array(tokens).fill(-1)); - let offset = 0; - for (let e = 0; e < experts; e++) { - const g: ExpertGroup = { expert: e, tokens: [], weights: [] }; - for (let t = 0; t < tokens; t++) { - // Top-4 never repeats an expert, so (e, t) matches at most one slot — - // each token contributes at most one row per expert group. - for (let s = 0; s < 4; s++) { - if (sel.idx[t]![s] === e) { - slotPos[s]![t] = offset + g.tokens.length; - g.tokens.push(t); - g.weights.push(sel.w[t]![s]!); - } - } - } - if (g.tokens.length > 0) { - groups.push(g); - offset += g.tokens.length; - } - } - return { groups, slotPos }; -} From 372f8c8de24aff658a947e3bfc2b8295fa526466 Mon Sep 17 00:00:00 2001 From: chmjkb Date: Tue, 27 Oct 2026 11:35:29 +0100 Subject: [PATCH 2/4] refactor(zoo): rename privacy-filter's EagerMoeMlp.forwardGather to forward --- packages/runntime/src/zoo/privacy-filter/model.ts | 2 +- packages/runntime/src/zoo/privacy-filter/moeMlp.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/runntime/src/zoo/privacy-filter/model.ts b/packages/runntime/src/zoo/privacy-filter/model.ts index 0c5e894..5c86c93 100644 --- a/packages/runntime/src/zoo/privacy-filter/model.ts +++ b/packages/runntime/src/zoo/privacy-filter/model.ts @@ -26,7 +26,7 @@ export class EagerBlock extends nn.Module<[Value, Value, Value], Value> { override forward(x: Value, cos: Value, sin: Value): Value { x = this.attn.forward(x, cos, sin); - x = this.mlp.forwardGather(x); + x = this.mlp.forward(x); return x; } } diff --git a/packages/runntime/src/zoo/privacy-filter/moeMlp.ts b/packages/runntime/src/zoo/privacy-filter/moeMlp.ts index 6dfc962..fd7919f 100644 --- a/packages/runntime/src/zoo/privacy-filter/moeMlp.ts +++ b/packages/runntime/src/zoo/privacy-filter/moeMlp.ts @@ -76,7 +76,7 @@ export class EagerMoeMlp extends nn.Module { return this.norm.scale.shape.dtype === 'f16' ? 'f16' : 'f32'; } - forwardGather(x: Value): Value { + override forward(x: Value): Value { const f = this.dims.weightFormat; const xn = this.norm.forward(x); // [T, H] From ffbe339876244ffa5815c10b4f559a32c8649023 Mon Sep 17 00:00:00 2001 From: chmjkb Date: Tue, 27 Oct 2026 13:44:49 +0100 Subject: [PATCH 3/4] test(zoo): check privacy-filter's MoE block against a CPU reference in f32, f16, int8 and int4 --- .../zoo/privacy-filter/moeMlp.gpu.test.ts | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts diff --git a/packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts b/packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts new file mode 100644 index 0000000..95a2a8d --- /dev/null +++ b/packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts @@ -0,0 +1,128 @@ +/** privacy-filter's MoE block against a CPU loop written from the formula: + * RMSNorm, router, top-4 softmax, clamped SwiGLU experts, residual. */ +import { describe, it } from 'vitest'; +import { toArray } from '../../../src/core/gpu/eval.ts'; +import { EagerMoeMlp } from '../../../src/zoo/privacy-filter/moeMlp.ts'; +import { randomQuantWeight } from '../../core/kernels/matmul/quantW.ts'; +import { expectClose, gpu, randomF32, requireFeature, roundF16, upload } from '../../gpu.ts'; + +const { root } = await gpu(); +const E = 8; +const H = 64; +const T = 5; +const G = 16; + +interface Experts { + glu: ArrayLike; + lin: ArrayLike; + out: ArrayLike; +} + +const clampedSwiglu = (glu: number, lin: number) => { + const g = Math.min(glu, 7); + const l = Math.min(Math.max(lin, -7), 7); + return (g / (1 + Math.exp(-1.702 * g))) * (l + 1); +}; + +function reference( + x: Float32Array, + scale: Float32Array, + gateW: Float32Array, + gateB: Float32Array, + w: Experts, + b: Experts, +): Float64Array { + const y = new Float64Array(T * H); + for (let t = 0; t < T; t++) { + const row = x.subarray(t * H, (t + 1) * H); + const ms = row.reduce((s, v) => s + v * v, 0) / H; + const xn = Array.from(row, (v, i) => (v / Math.sqrt(ms + 1e-5)) * scale[i]!); + const logits = Array.from({ length: E }, (_, e) => { + let s = gateB[e]!; + for (let i = 0; i < H; i++) s += xn[i]! * gateW[i * E + e]!; + return s; + }); + const ids = [...logits.keys()].sort((a, c) => logits[c]! - logits[a]! || a - c).slice(0, 4); + const top = ids.map((e) => logits[e]!); + const max = Math.max(...top); + const exps = top.map((v) => Math.exp(v - max)); + const sum = exps.reduce((a, c) => a + c, 0); + ids.forEach((e, slot) => { + const act = new Float64Array(H); + for (let n = 0; n < H; n++) { + let glu = b.glu[e * H + n]!; + let lin = b.lin[e * H + n]!; + for (let k = 0; k < H; k++) { + glu += xn[k]! * w.glu[(e * H + k) * H + n]!; + lin += xn[k]! * w.lin[(e * H + k) * H + n]!; + } + act[n] = clampedSwiglu(glu, lin); + } + for (let n = 0; n < H; n++) { + let o = b.out[e * H + n]!; + for (let k = 0; k < H; k++) o += act[k]! * w.out[(e * H + k) * H + n]!; + y[t * H + n]! += (exps[slot]! / sum) * o; + } + }); + for (let n = 0; n < H; n++) y[t * H + n]! += row[n]!; + } + return y; +} + +describe('privacy-filter EagerMoeMlp', () => { + it.for(['f32', 'f16', 'quantInt8', 'quantInt4'] as const)('%s', async (format, ctx) => { + const half = format === 'f16'; + if (half) requireFeature(ctx, root, 'shader-f16'); + const mlp = new EagerMoeMlp({ + experts: E, + hidden: H, + weightFormat: half ? 'f32' : format, + groupSize: G, + }); + if (half) mlp.half(root); + const dtype = half ? 'f16' : 'f32'; + const round = (a: Float32Array) => (half ? roundF16(a) : a); + const up = (data: Float32Array, dims: number[]) => + upload(root, data, { elems: data.length, dtype, dims }); + const x = round(randomF32(T * H, 1)); + const scale = round(randomF32(H, 2).map((v) => 1 + 0.2 * v)); + const gateW = round(randomF32(H * E, 3)); + const gateB = round(randomF32(E, 4).map((v) => 0.1 * v)); + const b = { + glu: round(randomF32(E * H, 5).map((v) => 0.1 * v)), + lin: round(randomF32(E * H, 6).map((v) => 0.1 * v)), + out: round(randomF32(E * H, 7).map((v) => 0.1 * v)), + }; + mlp.norm.scale.bind(up(scale, [H])); + mlp.gate.weight.bind(up(gateW, [H, E])); + mlp.gate.bias!.bind(up(gateB, [E])); + mlp.gluBias.bind(up(b.glu, [E, H])); + mlp.linBias.bind(up(b.lin, [E, H])); + mlp.outBias.bind(up(b.out, [E, H])); + let w: Experts; + if (format === 'f32' || half) { + const make = (seed: number) => round(randomF32(E * H * H, seed).map((v) => 0.2 * v)); + const [glu, lin, out] = [make(8), make(9), make(10)]; + mlp.gluWeight.bind(up(glu, [E * H, H])); + mlp.linWeight.bind(up(lin, [E * H, H])); + mlp.outWeight.bind(up(out, [E * H, H])); + w = { glu, lin, out }; + } else { + const bits = format === 'quantInt8' ? 8 : 4; + const make = (seed: number) => + randomQuantWeight(root, { experts: E, k: H, n: H, bits, groupSize: G }, seed); + const [glu, lin, out] = [make(8), make(9), make(10)]; + const scaleShape = [(E * H) / G, H] as const; + mlp.gluWeight.bind(glu.w); + mlp.linWeight.bind(lin.w); + mlp.outWeight.bind(out.w); + mlp.gluScales!.bind(upload(root, glu.scales, scaleShape)); + mlp.linScales!.bind(upload(root, lin.scales, scaleShape)); + mlp.outScales!.bind(upload(root, out.scales, scaleShape)); + w = { glu: glu.dequant, lin: lin.dequant, out: out.dequant }; + } + const got = await toArray(mlp.forward(up(x, [T, H]))); + const tol = half ? 3e-2 : 2e-3; + expectClose(got, reference(x, scale, gateW, gateB, w, b), { atol: tol, rtol: tol }); + }); +}); From 431f3e1bb8fe40de88f4c778a115cb75094e38f3 Mon Sep 17 00:00:00 2001 From: chmjkb Date: Tue, 27 Oct 2026 14:26:08 +0100 Subject: [PATCH 4/4] test(zoo): drop privacy-filter's MoE block test from this PR --- .../zoo/privacy-filter/moeMlp.gpu.test.ts | 128 ------------------ 1 file changed, 128 deletions(-) delete mode 100644 packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts diff --git a/packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts b/packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts deleted file mode 100644 index 95a2a8d..0000000 --- a/packages/runntime/test/zoo/privacy-filter/moeMlp.gpu.test.ts +++ /dev/null @@ -1,128 +0,0 @@ -/** privacy-filter's MoE block against a CPU loop written from the formula: - * RMSNorm, router, top-4 softmax, clamped SwiGLU experts, residual. */ -import { describe, it } from 'vitest'; -import { toArray } from '../../../src/core/gpu/eval.ts'; -import { EagerMoeMlp } from '../../../src/zoo/privacy-filter/moeMlp.ts'; -import { randomQuantWeight } from '../../core/kernels/matmul/quantW.ts'; -import { expectClose, gpu, randomF32, requireFeature, roundF16, upload } from '../../gpu.ts'; - -const { root } = await gpu(); -const E = 8; -const H = 64; -const T = 5; -const G = 16; - -interface Experts { - glu: ArrayLike; - lin: ArrayLike; - out: ArrayLike; -} - -const clampedSwiglu = (glu: number, lin: number) => { - const g = Math.min(glu, 7); - const l = Math.min(Math.max(lin, -7), 7); - return (g / (1 + Math.exp(-1.702 * g))) * (l + 1); -}; - -function reference( - x: Float32Array, - scale: Float32Array, - gateW: Float32Array, - gateB: Float32Array, - w: Experts, - b: Experts, -): Float64Array { - const y = new Float64Array(T * H); - for (let t = 0; t < T; t++) { - const row = x.subarray(t * H, (t + 1) * H); - const ms = row.reduce((s, v) => s + v * v, 0) / H; - const xn = Array.from(row, (v, i) => (v / Math.sqrt(ms + 1e-5)) * scale[i]!); - const logits = Array.from({ length: E }, (_, e) => { - let s = gateB[e]!; - for (let i = 0; i < H; i++) s += xn[i]! * gateW[i * E + e]!; - return s; - }); - const ids = [...logits.keys()].sort((a, c) => logits[c]! - logits[a]! || a - c).slice(0, 4); - const top = ids.map((e) => logits[e]!); - const max = Math.max(...top); - const exps = top.map((v) => Math.exp(v - max)); - const sum = exps.reduce((a, c) => a + c, 0); - ids.forEach((e, slot) => { - const act = new Float64Array(H); - for (let n = 0; n < H; n++) { - let glu = b.glu[e * H + n]!; - let lin = b.lin[e * H + n]!; - for (let k = 0; k < H; k++) { - glu += xn[k]! * w.glu[(e * H + k) * H + n]!; - lin += xn[k]! * w.lin[(e * H + k) * H + n]!; - } - act[n] = clampedSwiglu(glu, lin); - } - for (let n = 0; n < H; n++) { - let o = b.out[e * H + n]!; - for (let k = 0; k < H; k++) o += act[k]! * w.out[(e * H + k) * H + n]!; - y[t * H + n]! += (exps[slot]! / sum) * o; - } - }); - for (let n = 0; n < H; n++) y[t * H + n]! += row[n]!; - } - return y; -} - -describe('privacy-filter EagerMoeMlp', () => { - it.for(['f32', 'f16', 'quantInt8', 'quantInt4'] as const)('%s', async (format, ctx) => { - const half = format === 'f16'; - if (half) requireFeature(ctx, root, 'shader-f16'); - const mlp = new EagerMoeMlp({ - experts: E, - hidden: H, - weightFormat: half ? 'f32' : format, - groupSize: G, - }); - if (half) mlp.half(root); - const dtype = half ? 'f16' : 'f32'; - const round = (a: Float32Array) => (half ? roundF16(a) : a); - const up = (data: Float32Array, dims: number[]) => - upload(root, data, { elems: data.length, dtype, dims }); - const x = round(randomF32(T * H, 1)); - const scale = round(randomF32(H, 2).map((v) => 1 + 0.2 * v)); - const gateW = round(randomF32(H * E, 3)); - const gateB = round(randomF32(E, 4).map((v) => 0.1 * v)); - const b = { - glu: round(randomF32(E * H, 5).map((v) => 0.1 * v)), - lin: round(randomF32(E * H, 6).map((v) => 0.1 * v)), - out: round(randomF32(E * H, 7).map((v) => 0.1 * v)), - }; - mlp.norm.scale.bind(up(scale, [H])); - mlp.gate.weight.bind(up(gateW, [H, E])); - mlp.gate.bias!.bind(up(gateB, [E])); - mlp.gluBias.bind(up(b.glu, [E, H])); - mlp.linBias.bind(up(b.lin, [E, H])); - mlp.outBias.bind(up(b.out, [E, H])); - let w: Experts; - if (format === 'f32' || half) { - const make = (seed: number) => round(randomF32(E * H * H, seed).map((v) => 0.2 * v)); - const [glu, lin, out] = [make(8), make(9), make(10)]; - mlp.gluWeight.bind(up(glu, [E * H, H])); - mlp.linWeight.bind(up(lin, [E * H, H])); - mlp.outWeight.bind(up(out, [E * H, H])); - w = { glu, lin, out }; - } else { - const bits = format === 'quantInt8' ? 8 : 4; - const make = (seed: number) => - randomQuantWeight(root, { experts: E, k: H, n: H, bits, groupSize: G }, seed); - const [glu, lin, out] = [make(8), make(9), make(10)]; - const scaleShape = [(E * H) / G, H] as const; - mlp.gluWeight.bind(glu.w); - mlp.linWeight.bind(lin.w); - mlp.outWeight.bind(out.w); - mlp.gluScales!.bind(upload(root, glu.scales, scaleShape)); - mlp.linScales!.bind(upload(root, lin.scales, scaleShape)); - mlp.outScales!.bind(upload(root, out.scales, scaleShape)); - w = { glu: glu.dequant, lin: lin.dequant, out: out.dequant }; - } - const got = await toArray(mlp.forward(up(x, [T, H]))); - const tol = half ? 3e-2 : 2e-3; - expectClose(got, reference(x, scale, gateW, gateB, w, b), { atol: tol, rtol: tol }); - }); -});