mlx-bun 0.3.0 → 0.4.0

package/src/model/glm52-cache.ts

@@ -11,6 +11,8 @@
 // row's logical position without reconstructing per-head K/V.
 
 import { MlxArray } from "../mlx/array";
+import { FullPrefillPadding } from "./full-prefill-padding";
+import type { PrefillPadding } from "./gemma4-base";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import {
@@ -210,6 +212,8 @@ export class Glm52Cache implements BatchableCache {
   batchSize: number | null = null;
   rowOffsets: number[] = [];
   leftPad: number[] = [];
+  #beforeRound: number[] | undefined;
+  readonly #padding = new FullPrefillPadding(1);
 
   constructor(geometry: MLACacheGeometry) {
     positiveInteger("MLA kvLoraRank", geometry.kvLoraRank);
@@ -482,6 +486,56 @@ export class Glm52Cache implements BatchableCache {
     this.rowOffsets = this.rowOffsets.map((value) => value - n);
   }
 
+  preparePrefill(padding: PrefillPadding): void { this.#padding.prepare(this, padding); }
+  finalizePrefill(): void {
+    const next = this.#padding.finalize(this.state(), this);
+    if (!next) return;
+    this.latent?.dispose(); this.rope?.dispose();
+    this.latent = next[0] ?? null; this.rope = next[1] ?? null;
+    if (this.dsa && next[2]) this.dsa.restoreState(next[2], this.offset);
+  }
+
+  specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#beforeRound = undefined; }
+
+  /** Keep each accepted prefix in compressed form. Right alignment lets the
+   * next shared append reuse the existing MLA/DSA attention geometry. */
+  specRoundRollback(keep: number | readonly number[]): void {
+    const nextOffsets = this.#beforeRound!.map((before, row) => before +
+      (typeof keep === "number" ? keep : keep[row]!));
+    const rejected = this.rowOffsets.map((offset, row) => offset - nextOffsets[row]!);
+    if (rejected.every(count => count === rejected[0])) {
+      const count = this.batchSize!;
+      this.trim(rejected[0]!);
+      if (this.batchSize === null) {
+        const empty = this.makeEmptyBatch();
+        try { this.mergeRows(Array.from({ length: count }, () => empty)); }
+        finally { empty.dispose(); }
+      }
+      this.#beforeRound = undefined; return;
+    }
+    const width = this.offset;
+    const pads = this.leftPad.map((pad, row) => pad + rejected[row]!);
+    const removable = Math.min(...pads);
+    using indices = ops.fromInt32(rejected.flatMap(count =>
+      Array.from({ length: width }, (_, column) => (column - count + width) % width)),
+      [this.batchSize!, width, 1]);
+    const next: MlxArray[] = [];
+    try {
+      for (const array of this.state()) {
+        using shifted = ops.takeAlongAxis(array, indices, 1);
+        next.push(shifted.slice([0, removable, 0], [this.batchSize!, width, array.shape[2]!]));
+      }
+    } catch (error) { for (const array of next) array.dispose(); throw error; }
+    this.latent!.dispose(); this.rope!.dispose();
+    this.latent = next[0]!; this.rope = next[1]!;
+    if (this.dsa) this.dsa.restoreState(next[2]!, width - removable);
+    this.offset = width - removable;
+    this.rowOffsets = nextOffsets;
+    this.leftPad = pads.map(pad => pad - removable);
+    this.#beforeRound = undefined;
+  }
+
   /** Exact logical bytes of valid f32 compressed state. */
   get byteLength(): number {
     if (this.batchSize === null) return 0;
@@ -529,8 +583,16 @@ export class Glm52Cache implements BatchableCache {
         ) {
           throw new Error("MLA mergeRows cache geometry does not match");
         }
-        if (generic.batchSize === null || generic.offset === 0)
-          throw new Error("MLA mergeRows cannot merge an empty row");
+        if (generic.batchSize === null || generic.offset === 0) {
+          const count = generic.batchSize ?? 1;
+          for (let row = 0; row < count; row++) {
+            states.push({ latent: ops.zeros([1, 0, this.kvLoraRank], Dtype.float32),
+              rope: ops.zeros([1, 0, this.ropeHeadDim], Dtype.float32),
+              dsa: this.dsa ? ops.zeros([1, 0, this.dsa.headDim], Dtype.float32) : null });
+            offsets.push(0);
+          }
+          continue;
+        }
         for (let row = 0; row < generic.batchSize; row++) {
           states.push(generic.fetchRow(row));
           offsets.push(generic.rowOffset(row));
@@ -577,7 +639,10 @@ export class Glm52Cache implements BatchableCache {
       this.offset = width;
       this.rowOffsets = [...offsets];
       this.leftPad = offsets.map((value) => width - value);
-      if (this.dsa) this.dsa.restoreState(dsa!, width);
+      if (this.dsa) {
+        if (width) this.dsa.restoreState(dsa!, width);
+        else { this.dsa.data = dsa!; this.dsa.batchSize = offsets.length; }
+      }
     } finally {
       for (const state of states) {
         state.latent.dispose();
@@ -591,6 +656,7 @@ export class Glm52Cache implements BatchableCache {
   extractRow(row: number): Glm52Cache {
     if (this.batchSize === null) throw new Error("MLA cache is empty");
     validateRowIndex(row, this.batchSize, "MLA extract");
+    if (this.rowOffsets[row] === 0) return this.makeEmptyBatch();
     const state = this.fetchRow(row);
     const copy = (array: MlxArray): MlxArray => ops.mulScalar(array, 1);
     const latent = copy(state.latent);
@@ -666,7 +732,15 @@ export class Glm52Cache implements BatchableCache {
     this.offset = nextWidth;
     this.rowOffsets = nextOffsets;
     this.leftPad = selectedPad.map((value) => value - removablePad);
-    if (this.dsa) this.dsa.restoreState(dsa!, nextWidth);
+    this.#padding.filter(keep);
+    if (this.dsa) {
+      if (nextWidth) this.dsa.restoreState(dsa!, nextWidth);
+      else {
+        // A retained cold row still owns batch membership, even though the
+        // strict persisted-state restore contract requires nonempty tokens.
+        this.dsa.dispose(); this.dsa.data = dsa!; this.dsa.batchSize = keep.length;
+      }
+    }
   }
 
   projectedBytes(tokens: number): number {
@@ -709,6 +783,8 @@ export class Glm52Cache implements BatchableCache {
     this.batchSize = null;
     this.rowOffsets = [];
     this.leftPad = [];
+    this.#beforeRound = undefined;
+    this.#padding.clear();
   }
 }
 

package/src/model/glm52-mla.ts

@@ -484,7 +484,7 @@ export class Glm52Mla {
         }
         return this.attendAbsorbed(projection, state, selectedPositions);
       }
-      return this.attendReconstructed(projection, state);
+      return this.attendReconstructed(projection, state, cache.leftPad);
     } finally {
       disposeState(state);
       disposeProjection(projection);
@@ -495,6 +495,7 @@ export class Glm52Mla {
   attendReconstructed(
     projection: Pick<Glm52MlaProjection, "qNope" | "qRope">,
     state: Pick<MLACompressedState, "latent" | "rope">,
+    leftPadding: readonly number[] = [],
   ): MlxArray {
     const [batch, tokens] = projection.qNope.shape;
     const keys = state.latent.shape[1]!;
@@ -570,14 +571,19 @@ export class Glm52Mla {
 
     const prefixTokens = keys - tokens!;
     let maskedScores = scores;
-    if (tokens! > 1) {
-      const maskValues = new Float32Array(tokens! * keys);
-      for (let queryToken = 0; queryToken < tokens!; queryToken++) {
+    const padded = leftPadding.some(value => value > 0);
+    if (tokens! > 1 || padded) {
+      const maskRows = padded ? batch! : 1;
+      const maskValues = new Float32Array(maskRows * tokens! * keys);
+      for (let row = 0; row < maskRows; row++) for (let queryToken = 0; queryToken < tokens!; queryToken++) {
         const lastVisible = prefixTokens + queryToken;
+        const start = (row * tokens! + queryToken) * keys;
+        for (let keyToken = 0; keyToken < (leftPadding[row] ?? 0); keyToken++)
+          maskValues[start + keyToken] = -Infinity;
         for (let keyToken = lastVisible + 1; keyToken < keys; keyToken++)
-          maskValues[queryToken * keys + keyToken] = -Infinity;
+          maskValues[start + keyToken] = -Infinity;
       }
-      const mask = MlxArray.fromFloat32(maskValues, [1, 1, tokens!, keys]);
+      const mask = MlxArray.fromFloat32(maskValues, [maskRows, 1, tokens!, keys]);
       maskedScores = ops.add(scores, mask);
       mask.dispose();
       scores.dispose();

package/src/model/glm52.ts

@@ -144,8 +144,6 @@ class Glm52DsaIndexer {
     forceDense = false,
   ): DsaStep {
     const [batch, tokens, hidden] = input.shape;
-    if (batch! > 1 && tokens !== 1)
-      throw new Error("batched DSA supports ordinary single-token decode only");
     if (hidden !== this.config.hiddenSize)
       throw new Error(`DSA hidden width ${hidden} != ${this.config.hiddenSize}`);
 
@@ -228,7 +226,10 @@ class Glm52DsaIndexer {
     }
     const contextLength = allKeys.shape[1]!;
     let selectedPositions: MlxArray | Glm52MlaBatchedSelection | null;
-    if (batch! > 1) {
+    if (forceDense || contextLength <= this.config.indexTopk) {
+      state.selectFullDense(this.layer, contextLength);
+      selectedPositions = null;
+    } else if (batch! > 1) {
       const rowSelections: Array<MlxArray | null> = [];
       for (let row = 0; row < batch!; row++) {
         const rowContext = (cache.rowOffsets[row] ?? 0) + tokens!;
@@ -269,15 +270,6 @@ class Glm52DsaIndexer {
         scores.dispose();
       }
       selectedPositions = { rows: rowSelections };
-    } else if (forceDense) {
-      state.selectFullDense(this.layer, contextLength);
-      selectedPositions = null;
-    } else if (contextLength <= this.config.indexTopk) {
-      state.selectFull(
-        this.layer,
-        new Float32Array(contextLength),
-      );
-      selectedPositions = null;
     } else {
       if (tokens !== 1) {
         if (allKeys !== key) allKeys.dispose();
@@ -533,6 +525,7 @@ export class Glm52DecoderLayer {
       this.inputNorm,
       this.config.rmsNormEps,
     );
+    const batch = input.shape[0]!;
     const tokens = input.shape[1]!;
     const denseBenchmarkPrefill =
       this.denseBenchmarkPrefillTokens > 0 &&
@@ -552,10 +545,10 @@ export class Glm52DecoderLayer {
           for (let token = 0; token < tokens; token++) {
             const row = normalized.slice(
               [0, token, 0],
-              [1, token + 1, this.config.hiddenSize],
+              [batch, token + 1, this.config.hiddenSize],
             );
             let dsaKey: MlxArray | null = null;
-            let positions: MlxArray | null = null;
+            let positions: MlxArray | Glm52MlaBatchedSelection | null = null;
             try {
               if (this.dsa) {
                 const step = this.dsa.projectAndSelect(
@@ -564,18 +557,12 @@ export class Glm52DecoderLayer {
                   dsaState,
                 );
                 dsaKey = step.key;
-                if (
-                  step.selectedPositions !== null &&
-                  !(step.selectedPositions instanceof MlxArray)
-                ) {
-                  throw new Error("serial DSA verification received batched selection");
-                }
                 positions = step.selectedPositions;
               } else {
-                positions = dsaState.selectSharedPositions(
-                  this.layer,
-                  cache.offset + 1,
-                );
+                positions = batch > 1
+                  ? { rows: Array.from({ length: batch }, (_, row) =>
+                      dsaState.selectSharedPositions(this.layer, (cache.rowOffsets[row] ?? 0) + 1, row)) }
+                  : dsaState.selectSharedPositions(this.layer, cache.offset + 1);
               }
               rows.push(this.mla.forward(
                 row,
@@ -609,7 +596,9 @@ export class Glm52DecoderLayer {
           dsaKey = step.key;
           positions = step.selectedPositions;
         } else if (dsaState) {
-          positions = cache.batchSize !== null && cache.batchSize > 1
+          positions = denseBenchmarkPrefill || cache.offset + tokens <= this.config.indexTopk
+            ? null
+            : cache.batchSize !== null && cache.batchSize > 1
             ? {
                 rows: cache.rowOffsets.map((offset, row) =>
                   dsaState.selectSharedPositions(
@@ -871,18 +860,13 @@ export class Glm52Model {
       tokenCount > GLM52_DSA_SPARSE_DECODE_BATCH_MAX &&
       firstOffset + tokenCount > this.glmConfig.indexTopk
     ) {
-      if (ids.shape.length !== 2 || ids.shape[0] !== 1) {
-        throw new Error(
-          "G2 DSA sparse prefill requires serial [1,T] token ids",
-        );
-      }
-      // Colibri selects independently for every query row. G2 preserves that
-      // behavior with a deliberately serial correctness fallback once a
-      // multi-token prefill crosses index_topk; G6 may vectorize it.
+      // Bound sparse-selection state for long prefills while retaining all
+      // request rows in each forward. Short verify windows keep the residual
+      // and expert graph batched across both requests and query positions.
       const rows: MlxArray[] = [];
       try {
         for (let token = 0; token < tokenCount; token++) {
-          const row = ids.slice([0, token], [1, token + 1]);
+          const row = ids.slice([0, token], [ids.shape[0]!, token + 1]);
           try {
             rows.push(this.forwardHidden(row, cache));
           } finally {
@@ -922,15 +906,10 @@ export class Glm52Model {
       tokenCount > GLM52_DSA_SPARSE_DECODE_BATCH_MAX &&
       firstOffset + tokenCount > this.glmConfig.indexTopk
     ) {
-      if (ids.shape.length !== 2 || ids.shape[0] !== 1) {
-        throw new Error(
-          "G3 DSA sparse prefill requires serial [1,T] token ids",
-        );
-      }
       const rows: MlxArray[] = [];
       try {
         for (let token = 0; token < tokenCount; token++) {
-          const row = ids.slice([0, token], [1, token + 1]);
+          const row = ids.slice([0, token], [ids.shape[0]!, token + 1]);
           try {
             rows.push(await this.forwardHiddenAsync(row, cache));
           } finally {

package/src/model/kv-attention-view.ts

@@ -0,0 +1,67 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { Cache, KvAttentionView, KvDonorAttention } from "./gemma4-base";
+import { quantizedSdpa, disposeTriple } from "./gemma4-base";
+
+export function captureKvDonorAttention(cache: Cache): KvDonorAttention {
+  if (cache.captureDonorAttention) return cache.captureDonorAttention();
+  const view = cache.captureDonorRows!();
+  return {
+    width: view.keys.shape[2]!, dtype: view.keys.dtype,
+    offsets: view.offsets, starts: view.starts, ends: view.ends,
+    attend: (q, scale, mask) => ops.sdpa(q, view.keys, view.values, scale, mask.mode, mask.arr),
+    dispose() { view.keys.dispose(); view.values.dispose(); },
+  };
+}
+
+export function quantizedDonorAttention(keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
+  validity: Pick<KvDonorAttention, "offsets" | "starts" | "ends">,
+  groupSize: number, bits: number): KvDonorAttention {
+  return {
+    width: keys.packed.shape[2]!, dtype: keys.scales.dtype, ...validity,
+    attend: (q, scale, mask) => quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
+    dispose() { disposeTriple(keys); disposeTriple(values); },
+  };
+}
+
+export function captureKvAttention(cache: Cache, k: MlxArray, v: MlxArray): KvAttentionView {
+    if (cache.attentionState) return cache.attentionState.appendAndFetch(k, v);
+    const quantized = cache.quantizedAttention;
+    if (quantized) {
+      const { groupSize, bits } = quantized;
+      const [keys, values] = quantized.updateAndFetchQuantized(k, v);
+      return {
+        attend: (q, scale, mask) => quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
+        dispose() { disposeTriple(keys); disposeTriple(values); },
+      };
+    }
+    const [keys, values] = cache.updateAndFetch(k, v);
+    return {
+      attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+      dispose() { keys.dispose(); values.dispose(); },
+    };
+  }
+
+
+/** Takes ownership of already aligned per-row donor views. Each storage codec
+ * keeps its own attention arithmetic while the caller supplies one query batch. */
+export function combineKvDonorAttention(rows: readonly KvDonorAttention[]): KvDonorAttention {
+  return { width: rows[0]!.width, dtype: rows[0]!.dtype,
+    offsets: rows.flatMap(row => [...row.offsets]), starts: rows.flatMap(row => [...row.starts]),
+    ends: rows.flatMap(row => [...row.ends]),
+    attend(q, scale, mask) {
+      const outputs: MlxArray[] = [];
+      try {
+        for (const [index, row] of rows.entries()) {
+          using query = q.slice([index, 0, 0, 0], [index + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+          const from = mask.arr?.shape.map(() => 0), to = mask.arr ? [...mask.arr.shape] : undefined;
+          if (from && to && to.length === 4 && to[0]! > 1) { from[0] = index; to[0] = index + 1; }
+          using arr = mask.arr ? mask.arr.slice(from!, to!) : null;
+          outputs.push(row.attend(query, scale, { mode: mask.mode, arr }));
+        }
+        return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+      } finally { for (const output of outputs) output.dispose(); }
+    },
+    dispose() { for (const row of rows) row.dispose(); },
+  };
+}

package/src/model/minicpm5.ts

@@ -14,7 +14,6 @@ import {
   KVCache,
   LoraState,
   QuantizedKVCache,
-  TurboQuantKVCache,
   QuantizedEmbedding,
   QuantizedLinear,
   quantizedSdpa,
@@ -143,23 +142,28 @@ export class LlamaAttention {
     q = disposing(q, ropeStep(q));
     k = disposing(k, ropeStep(k));
     let attn: MlxArray;
-    if (cache instanceof QuantizedKVCache) {
-      const [keys, values] = cache.updateAndFetchQuantized(k, v);
+    if (cache.attentionState) {
+      const view = cache.attentionState.appendAndFetch(k, v);
+      k.dispose(); v.dispose();
+      try { attn = view.attend(q, this.scale, mask); } finally { view.dispose(); }
+    } else if (cache.quantizedAttention) {
+      const quantized = cache.quantizedAttention;
+      const [keys, values] = quantized.updateAndFetchQuantized(k, v);
       k.dispose();
       v.dispose();
-      attn = quantizedSdpa(q, keys, values, this.scale, mask, cache.groupSize, cache.bits);
+      attn = quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
       disposeTriple(keys);
       disposeTriple(values);
-    } else if (cache instanceof TurboQuantKVCache) {
+    } else if (cache.rotatedValueAttention) {
       // Deferred-V: sdpa runs on rotated values, then one InvFWHT on the
       // output (linear in V) — see TurboQuantKVCache.updateAndFetchDeferredV.
-      const [keys, values] = cache.updateAndFetchDeferredV(k, v);
+      const [keys, values] = cache.rotatedValueAttention.updateAndFetchDeferredV(k, v);
       k.dispose();
       v.dispose();
       const rotated = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr);
       keys.dispose();
       values.dispose();
-      attn = tqUnrotateValues(rotated);
+      attn = cache.rotatedValueAttention.captureValueTransform?.()(rotated) ?? tqUnrotateValues(rotated);
       rotated.dispose();
     } else {
       const [keys, values] = cache.updateAndFetch(k, v);

package/src/model/qwen3-delta-state.ts

@@ -0,0 +1,74 @@
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+import * as ops from "../mlx/ops";
+import { computeG } from "./qwen3-delta";
+
+// The same recurrent arithmetic as gated_delta_step, with no query/output
+// calculation. Each row replays only its retained prefix of the verify window.
+const SOURCE = String.raw`
+  const int T = Tin[0];
+  auto n = thread_position_in_grid.z;
+  auto b_idx = n / Hv;
+  auto hv_idx = n % Hv;
+  auto hk_idx = hv_idx / (Hv / Hk);
+  constexpr int n_per_t = Dk / 32;
+  auto dk_idx = thread_position_in_threadgroup.x;
+  auto dv_idx = thread_position_in_grid.y;
+  auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+  auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+  auto g_ = g + b_idx * T * Hv;
+  auto beta_ = beta + b_idx * T * Hv;
+  auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+  auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+  float state[n_per_t];
+  for (int i = 0; i < n_per_t; ++i)
+    state[i] = static_cast<float>(i_state[n_per_t * dk_idx + i]);
+  const int keep = lengths[b_idx];
+  for (int t = 0; t < keep; ++t) {
+    float kv_mem = 0.0f;
+    for (int i = 0; i < n_per_t; ++i) {
+      auto s_idx = n_per_t * dk_idx + i;
+      state[i] = state[i] * g_[hv_idx];
+      kv_mem += state[i] * k_[s_idx];
+    }
+    kv_mem = simd_sum(kv_mem);
+    auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+    for (int i = 0; i < n_per_t; ++i) {
+      auto s_idx = n_per_t * dk_idx + i;
+      state[i] = state[i] + k_[s_idx] * delta;
+    }
+    k_ += Hk * Dk;
+    v_ += Hv * Dv;
+    g_ += Hv;
+    beta_ += Hv;
+  }
+  for (int i = 0; i < n_per_t; ++i)
+    o_state[n_per_t * dk_idx + i] = state[i];
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** State-only prefix replay. Inputs keep their original [B,T,...] strides;
+ * lengths are request-owned accepted counts, never scheduling decisions. */
+export function gatedDeltaState(
+  k: MlxArray, v: MlxArray, a: MlxArray, b: MlxArray,
+  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null,
+  lengths: readonly number[],
+): MlxArray {
+  const [B, T, Hk, Dk] = k.shape as [number, number, number, number];
+  const [, , Hv, Dv] = v.shape as [number, number, number, number];
+  using beta = ops.sigmoid(b);
+  using g = computeG(aLog, a, dtBias);
+  using zero = state ? null : ops.zeros([B, Hv, Dv, Dk], Dtype.float32);
+  using t = ops.fromInt32([T], [1]);
+  using kept = ops.fromInt32([...lengths], [B]);
+  kernel ??= new MetalKernel({ name: "gated_delta_prefix_state",
+    inputNames: ["k", "v", "g", "beta", "state_in", "Tin", "lengths"],
+    outputNames: ["state_out"], source: SOURCE, ensureRowContiguous: true });
+  return kernel.apply([k, v, g, beta, state ?? zero!, t, kept], {
+    outputs: [{ shape: [B, Hv, Dv, Dk], dtype: Dtype.float32 }],
+    grid: [32, Dv, B * Hv], threadGroup: [32, 4, 1],
+    templateInts: { Dk, Dv, Hk, Hv },
+  })[0]!;
+}