mlx-bun 0.3.0 → 0.5.0

package/src/spec/assistant-rows.ts

@@ -0,0 +1,93 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { materializeCopy } from "../mlx/materialize";
+import { disposeResources, applyStateChanges } from "../engine/resources";
+import type { PreparedStateChange } from "../contracts/resources";
+import type { AssistantRowsTarget, DraftPrefillGroup, DraftRowGroup, DraftRowCheckpoint, GroupedDraftProvider } from "./source";
+import type { GemmaAssistantDrafter } from "./drafter";
+
+/** The assistant owns only its last true target hidden and coverage. Donor
+ * attention stays with target storage; checkpoints contain no duplicate KV. */
+class AssistantRows implements DraftPrefillGroup, DraftRowGroup {
+  readonly prefillMode = "tail-split" as const;
+  readonly tapLayers: readonly number[] = [];
+  #hidden: MlxArray | null = null;
+  #processed: number[] = [];
+  constructor(readonly target: AssistantRowsTarget, readonly drafter: GemmaAssistantDrafter,
+    readonly namespace: string, checkpoints: readonly (DraftRowCheckpoint | null)[]) { this.append(checkpoints); }
+  get rowCount(): number { return this.#processed.length; }
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void {
+    if (checkpoints.length) applyStateChanges([() => this.prepareAppend(checkpoints)]);
+  }
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange {
+    const processed = [...this.#processed, ...checkpoints.map(checkpoint => checkpoint?.processedTokens ?? 0)];
+    const example = this.#hidden ?? checkpoints.find(checkpoint => checkpoint !== null)?.attachment.tensors[0];
+    const held: MlxArray[] = []; let hidden: MlxArray | null = null;
+    try {
+      for (const checkpoint of checkpoints) if (checkpoint && (checkpoint.attachment.schema !== "gemma-assistant-v1" ||
+        checkpoint.attachment.metadata.processedTokens !== checkpoint.processedTokens))
+        throw new Error("invalid assistant checkpoint alignment");
+      if (example) {
+        const zero = (rows: number) => { const value = ops.zeros([rows,1,this.target.hiddenSize],example.dtype); held.push(value); return value; };
+        hidden = ops.concatAxis([...(this.rowCount ? [this.#hidden ?? zero(this.rowCount)] : []),
+          ...checkpoints.map(checkpoint => checkpoint?.attachment.tensors[0] ?? zero(1))],0);
+      }
+    } finally { disposeResources(held); }
+    return { commit: () => { const old = this.#hidden; this.#hidden=hidden; hidden=old; this.#processed=processed; },
+      dispose: () => { hidden?.dispose(); hidden=null; } };
+  }
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    using indices = ops.fromInt32([...keep],[keep.length]);
+    const next = this.#hidden && keep.length ? ops.takeAxis(this.#hidden,indices,0) : null;
+    this.#hidden?.dispose(); this.#hidden=next; this.#processed=keep.map(row=>this.#processed[row]!);
+  }
+  prefill(tokens: MlxArray, context?: MlxArray): void {
+    const [B,N] = tokens.shape as [number,number];
+    using tail = context!.slice([0,N-1,0],[B,N,context!.shape[2]!]);
+    const hidden = materializeCopy(tail);
+    this.#hidden?.dispose(); this.#hidden=hidden; this.#processed=this.#processed.map(count=>count+N);
+  }
+  materialize(): void { this.#hidden?.eval(); }
+  capture(row: number): DraftRowCheckpoint {
+    const processedTokens=this.#processed[row]!;
+    const hidden=this.#hidden!.slice([row,0,0],[row+1,1,this.target.hiddenSize]);
+    return { processedTokens, attachment: { schema:"gemma-assistant-v1",metadata:{processedTokens},tensors:[hidden] } };
+  }
+  draft(pending: readonly number[], depth: number, _steps: readonly number[]): number[][] {
+    if (!depth) return pending.map(()=>[]);
+    const B=pending.length, owned: MlxArray[]=[];
+    const donors=this.target.readDonors();
+    let ids=ops.fromInt32([...pending],[B,1]), hidden=this.#hidden!;
+    const tokens: MlxArray[]=[];
+    try {
+      for(let step=0;step<depth;step++) {
+        using emb=this.target.embed(ids);
+        const positions=donors.positions.map(position=>position+step);
+        const common=positions.every(position=>position===positions[0]);
+        using offsets=common ? null : ops.fromInt32(positions,[B]);
+        const next=this.drafter.forwardRows(emb,hidden,donors,offsets ?? positions[0]!);
+        owned.push(next.nextHidden); tokens.push(next.tokens); hidden=next.nextHidden;
+        ids.dispose(); ids=ops.reshape(next.tokens,[B,1]);
+      }
+      using packed=ops.concatAxis(tokens,0);
+      const values=packed.toIntTokens();
+      return pending.map((_,row)=>Array.from({length:depth},(_,step)=>values[step*B+row]!));
+    } finally { ids.dispose(); donors.dispose(); disposeResources([...owned,...tokens]); }
+  }
+  commit(accepted: readonly number[], context: MlxArray): void {
+    using indices=ops.fromInt32([...accepted],[accepted.length,1,1]);
+    const hidden=ops.takeAlongAxis(context,indices,1);
+    this.#hidden?.dispose(); this.#hidden=hidden;
+    this.#processed=this.#processed.map((count,row)=>count+accepted[row]!+1);
+  }
+  dispose(): void { this.#hidden?.dispose(); this.#hidden=null; this.#processed=[]; }
+}
+
+export function assistantGroups(drafter: GemmaAssistantDrafter, namespace: string): GroupedDraftProvider {
+  return {
+    checkpointNamespace: () => namespace,
+    openPrefill: ({target,checkpoints}) => new AssistantRows(target.assistantRows!,drafter,namespace,checkpoints),
+    open: ({target,checkpoints}) => new AssistantRows(target.assistantRows!,drafter,namespace,checkpoints),
+  };
+}

package/src/spec/assistant-source.ts

@@ -1,3 +1,6 @@
+import * as ops from "../mlx/ops";
+import { artifactIdentity } from "../model/artifact-identity";
+import { assistantGroups } from "./assistant-rows";
 // AssistantSource — the optiq KV-borrowing Gemma "-assistant" drafter
 // (src/spec/drafter.ts) behind the serve-time DraftSource seam. L2 oracle:
 // optiq spec_generate. This is the SAME drafter the standalone specGenerate
@@ -22,7 +25,7 @@ import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
 import { GemmaAssistantDrafter } from "./drafter";
-import type { DraftProvider, DraftSource, TargetView, AssistantTarget } from "./source";
+import type { DraftProvider, DraftSource, TargetView, AssistantRowsTarget } from "./source";
 
 /** Sum of the drafter's on-disk safetensors — admission accounting only. */
 function safetensorsBytes(dir: string): number {
@@ -35,20 +38,27 @@ function safetensorsBytes(dir: string): number {
 export class AssistantProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: import("./source").GroupedDraftProvider;
 
   private constructor(
     private readonly drafter: GemmaAssistantDrafter,
     id: string,
-    weightsBytes: number,
+    weightsBytes: number, namespace: string,
   ) {
     this.id = id;
     this.weightsBytes = weightsBytes;
+    this.grouped = assistantGroups(drafter,namespace);
   }
 
   static async load(modelDir: string): Promise<AssistantProvider> {
     const drafter = await GemmaAssistantDrafter.load(modelDir);
     const id = modelDir.split("/").filter(Boolean).at(-1)!;
-    return new AssistantProvider(drafter, id, safetensorsBytes(modelDir));
+    try {
+      const identity = await artifactIdentity(await Bun.file(`${modelDir}/config.json`).text(),
+        readdirSync(modelDir).filter(file => file.endsWith(".safetensors"))
+          .map(name => ({ name, path: join(modelDir, name) })));
+      return new AssistantProvider(drafter, id, safetensorsBytes(modelDir), `gemma-assistant-v1:${identity}`);
+    } catch (error) { drafter.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
@@ -56,18 +66,18 @@ export class AssistantProvider implements DraftProvider {
   }
 
   dispose(): void {
-    // The drafter's Weights are process-lifetime (pinned mmaps); nothing to free.
+    this.drafter.dispose();
   }
 }
 
 export class AssistantSource implements DraftSource {
   readonly weightsBytes = 0; // provider-owned weights; per-request adds nothing
-  private readonly target: AssistantTarget;
+  private readonly target: AssistantRowsTarget;
 
-  constructor(private readonly drafter: Pick<GemmaAssistantDrafter, "forward">, target: TargetView) {
-    if (!target.assistant)
+  constructor(private readonly drafter: Pick<GemmaAssistantDrafter, "forwardRows">, target: TargetView) {
+    if (!target.assistantRows)
       throw new Error("assistant drafter requires a Gemma4 target with donor views");
-    this.target = target.assistant;
+    this.target = target.assistantRows;
   }
 
   /** No own cache to prime — the drafter reads the target's live state. */
@@ -79,8 +89,8 @@ export class AssistantSource implements DraftSource {
    *  specGenerate's 2a block. anchorHidden is borrowed; never disposed here. */
   draft(feed: number[], n: number, _stepBase: number, anchorHidden?: MlxArray): number[] {
     if (!anchorHidden) throw new Error("assistant drafter needs the target anchor hidden");
-    const position = this.target.position();
     const shared = this.target.readDonors();
+    const position = shared.positions[0]!;
     const drafts: number[] = [];
     const ownedHiddens: MlxArray[] = [];
     let dTok = feed[feed.length - 1]!; // the pending token (anchorHidden's token)
@@ -88,22 +98,22 @@ export class AssistantSource implements DraftSource {
     let emb: MlxArray | null = null; // hoisted so a mid-loop throw still frees it
     try {
       for (let k = 0; k < n; k++) {
-        emb = this.target.embedScaled(dTok);
-        const step = this.drafter.forward(emb, dHid, shared, position + k);
+        using ids = ops.fromInt32([dTok], [1, 1]);
+        emb = this.target.embed(ids);
+        const step = this.drafter.forwardRows(emb, dHid, shared, position + k);
+        ownedHiddens.push(step.nextHidden);
+        let token: number;
+        try { token = ops.itemUint32(step.tokens); } finally { step.tokens.dispose(); }
         emb.dispose();
         emb = null;
-        drafts.push(step.token);
-        ownedHiddens.push(step.nextHidden);
-        dTok = step.token;
+        drafts.push(token);
+        dTok = token;
         dHid = step.nextHidden;
       }
     } finally {
       emb?.dispose();
       for (const a of ownedHiddens) a.dispose();
-      for (const [k, v] of [shared.sliding, shared.full]) {
-        k.dispose();
-        v.dispose();
-      }
+      shared.dispose();
     }
     return drafts;
   }

package/src/spec/cache-identity.ts

@@ -0,0 +1,22 @@
+import type { GenerateOptions } from "../generate";
+
+/** Snapshot numerical configuration when the request is bound. Caller-owned
+ * arrays must not change the graph or persistence identity after grouping. */
+export function captureSpeculativeOptions(options: GenerateOptions): GenerateOptions {
+  return { ...options,
+    ...(options.kvConfig ? { kvConfig: options.kvConfig.map(({ layerIdx, bits, groupSize }) =>
+      ({ layerIdx, bits, groupSize })).sort((a, b) => a.layerIdx - b.layerIdx) } : {}),
+    ...(options.turboQuant ? { turboQuant: { ...options.turboQuant } } : {}),
+  };
+}
+
+/** Existing namespaces remain byte-identical when no per-layer scheme is
+ * supplied. Execution and persistence consume the same captured settings. */
+export function speculativePrefixNamespace(method: string | undefined, adapters: string | string[],
+  options: GenerateOptions): string {
+  return JSON.stringify({ method, adapters, kvBits: options.kvBits ?? null,
+    kvGroupSize: options.kvGroupSize ?? 64, quantizedKvStart: options.quantizedKvStart ?? null,
+    turboQuant: options.turboQuant ?? null,
+    ...(options.kvConfig?.length ? { kvConfig: options.kvConfig } : {}),
+  });
+}

package/src/spec/deepspec-rows.ts

@@ -0,0 +1,29 @@
+import { projectedDraftGroups } from "./projected-draft-rows";
+import type { GroupedDraftProvider, TargetView } from "./source";
+import type { DeepspecDrafter } from "./dspark/deepspec-module";
+
+/** Existing pairing requirements are checked once when the provider opens. */
+export function bindDeepspecTarget(target: TargetView, drafter: Pick<DeepspecDrafter, "cfg">): void {
+  if (!target.gemmaTaps) throw new Error("DeepSpec drafter requires a Gemma4 target");
+  const layers = drafter.cfg.num_target_layers;
+  if (target.gemmaTaps.layerCount !== layers) throw new Error(
+    `DeepSpec drafter was trained for a ${layers}-layer target; ` +
+    `this model has ${target.gemmaTaps.layerCount} layers — wrong (target, drafter) pairing`);
+}
+
+export function deepspecGroups(drafter: DeepspecDrafter, namespace: string): GroupedDraftProvider {
+  return projectedDraftGroups(namespace, target => {
+    bindDeepspecTarget(target, drafter);
+    return {
+      namespace, schema: "deepspec-context-v1",
+      layers: drafter.cfg.num_hidden_layers, tapLayers: drafter.tapLayers,
+      project(hidden, positions) {
+        using projected = drafter.projectContext(hidden);
+        return drafter.projectContextKVRows(projected, positions);
+      },
+      draft(context, pending, positions, depth) {
+        return drafter.draftRows(context, pending, positions).tokens.map(row => row.slice(0, depth));
+      },
+    };
+  });
+}

package/src/spec/deepspec-source.ts

@@ -22,6 +22,8 @@
 // Context rows are cached post-norm-post-rope — bit-equivalent to their
 // in-round concat (per-row RMSNorm/RoPE; argued in projectContextKV's doc).
 
+import { artifactIdentity } from "../model/artifact-identity";
+import { deepspecGroups, bindDeepspecTarget } from "./deepspec-rows";
 import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
@@ -39,20 +41,27 @@ function safetensorsBytes(dir: string): number {
 export class DeepspecProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: import("./source").GroupedDraftProvider;
   /** Trained block width (config block_size, e.g. 7) — the server pins
    *  numDraftTokens to this. */
   readonly gamma: number;
 
-  private constructor(private readonly drafter: DeepspecDrafter, id: string, weightsBytes: number) {
+  private constructor(private readonly drafter: DeepspecDrafter, id: string, weightsBytes: number, namespace: string) {
     this.id = id;
     this.weightsBytes = weightsBytes;
     this.gamma = drafter.gamma;
+    this.grouped = deepspecGroups(drafter, namespace);
   }
 
   static async load(modelDir: string): Promise<DeepspecProvider> {
     const drafter = await DeepspecDrafter.load(modelDir);
     const id = modelDir.split("/").filter(Boolean).at(-1)!;
-    return new DeepspecProvider(drafter, id, safetensorsBytes(modelDir));
+    try {
+      const identity = await artifactIdentity(await Bun.file(`${modelDir}/config.json`).text(),
+        readdirSync(modelDir).filter(file => file.endsWith(".safetensors"))
+          .map(name => ({ name, path: join(modelDir, name) })));
+      return new DeepspecProvider(drafter, id, safetensorsBytes(modelDir), `deepspec-context-v1:${identity}`);
+    } catch (error) { drafter.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
@@ -72,14 +81,7 @@ export class DeepspecSource implements DraftSource {
 
   constructor(private readonly drafter: Pick<DeepspecDrafter,
     "cfg" | "tapLayers" | "projectContext" | "projectContextKV" | "draftBlock">, target: TargetView) {
-    if (!target.gemmaTaps)
-      throw new Error("DeepSpec drafter requires a Gemma4 target");
-    const nLayers = target.gemmaTaps.layerCount;
-    if (nLayers !== drafter.cfg.num_target_layers)
-      throw new Error(
-        `DeepSpec drafter was trained for a ${drafter.cfg.num_target_layers}-layer target; ` +
-          `this model has ${nLayers} layers — wrong (target, drafter) pairing`,
-      );
+    bindDeepspecTarget(target, drafter);
     this.tapLayers = drafter.tapLayers;
   }
 

package/src/spec/dflash-source.ts

@@ -18,6 +18,8 @@
 // target verify decides). The richer rejection-sampling verify stays in the
 // standalone dflashGenerate for the measure script.
 
+import { artifactIdentity } from "../model/artifact-identity";
+import { projectedDraftGroups } from "./projected-draft-rows";
 import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
@@ -37,20 +39,43 @@ function safetensorsBytes(dir: string): number {
 export class DflashProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: import("./source").GroupedDraftProvider;
   /** The trained block width — the server pins numDraftTokens to this so the
    *  serve loop never asks for more positions than the block was trained for. */
   readonly gamma: number;
 
-  private constructor(private readonly drafter: DflashDrafter, id: string, weightsBytes: number) {
+  private constructor(private readonly drafter: DflashDrafter, id: string, weightsBytes: number, namespace: string) {
     this.id = id;
     this.weightsBytes = weightsBytes;
     this.gamma = drafter.cfg.gamma;
+    this.grouped = projectedDraftGroups(namespace, target => {
+      if (!target.gemmaTaps) throw new Error("DSpark drafter requires a Gemma4 target");
+      const projection = target.gemmaTaps.projection;
+      const minConf = runtimeValue("MLX_BUN_DSPARK_MINCONF");
+      return {
+        namespace, schema: "dflash-context-v1",
+        layers: drafter.cfg.nLayers, tapLayers: drafter.cfg.tapLayers,
+        project: hidden => drafter.projectContextRows(hidden),
+        draft(context, pending, _positions, depth) {
+          return drafter.forwardRows(projection, null, pending, depth, {
+            collectLogits: false, collectConfidence: false,
+            thresholds: drafter.cfg.sts?.thresholds,
+            ...(minConf ? { minConf: Number(minConf) } : {}),
+          }, context).tokens;
+        },
+      };
+    });
   }
 
   static async load(modelDir: string): Promise<DflashProvider> {
     const drafter = loadDsparkDrafter(modelDir); // variant dispatch (dspark|legacy dflash)
     const id = modelDir.split("/").filter(Boolean).at(-1)!;
-    return new DflashProvider(drafter, id, safetensorsBytes(modelDir));
+    try {
+      const identity = await artifactIdentity(await Bun.file(`${modelDir}/dspark.json`).text(),
+        readdirSync(modelDir).filter(file => file.endsWith(".safetensors"))
+          .map(name => ({ name, path: join(modelDir, name) })));
+      return new DflashProvider(drafter, id, safetensorsBytes(modelDir), `dflash-context-v1:${identity}`);
+    } catch (error) { drafter.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
@@ -58,8 +83,7 @@ export class DflashProvider implements DraftProvider {
   }
 
   dispose(): void {
-    // Unlike the assistant drafter (process-pinned mmaps), DflashDrafter
-    // materializes owned MlxArrays — free them on teardown.
+    // Provider-owned weights are released when the mounted draft is closed.
     this.drafter.dispose();
   }
 }