mlx-bun 0.2.0 → 0.4.0

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

@@ -0,0 +1,54 @@
+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

@@ -0,0 +1,65 @@
+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

@@ -0,0 +1,130 @@
+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

@@ -0,0 +1,38 @@
+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

@@ -0,0 +1,9 @@
+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

@@ -0,0 +1,133 @@
+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

@@ -0,0 +1,212 @@
+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);
+  }
+}