mlx-bun 0.3.0 → 0.5.0

package/src/model/qwen3_5.ts

@@ -10,7 +10,11 @@
 // shared maybeQuantizeKv path).
 
 import type { ModelConfig } from "../config";
+import { quantizedAppendAttention } from "./quantized-append-attention";
+import { gatedDeltaState } from "./qwen3-delta-state";
 import type { Weights } from "../weights";
+import { runtimeFlag } from "../runtime-config";
+import { mapTokenGroups, type TokenGroup } from "./token-groups";
 import { MlxArray } from "../mlx/array";
 import { Dtype, deviceArchitecture } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
@@ -23,7 +27,6 @@ import {
   disposing,
   KVCache,
   LoraState,
-  QuantizedKVCache,
   QuantizedEmbedding,
   QuantizedLinear,
   quantizedSdpa,
@@ -36,8 +39,8 @@ import { materializeCopy } from "../mlx/materialize";
 import { gatedDeltaUpdate, SSMCache } from "./qwen3-delta";
 import type { QwenConvolution } from "./qwen-conv";
 import {
-  activeMrope, applyInterleavedRope, buildMropePositions, mropeInvFreq,
-  setActiveMrope, type MropeRequestState,
+  applyInterleavedRope, buildMropePositions, mropeInvFreq,
+  type MropeRequestState, type MropeForwardState,
 } from "./qwen3-mrope";
 
 const PREFIX = "language_model";
@@ -155,16 +158,20 @@ export class GatedDeltaNet {
     this.normWeight = weights.tensor(`${prefix}.norm.weight`);
   }
 
-  #convolve(qkv: MlxArray, state: MlxArray): [MlxArray, MlxArray] {
-    if (this.convolution) return this.convolution(qkv, state, this.convWeight);
-    const [B, S, D] = qkv.shape as [number, number, number];
-    const nKeep = this.convKernel - 1;
+  #convolve(qkv: MlxArray, state: MlxArray, rowLengths?: readonly number[]): [MlxArray, MlxArray] {
+    if (this.convolution) {
+      const result = this.convolution(qkv, state, this.convWeight);
+      if (!rowLengths) return result;
+      using input = ops.concatAxis([state, qkv], 1);
+      const tail = this.#convTail(input, qkv.shape[1]!, rowLengths);
+      result[1].dispose();
+      return [result[0], tail];
+    }
+    const [, S, D] = qkv.shape as [number, number, number];
     const input = ops.concatAxis([state, qkv], 1);
     // MLX copy and contiguous can both alias the whole prefill buffer.
     // Materialize only the tail so cache residency follows its logical size.
-    const view = input.slice([0, S, 0], [B, S + nKeep, D]);
-    const tail = materializeCopy(view);
-    view.dispose();
+    const tail = this.#convTail(input, S, rowLengths);
     const conv = ops.conv1d(input, this.convWeight, 1, 0, 1, D);
     input.dispose();
     const out = compiledSilu(conv);
@@ -172,7 +179,19 @@ export class GatedDeltaNet {
     return [out, tail];
   }
 
-  forward(x: MlxArray, cache: SSMCache, independentRows = false): MlxArray {
+  #convTail(input: MlxArray, processed: number, rowLengths?: readonly number[]): MlxArray {
+    const [B, , D] = input.shape as [number, number, number];
+    const nKeep = this.convKernel - 1;
+    if (rowLengths) {
+      using indices = ops.fromInt32(rowLengths.flatMap(length =>
+        Array.from({ length: nKeep }, (_, position) => length + position)), [B, nKeep, 1]);
+      return ops.takeAlongAxis(input, indices, 1);
+    }
+    using view = input.slice([0, processed, 0], [B, processed + nKeep, D]);
+    return materializeCopy(view);
+  }
+
+  forward(x: MlxArray, cache: SSMCache, independentRows = false, mask?: MlxArray | null): MlxArray {
     const [B, S] = x.shape as [number, number, number];
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
@@ -194,7 +213,7 @@ export class GatedDeltaNet {
     }
 
     const t0 = prof ? performance.now() : 0;
-    const qkv = this.inProjQkv.forward(x, independentRows); // [B,S,convDim]
+    let qkv = this.inProjQkv.forward(x, independentRows); // [B,S,convDim]
     let z = this.inProjZ.forward(x, independentRows);
     z = disposing(z, ops.reshape(z, [B, S, this.numVHeads, this.headVDim]));
     const b = this.inProjB.forward(x, independentRows); // [B,S,numVHeads]
@@ -202,10 +221,19 @@ export class GatedDeltaNet {
     if (prof) { ops.evalAll([qkv, z, b, a]); prof.proj = (prof.proj ?? 0) + performance.now() - t0; }
     const tc = prof ? performance.now() : 0; // [B,S,numVHeads]
 
-    // Causal depthwise conv with the conv-state prefix (B=1: no ssm mask).
+    // The model shares one padding mask across recurrent layers. Direct layer
+    // callers can request the same cache-owned mask without a model wrapper.
+    using localMask = mask === undefined ? cache.prefillPadding?.makeMask(S) ?? null : null;
+    const ssmMask = mask ?? localMask;
+    if (ssmMask) {
+      using expanded = ops.reshape(ssmMask, [B, S, 1]);
+      using zero = MlxArray.fromBytesCopy(new Uint8Array(qkv.dtype === Dtype.float32 ? 4 : 2), [], qkv.dtype);
+      qkv = disposing(qkv, ops.where(expanded, qkv, zero));
+    }
+    // Convolution keeps each row's last real-token tail, including empty rows.
     const convState =
       cache.conv ?? ops.zeros([B, nKeep, convDim], x.dtype);
-    const [convOut, newConv] = this.#convolve(qkv, convState);
+    const [convOut, newConv] = this.#convolve(qkv, convState, cache.prefillPadding?.convolutionLengths(S));
     if (!cache.conv) convState.dispose();
     if (spec) spec.qkv = qkv;
     else qkv.dispose();
@@ -236,7 +264,7 @@ export class GatedDeltaNet {
     const tk = prof ? performance.now() : 0;
 
     const [out, newState] = gatedDeltaUpdate(
-      q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent,
+      q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent, ssmMask,
     );
     if (prof) { ops.evalAll([out]); prof.kernel = (prof.kernel ?? 0) + performance.now() - tk; }
     const to = prof ? performance.now() : 0;
@@ -275,18 +303,21 @@ export class GatedDeltaNet {
    *  recurrent state is BIT-EXACTLY what a forward over only the accepted
    *  prefix would have produced. Output projections (z / out_proj) are state-
    *  free and skipped: only the states matter here. */
-  #replaySpecPrefix(cache: SSMCache, keep: number): void {
+  #replaySpecPrefix(cache: SSMCache, keep: number | readonly number[]): void {
     const r = cache.specRound;
-    if (!r || !r.qkv || !r.a || !r.b)      throw new Error("GatedDeltaNet replay without recorded round inputs");
+    if (!r || !r.qkv || !r.a || !r.b)
+      throw new Error("GatedDeltaNet replay without recorded round inputs");
     const B = r.qkv.shape[0]!;
+    const count = typeof keep === "number" ? keep : Math.max(...keep);
+    const rowLengths = typeof keep === "number" ? undefined : keep;
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
 
-    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, keep, convDim]);
+    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, count, convDim]);
     const qkvPfx = ops.contiguous(qkvPfxView);
     qkvPfxView.dispose();
     const convState = cache.conv ?? ops.zeros([B, nKeep, convDim], qkvPfx.dtype);
-    const [convOut, newConv] = this.#convolve(qkvPfx, convState);
+    const [convOut, newConv] = this.#convolve(qkvPfx, convState, rowLengths);
     if (!cache.conv) convState.dispose();
     qkvPfx.dispose();
     cache.conv?.dispose();
@@ -296,36 +327,45 @@ export class GatedDeltaNet {
       convOut, [this.keyDim, 2 * this.keyDim], -1,
     ) as [MlxArray, MlxArray, MlxArray];
     convOut.dispose();
-    let q = ops.reshape(qFlat, [B, keep, this.numKHeads, this.headKDim]);
+    let q = rowLengths ? null : ops.reshape(qFlat, [B, count, this.numKHeads, this.headKDim]);
     qFlat.dispose();
-    let k = ops.reshape(kFlat, [B, keep, this.numKHeads, this.headKDim]);
+    let k = ops.reshape(kFlat, [B, count, this.numKHeads, this.headKDim]);
     kFlat.dispose();
-    const v = disposing(vFlat, ops.reshape(vFlat, [B, keep, this.numVHeads, this.headVDim]));
+    const v = disposing(vFlat, ops.reshape(vFlat, [B, count, this.numVHeads, this.headVDim]));
     const invScale = Math.pow(this.headKDim, -0.5);
-    q = disposing(q, ops.rmsNorm(q, null, 1e-6));
-    q = disposing(q, ops.mulScalar(q, invScale * invScale));
+    if (q) {
+      q = disposing(q, ops.rmsNorm(q, null, 1e-6));
+      q = disposing(q, ops.mulScalar(q, invScale * invScale));
+    }
     k = disposing(k, ops.rmsNorm(k, null, 1e-6));
     k = disposing(k, ops.mulScalar(k, invScale));
 
-    const aPfxView = r.a.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const aPfxView = r.a.slice([0, 0, 0], [B, count, this.numVHeads]);
     const aPfx = ops.contiguous(aPfxView);
     aPfxView.dispose();
-    const bPfxView = r.b.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const bPfxView = r.b.slice([0, 0, 0], [B, count, this.numVHeads]);
     const bPfx = ops.contiguous(bPfxView);
     bPfxView.dispose();
 
-    const [y, newState] = gatedDeltaUpdate(
-      q, k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent,
-    );
-    q.dispose();
+    let newState: MlxArray;
+    if (rowLengths) {
+      // Unequal prefixes need a row-length-aware recurrence. The state-only
+      // kernel omits query/output work. Uniform replay retains its measured
+      // kernel until the separate optimization demonstrates a serving win.
+      newState = gatedDeltaState(k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent, rowLengths);
+    } else {
+      const [output, state] = gatedDeltaUpdate(q!, k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent);
+      output.dispose(); newState = state;
+    }
+    q?.dispose();
     k.dispose();
     v.dispose();
     aPfx.dispose();
     bPfx.dispose();
-    y.dispose(); // only the state advance matters on the rollback path
     cache.recurrent?.dispose();
     cache.recurrent = newState;
-    cache.advance(keep);
+    if (typeof keep === "number") cache.advance(keep);
+    else cache.advanceRows(keep);
   }
 
   private rmsNormGated(hidden: MlxArray, gate: MlxArray): MlxArray {
@@ -369,7 +409,8 @@ export class Qwen3Attention {
     this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
   }
 
-  forward(x: MlxArray, mask: Mask, cache: Cache, independentRows = false): MlxArray {
+  forward(x: MlxArray, mask: Mask, cache: Cache, independentRows = false,
+    mrope: MropeForwardState | null = null): MlxArray {
     const [B, L] = x.shape as [number, number, number];
 
     // q_proj emits 2× head_dim per head → split into queries + gate.
@@ -396,11 +437,10 @@ export class Qwen3Attention {
     // position as ropeOffsetArr (rows have different prompt lengths); the
     // dynamic-offset kernel is the same fast::rope, bit-exact vs the static
     // form (tests/unit/compile.test.ts). Serial lane: scalar offset, unchanged.
-    // Vision requests (serial lane only) install activeMrope for the current
-    // forward: 3D interleaved positions via the manual apply — the reference
+    // Media forwards supply their own 3D interleaved positions. The reference
     // never uses the fused fast-rope kernel when position_ids are supplied,
     // so this IS the oracle's own arithmetic. Text-only stays on ops.rope.
-    const mr = activeMrope;
+    const mr = mrope;
     const offArr = (cache as { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
     if (mr) {
       q = disposing(q, applyInterleavedRope(q, mr));
@@ -415,11 +455,18 @@ export class Qwen3Attention {
     }
 
     let attn: MlxArray;
-    if (cache instanceof QuantizedKVCache) {
-      const [keys, values] = cache.updateAndFetchQuantized(k, v);
+    const quantized = cache.quantizedAttention;
+    if (cache.attentionState) {
+      const view = cache.attentionState.appendAndFetch(k, v);
+      k.dispose(); v.dispose();
+      try { attn = view.attend(q, this.scale, mask, independentRows); } finally { view.dispose(); }
+    } else if (quantized) {
+      const [keys, values] = quantized.updateAndFetchQuantized(k, v);
       k.dispose();
       v.dispose();
-      attn = quantizedSdpa(q, keys, values, this.scale, mask, cache.groupSize, cache.bits);
+      attn = independentRows && L > 1
+        ? quantizedAppendAttention(q, keys, values, this.scale, quantized.groupSize, quantized.bits)
+        : quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
       disposeTriple(keys);
       disposeTriple(values);
     } else {
@@ -517,23 +564,27 @@ export class Qwen3Layer {
     this.postAttnNorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), config.text.rmsNormEps);
   }
 
-  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false): MlxArray {
-    const xn = this.inputNorm.forward(x);
-    const r = this.isLinear
-      ? this.linearAttn!.forward(xn, cache as SSMCache, independentRows)
-      : this.selfAttn!.forward(xn, faMask, cache, independentRows);
-    xn.dispose();
-    const h = ops.add(x, r);
-    r.dispose();
-    const hn = this.postAttnNorm.forward(h);
+  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false, ssmMask?: MlxArray | null,
+    mrope: MropeForwardState | null = null): MlxArray {
+    using h = this.forwardAttn(x, faMask, cache, independentRows, ssmMask, mrope);
+    return this.forwardMlp(h, independentRows);
+  }
+
+  forwardAttn(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false, ssmMask?: MlxArray | null,
+    mrope: MropeForwardState | null = null): MlxArray {
+    using xn = this.inputNorm.forward(x);
+    using r = this.isLinear
+      ? this.linearAttn!.forward(xn, cache as SSMCache, independentRows, ssmMask)
+      : this.selfAttn!.forward(xn, faMask, cache, independentRows, mrope);
+    return ops.add(x, r);
+  }
+
+  forwardMlp(h: MlxArray, independentRows = false): MlxArray {
+    using hn = this.postAttnNorm.forward(h);
     // RMSNorm allocates an aligned row-contiguous output. Pass that layout
     // proof so Trellis can match native small-prefill matmul without an eval.
-    const m = this.mlp.forward(hn, true, independentRows);
-    hn.dispose();
-    const out = ops.add(h, m);
-    h.dispose();
-    m.dispose();
-    return out;
+    using m = this.mlp.forward(hn, true, independentRows);
+    return ops.add(h, m);
   }
 }
 
@@ -627,10 +678,16 @@ export class Qwen35Model {
         return null;
     }
     return {
-      maxChunkSize: (state: readonly Cache[]) => qwenAppendChunkSize(state[0]!.offset),
+      affineKvBits: [4, 8],
+      turboQuantFormats: [{ kBits: 8, vBits: 3 }],
+      maxChunkSize: (state: readonly Cache[], rows = 1) => rows === 1 ? qwenAppendChunkSize(state[0]!.offset) : 1,
       forwardHidden: (ids: MlxArray, cache: Cache[]): MlxArray => {
+        // Wider cohorts retain their ordinary one-position numerical graph.
+        // The specialized multi-position append remains qualified at B=1.
+        if (ids.shape.length === 2 && ids.shape[0]! > 1 && ids.shape[1] === 1)
+          return this.forwardHidden(ids, cache);
         if (ids.shape.length !== 2 || ids.shape[0] !== 1 || ids.shape[1]! > 4)
-          throw new Error("Qwen committed-token append supports one row and at most four positions");
+          throw new Error("Qwen committed-token append supports one position per shared row or up to four positions at B=1");
         const hidden = this.embed.encode(ids);
         return this.forwardLayers(hidden, cache, true);
       },
@@ -642,13 +699,62 @@ export class Qwen35Model {
     return this.forwardLayers(h, cache);
   }
 
+  /** Borrowed request positions, including different positions for each row. */
+  forwardHiddenAtPositions(ids: MlxArray, cache: Cache[], positions: MlxArray): MlxArray {
+    return this.forwardLayers(this.embed.encode(ids), cache, false, positions);
+  }
+
+  forwardEmbeddingsAtPositions(embeds: MlxArray, cache: Cache[], positions: MlxArray): MlxArray {
+    return this.forwardLayers(ops.contiguous(embeds), cache, false, positions);
+  }
+
+  /** Attention and DeltaNet retain disjoint row state. Only tokenwise MLP
+   * work is eligible for packing; verification can preserve its geometry. */
+  forwardHiddenMixed(work: readonly TokenGroup[]): MlxArray[] {
+    if (work.length === 1 && !work[0]!.captureLayer) return [this.forwardHidden(work[0]!.ids, work[0]!.cache)];
+    const groups: Array<{ h: MlxArray; mask?: Mask; ssmMask?: MlxArray | null }> = [];
+    const results: MlxArray[] = [];
+    const pack = runtimeFlag("MLX_BUN_MIXED_PACKED_MLP", true);
+    const bounded = work.some(group => group.ids.shape[1]! > TRELLIS_MATVEC_MAX_M);
+    try {
+      for (const { ids, cache } of work) {
+        const group: typeof groups[number] = { h: this.embed.encode(ids) }; groups.push(group);
+        group.mask = cache[this.faIdx]!.makeMask(ids.shape[1]!, null);
+        group.ssmMask = (cache[0] as SSMCache).prefillPadding?.makeMask(ids.shape[1]!);
+      }
+      for (const [i, layer] of this.layers.entries()) {
+        const mids: MlxArray[] = [];
+        try {
+          for (const [row, group] of groups.entries())
+            mids.push(layer.forwardAttn(group.h, group.mask!, work[row]!.cache[i]!, false, group.ssmMask));
+          const outputs = mapTokenGroups(work, mids, hidden => layer.forwardMlp(hidden), pack);
+          for (const [row, group] of groups.entries()) { group.h.dispose(); group.h = outputs[row]!; }
+          // Materialize every layer's recurrent tail as well as the residual.
+          // It is not a dependency of h and otherwise pins prefill buffers.
+          if (bounded) {
+            const caches = work.map(group => group.cache[i]!);
+            const state = caches.map(cache => cache.state());
+            try { ops.evalAll([...groups.map(group => group.h), ...state.flat()]); }
+            finally { for (const [row, cache] of caches.entries())
+              if (cache.stateNeedsDispose) for (const value of state[row]!) value.dispose(); }
+          }
+          for (const [row, group] of groups.entries()) work[row]!.captureLayer?.(i, group.h);
+        } finally { for (const mid of mids) mid.dispose(); }
+      }
+      for (const [row, group] of groups.entries()) {
+        const hidden = this.finalNorm.forward(group.h); results.push(hidden);
+        work[row]!.captureLayer?.(this.layers.length, hidden);
+      }
+      return results;
+    } catch (error) { for (const result of results) result.dispose(); throw error; }
+    finally { for (const group of groups) { group.h.dispose(); group.mask?.arr?.dispose(); group.ssmMask?.dispose(); } }
+  }
+
   /** Active vision mRoPE request state (serial lane; set by the generation
    *  gateway around a vision request's run, null for text-only — which keeps
-   *  the bit-exact fast-rope path). While set, forwardLayers installs the
-   *  per-forward interleaved cos/sin consumed by every full-attn layer. */
+   *  the bit-exact fast-rope path). The explicit serial executor scopes this value.
+   *  Shared media forwards pass request-owned positions directly. */
   mrope: MropeRequestState | null = null;
-  /** Serving's memory-pressure check, scoped to the serial request. */
-  prefillMemoryGuard: (() => void) | null = null;
   #mropeInvFreq: MlxArray | null = null;
 
   /** Vision prefill: spliced input embeddings [1, L, H] (image features
@@ -682,23 +788,23 @@ export class Qwen35Model {
     tap.captured.set(i, copy);
   }
 
-  protected forwardLayers(h0: MlxArray, cache: Cache[], independentRows = false): MlxArray {
+  protected forwardLayers(h0: MlxArray, cache: Cache[], independentRows = false,
+    positions?: MlxArray): MlxArray {
     const L = h0.shape[1]!;
     // One full-attention mask shared by all full layers (same offset); linear
     // layers see no ssm mask at B=1.
     const faMask = cache[this.faIdx]!.makeMask(L, null);
+    using ssmMask = (cache[0] as SSMCache).prefillPadding?.makeMask(L) ?? null;
     // Vision requests: one interleaved-mRoPE cos/sin table per forward,
     // shared by all 12 full-attention layers (positions are layer-invariant).
     let mropeFwd: ReturnType<typeof buildMropePositions> | null = null;
-    if (this.mrope) {
+    if (positions || this.mrope) {
       const t = this.config.text;
       const ropeDims = Math.trunc(t.headDim * t.partialRotaryFactor);
       const base = t.ropeParameters.full_attention?.ropeTheta ?? 10000;
       this.#mropeInvFreq ??= mropeInvFreq(ropeDims, base);
-      mropeFwd = buildMropePositions(
-        this.mrope, cache[this.faIdx]!.offset, L, this.#mropeInvFreq, ropeDims,
-      );
-      setActiveMrope(mropeFwd);
+      mropeFwd = positions ? { posIds: positions, invFreq: this.#mropeInvFreq, rotaryDims: ropeDims }
+        : buildMropePositions(this.mrope!, cache[this.faIdx]!.offset, L, this.#mropeInvFreq, ropeDims);
     }
     let h: MlxArray | null = h0;
     const prof = (globalThis as Record<string, unknown>).__deltaProf as
@@ -706,14 +812,13 @@ export class Qwen35Model {
       | undefined;
     try {
       for (let i = 0; i < this.layers.length; i++) {
-        if (L > 1) this.prefillMemoryGuard?.();
         const tl = prof ? performance.now() : 0;
-        const next = this.layers[i]!.forward(h, faMask, cache[i]!, independentRows);
+        const next = this.layers[i]!.forward(h, faMask, cache[i]!, independentRows, ssmMask, mropeFwd);
         h.dispose();
         h = next;
-        const mlp = this.layers[i]!.mlp;
-        if (L > TRELLIS_MATVEC_MAX_M &&
-            [mlp.gate, mlp.up, mlp.down].some((p) => p instanceof TrellisLinear && !p.fallback)) {
+        if (L > TRELLIS_MATVEC_MAX_M) {
+          // Affine prefills also need bounded evaluation: a retained prefix
+          // history can leave too little room for the entire deferred graph.
           // Materialize the cache outputs too. In particular, the tiny
           // copied conv tail is not a dependency of `h`; leaving it lazy
           // pins the whole prefill conv buffer until the chunk ends.
@@ -732,10 +837,7 @@ export class Qwen35Model {
       h = null; // consumed — the finally must not double-free
       return out;
     } finally {
-      if (mropeFwd) {
-        setActiveMrope(null);
-        mropeFwd.posIds.dispose();
-      }
+      if (mropeFwd && !positions) mropeFwd.posIds.dispose();
       // A mid-loop layer throw must not strand the mask array or the
       // in-flight [1,L,H] residual (2026-08-18 review).
       faMask.arr?.dispose();

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

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