mlx-bun 0.2.0 → 0.4.0

package/src/model/gemma4.ts

@@ -23,7 +23,7 @@ import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
-import { isExpertTracing, recordRouting } from "../expert-trace";
+import { isExpertTracing, recordRouting } from "../lab/expert-trace/expert-trace";
 
 // The shared, config-independent machinery lives in gemma4-base.ts
 // (Phase B extraction); this file keeps the architecture-specific
@@ -45,18 +45,18 @@ import {
   RMSNorm,
   RotatingKVCache,
   RotatingQuantizedKVCache,
-  TurboQuantKVCache,
   quantizedSdpa,
   type Cache,
   type Mask,
   type SharedKv,
+  captureRopeOffsets,
   type LoraWeights,
 } from "./gemma4-base";
 import { Checkpoint } from "../mlx/checkpoint";
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 import { unrotateValues as tqUnrotateValues } from "../mlx/turboquant-ops";
 import { CompiledFunction } from "../mlx/compile";
-import { flagOn } from "../flags";
+import { flagOn } from "../runtime-config";
 
 /** Verbatim port of mlx_lm/models/gemma4_text.py:
  *    `@partial(mx.compile, shapeless=True) def geglu(gate, x): return nn.gelu_approx(gate) * x`
@@ -242,7 +242,7 @@ class Attention {
   /** A number offset takes the static fast::rope; an array offset (set by
    *  compiled-decode trace adapters) takes the dynamic variant — same
    *  kernel, offset read from the array (bit-exactness asserted in
-   *  tests/compile.test.ts). */
+   *  tests/unit/compile.test.ts). */
   rope(x: MlxArray, offset: number | MlxArray): MlxArray {
     return typeof offset === "number"
       ? ops.rope(x, this.headDim, this.ropeBase, offset, this.ropeFreqs)
@@ -304,15 +304,11 @@ class Attention {
       }
 
       const offset = cache.offset;
-      // Capture the RoPE offset array ONCE (like `offset` above): the
-      // updateAndFetch below advances cache.offset, so re-reading
-      // cache.ropeOffsetArr for Q *after* the write would hand K and Q
-      // different positions. Harmless today (real caches leave it unset;
-      // compiled-decode passes a constant trace input), but a hard
-      // prerequisite for batched decode, where ropeOffsetArr carries per-row
-      // positions derived from the pre-write offset and MUST be identical for
-      // this step's K and Q.
-      const offsetArr = cache.ropeOffsetArr;
+      // Own the pre-write positions through every query and shared-KV
+      // consumer. Advancing a row cache may replace and release its borrowed
+      // offset array before Q's RoPE is built. The fetched shared state owns
+      // this retained handle and releases it with K/V at the end of the pass.
+      const offsetArr = captureRopeOffsets(cache);
 
       const kNormed = this.kNorm!.forward(k);
       const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
@@ -328,23 +324,29 @@ class Attention {
       if (v !== k) v.dispose();
       k.dispose();
 
-      if (cache instanceof QuantizedKVCache || cache instanceof RotatingQuantizedKVCache) {
-        const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+      if (cache.attentionState) {
+        const attention = cache.attentionState.appendAndFetch(kRoped, vT);
+        kRoped.dispose(); vT.dispose();
+        shared = { kind: "view", attention, offset, offsetArr };
+      } else if (cache.quantizedAttention) {
+        const quantized = cache.quantizedAttention;
+        const [kq, vq] = quantized.updateAndFetchQuantized(kRoped, vT);
         kRoped.dispose();
         vT.dispose();
         shared = {
           kind: "quant", keys: kq, values: vq, offset,
-          groupSize: cache.groupSize, bits: cache.bits,
+          groupSize: quantized.groupSize, bits: quantized.bits,
           offsetArr,
         };
-      } else if (cache instanceof TurboQuantKVCache) {
+      } else if (cache.rotatedValueAttention) {
         // Deferred-V read: values stay rotated; the attention output is
         // un-rotated below (shared.vRotated) — including by KV-shared
         // consumer layers, which see the flag through sharedIn.
-        const [keys, values] = cache.updateAndFetchDeferredV(kRoped, vT);
+        const [keys, values] = cache.rotatedValueAttention.updateAndFetchDeferredV(kRoped, vT);
         kRoped.dispose();
         vT.dispose();
-        shared = { kind: "plain", keys, values, offset, offsetArr, vRotated: true };
+        shared = { kind: "plain", keys, values, offset, offsetArr, vRotated: true,
+          restoreValues: cache.rotatedValueAttention.captureValueTransform?.() };
       } else {
         const [keys, values] = cache.updateAndFetch(kRoped, vT);
         kRoped.dispose();
@@ -360,7 +362,9 @@ class Attention {
 
     let attn: MlxArray;
     const ta = getTrainingAttn();
-    if (shared.kind === "quant") {
+    if (shared.kind === "view") {
+      attn = shared.attention.attend(q, 1.0, mask);
+    } else if (shared.kind === "quant") {
       attn = quantizedSdpa(q, shared.keys, shared.values, 1.0, mask, shared.groupSize, shared.bits);
     } else if (
       ta === "flash" && flashSupported(q) && (mask.mode === "causal" || mask.mode === "array")
@@ -370,7 +374,7 @@ class Attention {
       // verified route: its dQ/dK/dV vjps are exact (finite-difference checked;
       // it's what mlx-lm's tuner differentiates). The two historical bugs were
       // in THIS flash kernel (spurious dK transpose, divergent dQ barrier),
-      // both fixed and gated by tests/flash-attention.test.ts vs the ops.sdpa
+      // both fixed and gated by tests/unit/flash-attention.test.ts vs the ops.sdpa
       // oracle — see flash-attention.ts's routing comment. Flash also remains
       // Gemma-guarded in the trainer (e4b >=2K SIGTRAP not yet re-validated).
       // Full layers → window 0 (pure causal); sliding layers → their window.
@@ -383,7 +387,7 @@ class Attention {
     if (shared.kind === "plain" && shared.vRotated) {
       // Attention is linear in V: one InvFWHT on the [B,H,L,D] output
       // replaces one per cached token (TurboQuant deferred-V).
-      attn = disposing(attn, tqUnrotateValues(attn));
+      attn = disposing(attn, shared.restoreValues ? shared.restoreValues(attn) : tqUnrotateValues(attn));
     }
     const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
     attn.dispose();
@@ -851,6 +855,13 @@ export class Gemma4Model {
     );
   }
 
+  /** Computed model constants must retain their constructor arithmetic when
+   * a restored request enters compiled decode without a prefill first. */
+  materializeGraphConstants(): void {
+    const constants = this.layers.flatMap(layer => layer.attn.ropeFreqs ? [layer.attn.ropeFreqs] : []);
+    if (constants.length) ops.evalAll(constants);
+  }
+
   /** ids [1, L] → final-norm hidden states [1, L, hidden]. */
   forwardHidden(ids: MlxArray, cache: Cache[]): MlxArray {
     let h = this.embed.encode(ids);
@@ -885,7 +896,7 @@ export class Gemma4Model {
    *  which attend bidirectionally among themselves (use_bidirectional_attention:
    *  "vision" — text stays causal; audio prompts pass NO bidir mask, and a
    *  mixed image+audio prompt drops the image overlay too, §3.3 Q1 of
-   *  docs/design/audio-input-plan.md). `ids` ([1, L], the spliced token ids)
+   *  docs/design/generic-model-support.md). `ids` ([1, L], the spliced token ids)
    *  is required for per-layer-input models (e2b/e4b): multimodal soft-token
    *  positions get token 0's per-layer embedding, and their tower content
    *  enters only through the projection term (the merged hidden). Matches
@@ -1000,13 +1011,7 @@ export class Gemma4Model {
     for (let i = 0; i < this.numDonors; i++) {
       const s = intermediates[i];
       if (!s) continue;
-      if (s.kind === "plain") {
-        s.keys.dispose();
-        s.values.dispose();
-      } else {
-        for (const t of [s.keys, s.values])
-          for (const a of [t.packed, t.scales, t.biases]) a.dispose();
-      }
+      Gemma4Model.disposeSharedKv(s);
     }
     for (const m of masks.values()) m.arr?.dispose();
 
@@ -1015,7 +1020,7 @@ export class Gemma4Model {
     return h;
   }
 
-  // --- Segmented-backward support (docs/design/segmented-backward-training.md
+  // --- Segmented-backward support (docs/design/orpo-training.md
   // §4 Phase B). These are ADDITIVE — forwardLayers is untouched. The segmented
   // driver builds masks + per-layer inputs ONCE, then drives runLayerRange per
   // segment, threading the KV-shared donor K/V across segment boundaries.
@@ -1056,7 +1061,10 @@ export class Gemma4Model {
   /** Dispose a fetched donor K/V (plain or quantized) — mirrors forwardLayers'
    *  donor cleanup. */
   private static disposeSharedKv(s: SharedKv): void {
-    if (s.kind === "plain") {
+    s.offsetArr?.dispose();
+    if (s.kind === "view") {
+      s.attention.dispose();
+    } else if (s.kind === "plain") {
       s.keys.dispose();
       s.values.dispose();
     } else {

package/src/model/generated/gemma4-12b.ts

@@ -4,8 +4,8 @@
 // fingerprint: 9f812d2eb461fcbe
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 
@@ -34,7 +34,10 @@ export const FINGERPRINT = "9f812d2eb461fcbe";
 export let generatedForwardUses = 0;
 
 function disposeSharedKv(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {

package/src/model/generated/gemma4-26b.ts

@@ -4,8 +4,8 @@
 // fingerprint: c9dd67ed5a525231
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 
@@ -34,7 +34,10 @@ export const FINGERPRINT = "c9dd67ed5a525231";
 export let generatedForwardUses = 0;
 
 function disposeSharedKv(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {

package/src/model/generated/gemma4-e4b.ts

@@ -4,8 +4,8 @@
 // fingerprint: 418e9adc386ea67c
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 
@@ -35,7 +35,10 @@ export const FINGERPRINT = "418e9adc386ea67c";
 export let generatedForwardUses = 0;
 
 function disposeSharedKv(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {

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-dsa.ts

@@ -291,7 +291,7 @@ function getDsaRankKeyKernel(): MetalKernel {
           ? ~bits
           : (bits ^ 0x80000000u);
         const uint descending = ~ascending;
-        keys[i] = (ulong(descending) << 32) | ulong(i);
+        keys[i] = descending;
       `,
     });
   }
@@ -308,7 +308,8 @@ function getDsaContractOrderKeyKernel(): MetalKernel {
         const uint i = thread_position_in_grid.x;
         const uint position = selected_positions[i];
         const uint threshold_class = selected_scores[i] > threshold ? 0u : 1u;
-        keys[i] = (ulong(threshold_class) << 32) | ulong(position);
+        // MLX axis lengths are signed int32; positions leave the high bit free.
+        keys[i] = (threshold_class << 31) | position;
       `,
     });
   }
@@ -318,8 +319,8 @@ function getDsaContractOrderKeyKernel(): MetalKernel {
 /**
  * Deterministic on-device equivalent of `selectDsaThresholdTiesF32`.
  *
- * The first uint64 key sorts by score descending and then position ascending,
- * making the top-k set deterministic even at the threshold. The second key
+ * Stable sorting of the first uint32 key orders scores descending and keeps
+ * equal-score positions ascending, including threshold ties. The second key
  * restores Colibri's observable two-scan order: all scores strictly above the
  * threshold in position order, followed by threshold ties in position order.
  * The caller owns the returned arrays and must call `dispose()` once the last
@@ -340,18 +341,21 @@ export function selectGlm52DsaDevice(
     ? scores
     : scores.astype(Dtype.float32);
   const [rankKeys] = getDsaRankKeyKernel().apply([scoresF32], {
-    outputs: [{ shape: [contextLength], dtype: Dtype.uint64 }],
+    outputs: [{ shape: [contextLength], dtype: Dtype.uint32 }],
     grid: [contextLength, 1, 1],
     threadGroup: [Math.min(256, contextLength), 1, 1],
   });
-  const partition = ops.argpartitionAxis(rankKeys!, topK - 1, 0);
-  const selectedUnordered = partition.slice([0], [topK]);
+  // Stable argsort retains ascending positions for equal score bits, so the
+  // ranking key needs no position field. MLX 0.31.2's GPU argpartition also
+  // sorts fully, but does not promise this tie order in its public contract.
+  const rankedPositions = ops.argsortAxis(rankKeys!, 0);
+  const selectedUnordered = rankedPositions.slice([0], [topK]);
   const selectedScores = ops.takeAlongAxis(scoresF32, selectedUnordered, 0);
   const threshold = ops.minAxis(selectedScores, 0, false);
   const [orderKeys] = getDsaContractOrderKeyKernel().apply(
     [selectedScores, selectedUnordered, threshold],
     {
-      outputs: [{ shape: [topK], dtype: Dtype.uint64 }],
+      outputs: [{ shape: [topK], dtype: Dtype.uint32 }],
       grid: [topK, 1, 1],
       threadGroup: [Math.min(256, topK), 1, 1],
     },
@@ -360,7 +364,7 @@ export function selectGlm52DsaDevice(
   const positions = ops.takeAlongAxis(selectedUnordered, order, 0);
 
   rankKeys!.dispose();
-  partition.dispose();
+  rankedPositions.dispose();
   selectedUnordered.dispose();
   selectedScores.dispose();
   orderKeys!.dispose();

package/src/model/glm52-memory.ts

@@ -4,7 +4,7 @@ import {
   type Glm52Config,
 } from "./glm52-config";
 import { ColibriGlm52Container } from "./glm52-container";
-import type { MemoryPlan } from "../memory-plan";
+import type { MemoryPlan } from "../admission-plan";
 
 export const GLM52_G5_GIB = 1024 ** 3;
 export const GLM52_G5_DEFAULT_MACHINE_BYTES = 32 * GLM52_G5_GIB;

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 {