mlx-bun 0.0.8 → 0.0.9

package/src/generate.ts

@@ -61,15 +61,43 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
    *  When set, `promptTokens` are the spliced (<|image|>-expanded) ids and the
    *  denoising engine prefills the merged vision features. Caller keeps ownership. */
   visionPixels?: MlxArray;
+  /** Capture each emitted token's log-probability (mlx_lm.server `logprobs`).
+   *  The distribution matches mlx-lm generate_step exactly: full-vocab
+   *  log-softmax of the logits AFTER logits processors, BEFORE the sampler's
+   *  temperature/top-p/top-k/min-p/XTC (generate.py L409-422). Off by default —
+   *  the hot path pays nothing when unset. */
+  logprobs?: boolean;
+  /** Capture the top-k (token id, logprob) pairs per emitted token from the
+   *  same distribution (mlx_lm.server `top_logprobs`). 0/unset = off. */
+  topLogprobs?: number;
 }
 
-/** Port of mlx-lm maybe_quantize_kv_cache + optiq serve's per-layer
- *  patched variant (incl. patch_rotating_to_quantized: rotating caches
- *  convert too — Phase 9). kvConfig overrides kvBits, matching optiq's
- *  --kv-config precedence; shipped kv_config.json files list
- *  full-attention layers only, so rotating quantization engages through
- *  uniform kvBits (like optiq --kv-bits) or a config that names sliding
- *  layers. */
+/** Per-token logprob info (only present when GenerateOptions requested it). */
+export interface TokenLogprobs {
+  /** logprob of the emitted token (requires options.logprobs). */
+  logprob?: number;
+  /** top-k (id, logprob) pairs, sorted by logprob descending (requires
+   *  options.topLogprobs > 0). mlx-lm leaves argpartition order unspecified;
+   *  sorting is the deterministic reading of the same set. */
+  top?: { id: number; logprob: number }[];
+}
+
+/** Port of mlx-lm maybe_quantize_kv_cache + BOTH halves of optiq serve's
+ *  per-layer patched variant (incl. patch_rotating_to_quantized: rotating
+ *  caches convert too — Phase 9):
+ *  - per-layer bits/group_size selection (kvConfig overrides kvBits,
+ *    matching optiq's --kv-config precedence; shipped kv_config.json
+ *    files list full-attention layers only, so rotating quantization
+ *    engages through uniform kvBits — like optiq --kv-bits — or a config
+ *    that names sliding layers), and
+ *  - STREAMING conversion (optiq streaming_kv_quant / serve.py
+ *    patched_maybe_quantize): eval each layer's quantized triples and
+ *    clear the buffer pool before building the next layer's conversion.
+ *    Lazily batching every layer's toQuantized into one eval pins ALL
+ *    layers' bf16 K/V as graph inputs alongside ALL quantized outputs —
+ *    the exact transient optiq's fix kills (16.35 → 7.60 GB at 32k on a
+ *    24 GB Mac). Numerics untouched: same quantize math, only the eval
+ *    ordering is forced (tests/kv-quant.test.ts, tests/rotating-kvq.test.ts). */
 export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void {
   const { kvBits, kvConfig } = options;
   if (!kvBits && !kvConfig?.length) return;
@@ -87,10 +115,16 @@ export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void
     if (c.offset === 0) continue;
     if (byLayer) {
       const e = byLayer.get(i);
-      if (e) cache[i] = c.toQuantized(e.groupSize, e.bits);
+      if (!e) continue;
+      cache[i] = c.toQuantized(e.groupSize, e.bits);
     } else {
       cache[i] = c.toQuantized(options.kvGroupSize ?? 64, kvBits!);
     }
+    // Streaming half: materialize THIS layer's conversion now so its bf16
+    // source (already released by toQuantized) frees before the next layer
+    // converts — the transient stays ~one layer, not the whole cache.
+    ops.evalAll(cache[i]!.state());
+    clearCache();
   }
 }
 
@@ -112,6 +146,8 @@ export interface GenerateStats {
 export interface GeneratedToken {
   token: number;
   index: number;
+  /** Present only when GenerateOptions.logprobs / topLogprobs requested it. */
+  logprobs?: TokenLogprobs;
 }
 
 export class Generation implements AsyncIterable<GeneratedToken> {
@@ -273,6 +309,14 @@ async function* generateInner(
   const processors = makeLogitsProcessors(options);
   const needsTokenHistory = processors.length > 0;
 
+  // logprobs capture (mlx_lm.server semantics — see GenerateOptions.logprobs).
+  // Everything below is gated: when neither flag is set, no extra ops, evals,
+  // or readbacks happen on the hot path.
+  const captureSelLp = options.logprobs === true;
+  const captureTopK =
+    options.topLogprobs && options.topLogprobs > 0 ? options.topLogprobs : 0;
+  const capture = captureSelLp || captureTopK > 0;
+
   const ownsCache = !options.cache;
   const cache = options.cache ?? model.makeCache();
   const cachedTokens = cache[0]!.offset;
@@ -283,16 +327,73 @@ async function* generateInner(
   /** device-side token history (only maintained when processors need it) */
   let history: MlxArray | null = null;
 
-  // logits [1,1,V] → sampled token array [1] (all on-device)
-  const sampleStep = (logits3d: MlxArray, step: number): MlxArray => {
+  /** Device-side logprob capture for one step, computed from the SAME lp the
+   *  sampler saw (post-processors, pre-truncation) — read back lazily with the
+   *  token so decode pipelining is preserved. */
+  interface StepExtras {
+    sel: MlxArray | null; // [1,1] emitted token's logprob
+    topIdx: MlxArray | null; // [1,k] top-k token ids
+    topVals: MlxArray | null; // [1,k] their logprobs
+  }
+  const extrasArrays = (e: StepExtras | null): MlxArray[] =>
+    e ? [e.sel, e.topIdx, e.topVals].filter((a): a is MlxArray => a !== null) : [];
+  const disposeExtras = (e: StepExtras | null): void => {
+    if (e) for (const a of [e.sel, e.topIdx, e.topVals]) a?.dispose();
+  };
+  /** Read extras back to JS (forces eval — they were async-dispatched with the
+   *  token, so this normally just copies) and dispose the device arrays. */
+  const readExtras = (e: StepExtras | null): TokenLogprobs | undefined => {
+    if (!e) return undefined;
+    const out: TokenLogprobs = {};
+    if (e.sel) out.logprob = e.sel.toFloat32()[0]!;
+    if (e.topIdx && e.topVals) {
+      // uint32 ids read via the f32 cast: exact for ids < 2^24 (vocab ≪ 16M)
+      const ids = e.topIdx.toFloat32();
+      const vals = e.topVals.toFloat32();
+      out.top = Array.from(ids, (id, i) => ({ id, logprob: vals[i]! })).sort(
+        (a, b) => b.logprob - a.logprob,
+      );
+    }
+    disposeExtras(e);
+    return out;
+  };
+
+  // logits [1,1,V] → sampled token array [1] (+ optional logprob capture,
+  // all on-device)
+  const sampleStep = (
+    logits3d: MlxArray,
+    step: number,
+  ): { tok: MlxArray; extras: StepExtras | null } => {
     const V = logits3d.shape[2]!;
     let logits = ops.reshape(logits3d, [1, V]);
     for (const p of processors) logits = disposing(logits, p(history, logits));
     const lp = toLogprobs(logits);
     logits.dispose();
     const tok = sampler(lp, step);
+    let extras: StepExtras | null = null;
+    if (capture) {
+      let sel: MlxArray | null = null;
+      let topIdx: MlxArray | null = null;
+      let topVals: MlxArray | null = null;
+      if (captureSelLp) {
+        const idx = ops.reshape(tok, [1, 1]);
+        sel = ops.takeAlongAxis(lp, idx, -1);
+        idx.dispose();
+      }
+      if (captureTopK > 0) {
+        // mlx-lm _format_top_logprobs: argpartition(-logprobs, kth=k-1)[:k]
+        const k = Math.min(captureTopK, V);
+        const negLp = ops.neg(lp);
+        const part = ops.argpartitionAxis(negLp, k - 1, -1);
+        const idxView = part.slice([0, 0], [1, k]);
+        topIdx = ops.contiguous(idxView); // slices are strided views; raw readback needs contiguous
+        topVals = ops.takeAlongAxis(lp, topIdx, -1);
+        for (const a of [negLp, part, idxView]) a.dispose();
+      }
+      extras = { sel, topIdx, topVals };
+    }
     lp.dispose();
-    return tok;
+    return { tok, extras };
   };
 
   const pushHistory = (tok: MlxArray) => {
@@ -317,6 +418,8 @@ async function* generateInner(
   const forwarded: number[] = [];
   let pending: MlxArray | null = null;
   let nextPending: MlxArray | null = null;
+  let pendingExtras: StepExtras | null = null;
+  let nextExtras: StepExtras | null = null;
   let finished = false;
   let threw = false;
   const makeStats = (): GenerateStats => ({
@@ -379,7 +482,9 @@ async function* generateInner(
     h0.dispose();
     const logits0 = model.logitsFromHidden(hLast);
     hLast.dispose();
-    pending = sampleStep(logits0, 0); // token array [1]
+    const s0 = sampleStep(logits0, 0); // token array [1] (+ optional extras)
+    pending = s0.tok;
+    pendingExtras = s0.extras;
     logits0.dispose();
     // mirror mlx-lm generate_step: async-dispatch the first token's
     // compute; the prefill clock keeps running until the token ARRIVES
@@ -389,7 +494,7 @@ async function* generateInner(
     // to prompt_time, not decode — replicated so cross-stack decode
     // tok/s measure the same quantity. The boundary cost is real and
     // scales with prompt length; it belongs to "having prefilled".
-    ops.asyncEvalAll([pending]);
+    ops.asyncEvalAll([pending, ...extrasArrays(pendingExtras)]);
 
     // ---- decode (pipelined) ----
     // Compiled decode (docs/design/optimization_plan.md Phase A): replay the per-step
@@ -413,8 +518,11 @@ async function* generateInner(
     let stop = false;
     while (!stop) {
       const cur = pending!;
+      const curExtras = pendingExtras;
+      pendingExtras = null;
       // build step n+1's graph from the *unread* pending token
       nextPending = null;
+      nextExtras = null;
       if (generated + 1 < maxTokens) {
         maybeQuantizeKv(cache, options);
         pushHistory(cur);
@@ -426,8 +534,11 @@ async function* generateInner(
             logits = r.logits;
             evalWith = r.evalWith;
           } catch (e) {
-            // growth done by a partial prepare is benign for the
-            // uncompiled path (updateAndFetch re-checks capacity)
+            // Safe to re-forward the SAME token below: a failed step is
+            // transactional — segmented mode stages ring adopts and rolls
+            // back committed js-layer writes on throw (see #stepSegmented),
+            // and buffer growth done by a partial prepare is benign for the
+            // uncompiled path (updateAndFetch re-checks capacity).
             compiled = null;
             console.warn(`compiled decode disabled for this generation: ${e}`);
           }
@@ -439,9 +550,11 @@ async function* generateInner(
           logits = model.logitsFromHidden(h);
           h.dispose();
         }
-        nextPending = sampleStep(logits, generated + 1);
+        const sn = sampleStep(logits, generated + 1);
+        nextPending = sn.tok;
+        nextExtras = sn.extras;
         logits.dispose();
-        ops.asyncEvalAll([nextPending, ...evalWith]);
+        ops.asyncEvalAll([nextPending, ...extrasArrays(nextExtras), ...evalWith]);
       }
 
       // sync-read step n's token while n+1 computes
@@ -460,11 +573,17 @@ async function* generateInner(
       if (nextPending !== null) forwarded.push(token);
 
       if (eosTokenIds.includes(token)) {
+        disposeExtras(curExtras);
         nextPending?.dispose();
         nextPending = null;
+        disposeExtras(nextExtras);
+        nextExtras = null;
         stop = true;
       } else {
-        yield { token, index: generated - 1 };
+        // readExtras before the yield: if the consumer breaks at this yield,
+        // the extras are already read and disposed.
+        const logprobs = readExtras(curExtras);
+        yield { token, index: generated - 1, ...(logprobs ? { logprobs } : {}) };
         // mlx-lm generate_step: clear_cache after token 0 (drops the
         // remaining prefill transients) and every 256 tokens after
         if ((generated - 1) % 256 === 0) clearCache();
@@ -473,6 +592,8 @@ async function* generateInner(
         } else {
           pending = nextPending;
           nextPending = null;
+          pendingExtras = nextExtras;
+          nextExtras = null;
         }
       }
     }
@@ -486,6 +607,8 @@ async function* generateInner(
     if (!finished) {
       pending?.dispose();
       nextPending?.dispose();
+      disposeExtras(pendingExtras);
+      disposeExtras(nextExtras);
     }
     if (ownsCache) for (const c of cache) c.dispose();
     history?.dispose();

package/src/harness-pi.ts

@@ -22,7 +22,10 @@ export const PI_PROVIDER_ID = "mlx-bun";
  *  server still reports the true id in /v1/models and responses. Gives a
  *  clean `mlx-bun/local` handle instead of `mlx-bun/mlx-community/…`. */
 export const PI_LOCAL_MODEL_ID = "local";
-export const DEFAULT_BASE_URL = "http://localhost:8090/v1";
+// 127.0.0.1 (the server's actual default bind), not "localhost": the generated
+// extension runs under the USER'S pi (node/jiti, not Bun), so don't depend on
+// how that runtime's resolver maps localhost.
+export const DEFAULT_BASE_URL = "http://127.0.0.1:8080/v1";
 
 export function defaultPiExtensionsDir(): string {
   return join(homedir(), ".pi", "agent", "extensions");
@@ -63,22 +66,34 @@ export interface ServerModel {
   contextWindow: number;
   maxTokens: number;
   reasoning: boolean;
+  vision: boolean;
 }
 
-/** Discover models from a running mlx-bun server; [] when unreachable. */
+/** Discover the SERVED model from a running mlx-bun server; [] when
+ *  unreachable. /v1/models lists the served model first (the only row
+ *  carrying the capability extras: context_window/reasoning/vision) and then
+ *  every other servable model in the registry — those are things `mlx-bun
+ *  serve` COULD run, not models this server answers for, so they must not be
+ *  registered with pi. */
 export async function fetchServerModels(baseUrl: string, timeoutMs = 3000): Promise<ServerModel[]> {
   try {
     const res = await fetch(`${baseUrl.replace(/\/$/, "")}/models`, {
       signal: AbortSignal.timeout(timeoutMs),
     });
     if (!res.ok) return [];
-    const body = (await res.json()) as { data?: Array<{ id: string; context_window?: number; reasoning?: boolean }> };
-    return (body.data ?? []).map((m) => ({
-      id: m.id,
-      contextWindow: m.context_window ?? 32768,
+    const body = (await res.json()) as {
+      data?: Array<{ id: string; context_window?: number; reasoning?: boolean; vision?: boolean }>;
+    };
+    const rows = body.data ?? [];
+    const served = rows.find((m) => m.context_window !== undefined) ?? rows[0];
+    if (!served) return [];
+    return [{
+      id: served.id,
+      contextWindow: served.context_window ?? 32768,
       maxTokens: 8192,
-      reasoning: !!m.reasoning,
-    }));
+      reasoning: !!served.reasoning,
+      vision: !!served.vision,
+    }];
   } catch {
     return [];
   }
@@ -113,16 +128,22 @@ export default async function (pi: any) {
     if (res.ok) {
       const { data } = await res.json();
       if (Array.isArray(data) && data.length > 0) {
-        models = data.map((m: any) => ({
+        // The SERVED model is the row with the capability extras (first);
+        // the rest are servable-but-not-served registry entries.
+        const served = data.find((m: any) => m.context_window !== undefined) ?? data[0];
+        models = [{
           id: ${JSON.stringify(PI_LOCAL_MODEL_ID)},
-          name: \`\${m.id} (mlx-bun local)\`,
-          reasoning: !!m.reasoning,
-          ...(m.reasoning ? { compat: { thinkingFormat: "qwen-chat-template" } } : {}),
-          input: ["text"],
+          name: \`\${served.id} (mlx-bun local)\`,
+          reasoning: !!served.reasoning,
+          compat: {
+            supportsDeveloperRole: false,
+            ...(served.reasoning ? { thinkingFormat: "qwen-chat-template" } : {}),
+          },
+          input: served.vision ? ["text", "image"] : ["text"],
           cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-          contextWindow: m.context_window ?? 32768,
+          contextWindow: served.context_window ?? 32768,
           maxTokens: 8192,
-        }));
+        }];
       }
     }
   } catch {}
@@ -142,8 +163,13 @@ function modelEntry(m: ServerModel) {
     id: PI_LOCAL_MODEL_ID,
     name: `${m.id} (mlx-bun local)`,
     reasoning: !!m.reasoning,
-    ...(m.reasoning ? { compat: { thinkingFormat: "qwen-chat-template" } } : {}),
-    input: ["text"],
+    // Kept aligned with buildPiProvider (src/pi-provider.ts): the system
+    // prompt must arrive as `system`, not `developer`.
+    compat: {
+      supportsDeveloperRole: false,
+      ...(m.reasoning ? { thinkingFormat: "qwen-chat-template" } : {}),
+    },
+    input: m.vision ? ["text", "image"] : ["text"],
     cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
     contextWindow: m.contextWindow,
     maxTokens: m.maxTokens,

package/src/index.ts

@@ -13,8 +13,10 @@ export { Weights } from "./weights";
 export { Gemma4Model } from "./model/gemma4";
 export { MiniCPM5Model } from "./model/minicpm5";
 export { createModel, type RuntimeModel } from "./model/factory";
+export { Qwen3Model } from "./model/qwen3";
 export { generate } from "./generate";
 export { loadTokenizer } from "./tokenizer";
+export { embedOne, embedMany, isEmbeddingModel, withInstruction, type EmbedResult } from "./embed";
 export { ChatTemplate } from "./chat-template";
 export { chooseAutoModel, COEXIST_FRACTION, DEFAULT_REPO_ID, fit, largestRecommendedRepoId, recommendedRepoId, skuMatrix, thisMachine } from "./fit";
 export { downloadModel } from "./download";

package/src/lora.ts

@@ -77,7 +77,7 @@ export function loadAdapterTensors(file: string): Map<string, MlxArray> {
  *  PEFT writes lora_alpha + r (scale = alpha / r). Reference: mount.py.
  *  `rsLora` (recorded at train time) means the effective per-layer scale is
  *  α/√rank — the caller divides by √(that layer's rank). */
-async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null; rsLora: boolean }> {
+export async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null; rsLora: boolean }> {
   for (const name of ["optiq_lora_config.json", "adapter_config.json"]) {
     const f = Bun.file(`${dir}/${name}`);
     if (await f.exists()) {
@@ -94,7 +94,7 @@ async function readAdapterScale(dir: string): Promise<{ scale: number; rank: num
   throw new Error(`no adapter_config.json in ${dir}`);
 }
 
-function adapterWeightsFile(dir: string): string {
+export function adapterWeightsFile(dir: string): string {
   for (const name of ["adapters.safetensors", "adapter_model.safetensors"]) {
     const p = `${dir}/${name}`;
     if (existsSync(p)) return p;