mlx-bun 0.5.0 → 1.0.0

package/src/backends/mlx/prefill-cohort.ts

@@ -1,100 +0,0 @@
-import type { Cache } from "../../model/gemma4-base";
-import type { RuntimeModel } from "../../model/factory";
-import type { MlxArray } from "../../mlx/array";
-import { nextPrefillStep } from "../../inference/prefill";
-import type { CacheCodecProvider } from "../../kv-store";
-import { disposeResources, cleanupFailure } from "../../engine/resources";
-import type { Row, RowPromptCache } from "./batch-group";
-import type { KvMaintenance } from "./kv-maintenance";
-import { MlxPrefillRows, type MlxPrefillState } from "./prefill-rows";
-
-export interface PrefillState extends MlxPrefillState {
-  snapAt: number | null;
-  continuation?: import("./continuation").OrdinaryContinuationState;
-  closePrefill?: () => void;
-}
-
-interface PrefillHost {
-  model: RuntimeModel;
-  chunkSize: number;
-  tailSplit: boolean;
-  promptCache?: RowPromptCache;
-  stateCodecs?: CacheCodecProvider;
-  maintain?: KvMaintenance;
-  forward(ids: MlxArray, caches: Cache[]): Promise<MlxArray>;
-  /** Target projection preserves the forward batch geometry before sampling. */
-  project(hidden: MlxArray, caches: Cache[], completed: readonly PrefillState[]): MlxArray;
-  /** Borrows [1,1,V] logits; takes ownership of the completed request caches. */
-  complete(state: PrefillState, logits: MlxArray): Promise<void>;
-  resume?(state: PrefillState): Promise<void>;
-  reject(row: Row, error: unknown): void;
-}
-
-/** Ordinary sampling/cache policy over the shared target preparation driver. */
-export class MlxPrefillCohort extends MlxPrefillRows<PrefillState> {
-  override get canAdmit(): boolean {
-    return super.canAdmit && this.supportsMixedWork;
-  }
-  get supportsMixedWork(): boolean { return !this.rows.some(row => row.req.promptInput); }
-  constructor(readonly host: PrefillHost) {
-    super({
-      stateCodecs: host.stateCodecs, maintain: host.maintain,
-      open(row) {
-        let owned: { caches: Cache[]; retain?: () => void } | undefined;
-        let closeCache: (() => void) | undefined;
-        let closePrefill: (() => void) | undefined;
-        try {
-          closeCache = row.req.trace?.begin("cache.lookup_restore", { mechanism: "continuous" });
-          const continuation = row.req.continuation?.restore(row.cacheNamespace ?? "");
-          if (continuation) {
-            owned = { caches: continuation.caches };
-            return { row, solo: continuation.caches, pos: row.req.promptIds.length,
-              snapAt: null, continuation };
-          }
-          const hit = host.promptCache?.take(row.req.promptIds, row.cacheNamespace, row.req.cacheSessionId) ?? null;
-          if (hit) owned = hit;
-          closeCache?.(); closeCache = undefined;
-          if (hit) row.cachedTokens = hit.tokens.length;
-          const length = row.req.promptIds.length;
-          const boundary = Math.min(row.req.snapshotAt ?? length, length - 1);
-          const snapAt = host.promptCache && boundary >= 256 && boundary > (hit?.tokens.length ?? 0) ? boundary : null;
-          closePrefill = row.req.trace?.begin("prefill.total", {
-            mechanism: "continuous", promptTokens: row.promptTokens, cachedTokens: row.cachedTokens,
-          });
-          const closeSetup = row.req.trace?.begin("prefill.batch_setup", { mechanism: "continuous" });
-          owned ??= { caches: row.req.statePolicy?.create() ?? host.model.makeCache() }; closeSetup?.();
-          if (row.req.promptInput) host.maintain?.(owned.caches);
-          return { row, solo: owned.caches, pos: hit?.tokens.length ?? 0, retain: owned.retain, snapAt, closePrefill };
-        } catch (error) {
-          return cleanupFailure(error, () => disposeResources([...(owned?.caches ?? []),
-            { dispose: () => owned?.retain?.() }, { dispose: () => closePrefill?.() }]));
-        } finally { closeCache?.(); }
-      },
-      ready: state => !!state.continuation,
-      plan(state) {
-        if (state.row.req.promptInput) return { kind: "final", start: state.pos,
-          end: state.row.promptTokens, snapshot: false, batchYield: false, atomic: true };
-        return nextPrefillStep({ length: state.row.promptTokens, position: state.pos,
-        chunkSize: state.row.req.prefillChunkSize ?? host.chunkSize,
-        tailSplit: host.tailSplit, snapshotAt: state.snapAt }); },
-      forward: (ids, caches, states, work) => states[0]?.row.req.promptInput
-        ? Promise.resolve(states[0].row.req.promptInput.forward(ids, caches, states[0].pos))
-        : work ? work(ids, caches) : host.forward(ids, caches),
-      project: host.project.bind(host),
-      complete: (state, logits) => state.continuation ? host.resume!(state) : host.complete(state, logits!), reject: host.reject.bind(host),
-      checkpoint(state, capture) {
-        let snapshot: Cache[] | undefined;
-        try {
-          snapshot = capture();
-          host.promptCache!.put(state.row.req.promptIds.slice(0, state.pos), snapshot, state.row.cacheNamespace, undefined, undefined, state.row.req.cacheSessionId);
-          snapshot = undefined;
-        } catch (error) {
-          if (snapshot) disposeResources(snapshot);
-          console.warn(`batch-lane boundary snapshot skipped: ${(error as Error).message}`);
-        }
-        state.snapAt = null;
-      },
-      close(state) { const close = state.closePrefill; state.closePrefill = undefined; close?.(); },
-    });
-  }
-}

package/src/backends/mlx/prefill-policy.ts

@@ -1,26 +0,0 @@
-import type { ModelConfig } from "../../config";
-import type { PrefillPolicy } from "../../inference/prefill";
-import { sdpaFallbackBytes } from "../../kv-scheme";
-import type { RuntimeConfig } from "../../runtime-config";
-
-/** Retain large chunks for short prompts, but bound the materialized attention
- * score workspace of recurrent-attention models at long context. This changes
- * work size, never request admission. Explicit chunk settings remain exact. */
-export function resolveMlxPrefillPolicy(config: ModelConfig, runtime: RuntimeConfig,
-  override?: number): PrefillPolicy {
-  const explicit = override ?? (runtime.value("MLX_BUN_RD_PREFILL_CHUNK") === undefined
-    ? undefined : runtime.number("MLX_BUN_RD_PREFILL_CHUNK", 2048));
-  if (explicit !== undefined) {
-    const chunk = Math.max(1, Math.floor(explicit));
-    return { chunkSize: () => chunk };
-  }
-  if (!config.text?.layerTypes?.includes("linear_attention")) return { chunkSize: () => 2048 };
-  // Capture geometry now, independently of later configuration mutation.
-  const geometry = { ...config, text: { ...config.text, layerTypes: [...config.text.layerTypes] } };
-  const workspaceBytes = 1024 ** 3;
-  return { chunkSize(promptTokens) {
-    let chunk = 2048;
-    while (chunk > 8 && sdpaFallbackBytes(geometry, chunk, promptTokens) > workspaceBytes) chunk /= 2;
-    return chunk;
-  } };
-}

package/src/backends/mlx/prefill-rows.ts

@@ -1,244 +0,0 @@
-import type { BatchableCache, Cache } from "../../model/gemma4-base";
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import { clearCache } from "../../mlx/ffi";
-import type { PrefillStep } from "../../inference/prefill";
-import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
-import { disposeResources, cleanupFailure } from "../../engine/resources";
-import type { MlxGroupPreparation, Row } from "./batch-group";
-import { evalCacheState } from "./prefill";
-import { prefillCacheLayout } from "./cache-layout";
-import { MlxStateRows } from "./state-rows";
-import type { P2RTracePhase, P2RTraceAttributes } from "../../serve/prompt-response-trace";
-import type { MlxForwardWork, MlxPreparationWork } from "./mixed-iteration";
-import type { KvMaintenance } from "./kv-maintenance";
-
-// Shared work is recorded on every participating request. workId identifies
-// duplicates when constructing a process timeline; row spans are not additive.
-let nextWorkId = 0;
-function traceRows(states: readonly MlxPrefillState[], phase: P2RTracePhase, attributes: P2RTraceAttributes) {
-  if (!states.some(state => state.row.req.trace)) return null;
-  const closes = states.map(state => state.row.req.trace?.begin(phase, attributes));
-  return { [Symbol.dispose]() { for (const close of closes) close?.(); } };
-}
-
-export interface MlxPrefillStep extends PrefillStep {
-  /** A non-causal input cannot be split at a scheduler token budget.
-   * Its preparation must decline joins until this step completes. */
-  readonly atomic?: boolean;
-  /** Some methods maintain their final target chunk before sampling. */
-  readonly maintain?: boolean;
-}
-
-export interface MlxPrefillState {
-  row: Row;
-  solo: Cache[];
-  pos: number;
-  retain?: () => void;
-  planned?: MlxPrefillStep;
-}
-
-/** The method owns initialization, planning, companion state, snapshots and
- * completion. This driver owns common target batching and work boundaries.
- * Checkpoint factories return owned cache state; close/dispose are idempotent. */
-export interface MlxPrefillRowsHost<State extends MlxPrefillState> {
-  stateCodecs?: CacheCodecProvider;
-  maintain?: KvMaintenance;
-  open(row: Row): State;
-  plan(state: State): MlxPrefillStep;
-  /** Already prepared state can complete without any target forward. */
-  ready?(state: State): boolean;
-  forward(ids: MlxArray, caches: Cache[], states: readonly State[], work?: MlxForwardWork): Promise<MlxArray>;
-  project(hidden: MlxArray, caches: Cache[], completed: readonly State[]): MlxArray | null;
-  /** Consume method context after target maintenance, before snapshots. */
-  afterForward?(ids: MlxArray, caches: Cache[], states: readonly State[], hidden: MlxArray): void | Promise<void>;
-  checkpoint(state: State, capture: () => Cache[], row: number): void;
-  complete(state: State, logits: MlxArray | null, row: number): Promise<void>;
-  reject(row: Row, error: unknown): void;
-  close(state: State): void;
-  filterRows?(keep: readonly number[]): void;
-  dispose?(): void;
-}
-
-/** One target preparation lifecycle at B1/B>1. Late joins retain each row's
- * planned chunk endpoint, precision transition and checkpoint boundary. */
-export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPreparation {
-  #states: State[] = [];
-  #rows: Row[] = [];
-  #stateRows?: MlxStateRows;
-  constructor(readonly operations: MlxPrefillRowsHost<State>) {}
-  get rows(): readonly Row[] { return this.#rows; }
-  get canAdmit(): boolean { return this.#rows.length > 0; }
-  /** Keep a request's prefill weight until retirement so a nearly completed
-   * request does not repeatedly attract another long prefill into its tail. */
-  get tokenWeight(): number { return this.#states.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); }
-
-  admit(row: Row): void {
-    this.#states.push(this.operations.open(row));
-    this.#rows.push(row);
-  }
-
-  #release(state: State): void {
-    const caches = state.solo; state.solo = [];
-    const retain = state.retain; state.retain = undefined;
-    disposeResources([...caches, { dispose: () => retain?.() }, { dispose: () => this.operations.close(state) }]);
-  }
-  #filter(keep: number[]): void {
-    if (keep.length) this.#stateRows?.filterRows(keep);
-    else { const rows = this.#stateRows; this.#stateRows = undefined; rows?.dispose(); }
-    this.operations.filterRows?.(keep);
-    this.#states = keep.map(row => this.#states[row]!);
-    this.#rows = keep.map(row => this.#rows[row]!);
-  }
-  /** Admission changes membership only between forwards. Existing planned
-   * endpoints survive a join, including partially consumed drain chunks. */
-  #synchronizeRows(): void {
-    if (!this.#stateRows && this.#states.length === 1) return;
-    if (!this.#stateRows) {
-      for (const state of this.#states) this.operations.maintain?.preparePrefill?.(state.solo);
-      const layouts: BatchableCache[] = [];
-      try { for (const cache of this.#states[0]!.solo) layouts.push(prefillCacheLayout(cache)); }
-      catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
-      this.#stateRows = new MlxStateRows(layouts);
-      this.#stateRows.mergeRows(this.#states.map(state => state.solo));
-    } else {
-      for (const state of this.#states.slice(this.#stateRows.rowCount)) {
-        this.operations.maintain?.preparePrefill?.(state.solo);
-        this.#stateRows.append(state.solo);
-      }
-    }
-    // Resolve input-backed copies before releasing SSD restore leases.
-    if (this.#states.some(state => state.solo.length)) {
-      evalCacheState(this.#stateRows.caches);
-      for (const state of this.#states) {
-        const caches = state.solo; state.solo = [];
-        const retain = state.retain; state.retain = undefined;
-        disposeResources([...caches, { dispose: () => retain?.() }]);
-      }
-    }
-    for (const cache of this.#stateRows.caches) cache.prefillMaintenance?.beginPrefill();
-  }
-
-  async advance(workLimit?: MlxPreparationWork): Promise<boolean> {
-    let remaining = workLimit?.maxTokens ?? Infinity;
-    // Complete restored inputs before merging: their cache coverage must not
-    // introduce padding or a zero-width forward into unrelated cold rows.
-    if (this.operations.ready && this.#states.some(state => this.operations.ready!(state))) {
-      const remaining: State[] = [];
-      for (const state of this.#states) {
-        if (!this.operations.ready(state)) { remaining.push(state); continue; }
-        try {
-          state.row.req.signal?.throwIfAborted();
-          await this.operations.complete(state, null, -1);
-          state.solo = []; state.retain = undefined;
-        } catch (error) { this.operations.reject(state.row, error); this.#release(state); }
-        finally { this.operations.close(state); }
-      }
-      this.#states = remaining; this.#rows = remaining.map(state => state.row);
-      if (!remaining.length) { this.dispose(); return true; }
-    }
-    {
-      using span = traceRows(this.#states, "prefill.row_sync", { workId: ++nextWorkId, batchSize: this.#states.length });
-      this.#synchronizeRows();
-    }
-    const keep: number[] = [];
-    for (const [index, state] of this.#states.entries()) {
-      if (!state.row.req.signal?.aborted) { keep.push(index); continue; }
-      this.#release(state); this.operations.reject(state.row, state.row.req.signal.reason);
-    }
-    if (keep.length !== this.#states.length) this.#filter(keep);
-    if (!this.#states.length) { this.dispose(); return true; }
-    while (this.#states.length) {
-      for (const state of this.#states) state.planned ??= this.operations.plan(state);
-      const count = Math.min(this.#states.some(state => state.planned!.atomic) ? Infinity
-        : Math.max(1, Math.floor(remaining / this.#states.length)),
-        ...this.#states.map(state => state.planned!.end - state.pos));
-      const work = { workId: ++nextWorkId, batchSize: this.#states.length, tokensPerRow: count };
-      const closes = this.#states.map(state => state.row.req.trace?.begin("prefill.chunk", {
-        mechanism: "continuous", startToken: state.pos, ...work,
-      }));
-      const caches = this.#stateRows?.caches ?? this.#states[0]!.solo;
-      try {
-        using ids = count ? ops.fromInt32(this.#states.flatMap(state => state.row.req.promptIds.slice(state.pos, state.pos + count)),
-          [this.#states.length, count]) : null;
-        let forwarded: MlxArray | null;
-        {
-          using span = traceRows(this.#states, "prefill.forward", work);
-          forwarded = ids ? await this.operations.forward(ids, caches, this.#states, workLimit?.forward) : null;
-        }
-        using hidden = forwarded;
-        const drains: number[] = [], finals: number[] = [];
-        let batchYield = false;
-        for (const [row, state] of this.#states.entries()) {
-          state.pos += count;
-          if (state.pos !== state.planned!.end) continue;
-          if (state.planned!.kind === "final") finals.push(row);
-          else drains.push(row);
-          batchYield ||= state.planned!.batchYield;
-        }
-        const maintained = [...drains, ...finals.filter(row => this.#states[row]!.planned!.maintain)];
-        if (count && (maintained.length || !finals.length)) {
-          {
-            using span = traceRows(this.#states, "prefill.evaluate", work);
-            evalCacheState(caches);
-          }
-          using span = traceRows(this.#states, "prefill.kv_maintenance", work);
-          if (this.#stateRows) {
-            for (const cache of caches) cache.prefillMaintenance?.commitPrefill(maintained);
-          } else if (maintained.length) this.operations.maintain?.(this.#states[0]!.solo);
-          // A scheduler budget can split one planned chunk into many pieces.
-          // Resolve state at each yield, but retain allocator reuse until the
-          // method's maintenance boundary instead of purging it every piece.
-          if (maintained.length) clearCache();
-        }
-        if (ids && this.operations.afterForward) {
-          using span = traceRows(this.#states, "prefill.companion", work);
-          const companion = this.operations.afterForward?.(ids, caches, this.#states, hidden!);
-          if (companion) await companion;
-        }
-        for (const row of [...drains, ...finals]) {
-          const state = this.#states[row]!;
-          if (state.planned!.snapshot) {
-            using span = traceRows([state], "prefill.checkpoint", { ...work, row });
-            this.operations.checkpoint(state, () =>
-              this.#stateRows ? this.#stateRows.extractRow(row) : cloneKvCaches(state.solo, this.operations.stateCodecs), row);
-          }
-          if (!finals.includes(row)) state.planned = undefined;
-        }
-        let projected: MlxArray | null = null;
-        if (hidden && finals.length) {
-          using span = traceRows(this.#states, "prefill.project", work);
-          projected = this.operations.project(hidden, caches, finals.map(row => this.#states[row]!));
-        }
-        using logits = projected;
-        for (const row of finals) {
-          const state = this.#states[row]!;
-          using span = traceRows([state], "prefill.complete", { ...work, row });
-          if (this.#stateRows) state.solo = this.#stateRows.extractRow(row);
-          using part = logits ? logits.slice([row, logits.shape[1]! - 1, 0], [row + 1, logits.shape[1]!, logits.shape[2]!]) : null;
-          try {
-            state.row.req.signal?.throwIfAborted();
-            await this.operations.complete(state, part, row);
-            state.solo = []; state.retain = undefined;
-          } catch (error) {
-            this.operations.reject(state.row, error);
-            this.#release(state);
-          } finally { this.operations.close(state); }
-        }
-        for (const cache of caches) (cache as { releaseRopeArr?: () => void }).releaseRopeArr?.();
-        if (finals.length) this.#filter(this.#states.flatMap((_, row) => finals.includes(row) ? [] : [row]));
-        if (!this.#states.length) { this.dispose(); return true; }
-        remaining -= count * work.batchSize;
-        if (batchYield || finals.length || remaining < this.#states.length) return false;
-      } finally { for (const close of closes) close?.(); }
-    }
-    return true;
-  }
-
-  dispose(): void {
-    const states = this.#states; this.#states = []; this.#rows = [];
-    const rows = this.#stateRows; this.#stateRows = undefined;
-    disposeResources([...(rows ? [rows] : []), ...states.map(state => ({ dispose: () => this.#release(state) })),
-      { dispose: () => this.operations.dispose?.() }]);
-  }
-}

package/src/backends/mlx/prefill.ts

@@ -1,33 +0,0 @@
-import type { PrefillStep } from "../../inference/prefill";
-import type { MlxArray } from "../../mlx/array";
-import type { Cache } from "../../model/gemma4";
-import * as ops from "../../mlx/ops";
-import { clearCache } from "../../mlx/ffi";
-import { leaseCacheStates } from "./state-views";
-import { withResource } from "../../engine/resources";
-
-/** State views with explicit ownership are released even on deferred errors.
- * Owned views precede borrowed state, preserving the serial eval ordering. */
-export function evalCacheState(cache: Cache[]): void {
-  withResource(leaseCacheStates(cache), (state) => { if (state.length) ops.evalAll([...state]); });
-}
-
-/** Forward once. Drains evaluate state, apply KV maintenance, then clear the
- * allocator cache. Finals return owned hidden state without forcing eval. */
-export async function executeMlxPrefillStep(
-  forward: (ids: MlxArray, cache: Cache[]) => MlxArray | Promise<MlxArray>,
-  cache: Cache[], prompt: readonly number[], step: PrefillStep,
-  maintain: () => void,
-): Promise<MlxArray | null> {
-  const chunk = prompt.slice(step.start, step.end);
-  const ids = ops.fromInt32(chunk, [1, chunk.length]);
-  let hidden: MlxArray;
-  try { hidden = await forward(ids, cache); }
-  finally { ids.dispose(); }
-  if (step.kind === "final") return hidden;
-  hidden.dispose();
-  evalCacheState(cache);
-  maintain();
-  clearCache();
-  return null;
-}

package/src/backends/mlx/projected-context-rows.ts

@@ -1,72 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import { Dtype } from "../../mlx/ffi";
-import * as ops from "../../mlx/ops";
-import { BatchedKVCache } from "../../model/batched-kv";
-import type { Cache } from "../../model/gemma4-base";
-import { MlxStateRows } from "./state-rows";
-import { disposeResources } from "../../engine/resources";
-
-export interface ProjectedContext { k: MlxArray; v: MlxArray; }
-export interface ContextProjection {
-  project(hidden: MlxArray, positions: number | MlxArray): ProjectedContext[];
-}
-
-/** Incremental projected context reuses ordinary full-attention row storage.
- * The projection graph owns numerics; this object owns positions and validity. */
-export class MlxProjectedContextRows {
-  readonly state: MlxStateRows<BatchedKVCache>;
-  constructor(layers: number, readonly graph: ContextProjection) {
-    this.state = new MlxStateRows(Array.from({ length: layers }, () => new BatchedKVCache()));
-  }
-  get positions(): readonly number[] { return this.state.caches[0]!.rowOffsets; }
-  prepareAppend(rows: readonly (readonly Cache[])[]) { return this.state.prepareAppendMany(rows); }
-  append(hidden: MlxArray, lengths: readonly number[]): void {
-    const width = Math.max(...lengths), B = lengths.length, positions = this.positions;
-    using context = hidden.slice([0, 0, 0], [B, width, hidden.shape[2]!]);
-    using offsets = positions.every(position => position === positions[0]) ? null : ops.fromInt32([...positions], [B]);
-    const projected = this.graph.project(context, offsets ?? positions[0]!);
-    try {
-      for (const [layer, cache] of this.state.caches.entries()) {
-        cache.specRoundBegin();
-        for (const array of cache.updateAndFetch(projected[layer]!.k, projected[layer]!.v)) array.dispose();
-        if (lengths.every(length => length === width)) cache.specRoundCommit();
-        else cache.specRoundRollback(lengths);
-      }
-    } finally { disposeResources(projected.flatMap(pair => [pair.k, pair.v])); }
-  }
-  materialize(): void { ops.evalAll(this.state.caches.flatMap(cache => cache.state())); }
-  readAttention() {
-    const views: import("../../model/gemma4-base").KvDonorRows[] = [];
-    try { for (const cache of this.state.caches) views.push(cache.captureDonorRows()); }
-    catch (error) { disposeResources(views.flatMap(view => [view.keys, view.values])); throw error; }
-    const masks = new Map<string, MlxArray | null>();
-    return {
-      attend(layer: number, query: MlxArray, blockKeys: MlxArray, blockValues: MlxArray, scale = 1): MlxArray {
-        const view = views[layer]!, width = view.keys.shape[2]!, block = blockKeys.shape[2]!, B = view.offsets.length;
-        // Membership and coverage advance together in every context layer.
-        const key = `${width}:${block}:${view.keys.dtype}`;
-        if (!masks.has(key)) {
-          let mask: MlxArray | null = null;
-          if (view.starts.some(start => start !== 0) || view.ends.some(end => end !== width)) {
-            using columns = ops.arange(0, width, 1, Dtype.int32);
-            using lower = ops.fromInt32([...view.starts], [B, 1, 1, 1]);
-            using upper = ops.fromInt32([...view.ends], [B, 1, 1, 1]);
-            using before = ops.less(columns, lower), after = ops.greaterEqual(columns, upper);
-            using invalid = ops.logicalOr(before, after);
-            using floats = invalid.astype(view.keys.dtype);
-            using prefix = ops.mulScalar(floats, -1e9);
-            using suffix = ops.zeros([B, 1, 1, block], view.keys.dtype);
-            mask = ops.concatAxis([prefix, suffix], 3);
-          }
-          masks.set(key, mask);
-        }
-        using keys = ops.concatAxis([view.keys, blockKeys], 2);
-        using values = ops.concatAxis([view.values, blockValues], 2);
-        const mask = masks.get(key)!;
-        return ops.sdpa(query, keys, values, scale, mask ? "array" : "", mask);
-      },
-      dispose() { disposeResources([...views.flatMap(view => [view.keys, view.values]), ...[...masks.values()].filter(value => value !== null)]); },
-    };
-  }
-  dispose(): void { this.state.dispose(); }
-}

package/src/backends/mlx/prompt-input.ts

@@ -1,29 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import type { Cache } from "../../model/gemma4-base";
-
-/** Prepared native input with an indivisible prefill. The preparation owner
- * retains its tensors; the model binding returns an owned final hidden row.
- * Decode, sampling and row retirement use the ordinary execution group. */
-export interface MlxPromptInput {
-  forward(ids: MlxArray, caches: Cache[], start?: number): MlxArray;
-  readonly decodeState?: MlxDecodeState;
-}
-
-/** Additional model-owned forward state. Equal keys declare compatible row
- * state; the backend supplies the current row order after joins/retirement. */
-export interface MlxDecodeState {
-  readonly key: string;
-  forward(ids: MlxArray, caches: Cache[], rows: readonly MlxDecodeState[]): MlxArray;
-}
-
-/** Preserve the ordinary generator's projection geometry for a prepared
- * embedding sequence: process the whole prompt, then project only its tip. */
-export function bindEmbeddingsInput(
-  forward: (ids: MlxArray, caches: Cache[], start: number) => MlxArray,
-): MlxPromptInput {
-  return { forward(ids, caches, start = 0) {
-    using hidden = forward(ids, caches, start);
-    const [batch, length, width] = hidden.shape;
-    return hidden.slice([0, length! - 1, 0], [batch!, length!, width!]);
-  } };
-}

package/src/backends/mlx/qwen-prompt-input.ts

@@ -1,40 +0,0 @@
-import type { Qwen35Model } from "../../model/qwen3_5";
-import { mropePositionIds, type MropeRequestState } from "../../model/qwen3-mrope";
-import type { MlxArray } from "../../mlx/array";
-import type { Cache } from "../../model/gemma4-base";
-import * as ops from "../../mlx/ops";
-import { Dtype } from "../../mlx/ffi";
-import { bindEmbeddingsInput, type MlxDecodeState, type MlxPromptInput } from "./prompt-input";
-
-class QwenMediaDecodeState implements MlxDecodeState {
-  readonly key = "qwen-interleaved-mrope";
-  constructor(readonly model: Qwen35Model, readonly delta: number) {}
-
-  forward(ids: MlxArray, caches: Cache[], rows: readonly MlxDecodeState[]): MlxArray {
-    const B = rows.length, L = ids.shape[1]!;
-    const cache = caches[this.model.faIdx]!;
-    // The cache supplies logical positions, including left padding, on device.
-    // No offset readback or request-global position mutation is needed.
-    const offsets = (cache as Cache & { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
-    using deltas = ops.fromInt32(rows.map(row => (row as QwenMediaDecodeState).delta), [B]);
-    using scalar = offsets ? null : ops.fromInt32([cache.offset], []);
-    using starts = ops.add(offsets ?? scalar!, deltas);
-    using startRows = ops.reshape(starts, [B, 1]);
-    using indices = ops.arange(0, L, 1, Dtype.int32);
-    using positions = ops.add(startRows, indices);
-    using axis = ops.reshape(positions, [1, B, L]);
-    using grid = ops.concatAxis([axis, axis, axis], 0);
-    return this.model.forwardHiddenAtPositions(ids, caches, grid);
-  }
-}
-
-/** The prompt owns its grid; only its continuation delta survives into decode. */
-export function bindQwenMediaInput(model: Qwen35Model, embeddings: MlxArray,
-  state: MropeRequestState): MlxPromptInput {
-  const input = bindEmbeddingsInput((ids, caches, start) => {
-    using positions = mropePositionIds(state, start, ids.shape[1]!);
-    return start > 0 ? model.forwardHiddenAtPositions(ids, caches, positions)
-      : model.forwardEmbeddingsAtPositions(embeddings, caches, positions);
-  });
-  return { ...input, decodeState: new QwenMediaDecodeState(model, state.delta) };
-}

package/src/backends/mlx/request-policy.ts

@@ -1,20 +0,0 @@
-import type { GenerateOptions } from "../../generate";
-
-/** Copy policy data without cloning native caches, controllers or callbacks. */
-export function snapshotGenerationPolicy<T extends GenerateOptions & { stopSequences?: readonly string[] }>(options: T): T {
-  const snapshot = { ...options };
-  const freeze = (value: unknown): void => {
-    if (!value || typeof value !== "object") return;
-    for (const item of Object.values(value)) freeze(item);
-    Object.freeze(value);
-  };
-  for (const key of ["stopSequences", "eosTokenIds", "adapters", "xtcSpecialTokens", "logitBias",
-    "kvConfig", "turboQuant", "pagedKv", "hlg", "curve", "decodePolicy"] as const) {
-    if (snapshot[key] === undefined) continue;
-    const value = structuredClone(snapshot[key]);
-    freeze(value);
-    Object.assign(snapshot, { [key]: value });
-  }
-  return Object.freeze(snapshot);
-}
-

package/src/backends/mlx/request-state-policy.ts

@@ -1,54 +0,0 @@
-import { namespacedCache } from "../../engine/namespaced-cache";
-import { runtimeConfig, runtimeFlag, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
-import type { GenerateOptions } from "../../generate";
-import type { RuntimeModel } from "../../model/factory";
-import { KVCache, type Cache } from "../../model/gemma4-base";
-import { PagedKVCache } from "../../lab/paged-kv/paged-kv";
-import type { RowPromptCache } from "./batch-group";
-import { disposeResources } from "../../engine/resources";
-
-/** Request state construction and reusable-prefix policy. Equal keys declare
- * compatible state layouts; the executor need not inspect their settings. */
-export interface MlxRequestStatePolicy {
-  readonly key: string;
-  readonly promptCache?: RowPromptCache;
-  create(): Cache[];
-}
-
-/** Replace fresh full-attention storage before prefill. Sliding layers retain
- * their layout; pre-warmed library caches retain their existing storage. */
-export function maybePageKv(
-  cache: Cache[], options: GenerateOptions, capacityTokens: number,
-): void {
-  if (!options.pagedKv) return;
-  if (cache.some((c) => c.offset > 0)) return;
-  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
-  for (let i = 0; i < cache.length; i++) {
-    if (cache[i] instanceof KVCache) {
-      cache[i]!.dispose(); // fresh (offset 0) — nothing stored yet
-      cache[i] = new PagedKVCache(capacityTokens, blockSize, undefined,
-        options.kvBits === 4 || options.kvBits === 8 ? { bits: options.kvBits, groupSize: options.kvGroupSize ?? 64 } : undefined);
-    }
-  }
-}
-
-export function pagedPrefixNamespace(options: GenerateOptions, base: string,
-  direct = runtimeFlag("MLX_BUN_PAGED_ATTN", false)): string {
-  return JSON.stringify(["paged-v1", options.pagedKv?.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE,
-    options.kvBits ?? 0, options.kvGroupSize ?? 64, direct, base]);
-}
-
-export function bindPagedRequestState(model: RuntimeModel, options: GenerateOptions,
-  capacityTokens: number, cache?: RowPromptCache, runtime: RuntimeConfig = runtimeConfig()): MlxRequestStatePolicy | undefined {
-  if (!options.pagedKv) return undefined;
-  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
-  const stateOptions = { pagedKv: { blockSize }, kvBits: options.kvBits, kvGroupSize: options.kvGroupSize };
-  const direct = runtime.flag("MLX_BUN_PAGED_ATTN", false);
-  const namespace = (base: string) => pagedPrefixNamespace(stateOptions, base, direct);
-  return { key: namespace(""),
-    promptCache: cache ? namespacedCache(cache, namespace) : undefined, create: () => withRuntimeConfig(runtime, () => {
-    const caches = model.makeCache();
-    try { maybePageKv(caches, stateOptions, capacityTokens); return caches; }
-    catch (error) { disposeResources(caches); throw error; }
-  }) };
-}