mlx-bun 0.4.0 → 1.0.0

package/src/model/rotating-kv-donor.ts

@@ -1,35 +0,0 @@
-import type { KvDonorAttention } from "./gemma4-base";
-import { BatchedRotatingState } from "./batched-rotating-state";
-import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
-import { plainRowStorage, quantizedRowStorage, temporalStorageView } from "./batched-row-storage";
-import { quantizedDonorAttention } from "./kv-attention-view";
-import { decodedKvDonorAttention } from "./decoded-kv-donor";
-import type { RotatingRing } from "./rotating-row-transaction";
-
-/** Capture the same chronological donor rectangle for plain and encoded rings. */
-export function captureRotatingDonorAttention(ring: RotatingRing): KvDonorAttention {
-  const position = new BatchedRotatingState(ring.maxSize, []);
-  position.restore(ring.positionSnapshot);
-  const offsets = position.offsets.map((_, row) => position.validOffset(row));
-  const ends = offsets.map((offset, row) => position.activeLength - (position.offsets[row]! - offset));
-  // Pending prefill keeps a shared physical rectangle, including when mixed
-  // rows are captured separately. Cropping from each row's valid end would
-  // give siblings different widths and can discard an earlier valid window.
-  const from = position.hasPendingPadding ? 0 : Math.max(0, position.activeLength - ring.maxSize);
-  const range = { from, to: position.activeLength };
-  const validity = { offsets,
-    starts: position.leftPad.map((pad, row) => Math.max(0, pad - from, ends[row]! - from - ring.maxSize)),
-    ends: ends.map(end => Math.max(0, end - from)) };
-  if (ring instanceof BatchedRotatingQuantCache) {
-    const keys = temporalStorageView(quantizedRowStorage, ring.keys!, position, range);
-    let values: import("../mlx/ops").QuantizedTensor;
-    try { values = temporalStorageView(quantizedRowStorage, ring.values!, position, range); }
-    catch (error) { quantizedRowStorage.dispose(keys); throw error; }
-    return quantizedDonorAttention(keys, values, validity, ring.groupSize, ring.bits);
-  }
-  const keys = temporalStorageView(plainRowStorage, ring.keys!, position, range);
-  let values: import("../mlx/array").MlxArray;
-  try { values = temporalStorageView(plainRowStorage, ring.values!, position, range); }
-  catch (error) { keys.dispose(); throw error; }
-  return decodedKvDonorAttention({ keys, values, ...validity });
-}

package/src/model/rotating-kv-layout.ts

@@ -1,262 +0,0 @@
-import { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { cloneKvCaches } from "../kv-store";
-import { type BatchableCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding, RotatingKVCache, RotatingQuantizedKVCache } from "./gemma4-base";
-import { BatchedRotatingCache, buildBatchedRotatingMask } from "./batched-rotating";
-import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
-import { BatchedRotatingState, type RotatingPositionSnapshot } from "./batched-rotating-state";
-import { plainRowStorage, quantizedRowStorage, temporalStorageView } from "./batched-row-storage";
-import type { TransitioningKvPositions, TransitionedKvLayout } from "../backends/mlx/transitioning-kv-rows";
-import { captureRotatingDonorAttention } from "./rotating-kv-donor";
-
-import { rollbackRotatingRing } from "./rotating-row-transaction";
-
-const fields = ["packed", "scales", "biases"] as const;
-type Ring = BatchedRotatingCache | BatchedRotatingQuantCache;
-
-function combinedPosition(rings: readonly Ring[]): RotatingPositionSnapshot {
-  const states = rings.map(ring => ring.positionSnapshot);
-  return { ...states[0]!, offsets: states.map(state => state.offsets[0]!),
-    leftPad: states.map(state => state.leftPad[0]!),
-    prefillEnds: states[0]!.prefillEnds
-      ? states.map(state => state.prefillEnds![0]!) : undefined };
-}
-
-/** One logical request with the group's physical ring geometry. This keeps
- * masks stable while different requests cross their precision boundary. */
-export class AlignedRotatingCache implements Cache {
-  minimumReusableOffset = 0;
-  constructor(public inner: Ring) {}
-  get offset(): number { return this.inner.validOffset(0); }
-  get affineConversion(): AlignedRotatingCache | undefined { return this.inner instanceof BatchedRotatingCache ? this : undefined; }
-  get quantizedAttention() { return this.inner instanceof BatchedRotatingQuantCache ? this.inner : undefined; }
-  signature(): string { return this.inner.signature(); }
-  captureDonorAttention() { return captureRotatingDonorAttention(this.inner); }
-  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.inner.updateAndFetch(k, v); }
-  makeMask(tokens: number, window: number | null): Mask { return this.inner.makeMask(tokens, window); }
-  isTrimmable(): boolean { return this.inner.positionSnapshot.totalOffset < this.inner.maxSize; }
-  trim(count: number): void { this.inner.trim(count); }
-  rollback(before: number, keep: number, preserveColumns: boolean): void {
-    this.inner = rollbackRotatingRing(this.inner, [before], keep, preserveColumns);
-  }
-  state(): MlxArray[] { return this.inner.state(); }
-  bytesPerToken(): number {
-    const state = this.state(), tokens = state[0]?.shape[2] ?? 0;
-    return tokens ? state.reduce((n, a) => n + a.nbytes, 0) / tokens : 0;
-  }
-  dispose(): void { this.inner.dispose(); }
-  extract(limit?: number): Cache {
-    const row = this.inner.extractRow(0, limit) ?? (this.inner instanceof BatchedRotatingQuantCache
-      ? new RotatingQuantizedKVCache(this.inner.maxSize, this.inner.groupSize, this.inner.bits)
-      : new RotatingKVCache(this.inner.maxSize));
-    row.minimumReusableOffset = this.minimumReusableOffset;
-    return row;
-  }
-  toQuantized(groupSize: number, bits: number): AlignedRotatingCache {
-    const plain = this.inner as BatchedRotatingCache;
-    const keys = ops.quantize(plain.keys!, groupSize, bits);
-    let values: ops.QuantizedTensor;
-    try { values = ops.quantize(plain.values!, groupSize, bits); }
-    catch (error) { quantizedRowStorage.dispose(keys); throw error; }
-    const next = new AlignedRotatingCache(BatchedRotatingQuantCache.adoptPhysical(keys, values, groupSize, bits, plain.positionSnapshot));
-    next.minimumReusableOffset = this.minimumReusableOffset;
-    this.dispose();
-    return next;
-  }
-}
-
-/** Borrow the newest valid window from a solo cache, including a temporary
- * oversized block. No dequantization or physical-ring normalization in place. */
-export function rotatingSourcePosition(row: RotatingKVCache | RotatingQuantizedKVCache): BatchedRotatingState {
-  const state = new BatchedRotatingState(row.maxSize, [0], [row.offset]);
-  const length = row instanceof RotatingKVCache ? row.keys?.shape[2] ?? 0 : row.keys?.packed.shape[2] ?? 0;
-  state.restore({ maxSize: row.maxSize, offsets: [row.offset], leftPad: [0], totalOffset: row.offset,
-    ringIndex: row.ringIdx, rotated: row.ringIdx < row.offset && row.ringIdx < length });
-  return state;
-}
-
-/** Take ownership of solo rows and align them once at membership changes. */
-export function alignRotatingRows(rows: Cache[]): AlignedRotatingCache[] {
-  const width = Math.max(0, ...rows.map(row => Math.min(row.offset, (row as RotatingKVCache).maxSize)));
-  const prototype = rows.find(row => row.state().length > 0) as RotatingKVCache | RotatingQuantizedKVCache | undefined;
-  const emptyPlane = (field: "keys" | "values"): MlxArray | null => {
-    if (!width || !prototype) return null;
-    const tensor = prototype[field]!;
-    if (tensor instanceof MlxArray) return ops.zeros([1, tensor.shape[1]!, width, tensor.shape[3]!], tensor.dtype);
-    return ops.zeros([1, tensor.packed.shape[1]!, width, tensor.packed.shape[3]! * 32 / (prototype as RotatingQuantizedKVCache).bits], tensor.scales.dtype);
-  };
-  const result: AlignedRotatingCache[] = [];
-  try {
-    for (const source of rows) {
-      if (!(source instanceof RotatingKVCache || source instanceof RotatingQuantizedKVCache))
-        throw new Error(`rotating layout cannot adopt ${source.signature()}`);
-      const length = Math.min(source.offset, source.maxSize), pad = width - length;
-      const position: RotatingPositionSnapshot = { maxSize: source.maxSize, offsets: [source.offset], leftPad: [pad],
-        totalOffset: width, ringIndex: width, rotated: false };
-      const state = rotatingSourcePosition(source);
-      const options = { from: Math.max(0, state.activeLength - length), to: state.activeLength };
-      let inner: Ring;
-      if (source instanceof RotatingKVCache) {
-        using k = source.keys ? temporalStorageView(plainRowStorage, source.keys, state, options) : null;
-        using v = source.values ? temporalStorageView(plainRowStorage, source.values, state, options) : null;
-        const keys = k ? plainRowStorage.padLeft(k, pad) : emptyPlane("keys");
-        let values: MlxArray | null;
-        try { values = v ? plainRowStorage.padLeft(v, pad) : emptyPlane("values"); }
-        catch (error) { keys?.dispose(); throw error; }
-        inner = BatchedRotatingCache.adoptPhysical(keys, values, position);
-      } else {
-        const k = source.keys ? temporalStorageView(quantizedRowStorage, source.keys, state, options) : null;
-        const v = source.values ? temporalStorageView(quantizedRowStorage, source.values, state, options) : null;
-        const empty = (field: "keys" | "values") => {
-          using tensor = emptyPlane(field);
-          return tensor ? ops.quantize(tensor, source.groupSize, source.bits) : null;
-        };
-        try {
-          const keys = k ? quantizedRowStorage.padLeft(k, pad) : empty("keys");
-          let values: ops.QuantizedTensor | null;
-          try { values = v ? quantizedRowStorage.padLeft(v, pad) : empty("values"); }
-          catch (error) { if (keys) quantizedRowStorage.dispose(keys); throw error; }
-          inner = BatchedRotatingQuantCache.adoptPhysical(keys, values, source.groupSize, source.bits, position);
-        } finally { if (k) quantizedRowStorage.dispose(k); if (v) quantizedRowStorage.dispose(v); }
-      }
-      const row = new AlignedRotatingCache(inner); row.minimumReusableOffset = source.minimumReusableOffset ?? 0;
-      result.push(row);
-    }
-  } catch (error) { for (const row of result) row.dispose(); throw error; }
-  for (const row of rows) row.dispose();
-  return result;
-}
-
-/** Position state has no codec or attention arithmetic. */
-export class RotatingKvPositions implements TransitioningKvPositions {
-  #state: BatchedRotatingState;
-  #rope?: MlxArray;
-  constructor(maxSize: number) { this.#state = new BatchedRotatingState(maxSize, []); }
-  get rowOffsets(): readonly number[] { return this.#state.offsets; }
-  get leftPad(): readonly number[] { return this.#state.leftPad; }
-  get offset(): number { return this.#state.totalOffset; }
-  get batchSize(): number | null { return this.#state.batchSize || null; }
-  get ropeOffsetArr(): MlxArray | undefined {
-    if (!this.batchSize) return undefined;
-    return this.#rope ??= ops.fromInt32([...this.rowOffsets], [this.batchSize]);
-  }
-  sync(rows: readonly Cache[]): void {
-    this.#rope?.dispose(); this.#rope = undefined;
-    if (!rows.length) { this.#state = new BatchedRotatingState(this.#state.maxSize, []); return; }
-    this.#state.restore(combinedPosition(rows.map(row => (row as AlignedRotatingCache).inner)));
-  }
-  filterRows(keep: readonly number[]): void { this.#state.filter(keep); this.#rope?.dispose(); this.#rope = undefined; }
-  makeMask(tokens: number, window: number | null): Mask {
-    return { mode: "array", arr: buildBatchedRotatingMask(this.#state.batchSize, tokens, this.#state.leftPad,
-      this.#state.maxSize, window ?? this.#state.maxSize, this.#state.ringIndex, this.#state.totalOffset, this.#state.rotated,
-      tokens > 1 || this.#state.hasPendingPadding) };
-  }
-  dispose(): void { this.#rope?.dispose(); this.#rope = undefined; this.#state = new BatchedRotatingState(this.#state.maxSize, []); }
-}
-
-/** Once all rows use affine storage, concatenate their already-aligned planes
- * without changing ring columns or rebuilding any attention mask. */
-export class RotatingAffineLayout implements TransitionedKvLayout, PaddedPrefillCache {
-  #inner?: BatchedRotatingQuantCache;
-  #before?: number[];
-  #minimum: number[] = [];
-  constructor(readonly maxSize: number, readonly groupSize: number, readonly bits: number) {}
-  get minimumReusableOffset(): number { return Math.max(0, ...this.#minimum); }
-  get quantizedAttention() { return this.#inner; }
-  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
-    return captureRotatingDonorAttention(this.#inner!);
-  }
-  get rowOffsets(): readonly number[] { return this.#inner?.offsetArr ?? []; }
-  get leftPad(): readonly number[] { return this.#inner?.leftPad ?? []; }
-  get offset(): number { return this.#inner?.offset ?? 0; }
-  get batchSize(): number | null { return this.#inner?.batchSize || null; }
-  get ropeOffsetArr(): MlxArray | undefined { return this.#inner?.ropeOffsetArr; }
-  get maxTokens(): number { return this.maxSize; }
-  signature(): string { return `kv:rotating-affine:${this.maxSize}:${this.groupSize}:${this.bits}`; }
-  makeEmptyBatch(): RotatingAffineLayout { return new RotatingAffineLayout(this.maxSize, this.groupSize, this.bits); }
-  makeMask(tokens: number, window: number | null): Mask { return this.#inner!.makeMask(tokens, window); }
-  preparePrefill(padding: PrefillPadding): void { this.#inner!.preparePrefill(padding); }
-  finalizePrefill(): void { this.#inner!.finalizePrefill(); }
-  bytesPerToken(): number {
-    const tokens = this.#inner?.keys?.packed.shape[2] ?? 0;
-    return tokens ? this.state().reduce((n, a) => n + a.nbytes, 0) / (tokens * this.#inner!.batchSize) : 0;
-  }
-  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
-  isTrimmable(): boolean { return this.#inner!.positionSnapshot.totalOffset < this.maxSize; }
-  trim(count: number): void { this.#inner!.trim(count); }
-  specRoundBegin(): void { this.#before = [...this.rowOffsets]; }
-  specRoundCommit(): void { this.#before = undefined; }
-  specRoundRollback(keep: number | readonly number[]): void {
-    this.#inner = rollbackRotatingRing(this.#inner!, this.#before!, keep);
-    this.#before = undefined;
-  }
-  updateAndFetch(): [MlxArray, MlxArray] { return this.#inner!.updateAndFetch(); }
-  releaseRopeArr(): void { this.#inner?.releaseRopeArr(); }
-  mergeRows(rows: readonly Cache[]): void {
-    const owned: Cache[] = [];
-    try {
-      for (const row of rows) {
-        if (row instanceof RotatingAffineLayout) {
-          for (let i = 0; i < (row.batchSize ?? 0); i++) owned.push(row.extractRow(i));
-        } else if (row instanceof AlignedRotatingCache) owned.push(row.extract());
-        else owned.push(...cloneKvCaches([row]));
-      }
-      if (!owned.length) { this.dispose(); return; }
-      const aligned = alignRotatingRows(owned); owned.length = 0;
-      try { this.adoptAlignedRows(aligned); } finally { for (const row of aligned) row.dispose(); }
-    } finally { for (const row of owned) row.dispose(); }
-  }
-  adoptAlignedRows(rows: readonly AlignedRotatingCache[]): void {
-    const rings = rows.map(row => (row as AlignedRotatingCache).inner as BatchedRotatingQuantCache);
-    const position = combinedPosition(rings);
-    if (rings.every(row => !row.keys)) {
-      const empty = BatchedRotatingQuantCache.adoptPhysical(null, null, this.groupSize, this.bits,
-        position);
-      this.#inner?.dispose(); this.#inner = empty; this.#minimum = rows.map(row => row.minimumReusableOffset ?? 0);
-      return;
-    }
-    const merge = (field: "keys" | "values") => {
-      const planes: MlxArray[] = [];
-      try {
-        for (const part of fields) planes.push(ops.concatAxis(rings.map(row => row[field]![part]), 0));
-        return { packed: planes[0]!, scales: planes[1]!, biases: planes[2]! };
-      } catch (error) { for (const plane of planes) plane.dispose(); throw error; }
-    };
-    const keys = merge("keys"); let values: ops.QuantizedTensor;
-    try { values = merge("values"); } catch (error) { quantizedRowStorage.dispose(keys); throw error; }
-    const inner = BatchedRotatingQuantCache.adoptPhysical(keys, values, this.groupSize, this.bits,
-      position);
-    this.#inner?.dispose(); this.#inner = inner; this.#minimum = rows.map(row => row.minimumReusableOffset ?? 0);
-  }
-  filterRows(keep: readonly number[]): void { this.#inner!.filterRows(keep); this.#minimum = keep.map(row => this.#minimum[row]!); }
-  extractRow(row: number): Cache {
-    return this.extract(row);
-  }
-  protected extract(row: number, limit?: number): Cache {
-    const result = this.#inner!.extractRow(row, limit) ?? new RotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits);
-    result.minimumReusableOffset = this.#minimum[row] ?? 0; return result;
-  }
-  state(): MlxArray[] { return this.#inner?.state() ?? []; }
-  dispose(): void { this.#inner?.dispose(); this.#inner = undefined; this.#minimum = []; this.#before = undefined; }
-}
-
-/** A committed speculative checkpoint contains only the newest valid window. */
-export class SpeculativeRotatingAffineLayout extends RotatingAffineLayout {
-  #ropeForOffset = -1;
-  override get ropeOffsetArr(): MlxArray | undefined {
-    if (this.#ropeForOffset !== this.offset) {
-      this.releaseRopeArr(); this.#ropeForOffset = this.offset;
-    }
-    return super.ropeOffsetArr;
-  }
-  override makeEmptyBatch(): SpeculativeRotatingAffineLayout {
-    return new SpeculativeRotatingAffineLayout(this.maxSize, this.groupSize, this.bits);
-  }
-  override makeMask(tokens: number, window: number | null): Mask {
-    if (this.leftPad.every(pad => pad === 0) &&
-        this.offset + tokens <= Math.min(this.maxSize, window ?? this.maxSize))
-      return { mode: tokens === 1 ? "" : "causal", arr: null };
-    return super.makeMask(tokens, window);
-  }
-  override extractRow(row: number): Cache { return this.extract(row, this.maxSize); }
-}

package/src/model/rotating-row-transaction.ts

@@ -1,35 +0,0 @@
-import * as ops from "../mlx/ops";
-import { BatchedRotatingCache } from "./batched-rotating";
-import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
-import { rotatingRollbackPlan } from "./batched-rotating-state";
-import { plainRowStorage, quantizedRowStorage, rowRollIndices } from "./batched-row-storage";
-
-export type RotatingRing = BatchedRotatingCache | BatchedRotatingQuantCache;
-
-/** Retain accepted columns without decoding or recomputing the ring. Ownership
- * moves to the returned ring; aligned members preserve their shared rectangle. */
-export function rollbackRotatingRing(ring: BatchedRotatingCache, before: readonly number[], keep: number | readonly number[], preserveColumns?: boolean): BatchedRotatingCache;
-export function rollbackRotatingRing(ring: BatchedRotatingQuantCache, before: readonly number[], keep: number | readonly number[], preserveColumns?: boolean): BatchedRotatingQuantCache;
-export function rollbackRotatingRing(ring: RotatingRing, before: readonly number[], keep: number | readonly number[], preserveColumns?: boolean): RotatingRing;
-export function rollbackRotatingRing(ring: RotatingRing, before: readonly number[],
-  keep: number | readonly number[], preserveColumns = false): RotatingRing {
-  const plan = rotatingRollbackPlan(ring.positionSnapshot, before, keep, preserveColumns);
-  if ("trim" in plan) { ring.trim(plan.trim); ring.releaseRopeArr(); return ring; }
-  let next: RotatingRing;
-  if (ring instanceof BatchedRotatingQuantCache) {
-    using indices = rowRollIndices(ring.keys!.packed.shape[2]!, plan.shifts);
-    const keys = quantizedRowStorage.rollRows(ring.keys!, indices);
-    let values: ops.QuantizedTensor;
-    try { values = quantizedRowStorage.rollRows(ring.values!, indices); }
-    catch (error) { quantizedRowStorage.dispose(keys); throw error; }
-    next = BatchedRotatingQuantCache.adoptPhysical(keys, values, ring.groupSize, ring.bits, plan.position);
-  } else {
-    using indices = rowRollIndices(ring.keys!.shape[2]!, plan.shifts);
-    const keys = plainRowStorage.rollRows(ring.keys!, indices);
-    let values: import("../mlx/array").MlxArray;
-    try { values = plainRowStorage.rollRows(ring.values!, indices); }
-    catch (error) { keys.dispose(); throw error; }
-    next = BatchedRotatingCache.adoptPhysical(keys, values, plan.position);
-  }
-  ring.dispose(); return next;
-}

package/src/model/speculative-rotating-kv.ts

@@ -1,101 +0,0 @@
-import { rotatingSourcePosition } from "./rotating-kv-layout";
-import { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { type BatchableCache, type Cache, type Mask, type PrefillPadding, RotatingKVCache } from "./gemma4-base";
-import { BatchedRotatingCache } from "./batched-rotating";
-import { rollbackRotatingRing } from "./rotating-row-transaction";
-import { BatchedRotatingState } from "./batched-rotating-state";
-import { plainRowStorage, temporalStorageView } from "./batched-row-storage";
-import { disposeResources } from "../engine/resources";
-
-/** Ring transactions retain accepted KV columns, including rounds that cross
- * the window boundary. The verify block already retains the complete history
- * needed by its earliest query; no target recomputation or scheduler policy. */
-export class SpeculativeRotatingKVCache implements BatchableCache {
-  #inner: BatchedRotatingCache;
-  #before?: number[];
-  constructor(readonly maxSize: number) { this.#inner = new BatchedRotatingCache(maxSize, []); }
-  signature(): string { return "kv:rotating-plain"; }
-  get offset(): number { return this.#inner.offset; }
-  get rowOffsets(): readonly number[] { return this.#inner.offsetArr; }
-  get leftPad(): readonly number[] { return this.#inner.leftPad; }
-  get batchSize(): number { return this.#inner.batchSize; }
-  get ropeOffsetArr(): MlxArray { return this.#inner.ropeOffsetArr; }
-  makeEmptyBatch(): SpeculativeRotatingKVCache { return new SpeculativeRotatingKVCache(this.maxSize); }
-  makeMask(tokens: number, window: number | null): Mask { return this.#inner.makeMask(tokens, window); }
-  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#inner.updateAndFetch(k, v); }
-  preparePrefill(padding: PrefillPadding): void { this.#inner.preparePrefill(padding); }
-  finalizePrefill(): void { this.#inner.finalizePrefill(); }
-  state(): MlxArray[] { return this.#inner.state(); }
-  isTrimmable(): boolean { return this.#inner.isTrimmable(); }
-  trim(count: number): void { this.#inner.trim(count); }
-  bytesPerToken(): number {
-    const state = this.state();
-    return state.length ? state.reduce((n, a) => n + a.nbytes, 0) / (this.batchSize * state[0]!.shape[2]!) : 0;
-  }
-  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
-  extractRow(row: number): Cache {
-    const result = new RotatingKVCache(this.maxSize);
-    if (!this.rowOffsets[row] || !this.#inner.keys || !this.#inner.values) return result;
-    const position = new BatchedRotatingState(this.maxSize, []); position.restore(this.#inner.positionSnapshot);
-    const range = { row, from: Math.max(0, position.leftPad[row]!, position.activeLength - this.maxSize),
-      to: position.activeLength, copy: true };
-    const keys = temporalStorageView(plainRowStorage, this.#inner.keys, position, range);
-    let values: MlxArray;
-    try { values = temporalStorageView(plainRowStorage, this.#inner.values, position, range); }
-    catch (error) { keys.dispose(); throw error; }
-    result.restoreState(keys, values, this.rowOffsets[row]!, keys.shape[2]!); return result;
-  }
-  captureDonorRows(): import("./gemma4-base").KvDonorRows {
-    const position = new BatchedRotatingState(this.maxSize, []); position.restore(this.#inner.positionSnapshot);
-    const from = Math.max(0, position.activeLength - this.maxSize);
-    const range = { from, to: position.activeLength };
-    const keys = temporalStorageView(plainRowStorage, this.#inner.keys!, position, range); let values: MlxArray;
-    try { values = temporalStorageView(plainRowStorage, this.#inner.values!, position, range); }
-    catch (error) { keys.dispose(); throw error; }
-    return { keys, values, offsets: [...this.rowOffsets], starts: this.leftPad.map(pad => Math.max(0, pad - from)),
-      ends: this.rowOffsets.map(() => position.activeLength - from) };
-  }
-
-  filterRows(keep: readonly number[]): void { this.#inner.filterRows(keep); }
-  mergeRows(sources: readonly Cache[]): void {
-    const only = sources.length === 1 ? sources[0] : undefined;
-    if (only instanceof SpeculativeRotatingKVCache) {
-      const keys = only.#inner.keys ? ops.copyOf(only.#inner.keys) : null;
-      let values: MlxArray | null;
-      try { values = only.#inner.values ? ops.copyOf(only.#inner.values) : null; }
-      catch (error) { keys?.dispose(); throw error; }
-      const next = BatchedRotatingCache.adoptPhysical(keys, values, only.#inner.positionSnapshot);
-      this.#inner.dispose(); this.#inner = next; return;
-    }
-    const held: Cache[] = [], planes: { keys: MlxArray; values: MlxArray }[] = [], offsets: number[] = [];
-    try {
-      const rows = sources.flatMap(source => source instanceof SpeculativeRotatingKVCache
-        ? Array.from({ length: source.batchSize }, (_, row) => { const cache = source.extractRow(row); held.push(cache); return cache; })
-        : [source]) as RotatingKVCache[];
-      const prototype = rows.find(row => row.keys !== null);
-      if (!prototype) {
-        this.#inner.dispose(); this.#inner = new BatchedRotatingCache(this.maxSize, rows.map(() => 0)); return;
-      }
-      for (const cache of rows) {
-        const state = rotatingSourcePosition(cache);
-        const range = { from: Math.max(0, state.activeLength - this.maxSize), to: state.activeLength };
-        const plane = (field: "keys" | "values") => cache[field]
-          ? temporalStorageView(plainRowStorage, cache[field]!, state, range)
-          : ops.zeros([1, prototype[field]!.shape[1]!, 0, prototype[field]!.shape[3]!], prototype[field]!.dtype);
-        const keys = plane("keys"); let values: MlxArray;
-        try { values = plane("values"); } catch (error) { keys.dispose(); throw error; }
-        planes.push({ keys, values }); offsets.push(cache.offset);
-      }
-      const next = BatchedRotatingCache.merge(planes, offsets, this.maxSize);
-      this.#inner.dispose(); this.#inner = next;
-    } finally { disposeResources([...held, ...planes.flatMap(row => [row.keys, row.values])]); }
-  }
-  specRoundBegin(): void { this.#before = [...this.rowOffsets]; }
-  specRoundCommit(): void { this.#before = undefined; }
-  specRoundRollback(keep: number | readonly number[]): void {
-    this.#inner = rollbackRotatingRing(this.#inner, this.#before!, keep);
-    this.#before = undefined;
-  }
-  dispose(): void { this.#before = undefined; this.#inner.dispose(); }
-}

package/src/model/ssm-prefill-padding.ts

@@ -1,39 +0,0 @@
-import type { PrefillPadding } from "./gemma4-base";
-import type { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
-
-/** ArraysCache's remaining lengths and padding from mlx-lm (MIT).
- * Recurrent state has no token axis to roll: masked tokens must not update it. */
-export class SsmPrefillPadding {
-  #left: number[];
-  #lengths: number[];
-  constructor(padding: PrefillPadding) {
-    this.#left = padding.leftPadding ? [...padding.leftPadding] : padding.lengths.map(() => 0);
-    this.#lengths = [...padding.lengths];
-  }
-  convolutionLengths(tokens: number): number[] {
-    return this.#lengths.map(length => Math.max(0, Math.min(tokens, length)));
-  }
-  makeMask(tokens: number): MlxArray {
-    const B = this.#lengths.length;
-    using arange = ops.arange(0, tokens, 1, Dtype.int32);
-    using positions = ops.reshape(arange, [1, tokens]);
-    using left = ops.fromInt32(this.#left, [B, 1]);
-    using ends = ops.fromInt32(this.#lengths, [B, 1]);
-    using after = ops.greaterEqual(positions, left);
-    using before = ops.less(positions, ends);
-    return ops.logicalAnd(after, before);
-  }
-  advance(tokens: number): number[] {
-    const counts = this.#lengths.map((length, row) => Math.max(0,
-      Math.min(tokens, length) - Math.max(0, Math.min(tokens, this.#left[row]!))));
-    this.#lengths = this.#lengths.map(length => length - tokens);
-    this.#left = this.#left.map(left => left - tokens);
-    return counts;
-  }
-  filter(keep: readonly number[]): void {
-    this.#left = keep.map(row => this.#left[row]!);
-    this.#lengths = keep.map(row => this.#lengths[row]!);
-  }
-}

package/src/model/support.ts

@@ -1,102 +0,0 @@
-import type { ModelConfig } from "../config";
-import { GENERIC_MODEL_TYPES, genericArgsFor, remapModelType } from "./universal/archs";
-
-export function isMiniCPM5Config(config: ModelConfig): boolean {
-  const t = config.text;
-  return config.modelType === "llama" &&
-    t.hiddenSize === 1536 &&
-    t.numHiddenLayers === 24 &&
-    t.numAttentionHeads === 16 &&
-    t.numKeyValueHeads === 2 &&
-    t.headDim === 128 &&
-    t.vocabSize === 130560 &&
-    t.tieWordEmbeddings === false;
-}
-
-/** Qwen3.5 hybrid gated-DeltaNet family (model_type qwen3_5 / qwen3_5_text).
- *  Dense MLP only for now — the MoE variant (qwen3_5_moe) is deferred. */
-export function isQwen35Config(config: ModelConfig): boolean {
-  return (config.modelType === "qwen3_5" || config.modelType === "qwen3_5_text") &&
-    !config.text.enableMoeBlock &&
-    config.text.numExperts === 0 &&
-    config.text.linearNumValueHeads > 0 &&
-    config.text.fullAttentionInterval > 0;
-}
-
-/** Plain Qwen3 (model_type `qwen3`, Qwen3ForCausalLM) — a standard dense decoder
- *  with per-head q/k norm and tied embeddings. Used here as the text-embedding
- *  backbone (mlx-community Qwen3-Embedding-*). Distinct from the qwen3_5 hybrid. */
-export function isQwen3Config(config: ModelConfig): boolean {
-  return config.modelType === "qwen3";
-}
-
-/** Qwen3-MoE (model_type `qwen3_moe`, Qwen3MoeForCausalLM) — plain-qwen3
- *  attention (per-head q/k norm, GQA, RoPE) with a sparse MoE FFN
- *  (gate Linear -> precise softmax -> top-k -> renormalize -> SwitchGLU experts
- *  with compiled swiglu). Distinct from the qwen3 embedding backbone and the
- *  qwen3_5 gated-DeltaNet hybrid. */
-export function isQwen3MoeConfig(config: ModelConfig): boolean {
-  return config.modelType === "qwen3_moe";
-}
-
-/** GLM-5.2 Colibri architecture. Its dedicated MLA/DSA/shared-MoE graph is
- *  intentionally distinct from the universal dense `glm4` descriptor. */
-export function isGlm52Config(config: ModelConfig): boolean {
-  return config.modelType === "glm_moe_dsa";
-}
-
-/** Speculative-decoding drafters (e.g. `gemma4_assistant`) are companion
- *  artifacts to a target model — Q-only, centroid-head, no standalone LM
- *  head. They are never servable/selectable on their own (the spec path
- *  loads them by explicit path), so they must be excluded from model
- *  resolution and the supported-model lists. See src/spec/drafter.ts. */
-export function isDrafterModelType(modelType: string): boolean {
-  return modelType.endsWith("_assistant");
-}
-
-/** DiffusionGemma (model_type `diffusion_gemma`): block/masked-diffusion canvas
- *  model. Non-autoregressive — routed through the diffusion engine, not the AR
- *  loop. See docs/design/generic-model-support.md. */
-export function isDiffusionGemmaConfig(config: ModelConfig): boolean {
-  return config.modelType === "diffusion_gemma";
-}
-
-/** Support tier for a registry record (docs/design/generic-model-support.md):
- *  "targeted" = dedicated/generated forward + L2/L3 paths;
- *  "generic"  = the Tier-0 universal module (L1 monolith only);
- *  null       = unsupported. Generic never shadows targeted. */
-export function supportTier(modelType: string, repoId = ""): "targeted" | "generic" | null {
-  if (isDrafterModelType(modelType)) return null;
-  if (modelType.startsWith("gemma4")) return "targeted";
-  if (modelType === "diffusion_gemma") return "targeted";
-  // qwen3_5 / qwen3_5_text (dense hybrid). The MoE variant is a separate type.
-  if (modelType === "qwen3_5" || modelType === "qwen3_5_text") return "targeted";
-  if (modelType === "qwen3") return "targeted"; // plain Qwen3 (embedding backbone)
-  if (modelType === "qwen3_moe") return "targeted"; // Qwen3-MoE (sparse experts)
-  if (modelType === "glm_moe_dsa") return "targeted"; // dedicated GLM-5.2 MLA/DSA port
-  if (modelType === "llama" && repoId.toLowerCase().includes("minicpm5-1b-optiq-4bit"))
-    return "targeted";
-  // Tier-0 generic fallback: the universal-dense descriptor table.
-  if (GENERIC_MODEL_TYPES.has(remapModelType(modelType))) return "generic";
-  return null;
-}
-
-export function isSupportedModelRecord(modelType: string, repoId = ""): boolean {
-  return supportTier(modelType, repoId) !== null;
-}
-
-export function isSupportedModelConfig(config: ModelConfig): boolean {
-  if (isDrafterModelType(config.modelType)) return false;
-  if (
-    config.modelType.startsWith("gemma4") || isMiniCPM5Config(config) ||
-    isQwen35Config(config) || isQwen3Config(config) || isQwen3MoeConfig(config) ||
-    isDiffusionGemmaConfig(config) || isGlm52Config(config)
-  ) return true;
-  // Generic tier: an arch descriptor exists for this model_type (and its
-  // config parses — a malformed config is unsupported, not a crash).
-  try {
-    return genericArgsFor(config) !== null;
-  } catch {
-    return false;
-  }
-}

package/src/model/trellis-balanced-scatter.ts

@@ -1,51 +0,0 @@
-// Experimental variant 8, k=3 / T=256 / L<=12 only.
-// The word-per-lane kernel uses 24 of 32 SIMD lanes for a 3-bit block.
-// Assign eight outputs to every lane, load adjacent packed words, and align
-// them once so each unrolled decoder has constant shifts. Row accumulation,
-// split boundaries and the separate partial reduction stay identical to v6.
-export const BALANCED_SCATTER_SOURCE = String.raw`
-  const uint lane = thread_index_in_simdgroup;
-  const uint sg = simdgroup_index_in_threadgroup;
-  const uint block = thread_position_in_grid.y * (uint)SG_TG + sg;
-  const uint nBlocks = (uint)C / (uint)BT;
-  const uint zi = thread_position_in_grid.z;
-  const uint sample = zi / (uint)SPLITS;
-  const uint split = zi - sample * (uint)SPLITS;
-  if (block >= nBlocks || sample >= (uint)M) return;
-  const uint wpb = (uint)(BT * K / 32);
-  const ulong rowWords = (ulong)(INTERLEAVE ? 2 : nBlocks) * wpb;
-  const uint pBase = lane * 8u * (uint)K;
-  const uint wi = pBase >> 5u;
-  const uint off = pBase & 31u;
-  const uint next = (wi + 1u == wpb) ? 0u : wi + 1u;
-  const device uint32_t* words = codes + (INTERLEAVE
-    ? (ulong)(block / 2) * R * 2 * wpb + (block % 2) * wpb : (ulong)block * wpb);
-  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
-  const uint r0 = split * rowsPer;
-  const uint r1 = metal::min((uint)R, r0 + rowsPer);
-  const device T* xs = x + (ulong)sample * (ulong)R;
-  float acc[8];
-  #pragma clang loop unroll(full)
-  for (uint i = 0; i < 8u; ++i) acc[i] = 0.0f;
-  for (uint r = r0; r < r1; ++r) {
-    const uint word = words[(ulong)r * rowWords + wi];
-    const uint nxt = words[(ulong)r * rowWords + next];
-    const uint lo = off == 0u ? word : (word >> off) | (nxt << (32u - off));
-    const uint hi = off == 0u ? nxt : nxt >> off;
-    const float xr = float(xs[r]);
-    const float sr = float(scales[r]);
-    #pragma clang loop unroll(full)
-    for (uint i = 0; i < 8u; ++i) {
-      const uint shift = i * (uint)K;
-      uint win = lo >> shift;
-      if (shift + (uint)L > 32u) win |= hi << (32u - shift);
-      win &= (1u << (uint)L) - 1u;
-      const float value = (float)trellis_y(win) * (1.0f / 147.800537109375f);
-      acc[i] = metal::fma(value * sr, xr, acc[i]);
-    }
-  }
-  const ulong outBase = ((ulong)sample * (ulong)SPLITS + split) * (ulong)C;
-  #pragma clang loop unroll(full)
-  for (uint i = 0; i < 8u; ++i)
-    partial[outBase + block * (uint)BT + (uint)BT - 1u - lane * 8u - i] = acc[i];
-`;