mlx-bun 0.4.0 → 1.0.0

package/src/backends/mlx/full-kv-row-donor.ts

@@ -1,54 +0,0 @@
-import { captureKvDonorAttention, combineKvDonorAttention } from "../../model/kv-attention-view";
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { Cache, KvDonorRows } from "../../model/gemma4-base";
-import { fullRowInner, fullRowPadding } from "../../model/full-prefill-row";
-import { disposeResources } from "../../engine/resources";
-
-/** Read full-attention rows in their existing physical columns. Valid token
- * coverage excludes prefill padding; no append, conversion or compaction. */
-export function captureFullKvDonorRows(rows: readonly Cache[], leftPad: readonly number[], width: number): KvDonorRows {
-  const parts: KvDonorRows[] = [], padded: MlxArray[] = [];
-  try {
-    for (const row of rows) parts.push(fullRowInner(row).captureDonorRows!());
-    const combine = (field: "keys" | "values") => {
-      const views = parts.map((part, index) => {
-        const a = part[field], before = leftPad[index]! - fullRowPadding(rows[index]!), after = width - before - a.shape[2]!;
-        if (!before && !after) return a;
-        using prefix = ops.zeros([1, a.shape[1]!, before, a.shape[3]!], a.dtype);
-        using suffix = ops.zeros([1, a.shape[1]!, after, a.shape[3]!], a.dtype);
-        const result = ops.concatAxis([prefix, a, suffix], 2); padded.push(result); return result;
-      });
-      return views.length === 1 ? ops.contiguous(views[0]!) : ops.concatAxis(views, 0);
-    };
-    const offsets = rows.map(row => row.offset), starts = [...leftPad];
-    const keys = combine("keys");
-    try { return { keys, values: combine("values"), offsets, starts, ends: offsets.map((offset, row) => starts[row]! + offset) }; }
-    catch (error) { keys.dispose(); throw error; }
-  } finally { disposeResources([...parts.flatMap(part => [part.keys, part.values]), ...padded]); }
-}
-
-
-/** Capture full rows in global columns while each encoded view retains its
- * own physical width and numerical attention operation. */
-export function captureFullKvDonorAttention(rows: readonly Cache[], leftPad: readonly number[], width: number): import("../../model/gemma4-base").KvDonorAttention {
-  const captured: import("../../model/gemma4-base").KvDonorAttention[] = [];
-  try {
-    for (const [index, row] of rows.entries()) {
-      const view = captureKvDonorAttention(fullRowInner(row));
-      const pad = leftPad[index]! - fullRowPadding(row);
-      captured.push({ width, dtype: view.dtype, offsets: [row.offset],
-        starts: [leftPad[index]!], ends: [leftPad[index]! + row.offset],
-        attend(q, scale, mask) {
-          if (!mask.arr) return view.attend(q, scale, mask);
-          const from = mask.arr.shape.map(() => 0), to = [...mask.arr.shape];
-          from[from.length - 1] = pad; to[to.length - 1] = pad + view.width;
-          using arr = mask.arr.slice(from, to);
-          return view.attend(q, scale, { mode: mask.mode, arr });
-        },
-        dispose: () => view.dispose(),
-      });
-    }
-    return combineKvDonorAttention(captured);
-  } catch (error) { disposeResources(captured); throw error; }
-}

package/src/backends/mlx/full-transitioning-kv-rows.ts

@@ -1,65 +0,0 @@
-import { KVCache, type Cache, type PaddedPrefillCache, type PrefillPadding } from "../../model/gemma4-base";
-import { FullPrefillRow, fullRowInner } from "../../model/full-prefill-row";
-import { cloneKvCaches } from "../../kv-store";
-import { FullKvPositions, SpeculativeTransitioningKvRows, type KvRowTransition, type SpeculativeTransitionedKvLayout } from "./transitioning-kv-rows";
-
-type FullLayout = SpeculativeTransitionedKvLayout & PaddedPrefillCache & {
-  restorePrefillEnds(ends: readonly number[] | undefined): void;
-}
-
-class PrefillPositions extends FullKvPositions {
-  override sync(rows: readonly Cache[], preserve: boolean): void {
-    const offsets = rows.map(row => row instanceof FullPrefillRow ? row.positionOffset : row.offset);
-    this.mergeRows(preserve
-      ? [{ planes: [], rowOffsets: offsets, leftPad: [...this.leftPad] }]
-      : rows.map((row, index) => ({ planes: [], rowOffsets: [offsets[index]!],
-        leftPad: [row instanceof FullPrefillRow ? row.leftPadding : 0] })));
-  }
-}
-
-/** Full-attention padding composes with the existing precision transition
- * and rollback lifecycle. Codec-specific attention stays in each adapter. */
-export abstract class FullTransitioningKvRows<Layout extends FullLayout>
-  extends SpeculativeTransitioningKvRows<Layout> implements PaddedPrefillCache {
-  constructor(transition: KvRowTransition<Layout>, row?: Cache) {
-    super({ ...transition,
-      converted: row => transition.converted(fullRowInner(row)),
-      extractRow: row => row instanceof FullPrefillRow ? row.extract() : cloneKvCaches([row])[0]!,
-      packRows(layout, rows) {
-        layout.mergeRows(rows.map(fullRowInner));
-        if (rows.some(row => row instanceof FullPrefillRow)) {
-          layout.preparePrefill({ lengths: rows.map(() => 0),
-            leftPadding: rows.map(row => row instanceof FullPrefillRow ? row.leftPadding : 0) });
-          const ends = rows.map(row => row instanceof FullPrefillRow ? row.padding.endOffsets?.[0] : undefined);
-          layout.restorePrefillEnds(ends.some(end => end !== undefined) ? ends.map(end => end ?? Infinity) : undefined);
-        }
-      },
-    }, row, new PrefillPositions());
-  }
-  preparePrefill(padding: PrefillPadding): void {
-    if (this.packed) { this.packed.preparePrefill(padding); return; }
-    if (!this.rows.length) this.rows = padding.lengths.map(() => new KVCache());
-    this.rows = this.rows.map(row => row instanceof FullPrefillRow ? row : new FullPrefillRow(row as KVCache));
-    for (const [index, row] of this.rows.entries()) (row as FullPrefillRow).preparePrefill({
-      lengths: [padding.lengths[index]!],
-      ...(padding.leftPadding ? { leftPadding: [padding.leftPadding[index]!] } : {}),
-      ...(padding.rightPadding ? { rightPadding: [padding.rightPadding[index]!] } : {}),
-    });
-    this.syncPositions();
-  }
-  finalizePrefill(): void {
-    if (this.packed) { this.packed.finalizePrefill(); return; }
-    for (const row of this.rows) if (row instanceof FullPrefillRow) row.finalizePrefill();
-    this.syncPositions();
-  }
-  override filterRows(keep: readonly number[]): void {
-    if (!this.packed && keep.length) {
-      const pads = [...this.leftPad], shared = Math.min(...keep.map(row => pads[row]!));
-      for (const index of keep) {
-        const row = this.rows[index]!;
-        if (row instanceof FullPrefillRow) row.removeLeftPadding(Math.max(0, row.leftPadding - (pads[index]! - shared)));
-      }
-    }
-    super.filterRows(keep);
-  }
-}

package/src/backends/mlx/gateway-binding.ts

@@ -1,130 +0,0 @@
-import { createOrdinaryContinuationRequest } from "./continuation-request";
-import type { MlxSerialServices } from "./serial-executor";
-import { bindPagedRequestState, type MlxRequestStatePolicy } from "./request-state-policy";
-import type { ExecutionContext } from "../../contracts/scheduling";
-import type { ModelConfig } from "../../config";
-import type { KvScheme } from "../../kv-scheme";
-import type { RuntimeModel } from "../../model/factory";
-import { DiffusionGemmaModel } from "../../model/diffusion-gemma";
-import { UniversalDenseModel } from "../../model/universal/dense";
-import { KVCache, RotatingKVCache, isBatchableCache, isPlainKvCache, isRotatingPlainCache } from "../../model/gemma4-base";
-import { SSMCache } from "../../model/qwen3-delta";
-import { Gemma4Model } from "../../model/gemma4";
-import { Qwen35Model } from "../../model/qwen3_5";
-import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
-import { disposeResources } from "../../engine/resources";
-import { legacyCompiledDecodeAvailable } from "./autoregressive";
-import { MlxBatchExecutionGroup, type MlxBatchExecutionGroupOptions, type MlxGroupMethodRequest } from "./batch-group";
-import type { DraftProvider } from "../../spec/source";
-import { targetRowLayoutFactory } from "./target-layout-capability";
-import { bindSpeculativeGroupRequests } from "./speculative-group";
-import type { GenerateOptions } from "../../generate";
-import type { ExecutionRequirements, ResolvedExecution } from "../../contracts/execution";
-import { resolveExecution } from "../../engine/execution-plan";
-
-export interface MlxBatchGroup extends Pick<MlxBatchExecutionGroup,
-  "activeRows" | "pendingRows" | "projectedKvBytes" | "kvBudgetBytes" | "submit" | "kick" | "close"> {}
-
-/** A model implementation owns capability checks and the execution group.
- * Scheduling never inspects concrete model/cache classes. */
-export interface MlxGatewayBinding {
-  configureContinuation?(services: MlxSerialServices): void;
-  continuationRequest?(execution: ResolvedExecution | undefined, options: GenerateOptions, prompt: number[],
-    onToken: Parameters<typeof createOrdinaryContinuationRequest>[0]["onToken"]): ReturnType<typeof createOrdinaryContinuationRequest> | undefined;
-  readonly config: ModelConfig;
-  readonly runtime: RuntimeConfig;
-  plan(request: ExecutionRequirements, options: GenerateOptions,
-    scheduling: { continuous: boolean; quantizedBatch: boolean; checkpoints: boolean }): ResolvedExecution;
-  bindAdapterContext?(adapters: string[], key: string): ExecutionContext;
-  cachesBatchable(): boolean;
-  kvBatchable(scheme: KvScheme): boolean;
-  statePolicy?(execution: ResolvedExecution | undefined, options: GenerateOptions, capacityTokens: number): MlxRequestStatePolicy | undefined;
-  methodRequest?(execution: ResolvedExecution | undefined, options: GenerateOptions): MlxGroupMethodRequest | undefined;
-  createBatchGroup(options: MlxBatchExecutionGroupOptions): MlxBatchGroup;
-}
-
-export function bindMlxGateway(model: RuntimeModel, draft?: { provider: DraftProvider; numDraftTokens: number }): MlxGatewayBinding {
-  const runtime = runtimeConfig();
-  let continuationServices: MlxSerialServices | undefined;
-  const kvBatchCapabilities = { delayedAffine: model instanceof Qwen35Model || model instanceof Gemma4Model };
-  const cachesBatchable = () => {
-    if (model instanceof UniversalDenseModel)
-      return !model.args.maskArray && !model.args.layerTypes?.includes("sliding_attention");
-    const caches = model.makeCache();
-    try {
-      const ssm = runtime.value("MLX_BUN_BATCH_SSM") !== "0";
-      return caches.every((cache) => cache instanceof KVCache || cache instanceof RotatingKVCache ||
-        isBatchableCache(cache) || (ssm && cache instanceof SSMCache));
-    } finally { disposeResources(caches); }
-  };
-  const supportsTargetRows = () => {
-    const caches = model.makeCache();
-    try { return caches.every(cache => targetRowLayoutFactory(cache) !== undefined); }
-    finally { disposeResources(caches); }
-  };
-  const speculative = draft?.provider.grouped && cachesBatchable() && supportsTargetRows()
-    ? bindSpeculativeGroupRequests(model, draft.provider, draft.numDraftTokens) : undefined;
-  const adapterState = "loraState" in model ? model.loraState : undefined;
-  return {
-    config: model.config, runtime,
-    configureContinuation: services => { continuationServices = services; },
-    continuationRequest(execution, options, prompt, onToken) {
-      if (!execution?.checkpoint || execution.mechanism !== "continuous") return undefined;
-      const services = continuationServices;
-      if (!services?.checkpoints || !services.checkpointPersistence)
-        throw new Error("qualified continuation requires bound persistence services");
-      return createOrdinaryContinuationRequest({ options, prompt, onToken, execution,
-        store: services.checkpoints, persistence: services.checkpointPersistence,
-        restore: entry => services.checkpoints!.restore(entry, model),
-        interval: services.checkpointEveryTokens!, identity: services.identity });
-    },
-    statePolicy: (execution, options, capacity) => execution?.pagedKv ? bindPagedRequestState(model, options, capacity) : undefined,
-    methodRequest: (execution, options) => execution?.method === "speculative" ? speculative?.(options) : undefined,
-    ...(adapterState ? { bindAdapterContext(adapters: string[], key: string): ExecutionContext {
-      const selected = [...adapters];
-      return { key, enter() {
-        const previous = adapterState.active;
-        adapterState.active = selected;
-        return () => { adapterState.active = previous; };
-      } };
-    } } : {}),
-    plan(request, options, scheduling) {
-      return resolveExecution(request, {
-        ...scheduling,
-        sharedCheckpoints: !!continuationServices?.checkpointPersistence &&
-          !request.hasDraft && !request.hasVision && !request.hasGrammar &&
-          !request.wantsLogprobs && !options.fill && !options.pagedKv,
-        adapterBatch: !!adapterState, pagedBatch: model instanceof Gemma4Model,
-        groupedMethods: speculative ? ["autoregressive", "speculative"] : ["autoregressive"],
-        speculativeLogprobs: scheduling.continuous && !!speculative,
-        sharedSpeculativeAdapters: scheduling.continuous && !!speculative && !!adapterState &&
-          draft?.provider.grouped?.supportsTargetAdapters === true,
-        turboQuantBatch: scheduling.quantizedBatch,
-        speculativeTurboQuant: scheduling.continuous && !!speculative && !!options.turboQuant,
-        method: model instanceof DiffusionGemmaModel ? "denoising" : "autoregressive",
-        compiledDecode: legacyCompiledDecodeAvailable(model),
-        grammarBatch: runtime.value("MLX_BUN_GRAMMAR_BATCH") !== "0",
-        speculativeKvQuant: (!(model instanceof Qwen35Model) || runtime.flag("MLX_BUN_QWEN_SPEC_KV4", true)) && (
-          (scheduling.continuous && !!speculative && (options.kvBits === 4 || options.kvBits === 8 || !!options.kvConfig?.length)) ||
-          (!options.kvConfig?.length && model instanceof Qwen35Model &&
-            (options.kvBits === 4 || (scheduling.continuous && !!speculative && options.kvBits === 8)) &&
-            (options.quantizedKvStart === 0 || (scheduling.continuous && !!speculative)))
-        ),
-      }, {
-        pagedKv: !!options.pagedKv, fill: !!options.fill,
-        compiledDecode: runtime.flag("MLX_BUN_COMPILED_DECODE", true),
-        grammarJump: runtime.flag("MLX_BUN_GRAMMAR_JUMP", false),
-      });
-    },
-    cachesBatchable,
-    kvBatchable(scheme) {
-      const caches = model.makeCache();
-      try {
-        return scheme.batchable(model.config,
-          (layer) => isPlainKvCache(caches[layer]) || isRotatingPlainCache(caches[layer]), caches.length,
-          kvBatchCapabilities);
-      } finally { disposeResources(caches); }
-    },
-    createBatchGroup: (options) => new MlxBatchExecutionGroup(model, { ...options, kvBatchCapabilities }),
-  };
-}

package/src/backends/mlx/graph.ts

@@ -1,38 +0,0 @@
-import type { AutoregressiveGraph, GraphDescriptor, LogitSelection } from "../../inference/graph";
-import type { MlxArray } from "../../mlx/array";
-
-/** Structural input: generated and hand-written graphs need no union membership. */
-export interface MlxGraphOperations<State> {
-  forwardHidden(ids: MlxArray, state: State): MlxArray;
-  forwardHiddenAsync?: (ids: MlxArray, state: State) => Promise<MlxArray>;
-  logitsFromHidden(hidden: MlxArray): MlxArray;
-}
-
-/** Bind once. No array wrapper, await, readback, or synchronization is added
- * to synchronous forward calls. Existing quant-specific kernels stay inside
- * the supplied operations. This adapter does not take ownership of weights. */
-export function bindMlxGraph<State>(
-  operations: MlxGraphOperations<State>,
-  descriptor: Omit<GraphDescriptor, "backend" | "graphAbi">,
-): AutoregressiveGraph<MlxArray, State, MlxArray> {
-  const forward = operations.forwardHiddenAsync
-    ? operations.forwardHiddenAsync.bind(operations)
-    : operations.forwardHidden.bind(operations);
-  const project = operations.logitsFromHidden.bind(operations);
-  return {
-    descriptor: Object.freeze({ ...descriptor, backend: "mlx", graphAbi: "mlx-hidden-bsh-v1" }),
-    forwardHidden: forward,
-    projectLogits(hidden: MlxArray, selection: LogitSelection) {
-      if (selection.type === "all") return project(hidden);
-      const shape = hidden.shape;
-      if (shape.length !== 3) throw new Error("MLX hidden state must have shape [batch, positions, width]");
-      const [batch, positions, width] = shape as [number, number, number];
-      const start = selection.type === "last" ? positions - 1 : selection.start;
-      const end = selection.type === "last" ? positions : selection.end;
-      if (!Number.isInteger(start) || !Number.isInteger(end) || start < 0 || end <= start || end > positions)
-        throw new Error("invalid logits position selection");
-      const selected = hidden.slice([0, start, 0], [batch, end, width]);
-      try { return project(selected); } finally { selected.dispose(); }
-    },
-  };
-}

package/src/backends/mlx/kv-codec.ts

@@ -1,9 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-
-/** A backend KV representation independent of row storage and scheduling.
- * Inputs are borrowed. Encoding and decoding return caller-owned arrays.
- * A codec may retain values in a rotated domain for attention to invert once. */
-export interface KvCodec<Encoded> {
-  encode(keys: MlxArray, values: MlxArray): Encoded;
-  decode(encoded: Encoded, length: number, headDim: number, deferValues?: boolean): [MlxArray, MlxArray];
-}

package/src/backends/mlx/kv-maintenance.ts

@@ -1,133 +0,0 @@
-import type { KvSchemeOptions } from "../../kv-scheme";
-import { KVCache, QuantizedKVCache, RotatingKVCache, RotatingQuantizedKVCache, TurboQuantKVCache, type Cache } from "../../model/gemma4-base";
-import type { KvQuantSpec, TurboQuantScheme } from "../../config";
-import * as ops from "../../mlx/ops";
-import { clearCache } from "../../mlx/ffi";
-
-import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
-import { DelayedQuantizedKVCache } from "../../model/delayed-quantized-kv";
-import { DelayedTurboQuantKVCache } from "../../model/delayed-turboquant-kv";
-
-export interface KvMaintenance {
-  (cache: Cache[]): void;
-  /** Bind state requiring row-local maintenance before shared decode. */
-  prepareBatch?(cache: Cache[]): void;
-  /** Bind all precision policies before a prefill cohort owns row boundaries. */
-  preparePrefill?(cache: Cache[]): void;
-}
-
-const unchanged = (_cache: Cache[]): void => {};
-let warnedTurboRotating = false;
-
-/** Port of mlx-lm maybe_quantize_kv_cache + BOTH halves of optiq serve's
- *  per-layer patched variant (incl. patch_rotating_to_quantized: rotating
- *  caches convert too — Phase 9):
- *  - per-layer bits/group_size selection (kvConfig overrides kvBits,
- *    matching optiq's --kv-config precedence; shipped kv_config.json
- *    files cover EVERY cache-owning layer, sliding ones included —
- *    verified 12B 48/48, 26B 30/30, e4b 24/24 distinct caches — so
- *    rotating quantization engages straight from the config; uniform
- *    kvBits — like optiq --kv-bits — reaches them too), and
- *  - STREAMING conversion (optiq streaming_kv_quant / serve.py
- *    patched_maybe_quantize): eval each layer's quantized triples and
- *    clear the buffer pool before building the next layer's conversion.
- *    Lazily batching every layer's toQuantized into one eval pins ALL
- *    layers' bf16 K/V as graph inputs alongside ALL quantized outputs —
- *    the exact transient optiq's fix kills (16.35 → 7.60 GB at 32k on a
- *    24 GB Mac). Numerics untouched: same quantize math, only the eval
- *    ordering is forced (tests/parity/kv-quant.test.ts, tests/parity/rotating-kvq.test.ts). */
-export function createKvMaintenance(options: Readonly<Omit<KvSchemeOptions, "kvConfig">> & {
-  readonly kvConfig?: readonly Readonly<KvQuantSpec>[];
-}): KvMaintenance {
-  const { kvBits, kvConfig, turboQuant } = options;
-  if (turboQuant) {
-    const start = options.quantizedKvStart ?? 0;
-    const scheme = { ...turboQuant };
-    const maintain: KvMaintenance = (cache) => maybeTurboQuantizeKv(cache, scheme, start);
-    maintain.preparePrefill = (cache) => {
-      for (let layer = 0; layer < cache.length; layer++) {
-        const row = cache[layer]!;
-        if (row instanceof KVCache || row instanceof TurboQuantKVCache)
-          cache[layer] = new DelayedTurboQuantKVCache(scheme.kBits, scheme.vBits, start, maintain, row);
-      }
-    };
-    if (start > 0) maintain.prepareBatch = maintain.preparePrefill;
-    return maintain;
-  }
-  if (!kvBits && !kvConfig?.length) return unchanged;
-  const start = options.quantizedKvStart ?? (kvConfig?.length ? 0 : 5000);
-  const groupSize = options.kvGroupSize ?? 64;
-  // Resolve layer policy once when composing execution, never per token.
-  const byLayer = kvConfig?.length
-    ? new Map(kvConfig.map((entry) => [entry.layerIdx, { ...entry }]))
-    : null;
-  const maintain: KvMaintenance = (cache) => {
-    for (let i = 0; i < cache.length; i++) {
-      const c = cache[i]!;
-      const conversion = c.affineConversion ?? (c instanceof KVCache || c instanceof RotatingKVCache ? c : undefined);
-      if (!conversion || conversion.offset < start || conversion.offset === 0) continue;
-      const conversionOffset = conversion.offset;
-      // The populated-cache boundary is part of oracle parity: quantizing an
-      // empty cache would also quantize the first prefill.
-      if (byLayer) {
-        const entry = byLayer.get(i);
-        if (!entry) continue;
-        cache[i] = conversion.toQuantized(entry.groupSize, entry.bits);
-      } else {
-        cache[i] = conversion.toQuantized(groupSize, kvBits!);
-      }
-      if (start > 0) cache[i]!.minimumReusableOffset = conversionOffset;
-      // Materialize one layer before converting the next to bound the live
-      // bf16 source plus quantized destination to one conversion at a time.
-      ops.evalAll(cache[i]!.state());
-      clearCache();
-    }
-  };
-  maintain.preparePrefill = (cache) => {
-    for (let layer = 0; layer < cache.length; layer++) {
-      const row = cache[layer]!;
-      if (!(row instanceof KVCache || row instanceof QuantizedKVCache || row instanceof RotatingKVCache || row instanceof RotatingQuantizedKVCache)) continue;
-      const spec = byLayer ? byLayer.get(layer) : { bits: kvBits!, groupSize };
-      if (!spec) continue;
-      const rowMaintenance = createKvMaintenance({ kvBits: spec.bits, kvGroupSize: spec.groupSize, quantizedKvStart: start });
-      cache[layer] = row instanceof RotatingKVCache || row instanceof RotatingQuantizedKVCache
-        ? new DelayedRotatingQuantizedKVCache(row.maxSize, spec.groupSize, spec.bits, start, rowMaintenance, row)
-        : new DelayedQuantizedKVCache(spec.groupSize, spec.bits, start, rowMaintenance, row);
-    }
-  };
-  if (start > 0) maintain.prepareBatch = maintain.preparePrefill;
-  return maintain;
-}
-
-function maybeTurboQuantizeKv(cache: Cache[], scheme: TurboQuantScheme, start: number): void {
-  for (let i = 0; i < cache.length; i++) {
-    const c = cache[i]!;
-    if (c instanceof RotatingKVCache) {
-      if (!warnedTurboRotating) {
-        warnedTurboRotating = true;
-        console.warn(
-          "[turbo-quant] sliding-window (RotatingKVCache) layers stay bf16 in v1 " +
-          "(full-attention only) — docs/design/turboquant.md.",
-        );
-      }
-      continue;
-    }
-    const conversion = c.turboConversion ?? (c instanceof KVCache ? c : undefined);
-    if (!conversion || conversion.offset < start || conversion.offset === 0) continue;
-    const conversionOffset = conversion.offset;
-    const tq = c.turboConversion ? c.turboConversion.toTurboQuantized(scheme.kBits, scheme.vBits)
-      : TurboQuantKVCache.fromKVCache(c as KVCache, scheme.kBits, scheme.vBits);
-    if (start > 0) tq.minimumReusableOffset = conversionOffset;
-    cache[i] = tq;
-    // state() allocates fresh trimmed slice views for this cache kind
-    // (see evalCacheState) — dispose after materializing (throw included),
-    // or they leak.
-    const state = tq.state();
-    try {
-      ops.evalAll(state);
-    } finally {
-      for (const a of state) a.dispose();
-    }
-    clearCache();
-  }
-}

package/src/backends/mlx/kv-tensor-rows.ts

@@ -1,212 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import { Dtype } from "../../mlx/ffi";
-import * as ops from "../../mlx/ops";
-import { materializeCopy } from "../../mlx/materialize";
-import { FullPrefillPadding } from "../../model/full-prefill-padding";
-import type { Mask, PrefillPadding } from "../../model/gemma4-base";
-import { cleanupFailure, disposeResources } from "../../engine/resources";
-
-/** Borrowed tensor planes with shared token positions. Plane shapes may differ
- * in head width and dtype. Every plane uses [row, head, token, field]. */
-export interface KvTensorRowView {
-  readonly planes: readonly MlxArray[];
-  readonly rowOffsets: readonly number[];
-  readonly leftPad: readonly number[];
-}
-
-/** Position and membership operations shared by packed KV representations.
- * A codec supplies tensor planes; this owner never encodes or decodes them. */
-export class KvTensorRows implements KvTensorRowView {
-  planes: MlxArray[] = [];
-  rowOffsets: number[] = [];
-  leftPad: number[] = [];
-  #rope: MlxArray | undefined;
-  #beforeRound: number[] | undefined;
-  readonly #padding = new FullPrefillPadding();
-
-  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#padding.restoreEnds(ends); }
-  preparePrefill(padding: PrefillPadding): void {
-    this.#padding.prepare(this, padding); this.#positionsChanged();
-  }
-  finalizePrefill(): void {
-    const next = this.#padding.finalize(this.planes, this);
-    if (next) this.#replace(next, this.rowOffsets, this.leftPad);
-  }
-
-  /** Snapshot logical validity without changing physical row alignment. */
-  captureDonorValidity(): Pick<import("../../model/gemma4-base").KvDonorRows, "offsets" | "starts" | "ends"> {
-    const offsets = this.rowOffsets.map((offset, row) => Math.max(0, this.#padding.validOffset(offset, row)));
-    return { offsets, starts: [...this.leftPad], ends: offsets.map((offset, row) => offset + this.leftPad[row]!) };
-  }
-  get offset(): number { return Math.max(0, ...this.#ends()); }
-  get batchSize(): number | null { return this.rowOffsets.length || null; }
-  #ends(): number[] { return this.rowOffsets.map((offset, row) => offset + this.leftPad[row]!); }
-  get ropeOffsetArr(): MlxArray | undefined {
-    if (this.leftPad.every(pad => pad === 0) && this.rowOffsets.every(offset => offset === this.rowOffsets[0])) return undefined;
-    return this.#rope ??= ops.fromInt32(this.rowOffsets, [this.rowOffsets.length]);
-  }
-  #positionsChanged(): void { this.#rope?.dispose(); this.#rope = undefined; }
-  bytesPerToken(): number {
-    return this.planes.length && this.rowOffsets.length ? this.planes.reduce((bytes, plane) => bytes + plane.nbytes, 0) /
-      (this.rowOffsets.length * this.planes[0]!.shape[2]!) : 0;
-  }
-  trim(count: number): void {
-    this.rowOffsets = this.rowOffsets.map(offset => offset - count); this.#positionsChanged();
-  }
-  specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
-  specRoundCommit(): void { this.#beforeRound = undefined; }
-  specRoundRollback(keep: number | readonly number[]): void {
-    this.rowOffsets = this.#beforeRound!.map((offset, row) => offset + (typeof keep === "number" ? keep : keep[row]!));
-    this.#beforeRound = undefined; this.#positionsChanged();
-  }
-  makeMask(N: number, window: number | null): Mask {
-    const B = this.rowOffsets.length, ends = this.#ends(), S = this.offset + N;
-    if (window === null && this.leftPad.every(pad => pad === 0) && ends.every(end => end === ends[0]))
-      return { mode: N === 1 ? "" : "causal", arr: null };
-    using starts = ops.fromInt32(ends, [B, 1, 1, 1]);
-    using relative = ops.arange(0, N, 1, Dtype.int32);
-    using queries = ops.reshape(relative, [1, 1, N, 1]);
-    using position = ops.add(starts, queries);
-    using indices = ops.arange(0, S, 1, Dtype.int32);
-    using keys = ops.reshape(indices, [1, 1, 1, S]);
-    using pads = ops.fromInt32(this.leftPad, [B, 1, 1, 1]);
-    using causal = ops.lessEqual(keys, position);
-    using valid = ops.greaterEqual(keys, pads);
-    const mask = ops.logicalAnd(causal, valid);
-    if (window === null) return { mode: "array", arr: mask };
-    using windowSize = ops.fromInt32([window], []);
-    using limit = ops.add(keys, windowSize);
-    using within = ops.less(position, limit);
-    try { return { mode: "array", arr: ops.logicalAnd(mask, within) }; }
-    finally { mask.dispose(); }
-  }
-  #replace(planes: MlxArray[], offsets: number[], pads: number[]): void {
-    const previous = this.planes;
-    this.planes = planes; this.rowOffsets = offsets; this.leftPad = pads;
-    this.#positionsChanged(); disposeResources(previous);
-  }
-  append(inputs: readonly MlxArray[]): void {
-    const B = this.rowOffsets.length, N = inputs[0]!.shape[2]!, ends = this.#ends();
-    const needed = Math.max(...ends) + N, capacity = Math.ceil(needed / 256) * 256;
-    const common = ends.every(end => end === ends[0]);
-    using positions = common ? null : ops.fromInt32(ends.flatMap(end =>
-      Array.from({ length: N }, (_, position) => end + position)), [B, 1, N, 1]);
-    const next: MlxArray[] = [];
-    try {
-      for (const [index, input] of inputs.entries()) {
-        const old = this.planes[index];
-        let base: MlxArray;
-        if (!old) base = ops.zeros([B, input.shape[1]!, capacity, input.shape[3]!], input.dtype);
-        else if (old.shape[2]! >= needed) base = old;
-        else {
-          using padding = ops.zeros([B, old.shape[1]!, capacity - old.shape[2]!, old.shape[3]!], old.dtype);
-          base = ops.concatAxis([old, padding], 2);
-        }
-        try {
-          const updated = common
-            ? ops.sliceUpdate(base, input, [0, 0, ends[0]!, 0], [B, base.shape[1]!, ends[0]! + N, base.shape[3]!])
-            : ops.putAlongAxis(base, positions!, input, 2);
-          next.push(updated);
-        } finally { if (base !== old) base.dispose(); }
-      }
-    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    this.#replace(next, this.rowOffsets.map(offset => offset + N), this.leftPad);
-  }
-  mergeRows(rows: readonly KvTensorRowView[]): void {
-    const next: MlxArray[] = [], held: MlxArray[] = [];
-    let offsets: number[], pads: number[];
-    try {
-      if (rows.length === 1) {
-        const row = rows[0]!;
-        for (const plane of row.planes) next.push(ops.copyOf(plane));
-        offsets = [...row.rowOffsets]; pads = [...row.leftPad];
-      } else {
-        const widths = rows.map(row => Math.max(0, ...row.rowOffsets.map((offset, i) => offset + row.leftPad[i]!)));
-        const width = Math.max(0, ...widths);
-        if (width) {
-          const prototype = rows.find((_, index) => widths[index]! > 0)!;
-          for (const [field, geometry] of prototype.planes.entries()) {
-            const pieces = rows.map((row, index) => {
-              const active = widths[index]!;
-              if (!active) {
-                const empty = ops.zeros([row.rowOffsets.length, geometry.shape[1]!, width, geometry.shape[3]!], geometry.dtype);
-                held.push(empty); return empty;
-              }
-              const plane = row.planes[field]!;
-              const view = plane.slice([0, 0, 0, 0], [row.rowOffsets.length, plane.shape[1]!, active, plane.shape[3]!]);
-              held.push(view);
-              if (active === width) return view;
-              using padding = ops.zeros([row.rowOffsets.length, plane.shape[1]!, width - active, plane.shape[3]!], plane.dtype);
-              const padded = ops.concatAxis([padding, view], 2); held.push(padded); return padded;
-            });
-            next.push(ops.concatAxis(pieces, 0));
-          }
-        }
-        offsets = rows.flatMap(row => [...row.rowOffsets]);
-        pads = rows.flatMap((row, index) => row.leftPad.map(pad => pad + width - widths[index]!));
-      }
-    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    finally { disposeResources(held); }
-    this.#replace(next, offsets, pads);
-  }
-  /** Preserve caller-owned physical positions when changing representation.
-   * Live tensor bytes are copied; padding never participates in a codec. */
-  alignRows(leftPad: readonly number[]): void {
-    if (leftPad.every((pad, row) => pad === this.leftPad[row])) return;
-    const width = Math.max(0, ...this.rowOffsets.map((offset, row) => offset + leftPad[row]!));
-    const rows: MlxArray[][] = [], planes: MlxArray[] = [];
-    try {
-      for (let row = 0; row < this.rowOffsets.length; row++) rows.push(this.extractRow(row));
-      for (let field = 0; field < this.planes.length; field++) {
-        const parts: MlxArray[] = [];
-        try {
-          for (let row = 0; row < rows.length; row++) {
-            const geometry = this.planes[field]!, live = rows[row]![field];
-            using before = ops.zeros([1, geometry.shape[1]!, leftPad[row]!, geometry.shape[3]!], geometry.dtype);
-            using after = ops.zeros([1, geometry.shape[1]!, width - leftPad[row]! - this.rowOffsets[row]!, geometry.shape[3]!], geometry.dtype);
-            parts.push(ops.concatAxis(live ? [before, live, after] : [before, after], 2));
-          }
-          planes.push(ops.concatAxis(parts, 0));
-        } finally { disposeResources(parts); }
-      }
-    } catch (error) { return cleanupFailure(error, () => disposeResources(planes)); }
-    finally { disposeResources(rows.flat()); }
-    this.#replace(planes, [...this.rowOffsets], [...leftPad]);
-  }
-  filterRows(keep: readonly number[]): void {
-    if (!keep.length) { this.dispose(); return; }
-    const next: MlxArray[] = [];
-    using indices = this.planes.length ? ops.fromInt32([...keep], [keep.length]) : null;
-    const pads = keep.map(row => this.leftPad[row]!);
-    const sharedPadding = Math.min(...pads);
-    try {
-      for (const plane of this.planes) {
-        const selected = ops.takeAxis(plane, indices!, 0);
-        if (!sharedPadding) next.push(selected);
-        else {
-          try { next.push(selected.slice([0, 0, sharedPadding, 0], [...selected.shape])); }
-          finally { selected.dispose(); }
-        }
-      }
-    }
-    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    this.#replace(next, keep.map(row => this.rowOffsets[row]!), pads.map(pad => pad - sharedPadding));
-    this.#padding.filter(keep);
-  }
-  extractRow(row: number): MlxArray[] {
-    const start = this.leftPad[row]!, count = this.rowOffsets[row]!, result: MlxArray[] = [];
-    if (!count) return result;
-    try {
-      for (const plane of this.planes) {
-        using view = plane.slice([row, 0, start, 0], [row + 1, plane.shape[1]!, start + count, plane.shape[3]!]);
-        result.push(materializeCopy(view));
-      }
-      return result;
-    } catch (error) { return cleanupFailure(error, () => disposeResources(result)); }
-  }
-  dispose(): void {
-    const previous = this.planes;
-    this.planes = []; this.rowOffsets = []; this.leftPad = []; this.#beforeRound = undefined; this.#padding.clear();
-    this.#positionsChanged(); disposeResources(previous);
-  }
-}