mlx-bun 0.4.0 → 1.0.0

package/src/audio/decode.ts

@@ -1,133 +0,0 @@
-// WAV decode for audio input (A1 of docs/design/generic-model-support.md).
-//
-// Scope: RIFF/WAVE containers — PCM 16/24/32-bit and IEEE float32, any
-// channel count (mean mixdown), any sample rate (linear resample to 16 kHz).
-// PCM16 scales by 1/32768 to match the oracle golden generator
-// (scripts/oracle/gen-e4b-audio-golden.py) exactly. Non-WAV containers (mp3/m4a/
-// flac/ogg) are transcoded to WAV by the serve layer via macOS `afconvert`
-// before reaching this parser.
-//
-// Resampling note: optiq's serve frontend never exposes audio, so there is
-// no oracle behavior for non-16 kHz input — linear interpolation is our
-// documented choice (the parity fixtures are native 16 kHz and skip it).
-
-export const AUDIO_SAMPLE_RATE = 16_000;
-export const AUDIO_MS_PER_TOKEN = 40;
-export const AUDIO_SEQ_LENGTH = 750; // 30 s cap, matches the oracle processor
-/** Raw-speech truncation bound — the oracle feature extractor's default
- *  (audio_feature_extractor.py __call__: truncation=True,
- *  max_length=480_000 → `w[:max_length]`, i.e. 30 s at 16 kHz) applied
- *  BEFORE feature extraction. Mirrored in decodeAudio so a >30 s clip
- *  yields exactly the capped 750 soft tokens AND the matching mel frame
- *  count — the prompt builder hard-checks that the tower's output length
- *  equals the splice's soft-token count. */
-export const AUDIO_MAX_SAMPLES = 30 * AUDIO_SAMPLE_RATE; // 480_000
-
-export interface DecodedWav {
-  samples: Float32Array; // mono
-  sampleRate: number;
-}
-
-/** Number of audio soft tokens for a clip — processing_gemma4.py
- *  `_compute_audio_num_tokens`: ceil(duration_ms / 40) capped at 750.
- *  MUST be computed from decoded samples, never from container metadata. */
-export function audioSoftTokenCount(
-  numSamples: number,
-  sampleRate: number = AUDIO_SAMPLE_RATE,
-): number {
-  const durationMs = (numSamples / sampleRate) * 1000;
-  return Math.min(Math.ceil(durationMs / AUDIO_MS_PER_TOKEN), AUDIO_SEQ_LENGTH);
-}
-
-const WAVE_FORMAT_PCM = 1;
-const WAVE_FORMAT_IEEE_FLOAT = 3;
-const WAVE_FORMAT_EXTENSIBLE = 0xfffe;
-
-/** Parse a RIFF/WAVE file into mono float samples at its native rate. */
-export function decodeWav(bytes: Uint8Array): DecodedWav {
-  const v = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
-  const tag = (o: number) => String.fromCharCode(bytes[o]!, bytes[o + 1]!, bytes[o + 2]!, bytes[o + 3]!);
-  if (bytes.length < 44 || tag(0) !== "RIFF" || tag(8) !== "WAVE")
-    throw new Error("not a RIFF/WAVE file");
-
-  let fmt: { format: number; channels: number; sampleRate: number; bits: number } | null = null;
-  let dataOff = -1;
-  let dataLen = 0;
-  for (let o = 12; o + 8 <= bytes.length; ) {
-    const id = tag(o);
-    const size = v.getUint32(o + 4, true);
-    const body = o + 8;
-    if (id === "fmt ") {
-      // the reads below reach body+16 (base fields) / body+26 (extensible)
-      if (size < 16 || body + 16 > bytes.length)
-        throw new Error("WAV: truncated fmt chunk");
-      let format = v.getUint16(body, true);
-      const channels = v.getUint16(body + 2, true);
-      const sampleRate = v.getUint32(body + 4, true);
-      const bits = v.getUint16(body + 14, true);
-      if (format === WAVE_FORMAT_EXTENSIBLE) {
-        if (size < 40 || body + 26 > bytes.length)
-          throw new Error("WAV: truncated WAVE_FORMAT_EXTENSIBLE fmt chunk");
-        format = v.getUint16(body + 24, true); // first 2 bytes of SubFormat GUID
-      }
-      fmt = { format, channels, sampleRate, bits };
-    } else if (id === "data") {
-      dataOff = body;
-      dataLen = Math.min(size, bytes.length - body);
-    }
-    o = body + size + (size & 1); // chunks are word-aligned
-  }
-  if (!fmt) throw new Error("WAV: missing fmt chunk");
-  if (dataOff < 0) throw new Error("WAV: missing data chunk");
-  const { format, channels, sampleRate, bits } = fmt;
-  if (channels < 1) throw new Error("WAV: zero channels");
-
-  const bytesPer = bits / 8;
-  const frames = Math.floor(dataLen / (bytesPer * channels));
-  const out = new Float32Array(frames);
-
-  const readSample = (off: number): number => {
-    if (format === WAVE_FORMAT_IEEE_FLOAT && bits === 32) return v.getFloat32(off, true);
-    if (format === WAVE_FORMAT_PCM && bits === 16) return v.getInt16(off, true) / 32768;
-    if (format === WAVE_FORMAT_PCM && bits === 24) {
-      const raw = bytes[off]! | (bytes[off + 1]! << 8) | (bytes[off + 2]! << 16);
-      return ((raw << 8) >> 8) / 8388608;
-    }
-    if (format === WAVE_FORMAT_PCM && bits === 32) return v.getInt32(off, true) / 2147483648;
-    throw new Error(`WAV: unsupported format ${format} / ${bits}-bit`);
-  };
-
-  for (let f = 0; f < frames; f++) {
-    const base = dataOff + f * bytesPer * channels;
-    let acc = 0;
-    for (let c = 0; c < channels; c++) acc += readSample(base + c * bytesPer);
-    out[f] = Math.fround(acc / channels);
-  }
-  return { samples: out, sampleRate };
-}
-
-/** Linear resample to the model rate (no-op passthrough at 16 kHz). */
-export function resampleTo16k(samples: Float32Array, sampleRate: number): Float32Array {
-  if (sampleRate === AUDIO_SAMPLE_RATE) return samples;
-  if (sampleRate <= 0) throw new Error(`bad sample rate ${sampleRate}`);
-  const outLen = Math.max(1, Math.round((samples.length * AUDIO_SAMPLE_RATE) / sampleRate));
-  const out = new Float32Array(outLen);
-  const step = samples.length / outLen;
-  for (let i = 0; i < outLen; i++) {
-    const pos = i * step;
-    const j = Math.floor(pos);
-    const fracPart = pos - j;
-    const a = samples[Math.min(j, samples.length - 1)]!;
-    const b = samples[Math.min(j + 1, samples.length - 1)]!;
-    out[i] = Math.fround(a + (b - a) * fracPart);
-  }
-  return out;
-}
-
-/** WAV bytes → mono 16 kHz float samples (the feature-extractor input),
- *  truncated at 30 s exactly like the oracle (AUDIO_MAX_SAMPLES). */
-export function decodeAudio(bytes: Uint8Array): Float32Array {
-  const { samples, sampleRate } = decodeWav(bytes);
-  const s = resampleTo16k(samples, sampleRate);
-  return s.length > AUDIO_MAX_SAMPLES ? s.subarray(0, AUDIO_MAX_SAMPLES) : s;
-}

package/src/audio/features.ts

@@ -1,202 +0,0 @@
-// USM log-mel feature extraction for gemma-4 audio — a verbatim-semantics
-// port of the oracle's Gemma4AudioFeatureExtractor (optiq/vlm/_mlxvlm/
-// models/gemma4/audio_feature_extractor.py, constructor DEFAULTS — §3.3 Q3
-// of docs/design/generic-model-support.md: no per-bin normalization, params are
-// fixed for all gemma-4 models). Gated bit-close against the T0 mel goldens
-// (goldens/e4b-audio-*-mel.bin) in tests/parity/audio-features.test.ts.
-//
-// Pipeline (single waveform, 16 kHz mono float32):
-//   pad to a multiple of 128 samples (validity mask over the original part)
-//   → semicausal left-pad of frame_length/2 = 160 zeros
-//   → unfold into 321-sample frames at hop 160, drop the 321st sample
-//     (preemphasis = 0 branch)
-//   → × periodic Hann window (numpy computes it in FLOAT32 — emulated with
-//     Math.fround so the products match the oracle's f32 rounding)
-//   → rfft n=512 (numpy promotes to float64; we FFT in f64)
-//   → |·| → @ mel filter bank (128 HTK bins over 0–8 kHz, float64, the
-//     transformers mel_filter_bank construction: linspace fft freqs,
-//     unguarded slope division, no norm)
-//   → log(x + 1e-3) → zero out invalid frames → float32 [T, 128].
-//
-// Frame validity: frame f is valid iff original-sample index
-// f·hop + frame_length lands inside the un-padded waveform.
-
-const SR = 16_000;
-const FRAME = 320; // 20 ms
-const HOP = 160; // 10 ms
-const FFT_N = 512; // 2^ceil(log2(320)), fft_overdrive=false
-const FREQ_BINS = FFT_N / 2 + 1; // 257
-const MEL_BINS = 128;
-const MEL_FMIN = 0.0;
-const MEL_FMAX = 8000.0;
-const MEL_FLOOR = 1e-3;
-const PAD_MULTIPLE = 128;
-
-export interface MelFeatures {
-  /** [frames, 128] row-major log-mel, float32 (invalid frames zeroed). */
-  features: Float32Array;
-  frames: number;
-  /** Per-frame validity (1 = real audio, 0 = padding). */
-  mask: Uint8Array;
-}
-
-// --- mel filter bank (float64, transformers semantics) -----------------------
-const hzToMel = (f: number): number => 2595.0 * Math.log10(1.0 + f / 700.0);
-const melToHz = (m: number): number => 700.0 * (10 ** (m / 2595.0) - 1.0);
-
-function buildMelFilters(): Float64Array {
-  // filter_freqs: 130 points linspace in mel space, mapped back to Hz
-  const melMin = hzToMel(MEL_FMIN);
-  const melMax = hzToMel(MEL_FMAX);
-  const nPts = MEL_BINS + 2;
-  const filterFreqs = new Float64Array(nPts);
-  for (let i = 0; i < nPts; i++)
-    filterFreqs[i] = melToHz(melMin + ((melMax - melMin) * i) / (nPts - 1));
-  // fft_freqs: linspace(0, sr//2, 257)
-  const fftFreqs = new Float64Array(FREQ_BINS);
-  for (let k = 0; k < FREQ_BINS; k++) fftFreqs[k] = (8000.0 * k) / (FREQ_BINS - 1);
-  // triangular filters: max(0, min(down, up)) — unguarded diffs like the oracle
-  const filters = new Float64Array(FREQ_BINS * MEL_BINS);
-  for (let m = 0; m < MEL_BINS; m++) {
-    const lower = filterFreqs[m]!;
-    const center = filterFreqs[m + 1]!;
-    const upper = filterFreqs[m + 2]!;
-    for (let k = 0; k < FREQ_BINS; k++) {
-      const down = (fftFreqs[k]! - lower) / (center - lower);
-      const up = (upper - fftFreqs[k]!) / (upper - center);
-      filters[k * MEL_BINS + m] = Math.max(0, Math.min(down, up));
-    }
-  }
-  return filters;
-}
-
-// --- periodic Hann window: the ORACLE'S EXACT float32 values -----------------
-// numpy builds this in float32 end-to-end (arange f32 → ×2π/320 → cos, all in
-// numpy's vectorized f32 routines) and those cos values differ from
-// fround(Math.cos(x)) by 1 ulp on some inputs. That 1 ulp is the ONLY
-// divergence in the whole pipeline: with these exact values our extractor
-// reproduces the oracle mel golden BIT-EXACTLY; with a recomputed window the
-// log-amplified cancellation error reaches ~4.5e-4 (verified 2026-07-07).
-// So the constant IS the spec. Regen (oracle venv):
-//   python -c "from optiq.vlm._mlxvlm.models.gemma4.audio_feature_extractor \
-//     import Gemma4AudioFeatureExtractor as F; import base64; \
-//     print(base64.b64encode(F().window.astype('<f4').tobytes()).decode())"
-const WINDOW_B64 =
-  "AAAAAAAgyjgAHMo5AFRjOgAHyjqAyR07ACJjOwCDmjvAtsk7wCf/O0BoHTyAVT48gFhiPNC2hDyQ" +
-  "yJk8UF+wPPB4yDwQE+I8ECv9PDDfDD0Q5Rs9qKUrPXAfPD3IUE09+DdfPUjTcT1okII9XI+MPYDl" +
-  "lj3QkaE9PJOsPbDotz0MkcM9LIvPPeDV2z30b+g9JFj1PZZGAT7iBgg+SOwOPhr2FT6mIx0+PHQk" +
-  "PhrnKz6OezM+zjA7Ph4GQz64+ko+0A1TPp4+Wz5OjGM+EvZrPhR7dD54Gn0+uemCPopShz5Ix4s+" +
-  "fEeQPrbSlD6GaJk+dQiePhCyoj7gZKc+biCsPkTksD7or7U+4oK6Prdcvz7uPMQ+CiPJPpEOzj4F" +
-  "/9I+6fPXPsTs3D4S6eE+W+jmPhvq6z7Y7fA+EvP1Pkr5+j4AAAA/XIMCP3cGBT8ViQc/8woKP9OL" +
-  "DD93Cw8/oIkRPwwGFD9+gBY/uPgYP3xuGz+K4R0/pVEgP5C+Ij8OKCU/340nP8rvKT+RTSw/+KYu" +
-  "P8b7MD++SzM/pZY1P0PcNz9dHDo/u1Y8PySLPj9iuUA/PeFCP3wCRT/sHEc/WjBJP448Sz9SQU0/" +
-  "eD5PP84zUT8dIVM/OgZVP/LiVj8Yt1g/eoJaP+5EXD9I/l0/Wq5fP/xUYT8C8mI/RIVkP5oOZj/e" +
-  "jWc/6gJpP5ltaj/GzWs/UCNtPxVubj/0rW8/zOJwP4EMcj/0KnM/Cj50P6ZFdT+wQXY/DTJ3P6gW" +
-  "eD9o73g/OLx5PwZ9ej+8MXs/Stp7P552fD+qBn0/YIp9P7ABfj+SbH4/+sp+P94cfz82Yn8//Zp/" +
-  "PyvHfz+85n8/r/l/PwAAgD+v+X8/vOZ/PyvHfz/9mn8/NmJ/P94cfz/6yn4/kmx+P7ABfj9gin0/" +
-  "qgZ9P552fD9J2ns/uzF7PwV9ej84vHk/aO94P6gWeD8NMnc/r0F2P6RFdT8IPnQ/8ipzP4AMcj/M" +
-  "4nA/861vPxRubj9QI20/xc1rP5dtaj/rAmk/3o1nP5oOZj9EhWQ/AfJiP/tUYT9Yrl8/SP5dP+9E" +
-  "XD95glo/FbdYP/DiVj84BlU/HCFTP80zUT93Pk8/UkFNP4o8Sz9ZMEk/6xxHP3gCRT894UI/YrlA" +
-  "PyKLPj+6Vjw/Whw6P0LcNz+mljU/vEszP8b7MD/4pi4/jk0sP8jvKT/cjSc/DCglP5C+Ij+kUSA/" +
-  "iuEdP3puGz+1+Bg/fYAWPwwGFD+diRE/dwsPP9KLDD/zCgo/EokHP3MGBT9dgwI/AAAAP0X5+j4R" +
-  "8/U+0+3wPhrq6z5U6OY+EenhPsXs3D7o89c+/v7SPo4Ozj4EI8k+6zzEPrhcvz7fgro+6a+1PkHk" +
-  "sD5oIKw+3WSnPgqyoj5zCJ4+iGiZPrTSlD58R5A+RMeLPoRShz616YI+ehp9Pgx7dD4Q9ms+Roxj" +
-  "Ppw+Wz7IDVM+qvpKPiQGQz7MMDs+hHszPhrnKz4ydCQ+piMdPhL2FT5E7A4+5AYIPpRGAT4UWPU9" +
-  "7G/oPdDV2z0oi889EJHDPajotz08k6w9yJGhPXDllj1Uj4w9WJCCPTjTcT34N189uFBNPXAfPD2Y" +
-  "pSs98OQbPTjfDD0QK/088BLiPPB4yDxAX7A8gMiZPMC2hDxgWGI8gFU+PCBoHTyAJ/87wLbJO8CC" +
-  "mjsAImM7gMkdOwAGyjoAVGM6ABzKOQAgyjg=";
-
-function buildWindow(): Float32Array {
-  const raw = Buffer.from(WINDOW_B64, "base64");
-  return new Float32Array(raw.buffer, raw.byteOffset, FRAME);
-}
-
-// --- iterative radix-2 complex FFT (float64), N = 512 ------------------------
-const FFT_COS = new Float64Array(FFT_N / 2);
-const FFT_SIN = new Float64Array(FFT_N / 2);
-for (let i = 0; i < FFT_N / 2; i++) {
-  FFT_COS[i] = Math.cos((-2 * Math.PI * i) / FFT_N);
-  FFT_SIN[i] = Math.sin((-2 * Math.PI * i) / FFT_N);
-}
-
-function fft512(re: Float64Array, im: Float64Array): void {
-  // bit-reversal permutation
-  for (let i = 1, j = 0; i < FFT_N; i++) {
-    let bit = FFT_N >> 1;
-    for (; j & bit; bit >>= 1) j ^= bit;
-    j ^= bit;
-    if (i < j) {
-      const tr = re[i]!; re[i] = re[j]!; re[j] = tr;
-      const ti = im[i]!; im[i] = im[j]!; im[j] = ti;
-    }
-  }
-  for (let len = 2; len <= FFT_N; len <<= 1) {
-    const half = len >> 1;
-    const step = FFT_N / len;
-    for (let i = 0; i < FFT_N; i += len) {
-      for (let k = 0; k < half; k++) {
-        const c = FFT_COS[k * step]!;
-        const s = FFT_SIN[k * step]!;
-        const xr = re[i + k + half]! * c - im[i + k + half]! * s;
-        const xi = re[i + k + half]! * s + im[i + k + half]! * c;
-        re[i + k + half] = re[i + k]! - xr;
-        im[i + k + half] = im[i + k]! - xi;
-        re[i + k] = re[i + k]! + xr;
-        im[i + k] = im[i + k]! + xi;
-      }
-    }
-  }
-}
-
-let melFilters: Float64Array | null = null;
-let hannWindow: Float32Array | null = null;
-
-/** 16 kHz mono float32 waveform → USM log-mel features [T, 128]. */
-export function extractMelFeatures(waveform: Float32Array): MelFeatures {
-  melFilters ??= buildMelFilters();
-  hannWindow ??= buildWindow();
-
-  const nReal = waveform.length;
-  const nPadded = Math.ceil(nReal / PAD_MULTIPLE) * PAD_MULTIPLE;
-  const PAD_LEFT = FRAME >> 1; // 160
-  const total = PAD_LEFT + nPadded;
-  const frames = Math.floor((total - (FRAME + 1)) / HOP) + 1;
-  if (frames <= 0) return { features: new Float32Array(0), frames: 0, mask: new Uint8Array(0) };
-
-  const features = new Float32Array(frames * MEL_BINS);
-  const mask = new Uint8Array(frames);
-  const re = new Float64Array(FFT_N);
-  const im = new Float64Array(FFT_N);
-  const mel = new Float64Array(MEL_BINS);
-
-  for (let f = 0; f < frames; f++) {
-    // frame f covers padded indices [f*HOP, f*HOP+FRAME) after the left pad;
-    // validity = the frame-END sample (index f*HOP+FRAME in the left-padded
-    // stream, i.e. f*HOP+FRAME-PAD_LEFT in the original) is real audio.
-    const endIdx = f * HOP + FRAME - PAD_LEFT;
-    const valid = endIdx >= 0 && endIdx < nReal;
-    mask[f] = valid ? 1 : 0;
-    if (!valid) continue; // oracle zeroes invalid frames; ours start zeroed
-
-    re.fill(0);
-    im.fill(0);
-    for (let n = 0; n < FRAME; n++) {
-      const srcIdx = f * HOP + n - PAD_LEFT; // <0 → semicausal zero pad
-      const sample = srcIdx >= 0 && srcIdx < nReal ? waveform[srcIdx]! : 0;
-      re[n] = Math.fround(sample * hannWindow[n]!); // numpy f32 product
-    }
-    fft512(re, im);
-
-    mel.fill(0);
-    for (let k = 0; k < FREQ_BINS; k++) {
-      const mag = Math.hypot(re[k]!, im[k]!);
-      if (mag === 0) continue;
-      const base = k * MEL_BINS;
-      for (let m = 0; m < MEL_BINS; m++) mel[m] = mel[m]! + mag * melFilters![base + m]!;
-    }
-    const out = f * MEL_BINS;
-    for (let m = 0; m < MEL_BINS; m++)
-      features[out + m] = Math.fround(Math.log(mel[m]! + MEL_FLOOR));
-  }
-  return { features, frames, mask };
-}

package/src/audio/transcode.ts

@@ -1,67 +0,0 @@
-// Non-WAV audio containers → 16 kHz WAV via macOS `afconvert` (CoreAudio —
-// built into darwin, our only platform; zero new dependencies). The serve
-// layer runs every clip through ensureWav(): RIFF/WAVE bytes pass through
-// untouched (decode.ts owns the parse), anything else (mp3/m4a/flac/ogg/
-// aiff/…) round-trips through a temp file under os.tmpdir().
-//
-// Two deliberate choices:
-//   - The WAV/non-WAV decision is CONTENT-based (RIFF magic), not
-//     metadata-based — a mislabelled `format` field must not break a valid
-//     WAV, and an mp3 labelled "wav" still transcodes.
-//   - Channel count is preserved (decodeWav does the mono mixdown); the
-//     sample rate converts here (`@16000`) because CoreAudio's resampler
-//     beats decode.ts's documented linear interpolation for real
-//     compressed sources.
-
-import { randomUUID } from "node:crypto";
-import { unlink } from "node:fs/promises";
-import { tmpdir } from "node:os";
-import { join } from "node:path";
-
-/** RIFF/WAVE magic sniff ("RIFF" at 0, "WAVE" at 8). */
-export function isRiffWave(bytes: Uint8Array): boolean {
-  return (
-    bytes.length >= 12 &&
-    bytes[0] === 0x52 && bytes[1] === 0x49 && bytes[2] === 0x46 && bytes[3] === 0x46 &&
-    bytes[8] === 0x57 && bytes[9] === 0x41 && bytes[10] === 0x56 && bytes[11] === 0x45
-  );
-}
-
-/** Transcode any CoreAudio-readable container to 16 kHz PCM16 WAV. Throws
- *  with afconvert's diagnostic on failure — the serve layer surfaces it as
- *  a 400 ("prompt build failed: audio transcode failed …"). Temp files are
- *  cleaned up on every path. */
-export async function transcodeToWav(bytes: Uint8Array): Promise<Uint8Array> {
-  const base = join(tmpdir(), `mlx-bun-audio-${randomUUID()}`);
-  const inPath = `${base}.in`;
-  const outPath = `${base}.wav`;
-  await Bun.write(inPath, bytes);
-  try {
-    const proc = Bun.spawn(
-      ["afconvert", "-f", "WAVE", "-d", "LEI16@16000", inPath, outPath],
-      { stdout: "ignore", stderr: "pipe" },
-    );
-    // ≤30 s clips convert in well under a second — a hung afconvert must
-    // not pin the request forever. kill() → nonzero exit → the 400 below.
-    const killer = setTimeout(() => proc.kill(), 30_000);
-    const exit = await proc.exited;
-    clearTimeout(killer);
-    if (exit !== 0) {
-      const err = (await new Response(proc.stderr).text()).trim();
-      throw new Error(
-        `audio transcode failed (afconvert exit ${exit}): ` +
-          `${err.split("\n").at(-1) || "unknown error"} — ` +
-          `supported containers: WAV (native), mp3/m4a/flac/ogg/aiff via CoreAudio`,
-      );
-    }
-    return new Uint8Array(await Bun.file(outPath).arrayBuffer());
-  } finally {
-    await Promise.allSettled([unlink(inPath), unlink(outPath)]);
-  }
-}
-
-/** WAV passthrough or afconvert transcode — every serve-side clip goes
- *  through here before decodeAudio. */
-export async function ensureWav(bytes: Uint8Array): Promise<Uint8Array> {
-  return isRiffWave(bytes) ? bytes : transcodeToWav(bytes);
-}

package/src/backends/mlx/assistant-target.ts

@@ -1,46 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import { Dtype } from "../../mlx/ffi";
-import * as ops from "../../mlx/ops";
-import type { Cache, KvDonorAttention } from "../../model/gemma4-base";
-import { captureKvDonorAttention } from "../../model/kv-attention-view";
-import type { AssistantAttention } from "../../spec/drafter";
-import { disposeResources } from "../../engine/resources";
-
-/** Bind captured cache validity to Q-only attention. Encoded layouts can
- * supply this same numerical port without changing the draft graph. */
-export function readAssistantDonors(sliding: Cache, full: Cache) {
-  const views: KvDonorAttention[] = [], masks: MlxArray[] = [];
-  const dispose = () => disposeResources([...masks, ...views]);
-  try {
-    for (const cache of [sliding, full]) views.push(captureKvDonorAttention(cache));
-    const bind = (view: KvDonorAttention): AssistantAttention => {
-      // All draft steps borrow the same target snapshot. Build each validity
-      // mask once, then reuse it across matching layers and the entire chain.
-      const byWindow = new Map<number | null, MlxArray | null>();
-      return { attend(query, scale, window) {
-        if (!byWindow.has(window)) {
-          const width = view.width, B = view.offsets.length;
-          const starts = view.starts.map((start, row) => Math.max(start,
-            window === null ? start : view.ends[row]! - window));
-          let mask: MlxArray | null = null;
-          if (starts.some(start => start !== 0) || view.ends.some(end => end !== width)) {
-            using columns = ops.arange(0, width, 1, Dtype.int32);
-            using lower = ops.fromInt32(starts, [B, 1, 1, 1]);
-            using upper = ops.fromInt32([...view.ends], [B, 1, 1, 1]);
-            using before = ops.less(columns, lower);
-            using after = ops.greaterEqual(columns, upper);
-            using invalid = ops.logicalOr(before, after);
-            using floats = invalid.astype(view.dtype);
-            mask = ops.mulScalar(floats, -1e9);
-            masks.push(mask);
-          }
-          byWindow.set(window, mask);
-        }
-        const mask = byWindow.get(window)!;
-        return view.attend(query, scale, { mode: mask ? "array" : "", arr: mask });
-      } };
-    };
-    return { positions: views[0]!.offsets.map(offset => offset - 1),
-      sliding: bind(views[0]!), full: bind(views[1]!), dispose };
-  } catch (error) { dispose(); throw error; }
-}

package/src/backends/mlx/autoregressive-draft-rows.ts

@@ -1,120 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { AutoregressiveGraph } from "../../inference/graph";
-import type { Cache, BatchableCache, PaddedPrefillCache } from "../../model/gemma4-base";
-import type { DraftRowSampling } from "../../spec/source";
-import { toLogprobs } from "../../sampler";
-import { disposeResources } from "../../engine/resources";
-import { MlxStateRows } from "./state-rows";
-import type { PreparedStateChange } from "../../contracts/resources";
-import { evalCacheState } from "./prefill";
-
-export interface MlxDraftLayout extends BatchableCache, PaddedPrefillCache {
-  makeEmptyBatch(): MlxDraftLayout;
-}
-
-/** The draft graph and sampler operate at B. State snapshots own rewind and
- * membership independently of the graph, target verifier and scheduler. */
-export class MlxAutoregressiveDraftRows {
-  #history: MlxStateRows[] = [];
-  #drafts: number[][] = [];
-  #depth = 0;
-  #pendingLast: Array<number | null>;
-  constructor(readonly graph: AutoregressiveGraph<MlxArray, Cache[], MlxArray>,
-    readonly state: MlxStateRows<MlxDraftLayout>, readonly sampling: DraftRowSampling | null) {
-    this.#pendingLast = Array.from({ length: state.rowCount }, () => null);
-  }
-  get pendingLast(): readonly (number | null)[] { return this.#pendingLast; }
-  prepareAppend(rows: readonly (readonly Cache[])[], pending: readonly (number | null)[]): PreparedStateChange {
-    const state = this.state.prepareAppendMany(rows), next = [...this.#pendingLast, ...pending];
-    return { commit: () => { state.commit(); this.#pendingLast = next; }, dispose: () => state.dispose() };
-  }
-  filterRows(keep: readonly number[]): void {
-    this.state.filterRows(keep); this.#pendingLast = keep.map(row => this.#pendingLast[row]!);
-  }
-
-  async #forward(ids: MlxArray, lengths?: readonly number[]): Promise<MlxArray> {
-    const width = ids.shape[1]!, padded = lengths?.some(length => length !== width);
-    if (padded) for (const cache of this.state.caches) cache.preparePrefill({
-      lengths: [...lengths!], rightPadding: lengths!.map(length => width - length),
-    });
-    const hidden = await this.graph.forwardHidden(ids, this.state.caches);
-    try {
-      if (padded) for (const cache of this.state.caches) cache.finalizePrefill();
-      return hidden;
-    } catch (error) { hidden.dispose(); throw error; }
-  }
-
-  /** A restored full-accept boundary may have one unprocessed draft token.
-   * Combine it with the incoming block on device, retaining per-row lengths. */
-  async prefill(tokens: MlxArray): Promise<void> {
-    const [B, N] = tokens.shape as [number, number];
-    if (this.#pendingLast.some(token => token !== null)) {
-      using previous = ops.fromInt32(this.#pendingLast.map(token => token ?? 0), [B, 1]);
-      using joined = ops.concatAxis([previous, tokens], 1);
-      const lengths = this.#pendingLast.map(token => N + Number(token !== null));
-      using indices = ops.fromInt32(lengths.flatMap(length => Array.from({ length: N + 1 },
-        (_, position) => length === N + 1 ? position : Math.min(position + 1, N))), [B, N + 1]);
-      using inputs = ops.takeAlongAxis(joined, indices, 1);
-      (await this.#forward(inputs, lengths)).dispose();
-    } else (await this.#forward(tokens)).dispose();
-    this.#pendingLast.fill(null); this.materialize();
-  }
-  materialize(): void { evalCacheState(this.state.caches); }
-
-  async draft(pending: readonly number[], depth: number, steps: readonly number[]): Promise<number[][]> {
-    this.#releaseHistory();
-    this.#depth = depth;
-    const tokens: MlxArray[] = [];
-    const feeds = pending.map((token, row) => this.#pendingLast[row] == null ? [token] : [this.#pendingLast[row]!, token]);
-    const lengths = feeds.map(feed => feed.length), width = Math.max(...lengths);
-    let ids = ops.fromInt32(feeds.flatMap(feed => [...feed, ...Array(width - feed.length).fill(0)]), [pending.length, width]);
-    try {
-      // With zero proposals, consume the pending input to keep companion
-      // coverage aligned with the target's one-token verification.
-      for (let position = 0; position < Math.max(1, depth); position++) {
-        const padded = position === 0 && lengths.some(length => length !== width);
-        using hidden = await this.#forward(ids, position === 0 ? lengths : undefined);
-        this.#history.push(this.state.clone());
-        if (depth === 0) break;
-        using selection = padded ? ops.fromInt32(lengths.map(length => length - 1), [pending.length, 1, 1]) : null;
-        using selected = selection ? ops.takeAlongAxis(hidden, selection, 1) : null;
-        using logits = this.graph.projectLogits(selected ?? hidden, { type: "last" });
-        using flat = ops.reshape(logits, [pending.length, logits.shape.at(-1)!]);
-        using logprobs = toLogprobs(flat);
-        const token = this.sampling!.sample(logprobs, steps.map(step => step + position));
-        tokens.push(token);
-        ids.dispose(); ids = ops.reshape(token, [pending.length, 1]);
-      }
-      if (depth === 0) this.#drafts = pending.map(() => []);
-      else {
-        using packed = ops.concatAxis(tokens, 0);
-        const values = packed.toIntTokens();
-        this.#drafts = pending.map((_, row) => Array.from({ length: depth },
-          (_, position) => values[position * pending.length + row]!));
-      }
-      return this.#drafts;
-    } catch (error) { this.#releaseHistory(); throw error; }
-    finally { ids.dispose(); disposeResources(tokens); }
-  }
-
-  /** Preserve mlx-lm's full-accept re-feed shape: the final draft remains
-   * pending and joins the next target token in one forward. Recurrent state
-   * uses saved boundaries, never trimming. The lag is companion state. */
-  async commit(accepted: readonly number[]): Promise<void> {
-    const selected: Cache[][] = [];
-    try {
-      if (!accepted.every(count => count >= this.#depth - 1)) {
-        for (const [row, count] of accepted.entries())
-          selected.push(this.#history[Math.min(count, Math.max(0, this.#depth - 1))]!.extractRow(row));
-        this.state.mergeRows(selected);
-      }
-    } finally { disposeResources(selected.flat()); this.#releaseHistory(); }
-
-    this.#pendingLast = accepted.map((count, row) => this.#depth > 0 && count === this.#depth ? this.#drafts[row]!.at(-1)! : null);
-    this.#drafts = []; this.#depth = 0;
-  }
-
-  #releaseHistory(): void { disposeResources(this.#history.splice(0)); }
-  dispose(): void { this.#releaseHistory(); this.#drafts = []; this.state.dispose(); }
-}