mlx-bun 0.0.13 → 0.3.0

package/src/backends/mlx/speculative.ts

@@ -0,0 +1,114 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cache } from "../../model/gemma4";
+import type { RuntimeModel } from "../../model/factory";
+import type { Sampler } from "../../sampler";
+import type { DraftProvider, DraftSource } from "../../spec/source";
+import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
+import type { MlxModelMemory } from "./autoregressive";
+import type { SpeculativeTransaction } from "../../inference/rollback";
+import { bindCacheRollback } from "./rollback";
+import type { GraphDescriptor } from "../../inference/graph";
+import { bindLegacyDraftTarget } from "./draft-target";
+
+/** Bound target operations and draft construction for one speculative run.
+ * A replacement graph supplies this entire port, including any hidden taps.
+ * The verifier never needs a RuntimeModel or an artifact-family check. */
+export interface MlxSpeculativeBinding {
+  readonly runtime?: RuntimeConfig;
+  readonly memory?: MlxModelMemory;
+  readonly descriptor: GraphDescriptor;
+  readonly eosTokenIds: readonly number[];
+  readonly prefillTailSplit: boolean;
+  makeCache(): Cache[];
+  openDraft(sampler: Sampler, caches: Cache[]): DraftSource;
+  bindRollback(caches: Cache[]): SpeculativeTransaction;
+  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[]):
+    Promise<{ hidden: MlxArray; ctxML: MlxArray | null }>;
+  projectLogits(hidden: MlxArray): MlxArray;
+  /** Establish the implementation's verify kernel context; restore on close. */
+  pinVerify?(): { close(): void };
+}
+
+/** Legacy mutable tap/kernel fields require the gateway's exclusive lease. */
+export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
+  const runtime = runtimeConfig();
+  return {
+    runtime,
+    memory: model,
+    descriptor: Object.freeze({ id: `legacy-spec:${model.config.modelType}`, backend: "mlx",
+      graphAbi: "mlx-hidden-bsh-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
+    eosTokenIds: model.config.eosTokenIds,
+    prefillTailSplit: runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
+    makeCache: model.makeCache.bind(model),
+    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }),
+    bindRollback: bindCacheRollback,
+    forward: (ids, caches, tapLayers) => legacyForwardWithTaps(model, ids, caches, tapLayers),
+    projectLogits: model.logitsFromHidden.bind(model),
+    ...("setSpecKernelPinned" in model ? {
+      pinVerify() {
+        model.setSpecKernelPinned(true);
+        return { close() { model.setSpecKernelPinned(false); } };
+      },
+    } : {}),
+  };
+}
+
+export function assertMlxSpeculativeBinding(binding: MlxSpeculativeBinding): void {
+  const descriptor = binding.descriptor;
+  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
+      descriptor.stateAbi !== "legacy-cache-array-v1")
+    throw new Error(`speculative binding ${descriptor.id} has an incompatible backend, graph, or state ABI`);
+}
+
+/** Run one target forward, and when `tapLayers` is set (⟹ a Gemma4 target — a
+ *  DSpark-style source), also return the captured multi-layer context
+ *  [1,L,m*H] (tapLayers concatenated on the feature axis; index nLayers is the
+ *  post-finalNorm sentinel). Non-tapping sources get ctxML=null and never
+ *  touch model.hiddenTap. Mirrors generate-dflash.ts forwardTapped. */
+async function legacyForwardWithTaps(
+  model: RuntimeModel,
+  ids: MlxArray,
+  caches: Cache[],
+  tapLayers: number[] | undefined,
+): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
+  if (!tapLayers) {
+    const asyncModel = model as RuntimeModel & {
+      forwardHiddenAsync?: (
+        ids: MlxArray,
+        caches: Cache[],
+      ) => Promise<MlxArray>;
+    };
+    const hidden = typeof asyncModel.forwardHiddenAsync === "function"
+      ? await asyncModel.forwardHiddenAsync(ids, caches)
+      : model.forwardHidden(ids, caches);
+    return { hidden, ctxML: null };
+  }
+  if (!("hiddenTap" in model)) throw new Error("target graph does not support hidden taps");
+  const m = model;
+  const previousTap = m.hiddenTap;
+  const cap = new Map<number, MlxArray>();
+  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
+  let hidden: MlxArray | null = null;
+  try {
+    hidden = model.forwardHidden(ids, caches);
+    const perLayer = tapLayers.map((li) => {
+      const a = cap.get(li);
+      if (!a) throw new Error(`spec tap: layer ${li} not captured`);
+      return a;
+    });
+    const ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
+    for (const [, a] of cap) a.dispose();
+    cap.clear(); // consumed — the finally must not double-dispose
+    const out = { hidden, ctxML };
+    hidden = null; // ownership returned to the caller
+    return out;
+  } finally {
+    // On any throw (forward mid-capture, missing layer, concat), free the
+    // partially-captured tap tensors and the orphaned hidden. On success both
+    // are already gone (cap cleared, hidden nulled).
+    hidden?.dispose();
+    for (const [, a] of cap) a.dispose();
+    m.hiddenTap = previousTap;
+  }
+}

package/src/backends/mlx/state-views.ts

@@ -0,0 +1,26 @@
+import type { StateView } from "../../contracts/resources";
+import { disposeResources, ownResource } from "../../engine/resources";
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+
+/** Compatibility boundary for legacy state(). Numerical code consumes a lease,
+ * never the legacy ownership marker. No copy, evaluation or fence is added. */
+export function leaseCacheState(cache: Cache): StateView<MlxArray> {
+  return ownResource<readonly MlxArray[]>(cache.state(), cache.stateNeedsDispose ? disposeResources : () => {});
+}
+
+/** Preserve the existing prefill evaluation order: temporary views first,
+ * borrowed state second. Partial acquisition releases every acquired view. */
+export function leaseCacheStates(caches: readonly Cache[]): StateView<MlxArray> {
+  const leases: StateView<MlxArray>[] = [];
+  const close = () => disposeResources(leases.map((lease) => ({ dispose: () => lease.close() })));
+  try {
+    for (const owned of [true, false]) {
+      for (const cache of caches) if (!!cache.stateNeedsDispose === owned) leases.push(leaseCacheState(cache));
+    }
+    return ownResource(leases.flatMap((lease) => [...lease.borrow()]), close);
+  } catch (error) {
+    try { close(); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "state view acquisition failed"); }
+    throw error;
+  }
+}

package/src/backends/mlx/text-engine.ts

@@ -0,0 +1,73 @@
+import { createCompletionMethod } from "../../serve/session-completion-engine";
+import type { ExecutionPlanner, InferenceEngine, Timer } from "../../contracts/generation";
+import { createInferenceEngine, type EngineOptions } from "../../engine/engine";
+import { throwIfCancelled } from "../../engine/cancellation";
+import type { GenerateOptions, GenerateStats } from "../../generate";
+import type { GenerationGateway, RequestShape } from "../../serve/generation-gateway";
+
+/** Data-only API for text requests whose policy/shape was already
+ * resolved by the caller. Native request resources still belong to the prepared
+ * completion API; they are deliberately absent from this token API. */
+export type TextGenerationOptions = Omit<GenerateOptions,
+  "signal" | "cache" | "grammar" | "fill" | "adapters" | "promptEmbeddings" |
+  "imageMask" | "multimodalMask" | "visionPixels" | "onPrefillDone" |
+  "onDecodeCheckpoint" | "initialPendingToken" | "initialGeneratedTokens" |
+  "originalPromptTokens" | "snapshotAt" | "checkpointEveryTokens"
+>;
+
+export interface TextGenerationRequest {
+  readonly promptIds: readonly number[];
+  readonly options: TextGenerationOptions;
+  readonly shape: RequestShape;
+}
+
+export const hostTimer: Timer = {
+  after(milliseconds, callback) {
+    const timer = setTimeout(callback, milliseconds);
+    return () => { clearTimeout(timer); };
+  },
+};
+
+/** Reuses the gateway's placement, exclusive lease, prompt cache, batching,
+ * and quant-specific execution. The portable session never imports them.
+ * Closing this engine cancels its sessions; it does not dispose a shared gateway. */
+export function createTextInferenceEngine(
+  gateway: Pick<GenerationGateway, "place" | "run">,
+  options: Omit<EngineOptions, "timer"> & { timer?: Timer } = {},
+): InferenceEngine<TextGenerationRequest, GenerateStats> {
+  const planner: ExecutionPlanner<TextGenerationRequest, GenerateStats> = {
+    async plan(request, cancellation) {
+      throwIfCancelled(cancellation);
+      if (request.shape.hasVision || request.shape.hasAdapters || request.shape.hasGrammar)
+        throw new Error("native request resources require the prepared completion API");
+      const placement = gateway.place(request.shape, request.options);
+      const maxTokens = request.options.maxTokens ?? 512;
+      return {
+        id: `text:${placement.mechanism}`, outputTokenLimit: maxTokens,
+        method: createCompletionMethod(gateway, {
+          prompt: [...request.promptIds], options: request.options,
+          shape: request.shape, placement,
+        }),
+      };
+    },
+  };
+  const engine = createInferenceEngine(planner, { ...options, timer: options.timer ?? hostTimer });
+  return {
+    open(request, control) {
+      // Snapshot before demand-start: callers may reuse/mutate request arrays.
+      // TextGenerationOptions contains data only; no native handles or callbacks.
+      const snapshot = structuredClone(request);
+      if (!snapshot.promptIds.length || snapshot.promptIds.some((id) => !Number.isInteger(id) || id < 0 || id > 0xffffffff))
+        return Promise.reject(new Error("prompt must contain valid token ids"));
+      const maxTokens = snapshot.options.maxTokens ?? 512;
+      if (!Number.isSafeInteger(maxTokens) || maxTokens <= 0)
+        return Promise.reject(new Error("text generation requires a positive token limit"));
+      if ((snapshot.options.topLogprobs ?? 0) > (options.maxTopLogprobs ?? 20))
+        return Promise.reject(new Error("top logprobs exceed session delivery capacity"));
+      if (control.output === "collect" && (snapshot.options.logprobs || snapshot.options.topLogprobs))
+        return Promise.reject(new Error("token logprobs require stream consumption"));
+      return engine.open(snapshot, control);
+    },
+    close: () => engine.close(),
+  };
+}