mlx-bun 0.3.0 → 0.4.0

package/src/preflight.ts

@@ -32,6 +32,9 @@ export interface MachineState {
   /** Foreign processes holding > bigRssMB resident (excludes kernel,
    *  WindowServer, and our own bun/python benchmark processes). */
   bigProcesses: { rssMB: number; command: string }[];
+  /** Explicitly permitted CPU-only background commands; RSS limits still apply. */
+  allowedCpuProcesses?: string[];
+  backgroundCpuProcesses?: { rssMB: number; command: string }[];
   ok: boolean;
   problems: string[];
   at: string;
@@ -41,6 +44,29 @@ export interface PreflightLimits {
   maxSwapUsedMB?: number;   // default 512
   minFreePercent?: number;  // default 35
   bigRssMB?: number;        // default 2048
+  allowCpuProcesses?: string[];
+}
+
+export function foreignProcessFindings(output: string, limits: PreflightLimits = {}) {
+  const bigProcesses: { rssMB: number; command: string }[] = [];
+  const backgroundCpuProcesses: { rssMB: number; command: string }[] = [];
+  for (const line of output.split("\n")) {
+    const m = line.match(/^\s*([\d.]+)\s+(\d+)\s+(.*)$/);
+    if (!m) continue;
+    const cpuPct = Number(m[1]);
+    const rssMB = Number(m[2]) / 1024;
+    const command = m[3]!;
+    if (rssMB < (limits.bigRssMB ?? 2048) && cpuPct < 50) continue;
+    if (/(bun|\.venv\/bin\/python|mlx_lm|optiq)/.test(command) ||
+        /(kernel_task|WindowServer|launchd|mds|spotlight)/i.test(command)) continue;
+    const finding = { rssMB: Math.round(rssMB),
+      command: command.slice(0, 80) + (cpuPct >= 50 ? ` (${cpuPct.toFixed(0)}% cpu)` : "") };
+    // Exact command match; permission to tolerate CPU activity never permits
+    // that process to consume the benchmark's memory headroom.
+    const allowed = rssMB < (limits.bigRssMB ?? 2048) && limits.allowCpuProcesses?.includes(command);
+    (allowed ? backgroundCpuProcesses : bigProcesses).push(finding);
+  }
+  return { bigProcesses, backgroundCpuProcesses };
 }
 
 function sh(cmd: string[]): string {
@@ -51,7 +77,6 @@ function sh(cmd: string[]): string {
 export function checkMachine(limits: PreflightLimits = {}): MachineState {
   const maxSwap = limits.maxSwapUsedMB ?? 512;
   const minFree = limits.minFreePercent ?? 35;
-  const bigRss = limits.bigRssMB ?? 2048;
   const problems: string[] = [];
 
   // swap: "vm.swapusage: total = 7168.00M  used = 6413.06M  free = ..."
@@ -95,26 +120,15 @@ export function checkMachine(limits: PreflightLimits = {}): MachineState {
   // one clean sample doesn't prove absence (the run-spread stability
   // policy in bench-h2h.ts is the direct signal); a hot sample here is
   // still a cheap early warning.
-  const SELF_PATTERN = /(bun|\.venv\/bin\/python|mlx_lm|optiq)/;
-  const SYSTEM_PATTERN = /(kernel_task|WindowServer|launchd|mds|spotlight)/i;
-  const BIG_CPU_PCT = 50;
-  const bigProcesses: { rssMB: number; command: string }[] = [];
-  for (const line of sh(["ps", "axo", "pcpu=,rss=,command="]).split("\n")) {
-    const m = line.match(/^\s*([\d.]+)\s+(\d+)\s+(.*)$/);
-    if (!m) continue;
-    const cpuPct = Number(m[1]);
-    const rssMB = Number(m[2]) / 1024;
-    const command = m[3]!.slice(0, 80);
-    if (rssMB < bigRss && cpuPct < BIG_CPU_PCT) continue;
-    if (SELF_PATTERN.test(command) || SYSTEM_PATTERN.test(command)) continue;
-    bigProcesses.push({ rssMB: Math.round(rssMB), command: cpuPct >= BIG_CPU_PCT ? `${command} (${cpuPct.toFixed(0)}% cpu)` : command });
-  }
+  const { bigProcesses, backgroundCpuProcesses } = foreignProcessFindings(
+    sh(["ps", "axo", "pcpu=,rss=,command="]), limits);
   for (const p of bigProcesses)
     problems.push(`big process: ${p.rssMB} MB — ${p.command}`);
 
   return {
     swapUsedMB, freePercent, cpuSpeedLimit, gpuWiredLimitMB, loadAvg1m,
     bigProcesses,
+    allowedCpuProcesses: limits.allowCpuProcesses ?? [], backgroundCpuProcesses,
     ok: problems.length === 0,
     problems,
     at: new Date().toISOString(),
@@ -136,6 +150,8 @@ export function machineStateJson(s: MachineState): string {
     gpu_wired_mb: s.gpuWiredLimitMB,
     load1m: s.loadAvg1m,
     big_procs: s.bigProcesses.length,
+    allowed_cpu_processes: s.allowedCpuProcesses ?? [],
+    background_cpu_processes: s.backgroundCpuProcesses ?? [],
     ok: s.ok,
     at: s.at,
   });

package/src/prompt-cache.ts

@@ -1,4 +1,5 @@
-import { leaseCacheState } from "./backends/mlx/state-views";
+import { type CheckpointAttachment, materializeCheckpoint, attachmentBytes, cloneAttachments, disposeAttachments } from "./backends/mlx/checkpoint-state";
+import { leaseCacheState, minimumReusableOffset } from "./backends/mlx/state-views";
 import { withResource, cleanupFailure, disposeResources, ownResource } from "./engine/resources";
 // Byte-capped LRU prompt cache — the RAM tier of the Layer-0 KV store.
 //
@@ -10,12 +11,9 @@ import { withResource, cleanupFailure, disposeResources, ownResource } from "./e
 //   const hit = cache.take(promptIds);   // longest-prefix match
 //   generate(model, promptIds, { cache: hit?.caches ?? fresh })
 //   cache.put([...promptIds, ...generated], caches, ns, hit?.retain);
-// take() is NON-CONSUMING for trimmable caches: it serves zero-copy clones
-// and leaves the donor in place, so N agents can share one prefill.
-// Qwen hybrid entries containing recurrent SSM state transfer ownership on
-// an exact hit, avoiding an extra resident recurrent snapshot. put() supersedes
-// same-ns prefix-ancestors (when the new entry is trimmable) so a conversation
-// stays one entry, not one per turn.
+// take() lends zero-copy views and leaves every RAM donor in place, including
+// recurrent and method state. SSD persistence runs independently of RAM reuse.
+// put() supersedes same-ns prefix-ancestors when the new entry is trimmable.
 //
 // TIERING (Layer 0, unified-engine plan): when a ColdTier is attached,
 // take() itself runs the two-tier dance — RAM peek vs cold find, restore
@@ -28,6 +26,7 @@ import { withResource, cleanupFailure, disposeResources, ownResource } from "./e
 // snapshot for both lanes.
 
 import type { Cache } from "./model/gemma4";
+import type { PrefixCache, PrefixCacheHit } from "./contracts/prefix-cache";
 import { cloneKvCaches } from "./kv-store";
 
 /** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
@@ -68,10 +67,10 @@ export interface ColdTier {
   /** Materialize a found entry as GPU-visible caches (zero-copy COW mmap;
    *  pages fault in lazily). `retain` must run after the caches are
    *  disposed (it unmaps the backing file). Null on any failure. */
-  restore(handle: unknown): { tokens: number[]; caches: Cache[]; retain: () => void } | null;
+  restore(handle: unknown): { tokens: number[]; caches: Cache[]; attachments?: CheckpointAttachment[]; retain: () => void } | null;
   /** Borrow live caches and finish persistence before returning. A false
-   *  result prevents a required demotion or recurrent ownership transfer. */
-  store(tokens: number[], caches: Cache[], ns: string): boolean | void;
+   *  result prevents a required demotion. */
+  store(tokens: number[], caches: Cache[], ns: string, attachments?: CheckpointAttachment[]): boolean | void;
 }
 
 /** Spill sink for eviction/demotion (write-behind, 2026-07-06).
@@ -95,9 +94,7 @@ export interface SpillSink {
   spillSync?: (entry: PromptCacheEntry) => void;
 }
 
-export interface PromptCacheEntry {
-  tokens: number[];
-  caches: Cache[];
+export interface PromptCacheEntry extends PrefixCacheHit<Cache[], CheckpointAttachment[]> {
   /** Namespace key — adapter spec for LoRA requests ("" = base model).
    *  KV computed under one adapter must never seed another's prefill. */
   ns: string;
@@ -134,7 +131,7 @@ interface EntryRecord {
   share: { acquire(): () => void };
 }
 
-export class PromptCache {
+export class PromptCache implements PrefixCache<Cache[], CheckpointAttachment[]> {
   readonly maxBytes: number;
   #entries: EntryRecord[] = [];
   #clock = 0;
@@ -161,6 +158,7 @@ export class PromptCache {
     spill: SpillSink | ((entry: PromptCacheEntry) => void) | null = null,
     cold: ColdTier | null = null,
     clone: (caches: Cache[]) => Cache[] = cloneKvCaches,
+    private readonly materialize: (caches: Cache[], attachments?: CheckpointAttachment[]) => void = materializeCheckpoint,
   ) {
     this.maxBytes = maxBytes;
     this.#spillOwned = typeof spill === "function" ? null : spill?.spillOwned ?? null;
@@ -185,23 +183,26 @@ export class PromptCache {
         // no backing mmap for retain to unmap, so nothing the clones
         // could alias goes away.
         let clones: Cache[] | null = null;
+        let attachments: CheckpointAttachment[] | undefined;
         try {
           clones = this.#clone(entry.caches);
-          this.#spillOwned({ tokens: [...entry.tokens], caches: clones, ns: entry.ns });
+          attachments = cloneAttachments(entry.attachments);
+          this.#spillOwned({ tokens: [...entry.tokens], caches: clones, ns: entry.ns, attachments });
         } catch {
           // Failure degrades to no-spill; eviction/demotion never unwinds.
           if (clones) for (const c of clones) try { c.dispose(); } catch {}
+          try { disposeAttachments(attachments); } catch {}
         }
       } else {
         const spill =
           this.#spillSync ??
           (this.#cold
-            ? (e: PromptCacheEntry) => this.#cold!.store(e.tokens, e.caches, e.ns)
+            ? (e: PromptCacheEntry) => this.#cold!.store(e.tokens, e.caches, e.ns, e.attachments)
             : null);
         if (spill) try { spill(entry); } catch {}
       }
     }
-    disposeResources(entry.caches);
+    disposeResources([...entry.caches, { dispose: () => disposeAttachments(entry.attachments) }]);
     entry.retain?.(); // backing release requires successful cache disposal
   }
 
@@ -227,7 +228,7 @@ export class PromptCache {
       // Keep the entry resident if persistence throws. A failed SSD write
       // must not be presented as a successful demotion.
       if (this.#cold) {
-        if (this.#cold.store(entry.tokens, entry.caches, entry.ns) === false) break;
+        if (this.#cold.store(entry.tokens, entry.caches, entry.ns, entry.attachments) === false) break;
       }
       else this.#spillSync?.(entry);
       this.#entries.splice(oldest, 1);
@@ -239,9 +240,9 @@ export class PromptCache {
   }
 
   /** Find the entry with the longest usable common prefix of `prompt`
-   *  across BOTH tiers. Trimmable RAM entries are served non-consumingly as
-   *  zero-copy clones; exact hits on Qwen hybrid entries containing recurrent
-   *  SSM state transfer the donor itself. The caller owns the returned caches —
+   *  across BOTH tiers. Every RAM entry lends zero-copy views and retains its
+   *  donor, including exact-boundary recurrent state and method attachments.
+   *  The caller owns the returned views —
    *  dispose them or put() an extended entry, honoring `retain` either way.
    *
    *  Usable prefix = common prefix capped at prompt.length - 1 (at least
@@ -262,9 +263,9 @@ export class PromptCache {
       const e = this.#entries[i]!.entry;
       if (e.ns !== ns) continue;
       const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
-      if (p <= bestLen) continue;
+      if (p <= bestLen || p < minimumReusableOffset(e.caches)) continue;
       const trimNeeded = e.tokens.length - p;
-      if (trimNeeded > 0 && !e.caches.every((c) => c.isTrimmable())) continue;
+      if (trimNeeded > 0 && (!!e.attachments?.length || !e.caches.every((c) => c.isTrimmable()))) continue;
       bestLen = p;
       bestIdx = i;
     }
@@ -283,7 +284,8 @@ export class PromptCache {
             this.#disposeEntry({ ...entry, ns }, false));
           const result = withResource(owner, (entry) => {
             const trimNeeded = entry.tokens.length - hit.prefixLen;
-            if (trimNeeded > 0 && !entry.caches.every((cache) => cache.isTrimmable())) {
+            if (hit.prefixLen < minimumReusableOffset(entry.caches)) return null;
+            if (trimNeeded > 0 && (!!entry.attachments?.length || !entry.caches.every((cache) => cache.isTrimmable()))) {
               console.warn(
                 `[prompt-cache] cold entry unusable: ${entry.tokens.length} tokens, ` +
                 `needs trim ${trimNeeded} but untrimmable (wrapped ring/SSM) — re-prefilling`,
@@ -312,28 +314,22 @@ export class PromptCache {
     const rec = this.#entries[bestIdx]!;
     rec.lastUsed = ++this.#clock;
     rec.lastUsedMs = Date.now();
-    // Qwen recurrent caches are untrimmable. Transfer the exact-boundary
-    // entry instead of retaining a second recurrent snapshot in RAM.
-    // This reduced memory pressure in the recorded multi-turn reproduction.
-    // Save the boundary before consumption: the extended recurrent state
-    // cannot be trimmed back, and a delayed writer can no longer find it
-    // in RAM after ownership moves. An already durable boundary is cheap.
-    if (rec.entry.caches.some((c) => c.signature() === "ssm")) {
-      if (this.#cold?.store(rec.entry.tokens, rec.entry.caches, ns) === false)
-        throw new Error("prompt-cache: failed to persist recurrent boundary before transfer");
-      this.#entries.splice(bestIdx, 1);
-      return rec.entry;
-    }
+    // Lend immutable views for every backend state. RAM retention and SSD
+    // durability remain cache responsibilities while the caller extends its
+    // own views. This keeps synchronous persistence out of a RAM hit.
     const clones = this.#clone(rec.entry.caches);
+    let attachments: CheckpointAttachment[] | undefined;
     const trimNeeded = rec.entry.tokens.length - bestLen;
     try {
+      attachments = cloneAttachments(rec.entry.attachments);
       if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
     } catch (error) {
-      cleanupFailure(error, () => disposeResources(clones));
+      cleanupFailure(error, () => disposeResources([...clones, { dispose: () => disposeAttachments(attachments) }]));
     }
     return {
       tokens: rec.entry.tokens.slice(0, bestLen),
       caches: clones,
+      attachments,
       ns,
       retain: rec.share.acquire(),
     };
@@ -346,9 +342,9 @@ export class PromptCache {
     for (const { entry: e } of this.#entries) {
       if (e.ns !== ns) continue;
       const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
-      if (p <= best) continue;
+      if (p <= best || p < minimumReusableOffset(e.caches)) continue;
       const trimNeeded = e.tokens.length - p; // same usability rule as take()
-      if (trimNeeded > 0 && !e.caches.every((c) => c.isTrimmable())) continue;
+      if (trimNeeded > 0 && (!!e.attachments?.length || !e.caches.every((c) => c.isTrimmable()))) continue;
       best = p;
     }
     return best;
@@ -368,10 +364,12 @@ export class PromptCache {
   /** Insert (or reinsert) an entry; evicts LRU entries over the byte cap
    *  (spilling each to the cold tier first, when one is attached). If the
    *  entry itself exceeds the cap it is spilled + disposed, not stored. */
-  put(tokens: number[], caches: Cache[], ns = "", retain?: () => void): void {
+  put(tokens: number[], caches: Cache[], ns = "", retain?: () => void, attachments?: CheckpointAttachment[]): void {
     // All cache-provided validation runs before ownership is adopted.
-    const bytes = cacheBytes(caches);
-    const trimmable = caches.every((cache) => cache.isTrimmable());
+    const bytes = cacheBytes(caches) + attachmentBytes(attachments);
+    const trimmable = !attachments?.length && caches.every((cache) => cache.isTrimmable());
+    // Storage takes resolved snapshots, independently of their producer.
+    this.materialize(caches, attachments);
     const discard = (entry: PromptCacheEntry, spill: boolean) => {
       try { this.#disposeEntry(entry, spill); }
       catch (error) {
@@ -380,7 +378,7 @@ export class PromptCache {
         console.warn(`[prompt-cache] entry cleanup failed: ${error}`);
       }
     };
-    const entry: PromptCacheEntry = { tokens, caches, ns, retain };
+    const entry: PromptCacheEntry = { tokens, caches, ns, retain, attachments };
     if (bytes > this.maxBytes) {
       discard(entry, true);
       return;
@@ -415,7 +413,7 @@ export class PromptCache {
     if (trimmable) {
       for (const old of [...this.#entries]) {
         if (old === rec || old.entry.ns !== ns) continue;
-        if (old.entry.tokens.length >= tokens.length) continue;
+        if (old.entry.tokens.length >= tokens.length || old.entry.tokens.length < minimumReusableOffset(caches)) continue;
         if (commonPrefixLength(old.entry.tokens, tokens) !== old.entry.tokens.length) continue;
         this.#entries = this.#entries.filter((r) => r !== old);
         discard(old.entry, false);

package/src/sampler.ts

@@ -11,6 +11,8 @@
 import { Dtype } from "./mlx/ffi";
 import { MlxArray } from "./mlx/array";
 import * as ops from "./mlx/ops";
+import type { TokenLogprobs } from "./contracts/generation";
+import type { SampledToken, SamplingSession } from "./contracts/sampling";
 import { applyCurve, type CurveParams } from "./lab/curve/curve-sampler";
 
 /** Resolved HLG sampling config (the user-facing knobs). The mid gain is NOT
@@ -120,10 +122,7 @@ export interface StepExtras {
   topVals: MlxArray | null;
 }
 
-export interface StepSample<T> {
-  token: T;
-  extras: StepExtras | null;
-}
+export type StepSample<T> = SampledToken<T, StepExtras>;
 
 interface StepGrammar {
   readonly isTerminated: boolean;
@@ -161,25 +160,8 @@ export interface NumberStepSamplerConfig extends StepSamplerConfigBase {
   grammarWait: "before-sample";
 }
 
-export interface DeviceStepSampler {
-  readonly isPlainGreedy: boolean;
-  readonly needsHistory: boolean;
-  sample(logits: MlxArray, step: number): StepSample<MlxArray>;
-  seedHistory(tokens: readonly number[]): void;
-  commitDevice(token: MlxArray): void;
-  commitNumbers(tokens: readonly number[]): void;
-  dispose(): void;
-}
-
-export interface NumberStepSampler {
-  readonly isPlainGreedy: boolean;
-  readonly needsHistory: boolean;
-  sample(logits: MlxArray, step: number): Promise<StepSample<number>>;
-  seedHistory(tokens: readonly number[]): void;
-  commitDevice(token: MlxArray): void;
-  commitNumbers(tokens: readonly number[]): void;
-  dispose(): void;
-}
+export type DeviceStepSampler = SamplingSession<MlxArray, MlxArray, StepSample<MlxArray>>;
+export type NumberStepSampler = SamplingSession<MlxArray, MlxArray, Promise<StepSample<number>>>;
 
 const GOLDEN = 0x9e3779b97f4a7c15n;
 
@@ -786,7 +768,7 @@ export function makeLogitsProcessors(opts: LogitsProcessorOptions = {}): LogitsP
   return out;
 }
 
-/** Whether a whole [B,V] row can use the scheduler's vectorized argmax path. */
+/** Whether scores can use stateless greedy sampling across independent rows. */
 export function isPlainGreedy(
   opts: StepSamplerOptions,
   hasProcessors = makeLogitsProcessors(opts).length > 0,
@@ -809,6 +791,23 @@ export function disposeStepExtras(extras: StepExtras | null): void {
   extras.topVals?.dispose();
 }
 
+/** Read captured probabilities with their token, without adding casts behind
+ * the next decode step. Consumes the capture on success and failure. */
+export function readStepExtras(extras: StepExtras | null): TokenLogprobs | undefined {
+  if (!extras) return undefined;
+  try {
+    const out: TokenLogprobs = {};
+    if (extras.sel) out.logprob = extras.sel.toFloat32Host()[0]!;
+    if (extras.topIdx && extras.topVals) {
+      const ids = extras.topIdx.toIntTokens();
+      const vals = extras.topVals.toFloat32Host();
+      out.top = Array.from(ids, (id, i) => ({ id, logprob: vals[i]! }))
+        .sort((a, b) => b.logprob - a.logprob);
+    }
+    return out;
+  } finally { disposeStepExtras(extras); }
+}
+
 export function makeStepSampler(
   options: StepSamplerOptions,
   config: DeviceStepSamplerConfig,
@@ -941,7 +940,10 @@ export function makeStepSampler(
   if (config.initialHistory) seedHistory(config.initialHistory);
 
   const common = {
+    independent: isPlainGreedy(options, processors.length > 0) && !capture && !config.sampler
+      ? independentGreedySampling : undefined,
     isPlainGreedy: isPlainGreedy(options, processors.length > 0),
+    capturesLogprobs: capture,
     needsHistory: processors.length > 0,
     seedHistory,
     commitDevice,
@@ -989,7 +991,16 @@ export function makeStepSampler(
   };
 }
 
-/** logits [1, V] → logprobs [1, V] (logits - logsumexp). */
+/** Stateless sampling across rows or verification positions. Preserve the
+ * normalized-score argmax, including rounding-created ties. */
+export const independentGreedySampling = Object.freeze({
+  sample(logits: MlxArray): MlxArray {
+    using logprobs = toLogprobs(logits);
+    return ops.argmaxAxis(logprobs, -1);
+  },
+});
+
+/** logits [..., V] → logprobs [..., V] (logits - logsumexp). */
 export function toLogprobs(logits: MlxArray): MlxArray {
   const lse = ops.logsumexpAxis(logits, -1, true);
   const out = ops.sub(logits, lse);

package/src/serve/chat-stage.ts

@@ -25,7 +25,7 @@ export class ChatStage {
     private readonly prep: RequestPrep,
     private readonly promptCache: Pick<PromptCache, "peekPrefixLen">,
     /** Admission ceiling (fit() / the GLM memory plan), resolved once. */
-    private readonly maxSafeContext: number,
+    private readonly contextLimit: number | null,
     /** `serve --adapter <dir>` startup default; a request's explicit
      *  `adapter` (incl. "none") wins. */
     private readonly defaultAdapter?: string,
@@ -165,7 +165,7 @@ export class ChatStage {
         promptIds,
         options,
         requestedMaxTokens,
-        maxSafeContext: this.maxSafeContext,
+        contextLimit: this.contextLimit,
         stream: request.stream,
         wantLogprobs,
         topLogprobs,

package/src/serve/generation-gateway.ts

@@ -29,7 +29,7 @@ import { cleanupFailure, disposeResources } from "../engine/resources";
 // server.py's all-caches-have-merge check). GLM's capability owns compressed
 // MLA/DSA merge/extract/filter and byte projection; Qwen3.5's SSM cache owns
 // its row state. The remaining request-shape exclusions include vision, LoRA,
-// serial-only KV modes, and user-fixed seed
+// serial-only KV modes
 // (reproducibility ⇒ solo, matching mlx-lm's _is_batchable). Temperature /
 // top-p / top-k DO batch (each row samples with its own seed). Full-attention
 // AND sliding-window (Gemma) models both batch — the scheduler assembles each
@@ -42,9 +42,9 @@ import { MlxArray } from "../mlx/array";
 import type { RuntimeModel } from "../model/factory";
 import type { GenerateOptions, GenerateStats, TokenLogprobs } from "../generate";
 import type { KvScheme } from "../kv-scheme";
+import { createRowSampling } from "../backends/mlx/row-sampling";
 import { makeStepSampler } from "../sampler";
 import type { RowPromptCache } from "../backends/mlx/batch-group";
-import type { MlxMemoryBudget } from "../backends/mlx/memory-guard";
 import { bindMlxGateway, type MlxGatewayBinding, type MlxBatchGroup } from "../backends/mlx/gateway-binding";
 import type { RuntimeConfig } from "../runtime-config";
 import type { ExecutionRequirements, ResolvedExecution } from "../contracts/execution";
@@ -114,7 +114,7 @@ export type Vision = {
 };
 
 /** Per-token sink: returning `false` halts this generation (stop sequence).
- *  `logprobs` is only populated on the serial lane when the request asked for
+ *  `logprobs` is populated on either lane when the request asked for
  *  logprobs capture (GenerateOptions.logprobs / topLogprobs). */
 export type OnToken = (
   token: number,
@@ -172,7 +172,6 @@ export class GenerationGateway {
     private readonly serialRun: SerialRun,
     private readonly opts: {
       kvBudgetBytes?: number;
-      memoryBudget?: MlxMemoryBudget;
       checkpoints?: boolean;
       stateCodecs?: import("../kv-store").CacheCodecProvider;
       /** The server-wide KV scheme (server.ts kvScheme) — threaded to the
@@ -185,9 +184,7 @@ export class GenerationGateway {
        *  back on finish. Safe to share with the serial lane — both use it
        *  only inside this gateway's mutual-exclusion domain. */
       promptCache?: RowPromptCache;
-      /** Finish cache writes before taking the generation lock. Writers
-       * need that same lock, so this hook must never run inside it. */
-      beforeSerial?: () => Promise<void>;
+      adapterNamespace?: (adapters: string[]) => string;
     } = {},
   ) {
     this.#binding = "createBatchGroup" in model ? model : bindMlxGateway(model);
@@ -235,16 +232,20 @@ export class GenerationGateway {
     await this.onIdle();
   }
 
-  /** Resolves once the engine is idle (poll-based, ~`pollMs` granularity —
-   *  cheap vs the multi-hundred-ms flushes it paces; resolves without a
-   *  timer when already idle). The SSD write-behind gates each per-tensor
-   *  flush step on this so durability work never steals blocking
-   *  GPU-sync + writeSync slices from an active decode (the 2026-07-07
-   *  decode@ctx contamination — see server.ts's write-behind block). */
+  /** Observe idleness. Call runWhenIdle for background work that must also
+   * acquire ownership atomically with the idle check. */
   async onIdle(pollMs = 20): Promise<void> {
     while (this.busy) await new Promise<void>((r) => setTimeout(r, pollMs));
   }
 
+  /** Background work waits without requesting a batch drain. The final busy
+   * check and mutex acquisition have no await between them, so a foreground
+   * request cannot enter between the check and ownership registration. */
+  async runWhenIdle<T>(fn: () => Promise<T>, pollMs = 20): Promise<T> {
+    while (this.busy) await new Promise<void>(resolve => setTimeout(resolve, pollMs));
+    return this.runExclusive(fn);
+  }
+
   /** Cumulative rows routed to the batch lane since server start. The serial
    *  lane never advances this — /stats' race-free lane-routing observable
    *  (active/pending are instantaneous and read 0 once a request finishes). */
@@ -272,7 +273,7 @@ export class GenerationGateway {
     const frozenShape = Object.freeze(shape);
     const execution = this.#binding.plan(frozenShape, options, {
       continuous: this.batchingEnabled,
-      quantizedBatch: shape.kvQuant && this.#kvBatchable(),
+      quantizedBatch: (shape.kvQuant || shape.turboQuant) && this.#kvBatchable(),
       checkpoints: this.opts.checkpoints === true,
     });
     return Object.freeze({ shape: frozenShape, mechanism: execution.mechanism, execution });
@@ -390,7 +391,7 @@ export class GenerationGateway {
         });
       };
       let started = false;
-      return Promise.resolve().then(() => this.opts.beforeSerial?.()).then(() => this.runExclusive(async () => {
+      return this.runExclusive(async () => {
         signal?.throwIfAborted();
         started = true;
         const stats = await this.serialRun(
@@ -398,7 +399,7 @@ export class GenerationGateway {
         );
         signal?.throwIfAborted();
         return stats;
-      }, trace, signal))
+      }, trace, signal)
         // Covers a serial waiter aborted before serialRun takes ownership.
         // generate() also disposes defensively; the operation is idempotent.
         .finally(() => {
@@ -407,16 +408,23 @@ export class GenerationGateway {
         });
     }
 
+    const adapters = options.adapters?.length ? [...options.adapters] : undefined;
+    const adapterKey = adapters ? JSON.stringify(adapters) : "";
+    const cacheNamespace = adapters
+      ? () => this.opts.adapterNamespace?.(adapters) ?? adapterKey : "";
+    const context = adapters
+      ? this.#binding.bindAdapterContext?.(adapters, `adapters:${adapterKey}`) : undefined;
+
     // The scheduler owns grammar ready/accept sequencing. StepSampler owns the
     // shared processors -> mask -> logprobs -> sample -> history contract.
-    const stepSampler = makeStepSampler(options, {
-      tokenRepresentation: "device",
-      grammarWait: "external",
-      historyUpdate: "after-sample",
-      initialHistory: promptIds,
-    });
-    const sample = (logits1V: MlxArray, step: number): MlxArray =>
-      stepSampler.sample(logits1V, step).token;
+    const method = this.#binding.methodRequest?.(placement.execution, options);
+    const continuation = method ? undefined : this.#binding.continuationRequest?.(placement.execution, options, promptIds, onToken);
+    const sampling = method ? undefined : continuation ?? createRowSampling(makeStepSampler(options, {
+      tokenRepresentation: "device", grammarWait: "external",
+      historyUpdate: "after-sample", initialHistory: promptIds,
+      captureSelectedLogprob: options.logprobs === true,
+      captureTopLogprobs: options.topLogprobs,
+    }), onToken);
 
     let st;
     this.#rowsSubmitted++;
@@ -425,13 +433,14 @@ export class GenerationGateway {
     });
     try {
       st = await this.#ensureScheduler().submit({
-        promptIds,
+        promptIds, context, cacheNamespace, prefillChunkSize: options.prefillChunkSize,
+        continuation: continuation?.continuation,
+        statePolicy: this.#binding.statePolicy?.(placement.execution, options, promptIds.length + (options.maxTokens ?? 512)),
         compiledDecode: placement.execution?.compiledDecode,
         maxTokens: options.maxTokens ?? 512,
         eosTokenIds: options.eosTokenIds ?? this.#binding.config.eosTokenIds,
-        sample,
-        plainGreedy: stepSampler.isPlainGreedy,
-        onToken,
+        ...(method ? { method, onToken } : { sample: sampling!.sample,
+          plainGreedy: sampling!.plainGreedy, onToken: sampling!.onToken }),
         onAdmitted: closeAdmission,
         trace,
         ...(signal ? { signal } : {}),
@@ -445,7 +454,7 @@ export class GenerationGateway {
       });
     } finally {
       closeAdmission?.();
-      stepSampler.dispose();
+      sampling?.dispose();
       // The scheduler never owns per-row grammar state.
       options.grammar?.dispose();
     }
@@ -460,6 +469,7 @@ export class GenerationGateway {
       prefillTps: st.prefillMs > 0 ? ((st.promptTokens - st.cachedTokens) / st.prefillMs) * 1000 : 0,
       decodeTps: st.decodeMs > 0 && st.generatedTokens > 1 ? ((st.generatedTokens - 1) / st.decodeMs) * 1000 : 0,
       cacheTokens: [],
+      ...(st.spec ? { spec: st.spec } : {}),
     };
   }
 
@@ -470,7 +480,6 @@ export class GenerationGateway {
         maxBatch: this.#batch,
         stateCodecs: this.opts.stateCodecs,
         kvBudgetBytes: this.opts.kvBudgetBytes,
-        memoryBudget: this.opts.memoryBudget,
         // Phase 3.1: the batchable kvConfig composition is applied by the
         // scheduler (solo rows convert at serial chunk boundaries, then
         // merge as quantized triples). Only threaded when the scheme

package/src/serve/model-host.ts

@@ -73,10 +73,9 @@ export interface ServerContext<Model = RuntimeModel> {
    *  optiq serve injects these as server defaults (gen_config.py);
    *  explicit request fields always win. */
   genDefaults: GenSamplingDefaults;
-  /** Speculative decoding (`serve --draft-model`, mlx_lm.server parity).
-   *  Server-level: when set, EVERY request routes to the serial lane
-   *  (upstream: is_batchable = draft is None) and spec-eligible ones decode
-   *  through src/spec/serve-loop.ts. null = no draft configured. */
+  /** Spec-eligible requests use src/spec/serve-loop.ts; ordinary fallbacks
+   *  use their own feature-compatible scheduling mechanism.
+   *  null = no draft configured. */
   draft?: {
     provider: import("../spec/source").DraftProvider;
     numDraftTokens: number;

package/src/serve/request-plan.ts

@@ -33,7 +33,7 @@ export interface RequestPlanInput {
   promptIds: number[];
   options: GenerateOptions & { stopSequences: string[] };
   requestedMaxTokens: number;
-  maxSafeContext: number;
+  contextLimit: number | null;
   stream: boolean;
   wantLogprobs: boolean;
   topLogprobs: number;
@@ -114,11 +114,12 @@ export type PlanRequestResult = RequestPlan | RequestRejection;
  * options, adapter selection, and lane shape are derived here once for chat
  * and raw completions. */
 export function planRequest(input: RequestPlanInput): PlanRequestResult {
-  const available = input.maxSafeContext - input.promptIds.length;
+  const available = input.contextLimit === null ? Infinity
+    : input.contextLimit - input.promptIds.length;
   if (available < 1) {
     return new RequestRejection(
       `prompt is ${input.promptIds.length} tokens but the memory budget caps ` +
-        `safe context at ${input.maxSafeContext} — no room to generate; ` +
+        `safe context at ${input.contextLimit} — no room to generate; ` +
         "shorten the prompt or raise --memory-budget",
       input.ownership,
     );