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 c4b4e45..fd7919f 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,38 +72,11 @@ 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'; } - forwardGather(x: Value): Value { + override forward(x: Value): Value { const f = this.dims.weightFormat; const xn = this.norm.forward(x); // [T, H] @@ -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 }; -}