mlx-bun 0.3.0 → 0.4.0

package/src/backends/mlx/target-layout-capability.ts

@@ -0,0 +1,35 @@
+import { Glm52Cache } from "../../model/glm52-cache";
+import { SpeculativeRotatingKVCache } from "../../model/speculative-rotating-kv";
+import { RotatingAffineLayout, SpeculativeRotatingAffineLayout } from "../../model/rotating-kv-layout";
+import type { Cache } from "../../model/gemma4-base";
+import { KVCache, QuantizedKVCache, TurboQuantKVCache, RotatingKVCache, RotatingQuantizedKVCache } from "../../model/gemma4-base";
+import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
+import { DelayedQuantizedKVCache } from "../../model/delayed-quantized-kv";
+import { DelayedTurboQuantKVCache } from "../../model/delayed-turboquant-kv";
+import { SSMCache } from "../../model/qwen3-delta";
+import { BatchedSSMCache } from "../../model/batched-ssm";
+import { BatchedKVCache } from "../../model/batched-kv";
+import { BatchedQuantizedKVCache } from "../../model/batched-quantized-kv";
+import { BatchedTurboQuantKVCache } from "../../model/batched-turboquant-kv";
+
+export type TargetRowLayout = Glm52Cache | DelayedRotatingQuantizedKVCache | RotatingAffineLayout | SpeculativeRotatingKVCache | BatchedKVCache | BatchedQuantizedKVCache | BatchedTurboQuantKVCache |
+  BatchedSSMCache | DelayedTurboQuantKVCache | DelayedQuantizedKVCache;
+
+/** The storage binding advertises a row transaction factory. The gateway can
+ * compose supported model graphs without identifying their model family. */
+export function targetRowLayoutFactory(cache: Cache): (() => TargetRowLayout) | undefined {
+  if (cache instanceof Glm52Cache) return () => cache.makeEmptyBatch();
+  if (cache instanceof DelayedRotatingQuantizedKVCache)
+    return () => new DelayedRotatingQuantizedKVCache(cache.maxSize, cache.groupSize, cache.bits, cache.start, cache.maintain, undefined, true);
+  if (cache instanceof RotatingAffineLayout || cache instanceof RotatingQuantizedKVCache)
+    return () => new SpeculativeRotatingAffineLayout(cache.maxSize, cache.groupSize, cache.bits);
+  if (cache instanceof RotatingKVCache || cache instanceof SpeculativeRotatingKVCache) return () => new SpeculativeRotatingKVCache(cache.maxSize);
+  if (cache instanceof DelayedTurboQuantKVCache || cache instanceof DelayedQuantizedKVCache ||
+      cache instanceof BatchedTurboQuantKVCache || cache instanceof BatchedKVCache ||
+      cache instanceof BatchedQuantizedKVCache || cache instanceof BatchedSSMCache) return () => cache.makeEmptyBatch();
+  if (cache instanceof TurboQuantKVCache) return () => new BatchedTurboQuantKVCache(cache.kBits, cache.vBits);
+  if (cache instanceof SSMCache) return () => new BatchedSSMCache();
+  if (cache instanceof QuantizedKVCache) return () => new BatchedQuantizedKVCache(cache.groupSize, cache.bits);
+  if (cache instanceof KVCache) return () => new BatchedKVCache();
+  return undefined;
+}

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

@@ -0,0 +1,165 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { BatchableCache, Cache, Mask } from "../../model/gemma4-base";
+import { KvTensorRows } from "./kv-tensor-rows";
+import { leaseCacheStates, minimumReusableOffset } from "./state-views";
+import { withResource, disposeResources } from "../../engine/resources";
+import { cloneKvCaches } from "../../kv-store";
+
+/** Encoded layouts retain physical row alignment across precision changes. */
+export interface TransitionedKvLayout extends BatchableCache {
+  alignRows?(leftPad: readonly number[]): void;
+}
+export interface SpeculativeTransitionedKvLayout extends TransitionedKvLayout {
+  specRoundBegin(): void;
+  specRoundCommit(): void;
+  specRoundRollback(keep: number | readonly number[]): void;
+}
+export interface KvRowTransition<Layout extends TransitionedKvLayout> {
+  readonly signature: string;
+  maintain(rows: Cache[]): void;
+  converted(row: Cache): boolean;
+  makeLayout(): Layout;
+  packRows?(layout: Layout, rows: readonly Cache[]): void;
+  prepareRows?(rows: Cache[]): Cache[];
+  extractRow?(row: Cache): Cache;
+  rollbackRow?(row: Cache, before: number, keep: number, preserveColumns: boolean): void;
+}
+
+
+export interface TransitioningKvPositions {
+  readonly rowOffsets: readonly number[];
+  readonly leftPad: readonly number[];
+  readonly offset: number;
+  readonly batchSize: number | null;
+  readonly ropeOffsetArr: MlxArray | undefined;
+  sync(rows: readonly Cache[], preserve: boolean): void;
+  filterRows(keep: readonly number[]): void;
+  makeMask(tokens: number, window: number | null): Mask;
+  dispose(): void;
+}
+export class FullKvPositions extends KvTensorRows implements TransitioningKvPositions {
+  sync(rows: readonly Cache[], preserve: boolean): void {
+    this.mergeRows(preserve
+      ? [{ planes: [], rowOffsets: rows.map(row => row.offset), leftPad: [...this.leftPad] }]
+      : rows.map(row => ({ planes: [], rowOffsets: [row.offset], leftPad: [0] })));
+  }
+}
+
+/** Owns membership, committed positions and precision transitions. Codec and
+ * attention adapters own tensor access; scheduling and persistence use the
+ * existing batch and checkpoint contracts. Row conversion never examines an
+ * uncommitted speculative suffix. */
+export abstract class TransitioningKvRows<Layout extends TransitionedKvLayout> implements BatchableCache {
+  protected readonly positions: TransitioningKvPositions;
+  protected rows: Cache[] = [];
+  protected packed?: Layout;
+  #prefilling = false;
+
+  readonly stateNeedsDispose = true;
+  constructor(readonly transition: KvRowTransition<Layout>, row?: Cache, positions: TransitioningKvPositions = new FullKvPositions()) {
+    this.positions = positions;
+    if (row) { this.rows = transition.prepareRows?.([row]) ?? [row]; this.syncPositions(); }
+  }
+  signature(): string { return this.transition.signature; }
+  get prefillMaintenance() { return this; }
+  beginPrefill(): void { this.#prefilling = true; }
+  commitPrefill(indices: readonly number[]): void {
+    if (this.packed) return;
+    for (const index of indices) {
+      const row = [this.rows[index]!];
+      try { this.transition.maintain(row); }
+      finally { this.rows[index] = row[0]!; }
+    }
+  }
+  endPrefill(): void { this.#prefilling = false; }
+  get minimumReusableOffset(): number { return this.packed?.minimumReusableOffset ?? minimumReusableOffset(this.rows); }
+  get rowOffsets(): readonly number[] { return this.packed?.rowOffsets ?? this.positions.rowOffsets; }
+  get leftPad(): readonly number[] { return this.packed?.leftPad ?? this.positions.leftPad; }
+  get offset(): number { return this.packed?.offset ?? this.positions.offset; }
+  get batchSize(): number | null { return this.packed?.batchSize ?? this.positions.batchSize; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.packed ? this.packed.ropeOffsetArr : this.positions.ropeOffsetArr; }
+  bytesPerToken(): number { return this.packed?.bytesPerToken?.() ?? Math.max(0, ...this.rows.map(row => row.bytesPerToken?.() ?? 0)); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return (this.packed ?? this.positions).makeMask(tokens, window); }
+  isTrimmable(): boolean { return this.packed?.isTrimmable() ?? this.rows.every(row => row.isTrimmable()); }
+  trim(count: number): void {
+    if (this.packed) this.packed.trim(count);
+    else { for (const row of this.rows) row.trim(count); this.syncPositions(true); }
+  }
+  protected syncPositions(preserve = false): void { this.positions.sync(this.rows, preserve); }
+  protected advance(): void {
+    if (this.packed) return;
+    if (!this.#prefilling) this.transition.maintain(this.rows);
+    if (this.rows.length && this.rows.every(row => this.transition.converted(row))) {
+      const packed = this.transition.makeLayout();
+      try {
+        if (this.transition.packRows) this.transition.packRows(packed, this.rows);
+        else packed.mergeRows(this.rows);
+        packed.alignRows?.(this.positions.leftPad);
+      }
+      catch (error) { packed.dispose(); throw error; }
+      const previous = this.rows; this.rows = []; this.packed = packed;
+      this.positions.dispose(); disposeResources(previous);
+    }
+  }
+  abstract updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray];
+  abstract makeEmptyBatch(): TransitioningKvRows<Layout>;
+  mergeRows(rows: readonly Cache[]): void {
+    const owned: Cache[] = [];
+    try {
+      for (const row of rows) {
+        if (row instanceof TransitioningKvRows) {
+          for (let index = 0; index < row.batchSize!; index++) owned.push(row.extractRow(index));
+        } else owned.push(...cloneKvCaches([row]));
+      }
+    } catch (error) { disposeResources(owned); throw error; }
+    let prepared: Cache[];
+    try { prepared = this.transition.prepareRows?.(owned) ?? owned; }
+    catch (error) { disposeResources(owned); throw error; }
+    this.dispose(); this.rows = prepared; this.syncPositions();
+  }
+  filterRows(keep: readonly number[]): void {
+    if (this.packed) { this.packed.filterRows(keep); return; }
+    const previous = this.rows; this.rows = keep.map(index => previous[index]!);
+    disposeResources(previous.filter((_, index) => !keep.includes(index))); this.positions.filterRows(keep);
+  }
+  extractRow(row: number): Cache { return this.packed?.extractRow(row) ?? (this.transition.extractRow?.(this.rows[row]!) ?? cloneKvCaches([this.rows[row]!])[0]!); }
+  state(): MlxArray[] {
+    return withResource(leaseCacheStates(this.packed ? [this.packed] : this.rows), arrays => {
+      const views: MlxArray[] = [];
+      try { for (const array of arrays) views.push(ops.contiguous(array)); return views; }
+      catch (error) { disposeResources(views); throw error; }
+    });
+  }
+  dispose(): void {
+    this.#prefilling = false;
+    const previous = this.rows; this.rows = [];
+    this.packed?.dispose(); this.packed = undefined; this.positions.dispose(); disposeResources(previous);
+  }
+}
+
+/** Rollback is a method capability, separate from ordinary cache membership. */
+export abstract class SpeculativeTransitioningKvRows<Layout extends SpeculativeTransitionedKvLayout> extends TransitioningKvRows<Layout> {
+  #beforeRound?: number[];
+  specRoundBegin(): void {
+    this.advance(); this.#beforeRound = [...this.rowOffsets]; this.packed?.specRoundBegin();
+  }
+  specRoundCommit(): void {
+    this.packed?.specRoundCommit(); this.#beforeRound = undefined; this.advance();
+  }
+  specRoundRollback(keep: number | readonly number[]): void {
+    if (this.packed) this.packed.specRoundRollback(keep);
+    else {
+      for (const [index, row] of this.rows.entries()) {
+        const count = typeof keep === "number" ? keep : keep[index]!;
+        const before = this.#beforeRound![index]!;
+        if (this.transition.rollbackRow) this.transition.rollbackRow(row, before, count, this.rows.length > 1);
+        else row.trim(row.offset - before - count);
+      }
+      this.syncPositions(true);
+    }
+    this.#beforeRound = undefined; this.advance();
+  }
+  override dispose(): void { this.#beforeRound = undefined; super.dispose(); }
+}

package/src/cli.ts

@@ -120,10 +120,11 @@ const SERVER_FLAGS = `Server options:
   --ssd-cache-verify        Verify tensor hashes on every restore (reads all
                             bytes eagerly — integrity paranoia only)
   --generation-checkpoint <tokens>
-                            Snapshot an in-flight serial generation to the SSD
+                            Queue an in-flight generation snapshot to the SSD
                             cache every N emitted tokens. Repeating the same
                             request after restart replays and resumes it.
-                            Requires --ssd-cache and --batch 1.
+                            Requires --ssd-cache; supports ordinary requests
+                            with shared cache layouts.
   --isolate                 Run the inference engine as a CHILD process and
                             keep this process a pure UI/API proxy — the UI
                             stays instant under GPU load and survives engine
@@ -169,8 +170,8 @@ Model & quality:
                             auto-detected — a full same-tokenizer model
                             (mlx_lm.server parity), a Gemma "-assistant"
                             KV-borrowing drafter, or a DSpark checkpoint.
-                            Serial lane only — with --batch N a mounted draft
-                            routes every request serial, like mlx_lm.server.
+                            Supported draft providers use shared batched
+                            execution; ordinary requests can batch too.
   --draft-kind <kind>       Override draft-artifact detection:
                             two-model | assistant | dspark | deepspec | mtp | ngram
                             (deepspec = DeepSeek's released Gemma4DSparkModel
@@ -205,9 +206,8 @@ Model & quality:
                             only; sliding-window layers stay bf16)]
   --paged-kv                OPTIONAL vLLM-style paged KV storage (block pool +
                             gather before the stock SDPA; env MLX_BUN_PAGED_KV=1).
-                            v1: serial only (pins --batch 1 unless --batch is
-                            given), Gemma4-family, bf16 — refuses --batch N>1,
-                            --kv-quant, --draft-model; bypasses the prompt
+                            Gemma4-family, bf16; supports --batch N.
+                            Refuses --kv-quant and --draft-model; bypasses the prompt
                             cache; runs uncompiled decode. Bit-exact with the
                             plain path; expect a small decode cost at batch=1
                             (the gather copy). docs/design/kv-cache.md
@@ -1012,16 +1012,11 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
   if (route.kvQuant !== undefined) serverOptions.kvQuant = route.kvQuant;
   if (route.turboQuant !== undefined) serverOptions.turboQuant = route.turboQuant;
   // --paged-kv (env: MLX_BUN_PAGED_KV=1): OPTIONAL vLLM-style block-pool KV
-  // storage, default off (docs/design/kv-cache.md). Serial-only in v1:
-  // with no explicit --batch, pin --batch 1 (the default is 8); an explicit
-  // --batch N>1 is refused loudly by createServer rather than downgraded.
+  // storage, default off (docs/design/kv-cache.md). The state policy owns
+  // block pools and cache reuse; paging preserves the selected batch size.
   if (flag("paged-kv") || runtimeValue("MLX_BUN_PAGED_KV") === "1") {
     const bsRaw = opt("paged-kv-block-size");
     serverOptions.pagedKv = bsRaw !== null ? { blockSize: Number(bsRaw) } : {};
-    if (batchRaw === null) {
-      serverOptions.batch = 1;
-      console.log("[paged-kv] serial-only in v1 — pinning --batch 1");
-    }
   }
   // Bind loopback unless asked otherwise (mlx_lm.server parity); --host
   // 0.0.0.0 is the explicit opt-in for LAN exposure. The chat-UI open and
@@ -1646,25 +1641,6 @@ switch (cmd) {
     const config = await loadModelConfig(m.path);
     const report = fit(config, m.sizeBytes, 8192, undefined, undefined, m.expertsBytes);
     sFit.done(`${style.bold(m.repoId)} ${style.dim(`· ${gb(m.sizeBytes)} · ~${report.predictedDecodeTps.toFixed(0)} tok/s predicted`)}${picked ? style.dim(" · auto-picked (override: mlx-bun serve <query>)") : ""}`);
-    // Near-ceiling advisory (admission doctrine: advise, never refuse). A
-    // model near the DEFAULT Metal wired ceiling (~75% of RAM) dies with an
-    // UNCATCHABLE async-GPU-OOM on long prefills (M4 24 GB × 17 GB model,
-    // 2026-08-20) — surface the standard remedy up front instead.
-    {
-      const { maxRecommendedWorkingSetSize } = await import("./mlx/ffi");
-      const ceiling = maxRecommendedWorkingSetSize();
-      const sysctlRaw = Bun.spawnSync(["sysctl", "-n", "iogpu.wired_limit_mb"]).stdout.toString().trim();
-      const limitIsDefault = sysctlRaw === "0" || sysctlRaw === "";
-      if (limitIsDefault && ceiling > 0 && m.sizeBytes > 0.8 * ceiling) {
-        const ramGb = Number(Bun.spawnSync(["sysctl", "-n", "hw.memsize"]).stdout.toString().trim()) / 2 ** 30;
-        const suggestMb = Math.floor((ramGb - 2.5) * 1024);
-        console.error(style.dim(
-          `  ! ${gb(m.sizeBytes)} model vs ~${gb(ceiling)} default GPU ceiling — long prompts can\n` +
-          `    hit an uncatchable Metal OOM. If that happens, raise the limit first:\n` +
-          `      sudo sysctl iogpu.wired_limit_mb=${suggestMb}   (resets on reboot)`,
-        ));
-      }
-    }
     const sNative = step("native runtime");
     await ensureNative(sNative);
     sNative.done("native runtime ready");
@@ -1684,9 +1660,8 @@ switch (cmd) {
     }
     // Speculative decoding (mlx_lm.server parity): --draft-model <path|query>
     // resolves like the main model; --num-draft-tokens defaults to 3 (the
-    // server's default upstream — mlx_lm.generate's is 2). Serial-lane-only:
-    // with --batch N, a mounted draft routes every request serial (upstream
-    // is_batchable = draft is None).
+    // server's default upstream — mlx_lm.generate's is 2). Qualified providers
+    // use shared execution; ordinary requests retain their own placement.
     const draftQuery = opt("draft-model");
     let draftModelDir: string | undefined;
     if (draftQuery !== null) {

package/src/contracts/execution.ts

@@ -17,9 +17,20 @@ export interface ExecutionCapabilities {
   readonly continuous: boolean;
   readonly quantizedBatch: boolean;
   readonly grammarBatch: boolean;
+  readonly adapterBatch?: boolean;
+  readonly pagedBatch?: boolean;
   readonly checkpoints: boolean;
+  /** Bound ordinary driver can restore/capture this request configuration. */
+  readonly sharedCheckpoints?: boolean;
   /** Model-qualified speculative execution can retain this request's KV codec. */
   readonly speculativeKvQuant?: boolean;
+  readonly speculativeTurboQuant?: boolean;
+  readonly turboQuantBatch?: boolean;
+  readonly speculativeLogprobs?: boolean;
+  /** The grouped provider supports the target's mounted adapter context. */
+  readonly sharedSpeculativeAdapters?: boolean;
+  /** Methods supplied by the model's shared execution binding. */
+  readonly groupedMethods?: readonly string[];
   /** A graph-owned compiled step exists; cache geometry can still decline it. */
   readonly compiledDecode?: boolean;
 }

package/src/contracts/prefix-cache.ts

@@ -0,0 +1,23 @@
+/** A reusable prefix and the backend/method state that covers its tokens.
+ * The caller owns a returned hit. After releasing or returning its state,
+ * it must release any retained backing through retain. */
+export interface PrefixCacheHit<State, Attachment = never> {
+  tokens: number[];
+  caches: State;
+  attachments?: Attachment;
+  retain?: () => void;
+}
+
+/** Storage is independent of request scheduling and the method producing
+ * state. Both prefill and decode can publish processed-token prefixes.
+ * State is opaque here; its backend/method defines layout and alignment.
+ * Implementations own prefix selection, retention and storage tiers. */
+export interface PrefixCache<State, Attachment = never> {
+  /** Return owned state for the longest usable prefix, or null on a miss.
+   * Namespace distinguishes incompatible artifacts, methods and adapters. */
+  take(prompt: number[], ns?: string): PrefixCacheHit<State, Attachment> | null;
+  /** Transfer state and its optional backing release to storage. If this
+   * throws, ownership stays with the caller; successful return transfers
+   * ownership even when storage immediately spills or disposes the entry. */
+  put(tokens: number[], caches: State, ns?: string, retain?: () => void, attachments?: Attachment): void;
+}

package/src/contracts/resources.ts

@@ -16,3 +16,10 @@ export interface StateView<T> {
 export interface DisposableResource {
   dispose(): void;
 }
+
+/** A prepared replacement for live state. Commit once between method rounds;
+ * it publishes already-built state without graph work or validation. Dispose
+ * releases the unpublished replacement, or the superseded state after commit. */
+export interface PreparedStateChange extends DisposableResource {
+  commit(): void;
+}

package/src/contracts/sampling.ts

@@ -0,0 +1,23 @@
+/** Sampling owns score processing, token selection, and request-local history.
+ * Score/token representations and sync/device execution belong to the backend;
+ * this contract has no scheduling, cache, or tensor-runtime dependency. */
+export interface SampledToken<Token, Metadata> {
+  token: Token;
+  extras: Metadata | null;
+}
+
+export interface SamplingSession<Scores, DeviceToken, Result> {
+  /** Optional stateless operation shared across compatible rows and positions.
+   * It borrows scores and returns owned device tokens, without advancing history. */
+  readonly independent?: { sample(scores: Scores): DeviceToken };
+  readonly isPlainGreedy: boolean;
+  readonly capturesLogprobs: boolean;
+  readonly needsHistory: boolean;
+  /** Borrows scores; transfers ownership of the result to the caller. */
+  sample(logits: Scores, step: number): Result;
+  seedHistory(tokens: readonly number[]): void;
+  /** Commit accepted tokens when configured for manual history updates. */
+  commitDevice(token: DeviceToken): void;
+  commitNumbers(tokens: readonly number[]): void;
+  dispose(): void;
+}

package/src/contracts/scheduling.ts

@@ -1,9 +1,27 @@
+/** Backend-owned context shared by compatible requests. A key identifies
+ * equivalent execution state; enter returns its scoped release. Scheduling
+ * observes group readiness without inspecting this context. */
+export interface ExecutionContext {
+  readonly key: string;
+  enter(): () => void;
+}
+
 /** A backend execution group exposes bounded work units and readiness. Queue
  * payloads, tensors, cache merging and sampling remain inside its adapter. */
 export interface ExecutionGroup {
   readonly active: number;
   readonly queued: number;
   readonly preparing: boolean;
+  readonly preparingRows?: number;
+  /** The preparation owner can accept another row before its next work unit. */
+  readonly canPrepareMore?: boolean;
+  /** Work estimates affect grouping only. A request larger than this budget
+   * still runs alone; it is never rejected or sent to another executor. */
+  readonly preparingTokens?: number;
+  readonly nextPreparationTokens?: number;
+  readonly maxPreparationTokens?: number;
+  /** The latest preparation published output that can be flushed by yielding. */
+  readonly preparationPublishedOutput?: boolean;
   readonly maxActive: number;
   readonly admissionHeld: boolean;
   readonly closed: boolean;

package/src/engine/execution-plan.ts

@@ -1,33 +1,35 @@
 import type { ExecutionCapabilities, ExecutionFeatures, ExecutionRequirements, ResolvedExecution } from "../contracts/execution";
 
-/** Preserve the shipped support matrix. A configured draft still forces serial
- * placement even when this request must use the ordinary AR method. */
+/** Select the request method, then place it using executor capabilities. */
 export function resolveExecution(
   request: ExecutionRequirements, capabilities: ExecutionCapabilities,
   features: ExecutionFeatures = { pagedKv: false, fill: false },
 ): ResolvedExecution {
   const reasons: string[] = [];
-  const continuousExclusions = [
+  const sharedRequestExclusions = [
     [!capabilities.continuous, "continuous-unavailable"],
-    [capabilities.method !== "autoregressive", "method-requires-serial"],
     [request.hasVision, "media-requires-serial"],
-    [request.hasAdapters, "adapters-require-serial"],
-    [request.wantsLogprobs, "logprobs-require-serial"],
-    [request.userSeed, "explicit-seed-requires-serial"],
+    [request.hasAdapters && !capabilities.adapterBatch, "adapters-require-serial"],
     [request.kvQuant && !capabilities.quantizedBatch, "kv-scheme-requires-serial"],
-    [request.turboQuant, "turbo-kv-requires-serial"],
-    [request.hasDraft, "configured-draft-requires-serial"],
+    [request.turboQuant && !capabilities.turboQuantBatch, "turbo-kv-requires-serial"],
     [request.hasGrammar && !capabilities.grammarBatch, "grammar-batching-disabled"],
-    [features.pagedKv, "paged-kv-requires-serial"],
+    [features.pagedKv && !capabilities.pagedBatch, "paged-kv-requires-serial"],
   ] as const;
-  for (const [excluded, reason] of continuousExclusions) if (excluded) reasons.push(reason);
-  const mechanism = reasons.length ? "serial" : "continuous";
+  const sharedAdapterMethod = capabilities.sharedSpeculativeAdapters === true &&
+    !sharedRequestExclusions.some(([excluded]) => excluded) &&
+    capabilities.groupedMethods?.includes("speculative") === true;
   const speculative = capabilities.method === "autoregressive" && request.hasDraft &&
-    !request.hasVision && !request.hasAdapters && !request.wantsLogprobs &&
+    !request.hasVision && (!request.hasAdapters || sharedAdapterMethod) && (!request.wantsLogprobs || capabilities.speculativeLogprobs === true) &&
     (!request.kvQuant || capabilities.speculativeKvQuant === true) &&
-    !request.turboQuant && !features.pagedKv;
-  if (request.hasDraft && !speculative) reasons.push("draft-incompatible-with-request");
+    (!request.turboQuant || capabilities.speculativeTurboQuant === true) && !features.pagedKv;
   const method = speculative ? "speculative" : capabilities.method;
+  const continuousExclusions = [sharedRequestExclusions[0],
+    [!(capabilities.groupedMethods ?? ["autoregressive"]).includes(method), "method-requires-serial"],
+    ...sharedRequestExclusions.slice(1),
+  ] as const;
+  for (const [excluded, reason] of continuousExclusions) if (excluded) reasons.push(reason);
+  const mechanism = reasons.length ? "serial" : "continuous";
+  if (request.hasDraft && !speculative) reasons.push("draft-incompatible-with-request");
   const pagedKv = features.pagedKv && !request.hasVision && !request.hasAdapters;
   if (features.pagedKv && !pagedKv) reasons.push("paged-kv-bypassed-for-media-or-adapters");
   const promptCache = method !== "speculative" && !request.hasVision && !pagedKv;
@@ -36,7 +38,7 @@ export function resolveExecution(
     !request.wantsLogprobs && !request.kvQuant && !request.turboQuant;
   if (features.fill && !fill) reasons.push("fill-incompatible-with-request");
   const checkpoint = capabilities.checkpoints && method === "autoregressive" &&
-    mechanism === "serial" && promptCache && !request.hasGrammar &&
+    (mechanism === "serial" || capabilities.sharedCheckpoints === true) && promptCache && !request.hasGrammar &&
     !features.fill && !request.wantsLogprobs;
   const compiledDecode = features.compiledDecode === true && capabilities.compiledDecode === true &&
     method === "autoregressive" && !request.hasAdapters && !pagedKv;

package/src/engine/resources.ts

@@ -1,4 +1,4 @@
-import type { DisposableResource, ResourceOwner } from "../contracts/resources";
+import type { DisposableResource, PreparedStateChange, ResourceOwner } from "../contracts/resources";
 
 export function ownResource<T>(value: T, release: (value: T) => void): ResourceOwner<T> {
   let held: { value: T } | undefined = { value };
@@ -35,6 +35,17 @@ export function cleanupFailure(error: unknown, cleanup: () => void): never {
   throw error;
 }
 
+/** Complete every preparation before publishing any owner's replacement.
+ * Old state is released only after all owners have committed. */
+export function applyStateChanges(prepare: readonly (() => PreparedStateChange)[]): void {
+  const changes: PreparedStateChange[] = [];
+  try {
+    for (const build of prepare) changes.push(build());
+    for (const change of changes) change.commit();
+  } catch (error) { return cleanupFailure(error, () => disposeResources(changes.splice(0))); }
+  disposeResources(changes);
+}
+
 export function withResource<T, R>(owner: Pick<ResourceOwner<T>, "borrow" | "close">, use: (value: T) => R): R {
   try { return use(owner.borrow()); }
   catch (error) { cleanupFailure(error, () => owner.close()); }

package/src/engine/scheduler.ts

@@ -30,15 +30,20 @@ export async function driveExecutionGroup(
 
       if (!group.preparing && !held && group.queued && group.active < group.maxActive)
         group.admitNext();
+      while (group.preparing && group.canPrepareMore && !held && group.queued &&
+          group.active + (group.preparingRows ?? 1) < group.maxActive && group.canBurst() &&
+          (group.preparingTokens ?? 0) + (group.nextPreparationTokens ?? 0) <= (group.maxPreparationTokens ?? Infinity)) {
+        if (!group.admitNext()) break;
+      }
       if (group.preparing) {
         const activeBefore = group.active;
         await group.advancePreparation();
-        // A backend may emit its first output while preparing the first row.
-        // Flush a lone admission before decode; queued short admissions still
-        // group together. Re-enter policy after yielding to observe new work,
+        // A bounded preparation may publish output before its state work ends.
+        // Flush it before the next work unit; queued short admissions still
+        // form a group across yields. Re-enter policy to observe new work,
         // cancellation, admission holds and shutdown.
-        if (yieldAfterPreparation && !activeBefore && !group.preparing &&
-            group.active === 1 && !group.queued) {
+        if (group.preparationPublishedOutput ||
+            (yieldAfterPreparation && !activeBefore && !group.preparing && group.active === 1 && !group.queued)) {
           lastYield = clock.now();
           await clock.yield();
           continue;

package/src/fit.ts

@@ -13,12 +13,12 @@ import type { ModelConfig } from "./config";
 import type { MemoryPlan } from "./admission-plan";
 import {
   kvBytesAt,
-  kvGeometry,
   kvQuantBytesPerElement,
+  sdpaFallbackBytes,
   type KvSchemeOptions,
 } from "./kv-scheme";
 
-export { kvBytesAt, kvQuantBytesPerElement } from "./kv-scheme";
+export { kvBytesAt, kvQuantBytesPerElement, sdpaFallbackBytes } from "./kv-scheme";
 export type FitKvScheme = KvSchemeOptions;
 
 /** Decode-efficiency vs theoretical bandwidth ceiling, measured on the
@@ -161,23 +161,37 @@ export function fit(
   kvScheme?: FitKvScheme,
 ): FitReport {
   const usable = usableBytes ?? machine.ramBytes * WIRED_FRACTION;
-  const transient = Math.min(chunk, ctx) * TRANSIENT_PER_TOKEN;
+  // Prefill transient = the calibrated per-chunk-token constant + the
+  // SDPA-fallback scores tensor (heads × min(chunk, ctx) × ctx × 2 B) for
+  // attention layers whose head dim MLX cannot fuse (kv-scheme.ts). The
+  // second term is what grows with context on Qwen3.x (head_dim 256) and
+  // Gemma 4 (256/512) and was unbilled before 2026-09-08.
+  const transientAt = (n: number) =>
+    Math.min(chunk, n) * TRANSIENT_PER_TOKEN + sdpaFallbackBytes(config, chunk, n);
+  const transient = transientAt(ctx);
   const kv = kvBytesAt(config, ctx, kvScheme);
   const total = weightsBytes + kv + transient;
 
-  // solve max context: weights + kv(ctx) + transient ≤ usable.
-  // Below the window both KV terms are linear in ctx; above it the
-  // sliding term saturates and only full-attention layers keep growing.
-  const g = kvGeometry(config, kvScheme);
-  const fixed = weightsBytes + chunk * TRANSIENT_PER_TOKEN + g.linearStateBytes;
+  // solve max context: weights + kv(n) + transient(n) ≤ usable. The same
+  // formula `fits` uses, so the solved ceiling always fits. It is monotone
+  // in n (KV linear, sliding saturating at the window, scores quadratic
+  // below the chunk then linear), so bisect over [0, maxPositionEmbeddings].
+  const totalAt = (n: number) => weightsBytes + kvBytesAt(config, n, kvScheme) + transientAt(n);
   let maxCtx = 0;
-  if (usable > fixed) {
-    const budget = usable - fixed;
-    const linear = Math.floor(budget / (g.fullBytesPerToken + g.slidingBytesPerToken));
-    maxCtx = linear <= g.window
-      ? linear
-      : Math.floor((budget - g.slidingBytesPerToken * g.window) / g.fullBytesPerToken);
-    maxCtx = Math.min(maxCtx, config.text.maxPositionEmbeddings);
+  {
+    let lo = 0;
+    let hi = Math.max(0, Math.floor(config.text.maxPositionEmbeddings));
+    if (totalAt(hi) <= usable) lo = hi;
+    else {
+      // invariant: totalAt(lo) ≤ usable (lo = 0 may itself not fit → 0)
+      if (totalAt(0) > usable) hi = 0;
+      while (hi - lo > 1) {
+        const mid = Math.floor((lo + hi) / 2);
+        if (totalAt(mid) <= usable) lo = mid;
+        else hi = mid;
+      }
+    }
+    maxCtx = lo;
   }
 
   // decode reads all weights + the KV cache once per token — except MoE