mlx-bun 0.2.0 → 0.3.0

package/src/grammar.ts

@@ -33,6 +33,7 @@ import xgrammar, {
 import { readFileSync } from "node:fs";
 import { MlxArray } from "./mlx/array";
 import * as ops from "./mlx/ops";
+import { applyTokenBitmask } from "./mlx/token-bitmask";
 import { runtimeValue } from "./runtime-config";
 import type { LoadedTokenizer } from "./tokenizer";
 
@@ -52,7 +53,7 @@ export interface GrammarRequest {
   guided_grammar?: string;
   /** vLLM/oMLX: regex. Compiled via `root ::= <regex>` EBNF embedding, so
    *  only the regex∩EBNF subset works (the F5 gap — real regex→EBNF or
-   *  structural-tag RegexFormat is the tracked fix, structured-output.md). */
+   *  structural-tag RegexFormat is the tracked fix, docs/reference/server-api.md). */
   guidedRegex?: string;
   guided_regex?: string;
   /** vLLM/oMLX: restrict output to one of these strings (enum). */
@@ -80,7 +81,7 @@ export interface CompiledGrammarResult {
 /** Module-level serializer for ALL xgrammar WASM calls. The WASM instance is
  *  single-threaded; CONCURRENT async calls corrupt emscripten's binding
  *  layer (BindingError: Expected null or instance of VectorInt, got an
- *  instance of VectorInt — reproduced in tests/grammar.test.ts B1 test, which
+ *  instance of VectorInt — reproduced in tests/parity/grammar.test.ts B1 test, which
  *  fires 4 overlapping compiles + fills). Under `--batch N` the scheduler
  *  fires N accept()s (each kicks an async fill) then awaits ready() on all;
  *  and the server compiles grammar for concurrent requests. Without
@@ -119,7 +120,7 @@ let cachedTokenizerInfo: {
  *  - decoder is Sequence w/ ByteFallback    → byte_fallback (Gemma/SP)
  *  - otherwise                              → raw
  *  prependSpace is true for byte_level (ByteLevel pretokenizers add a leading
- *  space). Validated in scripts/experiments/xgrammar-spike.ts across Gemma
+ *  space). Validated in xgrammar-spike.ts (deleted 2026-08-23; git history) across Gemma
  *  (byte_fallback, 256k), Llama-3 / Qwen2.5 (byte_level, 128k/151k). */
 function loadVocab(tokenizerJsonPath: string): {
   vocab: string[];
@@ -207,6 +208,7 @@ export class GrammarController {
   private pending: Promise<void> | null;
   private terminated = false;
   private disposed = false;
+  private readonly metalMask = runtimeValue("MLX_BUN_TOKEN_MASK") === "metal";
   /** Tokens emitted via jumpForward() this generation — telemetry/test hook. */
   jumpedTokens = 0;
 
@@ -238,6 +240,7 @@ export class GrammarController {
    *  Sync — the mask is already materialized. Returns a NEW array (caller
    *  disposes the input). Mirrors oMLX's apply_token_bitmask_mlx. */
   applyMask(logits: MlxArray): MlxArray {
+    if (this.metalMask) return applyTokenBitmask(logits, this.readyMask);
     // Build a -inf/0 additive float mask on device from the int32 bitmask.
     // Each int32 bit b (little-endian within the word) gates token
     // word*32 + b. Set = valid (xgrammar convention). Adding -inf → -inf,
@@ -255,7 +258,9 @@ export class GrammarController {
       const valid = (w >>> bit) & 1;
       maskArr[id] = valid ? 0 : -Infinity;
     }
-    const mask = MlxArray.fromFloat32(maskArr, [1, V]).astype(logits.dtype);
+    const raw = MlxArray.fromFloat32(maskArr, [1, V]);
+    const mask = raw.astype(logits.dtype);
+    raw.dispose();
     const out = ops.add(logits, mask);
     mask.dispose();
     return out;
@@ -285,7 +290,7 @@ export class GrammarController {
    *  time: a controller disposed while its fill is still queued (an exception
    *  between accept/jumpForward and ready()) must not call into the deleted
    *  WASM matcher — that BindingError poisons the module-wide wasmChain for
-   *  every later grammar request (found by tests/grammar-jump.test.ts). */
+   *  every later grammar request (found by tests/parity/grammar-jump.test.ts). */
   private fireFill(): Promise<void> {
     return wasmQueue(() =>
       this.disposed
@@ -458,7 +463,7 @@ export async function compileGrammarRequest(
     if (!degradeHint) degradeHint = resolved.degradeDescription ?? "grammar compile failed";
     // The compiler is the cached per-TokenizerInfo instance (F4) — a failed
     // compile must NOT dispose it; the WASM abort is catchable and the
-    // compiler state survives (verified in tests/grammar.test.ts).
+    // compiler state survives (verified in tests/parity/grammar.test.ts).
     // Return the hint (NOT bare null): it drives the degrade path's
     // system-prompt injection + Warning header. Dropping it here served
     // unconstrained output with no Warning at all (found 2026-07-07 sweep).
@@ -539,7 +544,7 @@ function resolveGrammarRequest(req: GrammarRequest):
           kind: "json_schema",
           schema,
           strict: js?.strict,
-          // Whitespace-stall escape hatch (structured-output.md known gaps):
+          // Whitespace-stall escape hatch (docs/reference/server-api.md known gaps):
           // any_whitespace=false compiles the schema with xgrammar's compact
           // separators — no whitespace choice points, so a whitespace-
           // degenerate base model can't greedily tab-loop to max_tokens.

package/src/harness-pi.ts

@@ -13,7 +13,9 @@ import { homedir } from "node:os";
 import { dirname, join } from "node:path";
 
 export const PI_EXTENSION_FILENAME = "mlx-bun-provider.ts";
-export const PI_PROVIDER_ID = "mlx-bun";
+import { PI_PROVIDER_ID, PI_LOCAL_MODEL_ID, PI_API_KEY, PI_API, piModelDefinition,
+  DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_TOKENS } from "./contracts/pi-provider";
+export { PI_PROVIDER_ID, PI_LOCAL_MODEL_ID } from "./contracts/pi-provider";
 /** Stable pi-facing model id. The server runs exactly one model at a time
  *  (memory: the 26B alone is 18 GB) and which one varies run-to-run, so
  *  advertising the real, changing repo id makes pi's persisted default go
@@ -21,7 +23,7 @@ export const PI_PROVIDER_ID = "mlx-bun";
  *  model is loaded; the real model is carried in the `name` field, and the
  *  server still reports the true id in /v1/models and responses. Gives a
  *  clean `mlx-bun/local` handle instead of `mlx-bun/mlx-community/…`. */
-export const PI_LOCAL_MODEL_ID = "local";
+
 // 127.0.0.1 (the server's actual default bind), not "localhost": the generated
 // extension runs under the USER'S pi (node/jiti, not Bun), so don't depend on
 // how that runtime's resolver maps localhost.
@@ -89,8 +91,8 @@ export async function fetchServerModels(baseUrl: string, timeoutMs = 3000): Prom
     if (!served) return [];
     return [{
       id: served.id,
-      contextWindow: served.context_window ?? 32768,
-      maxTokens: 8192,
+      contextWindow: served.context_window ?? DEFAULT_CONTEXT_WINDOW,
+      maxTokens: DEFAULT_MAX_TOKENS,
       reasoning: !!served.reasoning,
       vision: !!served.vision,
     }];
@@ -120,6 +122,7 @@ export function renderPiExtension(baseUrl: string, fallbackModels: ServerModel[]
 // is the fallback when the server isn't running.
 const BASE_URL = ${JSON.stringify(baseUrl)};
 const FALLBACK_MODELS = ${fallback};
+const modelEntry = ${piModelDefinition.toString()};
 
 export default async function (pi: any) {
   let models = FALLBACK_MODELS;
@@ -131,27 +134,19 @@ export default async function (pi: any) {
         // The SERVED model is the row with the capability extras (first);
         // the rest are servable-but-not-served registry entries.
         const served = data.find((m: any) => m.context_window !== undefined) ?? data[0];
-        models = [{
-          id: ${JSON.stringify(PI_LOCAL_MODEL_ID)},
+        models = [modelEntry({
           name: \`\${served.id} (mlx-bun local)\`,
-          reasoning: !!served.reasoning,
-          compat: {
-            supportsDeveloperRole: false,
-            ...(served.reasoning ? { thinkingFormat: "qwen-chat-template" } : {}),
-          },
-          input: served.vision ? ["text", "image"] : ["text"],
-          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-          contextWindow: served.context_window ?? 32768,
-          maxTokens: 8192,
-        }];
+          reasoning: !!served.reasoning, vision: !!served.vision,
+          contextWindow: served.context_window,
+        })];
       }
     }
   } catch {}
   if (models.length === 0) return; // nothing to register
   pi.registerProvider(${JSON.stringify(PI_PROVIDER_ID)}, {
     baseUrl: BASE_URL,
-    apiKey: "sk-mlx-bun-local",
-    api: "openai-completions",
+    apiKey: ${JSON.stringify(PI_API_KEY)},
+    api: ${JSON.stringify(PI_API)},
     models,
   });
 }
@@ -159,21 +154,7 @@ export default async function (pi: any) {
 }
 
 function modelEntry(m: ServerModel) {
-  return {
-    id: PI_LOCAL_MODEL_ID,
-    name: `${m.id} (mlx-bun local)`,
-    reasoning: !!m.reasoning,
-    // Kept aligned with buildPiProvider (src/pi-provider.ts): the system
-    // prompt must arrive as `system`, not `developer`.
-    compat: {
-      supportsDeveloperRole: false,
-      ...(m.reasoning ? { thinkingFormat: "qwen-chat-template" } : {}),
-    },
-    input: m.vision ? ["text", "image"] : ["text"],
-    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-    contextWindow: m.contextWindow,
-    maxTokens: m.maxTokens,
-  };
+  return piModelDefinition({ ...m, name: `${m.id} (mlx-bun local)` });
 }
 
 export interface InstallResult {

package/src/hub-rest.ts

@@ -3,7 +3,7 @@
 // Axis 3 "Hub" row). Same posture as src/memory/rest.ts: pure (request|url)
 // => Response handlers with no `ctx` dependency, dispatched from
 // src/server.ts by path+method, so they're unit-testable without booting a
-// model context (see tests/hub-rest.test.ts).
+// model context (see tests/serve/hub-rest.test.ts).
 //
 // Three sources, same as optiq Lab's Hub, plus the hardware-fit column optiq
 // lacks:
@@ -19,7 +19,7 @@
 //   - POST /api/hub/serve   — model-swap investigation (see the doc header
 //     above the handler): today this ALWAYS returns
 //     { ok:false, restart_required:true, command } — see
-//     docs/design/runtime-isolation.md for why no in-process swap seam is
+//     docs/reference/server-config.md for why no in-process swap seam is
 //     reachable from the web chat's serving path today.
 
 import { loadModelConfig } from "./config";
@@ -43,7 +43,7 @@ function jsonErr(error: string, status = 400): Response {
  *  fixed 8k context (the same "typical" context /fit and /library use, so
  *  the three surfaces never disagree). `reg`/`hubDir` are injectable so
  *  tests can point at a synthetic hub dir instead of the real
- *  ~/.cache/huggingface/hub (matching tests/registry-fit.test.ts's pattern);
+ *  ~/.cache/huggingface/hub (matching tests/unit/registry-fit.test.ts's pattern);
  *  the live server route omits both and gets the real registry. */
 export async function handleHubLocal(
   opts: { reg?: Registry; hubDir?: string } = {},
@@ -210,7 +210,7 @@ export async function handleHubDownload(request: Request): Promise<Response> {
 
 // ---- POST /api/hub/serve --------------------------------------------------
 
-/** Model-swap investigation (docs/design/runtime-isolation.md): the web
+/** Model-swap investigation (docs/reference/server-config.md): the web
  *  chat's `/ws/chat` is explicitly NOT proxied even under `--isolate`
  *  (src/serve/isolate.ts serves it a 501 with a pointer back to the direct,
  *  non-isolated server) and `--isolate` itself is opt-in, not the default —

package/src/index.ts

@@ -16,9 +16,15 @@ export {
   createModel,
   openGlm52RuntimeModel,
   openModel,
+  type ModelOpenOptions,
   type Glm52RuntimeOpenOptions,
   type RuntimeModel,
 } from "./model/factory";
+export { ModelImplementationRegistry, type ModelImplementation, type ModelImplementationProvider } from "./model/implementation";
+export type { ModelServingBinding, ModelPromptBuilder, BuiltPrompt, ServedModelInfo } from "./serve/model-binding";
+export type { ServingContext, ModelHostSource, AdapterService } from "./serve/model-host";
+export type { MlxGatewayBinding, MlxBatchGroup } from "./backends/mlx/gateway-binding";
+export type { MlxSerialServices, MlxSerialBinding } from "./backends/mlx/serial-executor";
 export { Qwen3Model } from "./model/qwen3";
 export { configFingerprint } from "./model/fingerprint";
 export {
@@ -49,6 +55,7 @@ export { downloadModel } from "./download";
 export { Registry } from "./registry";
 export {
   createServer,
+  shutdownServer,
   loadContext,
   type LoadContextOptions,
   type ServerContext,

package/src/inference/denoising.ts

@@ -0,0 +1,17 @@
+import type { GraphDescriptor } from "./graph";
+
+/** Denoising has canvas state and feedback, independent of AR token caches.
+ * Returned tensors/state are owned by the run. Inputs remain borrowed. A
+ * fused method may bypass this graph and implement InferenceMethod directly. */
+export interface DenoisingGraph<Tensor, State> {
+  readonly descriptor: GraphDescriptor;
+  readonly vocabSize: number;
+  readonly canvasLength: number;
+  readonly embedScale: number;
+  prefill(promptIds: number[], vision?: Tensor): State;
+  extendPrefill(tokens: Tensor, state: State): void;
+  decoderLogits(canvas: Tensor, state: State, feedback: Tensor | null): Tensor;
+  dequantEmbedWeight(): Tensor;
+  softEmbeddings(logits: Tensor, weight: Tensor): Tensor;
+  closeState(state: State): void;
+}

package/src/inference/graph.ts

@@ -0,0 +1,33 @@
+/** Describes an implementation's ABI, not a claim of oracle parity or speed. */
+export interface GraphDescriptor {
+  readonly id: string;
+  readonly backend: string;
+  readonly graphAbi: string;
+  readonly stateAbi: string;
+  readonly artifact: string;
+}
+
+export type LogitSelection =
+  | { readonly type: "all" }
+  | { readonly type: "last" }
+  | { readonly type: "range"; readonly start: number; readonly end: number };
+
+/** Backend-bound graph. Tensor/State/Inputs stay concrete within the binding.
+ * Inputs and hidden are borrowed until the returned lazy work completes;
+ * returned hidden/logits are owned by the caller. The graph appends to state.
+ * Hidden is [batch, positions, hidden-width]; logits are [batch, selected
+ * positions, vocabulary]. Selection precedes the vocabulary projection.
+ * Dtypes and evaluation/fence semantics belong to the declared backend ABI.
+ * Per-run masks, positions, media, taps, and adapters belong in Inputs or in
+ * a bound graph context; they are not mutable session fields. */
+export interface AutoregressiveGraph<Tensor, State, Inputs> {
+  readonly descriptor: GraphDescriptor;
+  forwardHidden(inputs: Inputs, state: State): Tensor | Promise<Tensor>;
+  projectLogits(hidden: Tensor, selection: LogitSelection): Tensor;
+}
+
+/** A fused InferenceMethod can bypass this graph interface entirely. It need
+ * only satisfy the method's output, cancellation, and ownership contracts. */
+export interface GraphFactory<Artifact, Binding, Graph> {
+  open(artifact: Artifact, binding: Binding): Promise<{ graph: Graph; close(): Promise<void> }>;
+}

package/src/inference/prefill.ts

@@ -0,0 +1,39 @@
+export interface PrefillPosition {
+  readonly length: number;
+  readonly position: number;
+  readonly chunkSize: number;
+  readonly tailSplit: boolean;
+  readonly snapshotAt?: number | null;
+}
+
+export interface PrefillStep {
+  readonly start: number;
+  readonly end: number;
+  readonly kind: "drain" | "final";
+  readonly snapshot: boolean;
+  /** Batch admission may finish a short last drain and token zero in one tick.
+   * Serial execution still yields between every drain and the next forward. */
+  readonly batchYield: boolean;
+}
+
+/** Shared serial/solo-batch prefill program. No tensors, runtime policy reads,
+ * or sampling. Advance position only after executing the returned step. */
+export function nextPrefillStep(input: PrefillPosition): PrefillStep {
+  const { length, position, chunkSize, tailSplit, snapshotAt } = input;
+  if (!Number.isSafeInteger(length) || !Number.isSafeInteger(position) ||
+      !Number.isSafeInteger(chunkSize) || chunkSize < 1 || position < 0 || position >= length)
+    throw new Error("invalid prefill position or chunk size");
+  if (snapshotAt != null && (!Number.isSafeInteger(snapshotAt) || snapshotAt < 0 || snapshotAt > length))
+    throw new Error("invalid prefill snapshot boundary");
+  if (snapshotAt != null && snapshotAt > position && snapshotAt < length) {
+    const end = Math.min(position + chunkSize, snapshotAt);
+    return { start: position, end, kind: "drain", snapshot: end === snapshotAt, batchYield: true };
+  }
+  const remaining = length - position;
+  if (remaining > (tailSplit ? 1 : chunkSize)) {
+    const end = position + (tailSplit ? Math.min(chunkSize, remaining - 1) : chunkSize);
+    return { start: position, end, kind: "drain", snapshot: false,
+      batchYield: !tailSplit || remaining - 1 > chunkSize };
+  }
+  return { start: position, end: length, kind: "final", snapshot: false, batchYield: false };
+}

package/src/inference/rollback.ts

@@ -0,0 +1,7 @@
+/** A complete transaction over every participating target state layer.
+ * A failed begin/resolve invalidates the run; its state must be discarded. */
+export interface SpeculativeTransaction {
+  canBegin(drafts: number): boolean;
+  begin(drafts: number): void;
+  resolve(accepted: number): void;
+}

package/src/jobs/db.ts

@@ -55,6 +55,10 @@ export class JobStore {
     this.dbPath = dbPath;
     this.logsDir = logsDir;
     this.db = new Database(dbPath);
+    // Parent submissions and child progress updates share this WAL database.
+    // Wait for a short writer transaction instead of failing a valid job with
+    // SQLITE_BUSY while the other process commits its terminal event.
+    this.db.exec("PRAGMA busy_timeout = 5000");
     this.db.exec("PRAGMA journal_mode = WAL");
     this.db.exec(SCHEMA);
     this.migrate();

package/src/jobs/index.ts

@@ -8,6 +8,8 @@ export {
   getRunner,
   submitInProcess,
   submitSubprocess,
+  closeSubprocessJobs,
+  closeInProcessJobs,
   makeEmit,
   isGpuBusy,
   currentGpuJob,

package/src/jobs/job-entry.ts

@@ -7,6 +7,7 @@
 // the runner, runs it with a log-appending emit, sets terminal status, and
 // exits 0 (done) / 1 (failed).
 
+import { createJobTaskClient } from "./task-client";
 import { JobStore } from "./db";
 import { getRunner, makeEmit } from "./runner";
 import type { JobRunner } from "./types";
@@ -16,9 +17,9 @@ import { runtimeValue } from "../runtime-config";
  *  kinds are built by parallel efforts; importing the module is expected to
  *  call `registerRunner(kind, ...)` at top level. Missing module ⇒ the job
  *  fails cleanly (see resolveRunner). */
-const KIND_MODULES: Record<string, string> = {
-  quantize: "../quantize/job.ts",
-  finetune: "../train/job.ts",
+const KIND_MODULES: Record<string, () => Promise<unknown>> = {
+  quantize: () => import("../quantize/job"),
+  finetune: () => import("../train/job"),
 };
 
 /** Resolve a runner for `kind`. Test kinds (noop/crash) are pre-registered in
@@ -28,19 +29,19 @@ async function resolveRunner(kind: string): Promise<JobRunner> {
   const existing = getRunner(kind);
   if (existing) return existing;
 
-  const modPath = KIND_MODULES[kind];
-  if (modPath) {
+  const load = KIND_MODULES[kind];
+  if (load) {
     try {
-      await import(modPath);
+      await load();
     } catch (e) {
       throw new Error(
-        `kind "${kind}": failed to load runner module ${modPath} ` +
+        `kind "${kind}": failed to load runner module ${kind} ` +
           `(${e instanceof Error ? e.message : String(e)})`,
       );
     }
     const loaded = getRunner(kind);
     if (loaded) return loaded;
-    throw new Error(`kind "${kind}": module ${modPath} did not register a runner`);
+    throw new Error(`kind "${kind}": module ${kind} did not register a runner`);
   }
   throw new Error(`no runner registered for kind "${kind}"`);
 }
@@ -49,8 +50,7 @@ function nowIso(): string {
   return new Date().toISOString().replace("T", " ").slice(0, 19);
 }
 
-async function main(): Promise<void> {
-  const jobId = process.argv[2];
+export async function runJobEntry(jobId = process.argv[2]): Promise<void> {
   if (!jobId) {
     console.error("usage: bun job-entry.ts <jobId>");
     process.exit(2);
@@ -87,7 +87,7 @@ async function main(): Promise<void> {
 
   try {
     const runner = await resolveRunner(row.kind);
-    const result = await runner(emit, config);
+    const result = await createJobTaskClient(runner).run(config, emit);
     const out = result?.outputPath;
     if (out) store.setOutputPath(jobId, out);
     store.setProgress(jobId, 1);
@@ -104,4 +104,4 @@ async function main(): Promise<void> {
   }
 }
 
-main();
+if (import.meta.main) void runJobEntry();