feat(ml): speculative decoding across the tier ladder

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,

src/ml/imf/model.ts

@@ -73,17 +73,120 @@ export class IMFModel {
   }
 
   async translate(text: string, maxLen = 256, opts: DecodeOptions = {}): Promise<string> {
-    // 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<Tensor | null> {
+    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<number[]> {
+    if (this.kv) {
+      const out: number[] = []
+      let current = [PAD_ID, ...prefix]
+      let present: ReadonlyMap<string, Tensor> | 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<void> {
     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

src/ml/imf/speculative.ts

@@ -0,0 +1,126 @@
+/**
+ * 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<string> {
+    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)
+  }
+}

test/speculative.test.ts

@@ -0,0 +1,93 @@
+/**
+ * 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)
+})