mlx-bun 0.5.0 → 1.0.0

package/src/backends/mlx/fill-group.ts

@@ -1,497 +0,0 @@
-import type { GenerateOptions } from "../../generate";
-import type { RuntimeModel } from "../../model/factory";
-import type { Cache } from "../../model/gemma4-base";
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import { clearCache } from "../../mlx/ffi";
-import { makeStepSampler, type DeviceStepSampler, type NumberStepSampler } from "../../sampler";
-import type { FillSession, Proposal } from "../../fill/fill-session";
-import { disposeResources, cleanupFailure, withResource } from "../../engine/resources";
-import { snapshotGenerationPolicy } from "./request-policy";
-import { MlxPrefillCohort } from "./prefill-cohort";
-import { MlxStateRows } from "./state-rows";
-import { targetCacheLayout } from "./cache-layout";
-import { leaseCacheStates } from "./state-views";
-import { createKvMaintenance, type KvMaintenance } from "./kv-maintenance";
-import { bindLegacyAutoregressiveModel, supportsCommittedAppendCache, type MlxAutoregressiveBinding, type MlxTokenAppend } from "./autoregressive";
-import { appendHiddenRows } from "./fill-append";
-import { advanceSpeculativeOutputs } from "./speculative-round";
-import { bindRowCacheRollback } from "./rollback";
-import { bindSpeculativeTargetModel, type MlxSpeculativeTargetBinding } from "./speculative";
-import type { MlxForwardWork } from "./mixed-iteration";
-import type { MlxGroupedMethod, MlxGroupMethodHost, MlxGroupMethodRequest,
-  MlxGroupPreparation, Row } from "./batch-group";
-
-interface RequestState {
-  sampling: DeviceStepSampler;
-  fill: FillSession;
-  proposal?: { value: Proposal; emitted: number };
-}
-
-/** Strict continuations share preparation, row layouts and queueing with other
- * methods. Sampled tokens remain pipelined; committed tokens bypass sampling.
- * This binding is qualified independently before serving advertises it. */
-export function bindFillGroupRequests(model: RuntimeModel) {
-  return (input: GenerateOptions): MlxGroupMethodRequest => {
-    const options = snapshotGenerationPolicy(input);
-    return {
-      key: JSON.stringify(["fill", options.kvBits ?? null, options.kvGroupSize ?? 64,
-        options.quantizedKvStart ?? null, options.turboQuant ?? null, options.kvConfig ?? null]),
-      data: options,
-      open: host => new FillGroup(host, model, options),
-    };
-  };
-}
-
-class FillGroup implements MlxGroupedMethod {
-  readonly #binding: MlxAutoregressiveBinding;
-  readonly #maintain: KvMaintenance;
-  readonly #append: MlxTokenAppend | null | undefined;
-  readonly #verify: MlxSpeculativeTargetBinding;
-  readonly #requests = new Map<Row, RequestState>();
-  #target: MlxStateRows<ReturnType<typeof targetCacheLayout>> | null = null;
-  #pending: MlxArray | null = null;
-  #known: boolean[] = [];
-  #real: boolean[] = [];
-  #steps = 0;
-
-  constructor(readonly host: MlxGroupMethodHost, readonly model: RuntimeModel, options: GenerateOptions) {
-    this.#binding = bindLegacyAutoregressiveModel(model);
-    this.#verify = bindSpeculativeTargetModel(model);
-    this.#maintain = createKvMaintenance(options);
-    const append = this.#binding.createAppend?.({ hasAdapters: !!this.#binding.adapters?.active.length, pagedKv: false });
-    this.#append = supportsCommittedAppendCache(append, options) ? append : null;
-  }
-
-  get runningTokens(): number { return this.host.rows.length * Math.max(1, this.#burstLength(), this.#verifyLength()); }
-
-  #verifyLength(): number {
-    if (!this.#pending || !this.host.rows.length ||
-        this.host.rows.some(row => this.#requests.get(row)!.proposal?.value.policy === "assert")) return 0;
-    return Math.max(...this.host.rows.map(row => {
-      const proposal = this.#requests.get(row)!.proposal;
-      return proposal ? proposal.value.ids.length - proposal.emitted : 0;
-    }));
-  }
-
-  #burstLength(): number {
-    const rows = this.host.rows;
-    if (!rows.length || !this.#pending || this.#known.some(known => !known) ||
-        (this.#append?.maxChunkSize(this.#target!.caches, rows.length) ?? 1) <= 1) return 0;
-    return Math.min(...rows.map(row => {
-      const proposal = this.#requests.get(row)!.proposal;
-      return proposal?.value.policy === "assert" ? proposal.value.ids.length - proposal.emitted : 0;
-    }));
-  }
-
-  prepare(first: Row): MlxGroupPreparation {
-    const ready: { row: Row; caches: Cache[]; retain?: () => void }[] = [];
-    const target = new MlxPrefillCohort({
-      model: this.model, chunkSize: this.host.prefillChunkSize,
-      tailSplit: this.host.runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
-      promptCache: this.host.promptCache, maintain: this.#maintain,
-      forward: async (ids, caches) => this.#binding.graph.forwardHidden(ids, caches),
-      project: hidden => this.#binding.graph.projectLogits(hidden, { type: "all" }),
-      complete: async (state, logits) => {
-        const row = state.row;
-        const request = this.#requests.get(row)!;
-        using scores = ops.reshape(logits, [1, logits.shape.at(-1)!]);
-        using token = request.sampling.sample(scores, 0).token;
-        row.generated = row.sampled = 1;
-        try {
-          const reason = await this.#emit(row, token.toIntTokens()[0]!, false);
-          this.#maintain(state.solo);
-          if (reason) {
-            this.#put(row, state.solo, state.retain);
-            state.solo = []; state.retain = undefined;
-            this.#closeRequest(row); this.host.finish(row, reason);
-          } else {
-            ready.push({ row, caches: state.solo, retain: state.retain });
-          }
-        } catch (error) { this.#closeRequest(row); throw error; }
-      },
-      reject: (row, error) => { this.#closeRequest(row); row.reject(error); },
-    });
-    const admit = (row: Row) => {
-      const options = row.req.method!.data as GenerateOptions;
-      const request: RequestState = { fill: options.fill!,
-        sampling: makeStepSampler(options, { tokenRepresentation: "device", grammarWait: "external",
-          historyUpdate: "manual", initialHistory: row.req.promptIds }) };
-      row.fill = request.fill.stats;
-      this.#requests.set(row, request);
-      try { target.admit(row); }
-      catch (error) { this.#closeRequest(row); throw error; }
-    };
-    const preparation: MlxGroupPreparation = {
-      get rows() { return [...target.rows, ...ready.map(state => state.row)]; },
-      get canAdmit() { return target.canAdmit; },
-      get tokenWeight() { return target.tokenWeight + ready.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); },
-      admit,
-      advance: async work => {
-        if (!ready.length) {
-          const done = await target.advance(work);
-          return done && !ready.length;
-        }
-        // Flush before membership changes. Prepared output was already
-        // published on the preceding scheduler iteration.
-        await this.#flush();
-        while (ready.length) {
-          const state = ready.shift()!;
-          try {
-            if (state.row.req.signal?.aborted) {
-              this.#closeRequest(state.row); state.row.reject(state.row.req.signal.reason); continue;
-            }
-            this.#maintain.prepareBatch?.(state.caches);
-            this.#target ??= new MlxStateRows(state.caches.map(targetCacheLayout));
-            this.#target.append(state.caches);
-            // Copies must settle before releasing a restored SSD lease.
-            if (state.retain) withResource(leaseCacheStates(this.#target.caches), arrays => ops.evalAll([...arrays]));
-            this.host.join(state.row);
-          } finally { disposeResources([...state.caches, { dispose: () => state.retain?.() }]); }
-        }
-        return !target.rows.length;
-      },
-      dispose: () => {
-        const rows = preparation.rows;
-        disposeResources([target, ...ready.splice(0).flatMap(state => [...state.caches,
-          { dispose: () => state.retain?.() }]), { dispose: () => {
-          for (const row of rows) if (!this.host.rows.includes(row)) this.#closeRequest(row);
-        } }]);
-      },
-    };
-    try { admit(first); return preparation; }
-    catch (error) { return cleanupFailure(error, () => preparation.dispose()); }
-  }
-
-  async advance(work?: MlxForwardWork): Promise<void> {
-    const live = this.host.rows.flatMap((row, index) => row.req.signal?.aborted ? [] : [index]);
-    if (live.length !== this.host.rows.length) {
-      for (const row of this.host.rows) if (row.req.signal?.aborted) row.reject(row.req.signal.reason);
-      this.host.filterRows(live);
-    }
-    const rows = [...this.host.rows], B = rows.length;
-    if (!B) return;
-    const verifyLength = this.#verifyLength();
-    if (!work && verifyLength > 1) { await this.#verifyKnown(verifyLength); return; }
-    const burst = this.#burstLength();
-    if (!work && burst > 1) { await this.#appendKnown(burst); return; }
-    const previous = this.#pending, known = this.#known, real = this.#real;
-    const anyLive = rows.some(row => row.sampled < row.req.maxTokens);
-    let next: MlxArray | null = null;
-    const nextKnown: boolean[] = [], nextReal: boolean[] = [];
-    if (anyLive) {
-      using input = previous ? ops.reshape(previous, [B, 1]) : ops.fromInt32(rows.map(row => row.current), [B, 1]);
-      using hidden = await (work ? work(input, this.#target!.caches) : this.#binding.graph.forwardHidden(input, this.#target!.caches));
-      const selected = rows.map((row, index) => {
-        const request = this.#requests.get(row)!;
-        // The next selected token is not committed until it becomes input.
-        // Discarding a speculative sample for an assert cannot pollute history.
-        if (request.sampling.needsHistory) {
-          using current = input.slice([index, 0], [index + 1, 1]);
-          request.sampling.commitDevice(current);
-        }
-        const value = request.proposal?.value.policy === "assert"
-          ? request.proposal.value.ids[request.proposal.emitted + Number(!!previous && known[index])] : undefined;
-        nextKnown.push(value !== undefined); nextReal.push(row.sampled < row.req.maxTokens);
-        return value;
-      });
-      const needScores = rows.some((_row, index) => nextReal[index] && !nextKnown[index]);
-      using logits = needScores ? this.#binding.graph.projectLogits(hidden, { type: "last" }) : null;
-      const tokens: MlxArray[] = [];
-      try {
-        for (let index = 0; index < B; index++) {
-          const row = rows[index]!;
-          if (!nextReal[index]) { tokens.push(ops.fromInt32([0], [1])); continue; }
-          if (nextKnown[index]) tokens.push(ops.fromInt32([selected[index]!], [1]));
-          else {
-            using slice = logits!.slice([index, 0, 0], [index + 1, 1, logits!.shape[2]!]);
-            using scores = ops.reshape(slice, [1, slice.shape[2]!]);
-            tokens.push(this.#requests.get(row)!.sampling.sample(scores, row.sampled).token);
-          }
-          row.sampled++;
-        }
-        next = ops.concatAxis(tokens, 0);
-        // Constant known tokens have no dependency on the model graph.
-        // Submit state explicitly when the vocabulary projection is omitted.
-        if (needScores) ops.asyncEvalAll([next]);
-        else withResource(leaseCacheStates(this.#target!.caches), state => ops.asyncEvalAll([next!, hidden, ...state]));
-      } catch (error) { next?.dispose(); next = null; throw error; }
-      finally { disposeResources(tokens); }
-    }
-    this.#pending = next; this.#known = nextKnown; this.#real = nextReal;
-    if (previous) {
-      let values: number[];
-      try { values = previous.toIntTokens(); } finally { previous.dispose(); }
-      if (anyLive) for (let index = 0; index < B; index++) {
-        if (real[index]) rows[index]!.fed.push(values[index]!);
-        else rows[index]!.fedTainted = true;
-      }
-      await this.#emitRows(values, known);
-    } else if (anyLive) for (const row of rows) row.fed.push(row.current);
-    this.#checkPendingProposals();
-    if (++this.#steps % 256 === 0) clearCache();
-  }
-
-  /** The pending sample already verifies an echo's first token. Asserted
-   * continuations keep their append policy; other rows can join a verified
-   * span with an empty proposal through the existing row verifier. */
-  #checkPendingProposals(): void {
-    if (!this.#pending) return;
-    const rows = this.host.rows;
-    if (!rows.some(row => this.#requests.get(row)!.proposal?.value.policy === "verify")) return;
-    const tokens = this.#pending.toIntTokens();
-    for (const [index, row] of rows.entries()) {
-      const request = this.#requests.get(row)!, proposal = request.proposal;
-      if (proposal?.value.policy !== "verify") continue;
-      if (tokens[index] === proposal.value.ids[proposal.emitted]) this.#known[index] = true;
-      else {
-        request.fill.commit(proposal.value, proposal.emitted);
-        request.proposal = undefined;
-      }
-    }
-  }
-
-  async #verifyKnown(width: number): Promise<void> {
-    // Publish the already verified pending token before changing geometry.
-    await this.#flush();
-    const rows = [...this.host.rows];
-    if (!rows.length) return;
-    // Publishing an ordinary row can discover a strict template span. Let
-    // its append policy choose the next operation before verifying echoes.
-    if (rows.some(row => this.#requests.get(row)!.proposal?.value.policy === "assert")) return;
-    const pending = rows.map(row => row.current);
-    const proposals = rows.map(row => {
-      const proposal = this.#requests.get(row)!.proposal;
-      // A newly discovered proposal cannot exceed this iteration's declared
-      // token work. Its remaining continuation stays request-owned.
-      return proposal?.value.policy === "verify"
-        ? proposal.value.ids.slice(proposal.emitted, proposal.emitted + width - 1) : [];
-    });
-    const depth = Math.max(...proposals.map(ids => ids.length));
-    const samplers = rows.map(row => {
-      const sampling = this.#requests.get(row)!.sampling;
-      sampling.commitNumbers([row.current]);
-      let previous: number | undefined;
-      // A second sample is requested only after the preceding candidate was
-      // accepted. Leave the final correction out of manual processor history
-      // until the next forward actually consumes it.
-      return { ...sampling,
-        async sample(scores, step) {
-          if (previous !== undefined) sampling.commitNumbers([previous]);
-          const result = sampling.sample(scores, step);
-          try { previous = result.token.toIntTokens()[0]!; return { token: previous, extras: result.extras }; }
-          finally { result.token.dispose(); }
-        },
-      } satisfies NumberStepSampler;
-    });
-    const halted = new Set<Row>();
-    const rollback = bindRowCacheRollback(this.#target!.caches, rows.length);
-    const verifying = rows.filter(row => this.#requests.get(row)!.proposal?.value.policy === "verify")
-      .map(row => this.#requests.get(row)!.fill);
-    const pin = this.#verify.pinVerify?.();
-    try {
-      const completed = await advanceSpeculativeOutputs(rows.map((row, index) => ({
-        pending: pending[index]!, step: row.sampled, remaining: row.req.maxTokens - row.generated,
-        eosTokenIds: row.req.eosTokenIds, sampling: samplers[index]!,
-        output: { commit: async ids => {
-          const proposal = this.#requests.get(row)!.proposal;
-          const known = proposal?.value.policy === "verify" && ids[0] === proposal.value.ids[proposal.emitted];
-          if (!known) this.#settleProposal(row);
-          row.generated++;
-          const reason = await this.#emit(row, ids[0]!, known);
-          if (reason === "stop") { halted.add(row); return false; }
-        } },
-      })), depth, { draft: () => proposals, commit: () => {} }, {
-        transaction: { ...rollback,
-          begin(depth) {
-            const start = performance.now(); rollback.begin(depth);
-            const elapsed = performance.now() - start;
-            for (const fill of verifying) fill.noteVerifyEvent(elapsed);
-          },
-          resolve(accepted) {
-            const start = performance.now(); rollback.resolve(accepted);
-            const elapsed = performance.now() - start;
-            for (const fill of verifying) fill.stats.checkpointMs += elapsed;
-          },
-        },
-        forward: async ids => {
-          const result = await this.#verify.forward(ids, this.#target!.caches);
-          const context = result.ctxML ?? result.hidden;
-          try { return { logits: this.#verify.projectLogits(result.hidden), context }; }
-          catch (error) { context.dispose(); throw error; }
-          finally { if (result.ctxML) result.hidden.dispose(); }
-        },
-      });
-      const keep: number[] = [];
-      for (const [index, row] of rows.entries()) {
-        const output = completed.outputs[index]!, round = completed.rounds[index]!;
-        if (output.kind !== "failed" && output.kind !== "cancelled")
-          row.fed.push(pending[index]!, ...round.drafts.slice(0, output.accepted));
-        row.sampled += round.acceptance.emitted.length + Number(round.acceptance.sawEos);
-        if (output.kind === "continue") { row.current = output.pending; keep.push(index); }
-        else if (output.kind === "failed") row.reject(output.error);
-        else if (output.kind === "cancelled") row.reject(row.req.signal?.reason ?? new Error(output.reason));
-        else {
-          if (!halted.has(row) && round.acceptance.sawEos && row.generated < row.req.maxTokens) row.generated++;
-          this.#put(row, this.#target!.extractRow(index));
-          this.#settleProposal(row); this.host.finish(row, output.kind);
-        }
-      }
-      if (keep.length !== rows.length) this.host.filterRows(keep);
-      if (++this.#steps % 256 === 0) clearCache();
-    } finally { pin?.close(); }
-  }
-
-  /** Multi-position arithmetic is model-owned. Wider cohorts whose model
-   * declares one position keep the pipelined single-position implementation. */
-  async #appendKnown(length: number): Promise<void> {
-    const rows = [...this.host.rows], B = rows.length;
-    const spans = new Map(rows.map(row => {
-      const proposal = this.#requests.get(row)!.proposal!;
-      return [row, proposal.value.ids.slice(proposal.emitted, proposal.emitted + length)] as const;
-    }));
-    using input = ops.fromInt32(rows.flatMap(row => spans.get(row)!), [B, length]);
-    using hidden = await appendHiddenRows(this.#append!.forwardHidden.bind(this.#append), this.#target!.caches, input,
-      (caches, count) => Math.min(this.#append!.maxChunkSize(caches, count),
-        this.#maintain.maxAppendTokens?.(caches) ?? Infinity,
-        ...rows.map(row => this.#requests.get(row)!.fill.appendChunkSize || Infinity)));
-    const nextKnown: boolean[] = [], nextReal: boolean[] = [];
-    const selected = rows.map(row => {
-      const request = this.#requests.get(row)!, proposal = request.proposal!;
-      request.sampling.commitNumbers(spans.get(row)!);
-      row.fed.push(...spans.get(row)!);
-      const next = proposal.value.ids[proposal.emitted + length];
-      const real = row.generated + length < row.req.maxTokens;
-      nextKnown.push(next !== undefined); nextReal.push(real);
-      row.sampled = row.generated + length + Number(real);
-      return next;
-    });
-    const needScores = rows.some((_row, index) => nextReal[index] && !nextKnown[index]);
-    using logits = needScores ? this.#binding.graph.projectLogits(hidden, { type: "last" }) : null;
-    const tokens: MlxArray[] = [];
-    let next: MlxArray | null = null;
-    try {
-      for (const [index, row] of rows.entries()) {
-        if (!nextReal[index] || nextKnown[index]) tokens.push(ops.fromInt32([selected[index] ?? 0], [1]));
-        else {
-          using slice = logits!.slice([index, 0, 0], [index + 1, 1, logits!.shape[2]!]);
-          using scores = ops.reshape(slice, [1, slice.shape[2]!]);
-          tokens.push(this.#requests.get(row)!.sampling.sample(scores, row.generated + length).token);
-        }
-      }
-      next = ops.concatAxis(tokens, 0);
-      if (needScores) ops.asyncEvalAll([next]);
-      else withResource(leaseCacheStates(this.#target!.caches), state => ops.asyncEvalAll([hidden, ...state]));
-      this.#pending!.dispose(); this.#pending = next; next = null;
-      this.#known = nextKnown; this.#real = nextReal;
-    } finally { next?.dispose(); disposeResources(tokens); }
-    for (let column = 0; column < length && this.host.rows.length; column++) {
-      const active = [...this.host.rows];
-      await this.#emitRows(active.map(row => spans.get(row)![column]!), active.map(() => true));
-    }
-    if (++this.#steps % 256 === 0) clearCache();
-  }
-
-  async #emit(row: Row, token: number, known: boolean): Promise<"stop" | "length" | undefined> {
-    row.req.signal?.throwIfAborted();
-    const request = this.#requests.get(row)!;
-    if (row.req.eosTokenIds.includes(token)) return "stop";
-    // Account a known token even when the consumer stops on this token.
-    if (known) {
-      const proposal = request.proposal!;
-      if (++proposal.emitted === proposal.value.ids.length) {
-        request.fill.commit(proposal.value, proposal.emitted); request.proposal = undefined;
-      }
-    }
-    if (!known) {
-      const proposal = request.fill.push(token, row.req.maxTokens - row.generated);
-      if (proposal) {
-        request.proposal = { value: proposal, emitted: 0 };
-      }
-    }
-    const more = await this.host.publish(row, token);
-    if (more === false) return "stop";
-    if (row.generated >= row.req.maxTokens) return "length";
-    row.current = token;
-  }
-
-  async #emitRows(values: number[], known: readonly boolean[]): Promise<void> {
-    const rows = [...this.host.rows], keep: number[] = [];
-    const replacement = new Map<number, number>();
-    for (let index = 0; index < rows.length; index++) {
-      const row = rows[index]!;
-      row.generated++;
-      try {
-        const reason = await this.#emit(row, values[index]!, known[index] ?? false);
-        if (reason) {
-          this.#put(row, this.#target!.extractRow(index));
-          this.#settleProposal(row); this.host.finish(row, reason);
-        } else {
-          keep.push(index);
-          const request = this.#requests.get(row)!;
-          if (!known[index] && request.proposal?.value.policy === "assert" && this.#pending) {
-            replacement.set(index, request.proposal.value.ids[0]!);
-            request.fill.noteWastedSample(); this.#known[index] = true;
-          }
-        }
-      } catch (error) { row.reject(error); }
-    }
-    if (replacement.size) {
-      const parts: MlxArray[] = [];
-      try {
-        for (let index = 0; index < rows.length; index++) parts.push(replacement.has(index)
-          ? ops.fromInt32([replacement.get(index)!], [1]) : this.#pending!.slice([index], [index + 1]));
-        const next = ops.concatAxis(parts, 0); this.#pending!.dispose(); this.#pending = next;
-      } finally { disposeResources(parts); }
-    }
-    if (keep.length !== rows.length) this.host.filterRows(keep);
-  }
-
-  async #flush(): Promise<void> {
-    const pending = this.#pending, known = this.#known;
-    if (!pending) return;
-    this.#pending = null; this.#known = []; this.#real = [];
-    let values: number[];
-    try { values = pending.toIntTokens(); } finally { pending.dispose(); }
-    await this.#emitRows(values, known);
-  }
-
-  #put(row: Row, caches: Cache[], retain?: () => void): void {
-    try {
-      if (this.host.promptCache && !row.fedTainted && !row.req.signal?.aborted) {
-        this.host.promptCache.put([...row.req.promptIds, ...row.fed], caches,
-          row.cacheNamespace, retain, undefined, row.req.cacheSessionId);
-        caches = []; retain = undefined;
-      }
-    } finally { disposeResources([...caches, { dispose: () => retain?.() }]); }
-  }
-
-  #settleProposal(row: Row): void {
-    const request = this.#requests.get(row);
-    if (request?.proposal) {
-      const proposal = request.proposal; request.proposal = undefined;
-      request.fill.commit(proposal.value, proposal.emitted);
-    }
-  }
-  #closeRequest(row: Row): void {
-    this.#settleProposal(row);
-    const request = this.#requests.get(row); this.#requests.delete(row); request?.sampling.dispose();
-  }
-  filterRows(keep: readonly number[], _discard: boolean): void {
-    if (this.#pending) {
-      using indices = ops.fromInt32([...keep], [keep.length]);
-      const next = keep.length ? ops.takeAxis(this.#pending, indices, 0) : null;
-      this.#pending.dispose(); this.#pending = next;
-    }
-    this.#known = keep.map(index => this.#known[index]!); this.#real = keep.map(index => this.#real[index]!);
-    this.#target?.filterRows(keep);
-    const retained = new Set(keep.map(index => this.host.rows[index]!));
-    for (const row of this.host.rows) if (!retained.has(row)) this.#closeRequest(row);
-  }
-  dispose(): void {
-    const pending = this.#pending, target = this.#target;
-    this.#pending = null; this.#target = null;
-    disposeResources([...(pending ? [pending] : []), ...(target ? [target] : []),
-      { dispose: () => { for (const row of this.#requests.keys()) this.#closeRequest(row); } }]);
-  }
-}

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

@@ -1,32 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { Cache } from "../../model/gemma4-base";
-import { fullRowPadding } from "../../model/full-prefill-row";
-import { disposeResources } from "../../engine/resources";
-
-/** Assemble full-attention tensors at the model's B while storage retains
- * per-row formats and positions. Returned arrays own their captured state. */
-export function appendFullKvRows(rows: readonly Cache[], leftPad: readonly number[], width: number,
-  k: MlxArray, v: MlxArray, deferredValues = false): [MlxArray, MlxArray] {
-  const parts: [MlxArray, MlxArray][] = [], padded: MlxArray[] = [];
-  try {
-    for (const [index, row] of rows.entries()) {
-      using kr = k.slice([index, 0, 0, 0], [index + 1, k.shape[1]!, k.shape[2]!, k.shape[3]!]);
-      using vr = v.slice([index, 0, 0, 0], [index + 1, v.shape[1]!, v.shape[2]!, v.shape[3]!]);
-      parts.push(deferredValues && row.rotatedValueAttention
-        ? row.rotatedValueAttention.updateAndFetchDeferredV(kr, vr) : row.updateAndFetch(kr, vr));
-    }
-    const combine = (field: 0 | 1) => {
-      const views = parts.map((part, row) => {
-        const a = part[field], before = leftPad[row]! - fullRowPadding(rows[row]!), after = width - before - a.shape[2]!;
-        if (!before && !after) return a;
-        using prefix = ops.zeros([1, a.shape[1]!, before, a.shape[3]!], a.dtype);
-        using suffix = ops.zeros([1, a.shape[1]!, after, a.shape[3]!], a.dtype);
-        const result = ops.concatAxis([prefix, a, suffix], 2); padded.push(result); return result;
-      });
-      return views.length === 1 ? ops.contiguous(views[0]!) : ops.concatAxis(views, 0);
-    };
-    const keys = combine(0);
-    try { return [keys, combine(1)]; } catch (error) { keys.dispose(); throw error; }
-  } finally { disposeResources([...parts.flat(), ...padded]); }
-}

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

@@ -1,54 +0,0 @@
-import { captureKvDonorAttention, combineKvDonorAttention } from "../../model/kv-attention-view";
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { Cache, KvDonorRows } from "../../model/gemma4-base";
-import { fullRowInner, fullRowPadding } from "../../model/full-prefill-row";
-import { disposeResources } from "../../engine/resources";
-
-/** Read full-attention rows in their existing physical columns. Valid token
- * coverage excludes prefill padding; no append, conversion or compaction. */
-export function captureFullKvDonorRows(rows: readonly Cache[], leftPad: readonly number[], width: number): KvDonorRows {
-  const parts: KvDonorRows[] = [], padded: MlxArray[] = [];
-  try {
-    for (const row of rows) parts.push(fullRowInner(row).captureDonorRows!());
-    const combine = (field: "keys" | "values") => {
-      const views = parts.map((part, index) => {
-        const a = part[field], before = leftPad[index]! - fullRowPadding(rows[index]!), after = width - before - a.shape[2]!;
-        if (!before && !after) return a;
-        using prefix = ops.zeros([1, a.shape[1]!, before, a.shape[3]!], a.dtype);
-        using suffix = ops.zeros([1, a.shape[1]!, after, a.shape[3]!], a.dtype);
-        const result = ops.concatAxis([prefix, a, suffix], 2); padded.push(result); return result;
-      });
-      return views.length === 1 ? ops.contiguous(views[0]!) : ops.concatAxis(views, 0);
-    };
-    const offsets = rows.map(row => row.offset), starts = [...leftPad];
-    const keys = combine("keys");
-    try { return { keys, values: combine("values"), offsets, starts, ends: offsets.map((offset, row) => starts[row]! + offset) }; }
-    catch (error) { keys.dispose(); throw error; }
-  } finally { disposeResources([...parts.flatMap(part => [part.keys, part.values]), ...padded]); }
-}
-
-
-/** Capture full rows in global columns while each encoded view retains its
- * own physical width and numerical attention operation. */
-export function captureFullKvDonorAttention(rows: readonly Cache[], leftPad: readonly number[], width: number): import("../../model/gemma4-base").KvDonorAttention {
-  const captured: import("../../model/gemma4-base").KvDonorAttention[] = [];
-  try {
-    for (const [index, row] of rows.entries()) {
-      const view = captureKvDonorAttention(fullRowInner(row));
-      const pad = leftPad[index]! - fullRowPadding(row);
-      captured.push({ width, dtype: view.dtype, offsets: [row.offset],
-        starts: [leftPad[index]!], ends: [leftPad[index]! + row.offset],
-        attend(q, scale, mask) {
-          if (!mask.arr) return view.attend(q, scale, mask);
-          const from = mask.arr.shape.map(() => 0), to = [...mask.arr.shape];
-          from[from.length - 1] = pad; to[to.length - 1] = pad + view.width;
-          using arr = mask.arr.slice(from, to);
-          return view.attend(q, scale, { mode: mask.mode, arr });
-        },
-        dispose: () => view.dispose(),
-      });
-    }
-    return combineKvDonorAttention(captured);
-  } catch (error) { disposeResources(captured); throw error; }
-}

package/src/backends/mlx/full-transitioning-kv-rows.ts

@@ -1,65 +0,0 @@
-import { KVCache, type Cache, type PaddedPrefillCache, type PrefillPadding } from "../../model/gemma4-base";
-import { FullPrefillRow, fullRowInner } from "../../model/full-prefill-row";
-import { cloneKvCaches } from "../../kv-store";
-import { FullKvPositions, SpeculativeTransitioningKvRows, type KvRowTransition, type SpeculativeTransitionedKvLayout } from "./transitioning-kv-rows";
-
-type FullLayout = SpeculativeTransitionedKvLayout & PaddedPrefillCache & {
-  restorePrefillEnds(ends: readonly number[] | undefined): void;
-}
-
-class PrefillPositions extends FullKvPositions {
-  override sync(rows: readonly Cache[], preserve: boolean): void {
-    const offsets = rows.map(row => row instanceof FullPrefillRow ? row.positionOffset : row.offset);
-    this.mergeRows(preserve
-      ? [{ planes: [], rowOffsets: offsets, leftPad: [...this.leftPad] }]
-      : rows.map((row, index) => ({ planes: [], rowOffsets: [offsets[index]!],
-        leftPad: [row instanceof FullPrefillRow ? row.leftPadding : 0] })));
-  }
-}
-
-/** Full-attention padding composes with the existing precision transition
- * and rollback lifecycle. Codec-specific attention stays in each adapter. */
-export abstract class FullTransitioningKvRows<Layout extends FullLayout>
-  extends SpeculativeTransitioningKvRows<Layout> implements PaddedPrefillCache {
-  constructor(transition: KvRowTransition<Layout>, row?: Cache) {
-    super({ ...transition,
-      converted: row => transition.converted(fullRowInner(row)),
-      extractRow: row => row instanceof FullPrefillRow ? row.extract() : cloneKvCaches([row])[0]!,
-      packRows(layout, rows) {
-        layout.mergeRows(rows.map(fullRowInner));
-        if (rows.some(row => row instanceof FullPrefillRow)) {
-          layout.preparePrefill({ lengths: rows.map(() => 0),
-            leftPadding: rows.map(row => row instanceof FullPrefillRow ? row.leftPadding : 0) });
-          const ends = rows.map(row => row instanceof FullPrefillRow ? row.padding.endOffsets?.[0] : undefined);
-          layout.restorePrefillEnds(ends.some(end => end !== undefined) ? ends.map(end => end ?? Infinity) : undefined);
-        }
-      },
-    }, row, new PrefillPositions());
-  }
-  preparePrefill(padding: PrefillPadding): void {
-    if (this.packed) { this.packed.preparePrefill(padding); return; }
-    if (!this.rows.length) this.rows = padding.lengths.map(() => new KVCache());
-    this.rows = this.rows.map(row => row instanceof FullPrefillRow ? row : new FullPrefillRow(row as KVCache));
-    for (const [index, row] of this.rows.entries()) (row as FullPrefillRow).preparePrefill({
-      lengths: [padding.lengths[index]!],
-      ...(padding.leftPadding ? { leftPadding: [padding.leftPadding[index]!] } : {}),
-      ...(padding.rightPadding ? { rightPadding: [padding.rightPadding[index]!] } : {}),
-    });
-    this.syncPositions();
-  }
-  finalizePrefill(): void {
-    if (this.packed) { this.packed.finalizePrefill(); return; }
-    for (const row of this.rows) if (row instanceof FullPrefillRow) row.finalizePrefill();
-    this.syncPositions();
-  }
-  override filterRows(keep: readonly number[]): void {
-    if (!this.packed && keep.length) {
-      const pads = [...this.leftPad], shared = Math.min(...keep.map(row => pads[row]!));
-      for (const index of keep) {
-        const row = this.rows[index]!;
-        if (row instanceof FullPrefillRow) row.removeLeftPadding(Math.max(0, row.leftPadding - (pads[index]! - shared)));
-      }
-    }
-    super.filterRows(keep);
-  }
-}