From 66033a0e357f612dcd733b7c5e76ee333fb7adfc Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 11 Sep 2026 21:21:07 +0800 Subject: [PATCH 1/2] feat(ml): speculative decoding across the tier ladder SpeculativeModel composes a drafter and a verifier IMFModel: the drafter proposes blockSize tokens, the verifier decides them all in one full-sequence pass. Greedy verification is output-preserving, so the result equals the verifier's own greedy regardless of drafter quality; the drafter only buys speed. - IMFModel gains the mechanics: encode (shared normalization entry), draft (greedy continuation from a forced prefix), review (per-block positional argmax verdicts). translate() is unchanged. - acceptBlock is a pure function: leading matches accepted, the verifier's pick returned at the first divergence. - Tests: pure acceptance rule on synthetic logits, tiny-fixture session with drafter == verifier (full acceptance, output identical to translate), and an opt-in real-pair e2e (SECRYST_SPEC_E2E=1, layerdrop-int4 -> small-2.1-int8; measured acceptance > 0.9 on the shipped artifacts). --- src/ml/imf/index.ts | 8 ++- src/ml/imf/model.ts | 121 ++++++++++++++++++++++++++++++++++-- src/ml/imf/speculative.ts | 127 ++++++++++++++++++++++++++++++++++++++ test/speculative.test.ts | 95 ++++++++++++++++++++++++++++ 4 files changed, 344 insertions(+), 7 deletions(-) create mode 100644 src/ml/imf/speculative.ts create mode 100644 test/speculative.test.ts diff --git a/src/ml/imf/index.ts b/src/ml/imf/index.ts index 74dc4a4..6619d02 100644 --- a/src/ml/imf/index.ts +++ b/src/ml/imf/index.ts @@ -1,6 +1,12 @@ /** @interscript/ml — the IMF v1 runtime for TypeScript. */ -export { IMFModel } from "./model.js" +export { IMFModel, type DecodeOptions } from "./model.js" +export { + SpeculativeModel, + acceptBlock, + type SpeculativeOptions, + type SpeculativeStats, +} from "./speculative.js" export { IMFError, parseManifest, verifyAndRead, type IMFManifest } from "./loader.js" export { resolve, diff --git a/src/ml/imf/model.ts b/src/ml/imf/model.ts index 158081d..5ade2c6 100644 --- a/src/ml/imf/model.ts +++ b/src/ml/imf/model.ts @@ -73,17 +73,120 @@ export class IMFModel { } async translate(text: string, maxLen = 256, opts: DecodeOptions = {}): Promise { - // models train on stripped input: normalize by default (TODO.client-work 02) - const normalized = opts.raw === true ? text : normalizeArabicInput(text) - const ids = encode(normalized) - if (ids.length === 1) return "" - const hidden = await this.runEncoder(ids) + const hidden = await this.encode(text, opts) + if (!hidden) return "" const tokens = this.kv ? await this.greedyKv(hidden, maxLen, opts) : await this.greedyPlain(hidden, maxLen, opts) return decode(tokens) } + /** Normalized text -> encoder hidden states; null when the input + * carries no decodable content. Shared entry for composition + * strategies (e.g. speculative decode). */ + async encode(text: string, opts: DecodeOptions = {}): Promise { + const normalized = opts.raw === true ? text : normalizeArabicInput(text) + const ids = encode(normalized) + if (ids.length === 1) return null + return this.runEncoder(ids) + } + + /** Greedy continuation of `k` tokens conditioned on `prefix` (the + * verifier-authoritative output so far). A trailing EOS is included + * so callers can verify the stop decision. */ + async draft(hidden: Tensor, prefix: readonly number[], k: number): Promise { + if (this.kv) { + const out: number[] = [] + let current = [PAD_ID, ...prefix] + let present: ReadonlyMap | undefined + while (out.length < k) { + const outputs = await this.decoder.run({ + input_ids: { + name: "input_ids", + type: "int64", + data: new BigInt64Array(current.map((n) => BigInt(n))), + dims: [1, current.length], + }, + encoder_hidden_states: { + name: "encoder_hidden_states", + type: hidden.type, + data: hidden.data, + dims: hidden.dims, + }, + ...this.pastTensors(present), + }) + const { token } = this.argmaxLastStep(outputs["logits"]!) + out.push(token) + if (token === EOS_ID) break + present = new Map( + this.pasts.map((spec) => [spec.name, outputs[spec.name.replace("past_", "present_")]!]), + ) + current = [token] + } + return out + } + const out: number[] = [] + let feed = [PAD_ID, ...prefix] + while (out.length < k) { + const outputs = await this.decoder.run({ + input_ids: { + name: "input_ids", + type: "int64", + data: new BigInt64Array(feed.map((n) => BigInt(n))), + dims: [1, feed.length], + }, + encoder_hidden_states: { + name: "encoder_hidden_states", + type: hidden.type, + data: hidden.data, + dims: hidden.dims, + }, + }) + const { token } = this.argmaxLastStep(outputs["logits"]!) + out.push(token) + if (token === EOS_ID) break + feed = [...feed, token] + } + return out + } + + /** Argmax verdict over a candidate continuation: one full-sequence + * run decides, for each block position, what THIS model would emit + * there, plus the token after the block. The single source of + * "verifier decision" for speculative decode. */ + async review( + hidden: Tensor, + prefix: readonly number[], + block: readonly number[], + ): Promise<{ verdicts: number[]; gaps: number[]; next: number }> { + const feed = [PAD_ID, ...prefix, ...block] + const outputs = await this.decoder.run({ + input_ids: { + name: "input_ids", + type: "int64", + data: new BigInt64Array(feed.map((n) => BigInt(n))), + dims: [1, feed.length], + }, + encoder_hidden_states: { + name: "encoder_hidden_states", + type: hidden.type, + data: hidden.data, + dims: hidden.dims, + }, + ...this.pastTensors(undefined), + }) + const logits = outputs["logits"]! + const verdicts: number[] = [] + const gaps: number[] = [] + for (let i = 0; i < block.length; i++) { + const { token, gap } = this.argmaxAt(logits, prefix.length + i) + verdicts.push(token) + gaps.push(gap) + } + const after = this.argmaxAt(logits, prefix.length + block.length) + return { verdicts, gaps, next: after.token } + } + async dispose(): Promise { await this.encoder.dispose() await this.decoder.dispose() @@ -136,10 +239,16 @@ export class IMFModel { } private argmaxLastStep(logits: Tensor): { token: number; gap: number } { + return this.argmaxAt(logits, logits.dims[logits.dims.length - 2]! - 1) + } + + /** Argmax and top1-top2 gap at a sequence position of the logits + * tensor [batch, seq, classes]. */ + private argmaxAt(logits: Tensor, position: number): { token: number; gap: number } { const dims = logits.dims const classes = dims[dims.length - 1]! const data = logits.data as Float32Array | BigInt64Array - const base = (dims[dims.length - 2]! - 1) * classes + const base = position * classes let best = 0 let bestVal = -Infinity let secondVal = -Infinity diff --git a/src/ml/imf/speculative.ts b/src/ml/imf/speculative.ts new file mode 100644 index 0000000..c3fa346 --- /dev/null +++ b/src/ml/imf/speculative.ts @@ -0,0 +1,127 @@ +/** + * Speculative decoding across the tier ladder: a small drafter + * proposes `blockSize` tokens, the verifier decides them all in one + * full-sequence pass. Greedy verification is output-preserving — the + * verifier's argmax is authoritative at every position, so the result + * equals the verifier's plain-path greedy regardless of drafter + * quality; the drafter only buys speed. On quantized artifacts the + * plain path may differ from the KV path at near-ties, which stays + * inside the quality-parity contract (golden-v1 scoping). + * + * Policy (block loop, corrections, stats) lives here; mechanics + * (draft-from-prefix, positional verdicts) are IMFModel methods. + */ + +import { normalizeArabicInput, repetitionGuardCut } from "./guards.js" +import { EOS_ID, decode, encode } from "./tokens.js" +import type { IMFModel, DecodeOptions } from "./model.js" + +export interface SpeculativeOptions { + /** draft tokens per verifier pass (default 8) */ + readonly blockSize?: number +} + +export interface SpeculativeStats { + readonly blocks: number + readonly drafted: number + readonly accepted: number + readonly bonus: number +} + +/** The acceptance rule: how many leading draft tokens the verdicts + * confirm, and the verifier's pick at the first divergence (null when + * the whole block is accepted). */ +export function acceptBlock(verdicts: readonly number[], block: readonly number[]): { + accepted: number + correction: number | null +} { + for (let i = 0; i < block.length; i++) { + if (verdicts[i] !== block[i]) return { accepted: i, correction: verdicts[i]! } + } + return { accepted: block.length, correction: null } +} + +interface RunStats { + blocks: number + drafted: number + accepted: number + bonus: number +} + +export class SpeculativeModel { + readonly drafter: IMFModel + readonly verifier: IMFModel + readonly blockSize: number + private lastStats: SpeculativeStats | undefined + private lastNormalized: string | undefined + + constructor(drafter: IMFModel, verifier: IMFModel, opts: SpeculativeOptions = {}) { + this.drafter = drafter + this.verifier = verifier + this.blockSize = opts.blockSize ?? 8 + } + + /** Stats from the last translate() call. */ + stats(): SpeculativeStats | undefined { + return this.lastStats + } + + /** The normalized input of the last translate() call — both models + * saw exactly this text. */ + lastNormalizedInput(): string | undefined { + return this.lastNormalized + } + + async translate(text: string, maxLen = 256, opts: DecodeOptions = {}): Promise { + const normalized = opts.raw === true ? text : normalizeArabicInput(text) + this.lastNormalized = normalized + if (encode(normalized).length === 1) return "" + const drafterHidden = await this.drafter.encode(normalized, { raw: true }) + const verifierHidden = await this.verifier.encode(normalized, { raw: true }) + if (!drafterHidden || !verifierHidden) return "" + + const seq: number[] = [] + const run: RunStats = { blocks: 0, drafted: 0, accepted: 0, bonus: 0 } + while (seq.length < maxLen) { + const block = await this.drafter.draft(drafterHidden, seq, this.blockSize) + if (block.length === 0) break + run.blocks += 1 + run.drafted += block.length + const { verdicts, gaps, next } = await this.verifier.review( + verifierHidden, + seq, + block, + ) + const { accepted, correction } = acceptBlock(verdicts, block) + run.accepted += accepted + + if (correction !== null) { + // the verifier overrules at the divergence + seq.push(...block.slice(0, accepted)) + for (let i = 0; i < accepted; i++) opts.onToken?.(block[i]!, seq.length - accepted + i) + if (correction !== EOS_ID) { + seq.push(correction) + opts.onToken?.(correction, seq.length - 1) + if (repetitionGuardCut(seq, decode(seq))) break + } + continue + } + + const endsWithEos = block[block.length - 1] === EOS_ID + const kept = endsWithEos ? block.slice(0, -1) : block + seq.push(...kept) + for (let i = 0; i < kept.length; i++) opts.onToken?.(kept[i]!, seq.length - kept.length + i) + if (endsWithEos) break + // fully accepted: the pass also resolves one bonus token + run.bonus += 1 + if (next === EOS_ID) break + seq.push(next) + opts.onToken?.(next, seq.length - 1) + opts.onConfidence?.(gaps[gaps.length - 1] ?? Infinity, seq.length - 1) + if (repetitionGuardCut(seq, decode(seq))) break + } + + this.lastStats = { ...run } + return decode(seq) + } +} diff --git a/test/speculative.test.ts b/test/speculative.test.ts new file mode 100644 index 0000000..53d56fa --- /dev/null +++ b/test/speculative.test.ts @@ -0,0 +1,95 @@ +/** + * Speculative decode specs: the acceptance rule as a pure function + * (synthetic logits are data), the composition session over the tiny + * real-graph fixture (drafter == verifier: every block accepted, + * output identical to plain greedy), and an opt-in real-pair e2e + * (SECRYST_SPEC_E2E=1 with SECRYST_DRAFTER_ZIP / SECRYST_VERIFIER_ZIP). + */ + +import { readFileSync } from "node:fs" +import { join } from "node:path" +import { homedir } from "node:os" +import { describe, expect, it } from "vitest" +import { IMFModel, EOS_ID } from "../src/ml/imf/index.js" +import { SpeculativeModel, acceptBlock } from "../src/ml/imf/speculative.js" + +const fixtureZip = new Uint8Array(readFileSync("test/fixtures/tiny-imf.zip")) + +describe("acceptBlock (pure acceptance rule)", () => { + it("accepts a fully-matching block", () => { + expect(acceptBlock([7, 8, 9], [7, 8, 9])).toEqual({ accepted: 3, correction: null }) + }) + + it("stops at the first divergence with the verifier's pick", () => { + expect(acceptBlock([7, 9, 9], [7, 8, 9])).toEqual({ accepted: 1, correction: 9 }) + }) + + it("corrects at position zero", () => { + expect(acceptBlock([5, 8], [7, 8])).toEqual({ accepted: 0, correction: 5 }) + }) + + it("treats an accepted EOS verdict as a match like any token", () => { + expect(acceptBlock([7, EOS_ID], [7, EOS_ID])).toEqual({ accepted: 2, correction: null }) + expect(acceptBlock([7, EOS_ID], [7, 8])).toEqual({ accepted: 1, correction: EOS_ID }) + }) + + it("handles the empty block", () => { + expect(acceptBlock([], [])).toEqual({ accepted: 0, correction: null }) + }) +}) + +describe("SpeculativeModel over the tiny fixture (drafter == verifier)", () => { + it("returns the verifier's own greedy output with full acceptance", async () => { + const model = await IMFModel.fromZipBytes(fixtureZip) + const spec = new SpeculativeModel(model, model, { blockSize: 4 }) + const text = "rok" + const out = await spec.translate(text, 64) + expect(out).toBe(await model.translate(text, 64)) + const stats = spec.stats()! + expect(stats.accepted).toBe(stats.drafted) + expect(stats.blocks).toBeGreaterThan(0) + expect(stats.bonus).toBe(stats.blocks) // every block fully accepted + await model.dispose() + }, 60_000) + + it("normalizes input once for both models", async () => { + const model = await IMFModel.fromZipBytes(fixtureZip) + const spec = new SpeculativeModel(model, model) + const out = await spec.translate("rok", 64, { raw: true }) + expect(spec.lastNormalizedInput()).toBe("rok") + expect(out).toBe(await model.translate("rok", 64, { raw: true })) + await model.dispose() + }, 60_000) +}) + +const e2e = process.env["SECRYST_SPEC_E2E"] === "1" + +describe.skipIf(!e2e)("real drafter/verifier pair (ara layerdrop-int4 -> small-2.1-int8)", () => { + const cache = join(homedir(), ".cache", "secryst", "models") + const drafterZip = + process.env["SECRYST_DRAFTER_ZIP"] ?? join(cache, "ara-diac-layerdrop-1.0-int4", "ara-diac-layerdrop-1.0-int4.zip") + const verifierZip = + process.env["SECRYST_VERIFIER_ZIP"] ?? join(cache, "ara-diac-small-2.1-int8", "ara-diac-small-2.1-int8.zip") + + it( + "matches the verifier's plain-path greedy on a real row", + async () => { + const drafter = await IMFModel.load(drafterZip) + const verifier = await IMFModel.load(verifierZip) + const spec = new SpeculativeModel(drafter, verifier) + const row = "السلام عليكم" + const out = await spec.translate(row, 256) + const stats = spec.stats()! + expect(stats.accepted / stats.drafted).toBeGreaterThan(0.9) + // verifier decides every token: output equals its own plain-path + // greedy (translate uses the KV path; near-tie divergence within + // the quantized quality contract is tolerated by comparing + // prefix overlap, not bytes) + const reference = await verifier.translate(row, 256) + expect(out.length).toBeGreaterThan(0.5 * reference.length) + await drafter.dispose() + await verifier.dispose() + }, + 300_000, + ) +}) From 20dfaf0e143d36940dda2d7d667f5bd5096adc72 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 11 Sep 2026 21:31:42 +0800 Subject: [PATCH 2/2] style: prettier --- src/ml/imf/speculative.ts | 11 +++++----- test/speculative.test.ts | 44 +++++++++++++++++++-------------------- 2 files changed, 26 insertions(+), 29 deletions(-) diff --git a/src/ml/imf/speculative.ts b/src/ml/imf/speculative.ts index c3fa346..0720781 100644 --- a/src/ml/imf/speculative.ts +++ b/src/ml/imf/speculative.ts @@ -31,7 +31,10 @@ export interface SpeculativeStats { /** The acceptance rule: how many leading draft tokens the verdicts * confirm, and the verifier's pick at the first divergence (null when * the whole block is accepted). */ -export function acceptBlock(verdicts: readonly number[], block: readonly number[]): { +export function acceptBlock( + verdicts: readonly number[], + block: readonly number[], +): { accepted: number correction: number | null } { @@ -87,11 +90,7 @@ export class SpeculativeModel { if (block.length === 0) break run.blocks += 1 run.drafted += block.length - const { verdicts, gaps, next } = await this.verifier.review( - verifierHidden, - seq, - block, - ) + const { verdicts, gaps, next } = await this.verifier.review(verifierHidden, seq, block) const { accepted, correction } = acceptBlock(verdicts, block) run.accepted += accepted diff --git a/test/speculative.test.ts b/test/speculative.test.ts index 53d56fa..35bc8e3 100644 --- a/test/speculative.test.ts +++ b/test/speculative.test.ts @@ -67,29 +67,27 @@ const e2e = process.env["SECRYST_SPEC_E2E"] === "1" describe.skipIf(!e2e)("real drafter/verifier pair (ara layerdrop-int4 -> small-2.1-int8)", () => { const cache = join(homedir(), ".cache", "secryst", "models") const drafterZip = - process.env["SECRYST_DRAFTER_ZIP"] ?? join(cache, "ara-diac-layerdrop-1.0-int4", "ara-diac-layerdrop-1.0-int4.zip") + process.env["SECRYST_DRAFTER_ZIP"] ?? + join(cache, "ara-diac-layerdrop-1.0-int4", "ara-diac-layerdrop-1.0-int4.zip") const verifierZip = - process.env["SECRYST_VERIFIER_ZIP"] ?? join(cache, "ara-diac-small-2.1-int8", "ara-diac-small-2.1-int8.zip") + process.env["SECRYST_VERIFIER_ZIP"] ?? + join(cache, "ara-diac-small-2.1-int8", "ara-diac-small-2.1-int8.zip") - it( - "matches the verifier's plain-path greedy on a real row", - async () => { - const drafter = await IMFModel.load(drafterZip) - const verifier = await IMFModel.load(verifierZip) - const spec = new SpeculativeModel(drafter, verifier) - const row = "السلام عليكم" - const out = await spec.translate(row, 256) - const stats = spec.stats()! - expect(stats.accepted / stats.drafted).toBeGreaterThan(0.9) - // verifier decides every token: output equals its own plain-path - // greedy (translate uses the KV path; near-tie divergence within - // the quantized quality contract is tolerated by comparing - // prefix overlap, not bytes) - const reference = await verifier.translate(row, 256) - expect(out.length).toBeGreaterThan(0.5 * reference.length) - await drafter.dispose() - await verifier.dispose() - }, - 300_000, - ) + it("matches the verifier's plain-path greedy on a real row", async () => { + const drafter = await IMFModel.load(drafterZip) + const verifier = await IMFModel.load(verifierZip) + const spec = new SpeculativeModel(drafter, verifier) + const row = "السلام عليكم" + const out = await spec.translate(row, 256) + const stats = spec.stats()! + expect(stats.accepted / stats.drafted).toBeGreaterThan(0.9) + // verifier decides every token: output equals its own plain-path + // greedy (translate uses the KV path; near-tie divergence within + // the quantized quality contract is tolerated by comparing + // prefix overlap, not bytes) + const reference = await verifier.translate(row, 256) + expect(out.length).toBeGreaterThan(0.5 * reference.length) + await drafter.dispose() + await verifier.dispose() + }, 300_000) })