mlx-bun 0.0.8 → 0.0.9

package/src/registry.ts

@@ -6,7 +6,8 @@
 
 import { Database } from "bun:sqlite";
 import {
-  closeSync, existsSync, openSync, readdirSync, readFileSync, readSync, statSync,
+  closeSync, existsSync, lstatSync, openSync, readdirSync, readFileSync,
+  readlinkSync, readSync, rmSync, statSync,
 } from "node:fs";
 import { join } from "node:path";
 import { isDrafterModelType } from "./model/support";
@@ -27,6 +28,12 @@ export interface ModelRecord {
   quantGroupSize: number | null;
   quantMode: string | null;
   hasVisionSidecar: boolean;
+  /** config.json's vision_config.model_type when it names a vision tower
+   *  (e.g. "gemma4_vision" — SigLIP, needs the bf16 sidecar — or
+   *  "gemma4_unified_vision" — encoder-free). null for text-only models
+   *  AND for configs whose nested vision_config is not a vision tower
+   *  (Qwen3.5 nests a copy of its own text config there). */
+  visionConfigType: string | null;
   hasKvConfig: boolean;
   hasToolTemplate: boolean;
   numLayers: number | null;
@@ -52,6 +59,7 @@ CREATE TABLE IF NOT EXISTS models (
   quant_group_size INTEGER,
   quant_mode TEXT,
   has_vision_sidecar INTEGER NOT NULL,
+  vision_config_type TEXT,
   has_kv_config INTEGER NOT NULL,
   has_tool_template INTEGER NOT NULL,
   num_layers INTEGER,
@@ -62,7 +70,14 @@ CREATE TABLE IF NOT EXISTS models (
 );
 `;
 
-export const DEFAULT_HUB = `${process.env.HOME}/.cache/huggingface/hub`;
+/** HF hub cache root, honoring the standard env overrides the same way
+ *  huggingface_hub (and our server.ts) do: HF_HUB_CACHE > HF_HOME/hub >
+ *  ~/.cache/huggingface/hub. */
+export const DEFAULT_HUB =
+  process.env.HF_HUB_CACHE ??
+  (process.env.HF_HOME
+    ? join(process.env.HF_HOME, "hub")
+    : `${process.env.HOME}/.cache/huggingface/hub`);
 export const DEFAULT_DB = `${process.env.HOME}/.cache/mlx-bun/registry.sqlite`;
 
 export class Registry {
@@ -81,7 +96,8 @@ export class Registry {
     if (
       !cols.includes("sidecar_bytes") ||
       !cols.includes("experts_bytes") ||
-      !cols.includes("license")
+      !cols.includes("license") ||
+      !cols.includes("vision_config_type")
     ) {
       this.db.exec("DROP TABLE models");
       this.db.exec(SCHEMA);
@@ -99,7 +115,7 @@ export class Registry {
     const upsert = this.db.prepare(`
       INSERT OR REPLACE INTO models VALUES
       ($path, $repo, $type, $params, $size, $sidecar, $experts, $bits, $gs, $mode,
-       $vision, $kv, $tools, $layers, $hidden, $vocab, $license, $at)
+       $vision, $vtype, $kv, $tools, $layers, $hidden, $vocab, $license, $at)
     `);
     for (const entry of readdirSync(hubDir)) {
       if (!entry.startsWith("models--")) continue;
@@ -116,6 +132,7 @@ export class Registry {
           $sidecar: rec.sidecarBytes, $experts: rec.expertsBytes,
           $bits: rec.quantBits, $gs: rec.quantGroupSize, $mode: rec.quantMode,
           $vision: rec.hasVisionSidecar ? 1 : 0,
+          $vtype: rec.visionConfigType,
           $kv: rec.hasKvConfig ? 1 : 0,
           $tools: rec.hasToolTemplate ? 1 : 0,
           $layers: rec.numLayers, $hidden: rec.hiddenSize, $vocab: rec.vocabSize,
@@ -131,8 +148,12 @@ export class Registry {
     const clauses: string[] = [];
     const params: Record<string, unknown> = {};
     if (filter.vision !== undefined) {
-      clauses.push("has_vision_sidecar = $vision");
-      params.$vision = filter.vision ? 1 : 0;
+      // Vision-capable = SigLIP sidecar present OR the encoder-free unified
+      // tower declared in config.json (12B ships no sidecar upstream anymore).
+      // COALESCE: NULL vision_config_type must read as "no", not SQL-NULL
+      // (a bare NULL LIKE poisons the NOT branch of the filter).
+      const cap = "(has_vision_sidecar = 1 OR COALESCE(vision_config_type,'') LIKE '%unified_vision')";
+      clauses.push(filter.vision ? cap : `NOT ${cap}`);
     }
     if (filter.maxBytes !== undefined) {
       clauses.push("size_bytes <= $max");
@@ -152,6 +173,24 @@ export class Registry {
     return rows.map(rowToRecord).filter((r) => existsSync(r.path));
   }
 
+  /** list(), collapsed to ONE record per repo: the canonical snapshot
+   *  (refs/main, else most recently scanned). Upstream pushes create a new
+   *  snapshots/<commit> dir per revision and the old ones are never deleted,
+   *  so raw list() shows a row per revision — this is the user-facing view
+   *  (`ls`, /v1/models, /library). Order follows list() (size ascending). */
+  listCanonical(filter: Parameters<Registry["list"]>[0] = {}): ModelRecord[] {
+    const byRepo = new Map<string, ModelRecord[]>();
+    const order: string[] = [];
+    for (const m of this.list(filter)) {
+      if (!byRepo.has(m.repoId)) {
+        byRepo.set(m.repoId, []);
+        order.push(m.repoId);
+      }
+      byRepo.get(m.repoId)!.push(m);
+    }
+    return order.map((repo) => pickCanonicalRevision(byRepo.get(repo)!));
+  }
+
   /** Resolve a fuzzy query to exactly one model (error listing candidates otherwise).
    *  Speculative-decoding drafters are companion artifacts, never selectable on
    *  their own, so they never count as candidates here. */
@@ -172,10 +211,17 @@ export class Registry {
   }
 }
 
+/** True when the model can answer vision requests: it either ships the bf16
+ *  SigLIP sidecar (gemma4_vision: e2b/e4b/26B/31B) or declares the
+ *  encoder-free unified tower in config.json (gemma4_unified_vision: 12B). */
+export function visionCapable(m: Pick<ModelRecord, "hasVisionSidecar" | "visionConfigType">): boolean {
+  return m.hasVisionSidecar || (m.visionConfigType?.endsWith("unified_vision") ?? false);
+}
+
 /** Of several cached revisions of ONE repo, pick the canonical snapshot: the
  *  revision refs/main points at, else the most recently scanned (stable
  *  tie-break on path). Paths are `<hub>/models--<repo>/snapshots/<hash>`. */
-function pickCanonicalRevision(matches: ModelRecord[]): ModelRecord {
+export function pickCanonicalRevision(matches: ModelRecord[]): ModelRecord {
   if (matches.length === 1) return matches[0]!;
   const root = matches[0]!.path.split("/snapshots/")[0]!;
   try {
@@ -203,6 +249,7 @@ function rowToRecord(r: Record<string, unknown>): ModelRecord {
     quantGroupSize: r.quant_group_size as number | null,
     quantMode: r.quant_mode as string | null,
     hasVisionSidecar: !!r.has_vision_sidecar,
+    visionConfigType: (r.vision_config_type as string | null) ?? null,
     hasKvConfig: !!r.has_kv_config,
     hasToolTemplate: !!r.has_tool_template,
     numLayers: r.num_layers as number | null,
@@ -258,6 +305,166 @@ function expertTensorBytes(path: string): number {
   }
 }
 
+// ---------------------------------------------------------------------------
+// gc: reclaim superseded snapshots + dead blobs.
+//
+// The HF cache layout keeps one snapshots/<commit> dir per downloaded
+// revision and NEVER deletes old ones — every upstream push a `get` follows
+// strands the previous snapshot (and any blobs only it references). gc keeps
+// the snapshots refs/* point at, deletes the rest, then deletes blobs no
+// surviving snapshot symlink targets.
+
+export interface GcSkippedSnapshot {
+  path: string;
+  /** Files this snapshot has that NO kept snapshot has (relative names).
+   *  Deleting it would delete the only copy — needs an explicit --force. */
+  extraFiles: string[];
+}
+
+export interface GcRepoPlan {
+  repoId: string;
+  repoDir: string;
+  /** Snapshots referenced by refs/* — never deleted. */
+  keepSnapshots: string[];
+  /** Unreferenced snapshots safe to delete (file set ⊆ kept snapshots'). */
+  pruneSnapshots: string[];
+  /** Unreferenced snapshots with files the kept set lacks — skipped unless
+   *  force (their blobs count as live). */
+  skippedSnapshots: GcSkippedSnapshot[];
+  /** blobs/ entries no surviving snapshot links to. */
+  deadBlobs: string[];
+  reclaimBytes: number;
+}
+
+/** Relative paths of all files/symlinks under a snapshot dir. */
+function snapshotFiles(snapDir: string): string[] {
+  const out: string[] = [];
+  const walk = (dir: string, prefix: string) => {
+    for (const name of readdirSync(dir)) {
+      const p = join(dir, name);
+      const st = lstatSync(p);
+      if (st.isDirectory()) walk(p, `${prefix}${name}/`);
+      else out.push(`${prefix}${name}`);
+    }
+  };
+  try { walk(snapDir, ""); } catch {}
+  return out;
+}
+
+/** Blob basenames a snapshot's symlinks resolve to (within this repo). */
+function snapshotBlobTargets(snapDir: string): Set<string> {
+  const targets = new Set<string>();
+  for (const rel of snapshotFiles(snapDir)) {
+    const p = join(snapDir, rel);
+    try {
+      const t = readlinkSync(p); // throws for regular files — those aren't blobs
+      const base = t.split("/").pop();
+      if (base && t.includes("blobs")) targets.add(base);
+    } catch {}
+  }
+  return targets;
+}
+
+/** Plan gc for one models--* repo dir. Conservative by construction:
+ *  - no refs/ (or refs pointing at missing snapshots) → nothing is pruned;
+ *  - an unreferenced snapshot carrying files the kept set lacks is only
+ *    prunable under `force` (the live example: a stale snapshot whose
+ *    optiq_vision.safetensors the canonical revision dropped);
+ *  - blobs are dead only when NO surviving snapshot (kept or skipped)
+ *    links to them. `.incomplete`/`.lock` resume artifacts are never touched. */
+export function planRepoGc(repoDir: string, opts: { force?: boolean } = {}): GcRepoPlan {
+  const entry = repoDir.split("/").pop() ?? "";
+  const repoId = entry.slice("models--".length).replaceAll("--", "/");
+  const plan: GcRepoPlan = {
+    repoId, repoDir,
+    keepSnapshots: [], pruneSnapshots: [], skippedSnapshots: [],
+    deadBlobs: [], reclaimBytes: 0,
+  };
+  const snapsDir = join(repoDir, "snapshots");
+  const refsDir = join(repoDir, "refs");
+  if (!existsSync(snapsDir)) return plan;
+
+  const refCommits = new Set<string>();
+  if (existsSync(refsDir)) {
+    for (const r of readdirSync(refsDir)) {
+      try { refCommits.add(readFileSync(join(refsDir, r), "utf8").trim()); } catch {}
+    }
+  }
+
+  const candidates: string[] = [];
+  for (const snap of readdirSync(snapsDir)) {
+    const p = join(snapsDir, snap);
+    if (!lstatSync(p).isDirectory()) continue;
+    (refCommits.has(snap) ? plan.keepSnapshots : candidates).push(p);
+  }
+  // No ref resolves to an existing snapshot → we cannot tell what's live.
+  if (plan.keepSnapshots.length === 0) return plan;
+
+  const keptFiles = new Set<string>();
+  for (const k of plan.keepSnapshots) for (const f of snapshotFiles(k)) keptFiles.add(f);
+  for (const c of candidates) {
+    const extra = snapshotFiles(c).filter((f) => !keptFiles.has(f));
+    if (extra.length > 0 && !opts.force) plan.skippedSnapshots.push({ path: c, extraFiles: extra });
+    else plan.pruneSnapshots.push(c);
+  }
+
+  const live = new Set<string>();
+  for (const s of [...plan.keepSnapshots, ...plan.skippedSnapshots.map((x) => x.path)])
+    for (const b of snapshotBlobTargets(s)) live.add(b);
+  const blobsDir = join(repoDir, "blobs");
+  if (existsSync(blobsDir)) {
+    for (const b of readdirSync(blobsDir)) {
+      if (b.endsWith(".incomplete") || b.endsWith(".lock")) continue;
+      if (live.has(b)) continue;
+      const p = join(blobsDir, b);
+      plan.deadBlobs.push(p);
+      try { plan.reclaimBytes += lstatSync(p).size; } catch {}
+    }
+  }
+  return plan;
+}
+
+export function planGc(hubDir: string = DEFAULT_HUB, opts: { force?: boolean } = {}): GcRepoPlan[] {
+  if (!existsSync(hubDir)) return [];
+  const plans: GcRepoPlan[] = [];
+  for (const entry of readdirSync(hubDir)) {
+    if (!entry.startsWith("models--")) continue;
+    plans.push(planRepoGc(join(hubDir, entry), opts));
+  }
+  return plans;
+}
+
+/** Execute a plan: snapshot dirs first, then dead blobs. Returns totals. */
+export function executeGc(plans: GcRepoPlan[]): {
+  snapshots: number; blobs: number; reclaimedBytes: number;
+} {
+  let snapshots = 0, blobs = 0, reclaimedBytes = 0;
+  for (const p of plans) {
+    for (const s of p.pruneSnapshots) {
+      rmSync(s, { recursive: true, force: true });
+      snapshots++;
+    }
+    for (const b of p.deadBlobs) {
+      try {
+        const size = lstatSync(b).size;
+        rmSync(b, { force: true });
+        blobs++;
+        reclaimedBytes += size;
+      } catch {}
+    }
+  }
+  return { snapshots, blobs, reclaimedBytes };
+}
+
+/** vision_config.model_type, but only when it actually names a vision tower
+ *  (`*_vision`). Some text-only configs nest a serialized copy of themselves
+ *  under vision_config (Qwen3.5 puts model_type "qwen3_5" there) — presence
+ *  of the key alone is NOT a vision signal. */
+function visionConfigTypeOf(config: Record<string, any>): string | null {
+  const vt = (config.vision_config as Record<string, any> | undefined)?.model_type;
+  return typeof vt === "string" && vt.endsWith("_vision") ? vt : null;
+}
+
 async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
   const configPath = join(dir, "config.json");
   if (!existsSync(configPath)) return null;
@@ -325,6 +532,7 @@ async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord |
     quantGroupSize: quant?.group_size ?? null,
     quantMode: quant?.mode ?? (quant ? "affine" : null),
     hasVisionSidecar: existsSync(join(dir, "optiq_vision.safetensors")),
+    visionConfigType: visionConfigTypeOf(config),
     hasKvConfig: existsSync(join(dir, "kv_config.json")),
     hasToolTemplate,
     numLayers: text.num_hidden_layers ?? null,

package/src/responses.ts

@@ -121,6 +121,19 @@ export interface ResponsesRequest {
   top_p?: number;
   top_k?: number;
   stream?: boolean;
+  /** mlx-lm sampler/penalty extensions + logit_bias — not part of the
+   *  Responses protocol proper; accepted as pass-through extras with the
+   *  same names as our /v1 chat surface. */
+  min_p?: number;
+  xtc_probability?: number;
+  xtc_threshold?: number;
+  logit_bias?: Record<string, number>;
+  repetition_penalty?: number;
+  repetition_context_size?: number;
+  presence_penalty?: number;
+  presence_context_size?: number;
+  frequency_penalty?: number;
+  frequency_context_size?: number;
   tools?: Array<Record<string, unknown>>;
   tool_choice?: unknown;
 }
@@ -235,6 +248,16 @@ export function responsesToChatBody(body: ResponsesRequest): Record<string, unkn
   if (body.top_p != null) oai.top_p = body.top_p;
   if (body.top_k != null) oai.top_k = body.top_k;
   if (body.stream) oai.stream = true;
+  // mlx-lm sampler/penalty extensions (+ OpenAI logit_bias): straight
+  // pass-through — same wire names on the chat surface.
+  for (const k of [
+    "min_p", "xtc_probability", "xtc_threshold", "logit_bias",
+    "repetition_penalty", "repetition_context_size",
+    "presence_penalty", "presence_context_size",
+    "frequency_penalty", "frequency_context_size",
+  ] as const) {
+    if (body[k] != null) oai[k] = body[k];
+  }
 
   // Flat Responses tool shape → nested chat shape; built-ins
   // (web_search, file_search, mcp, computer_use) dropped silently.

package/src/sampler.ts

@@ -1,8 +1,12 @@
 // Sampling — port of mlx-lm's sample_utils.make_sampler /
-// make_logits_processors (temperature, top-p, top-k, repetition penalty).
+// make_logits_processors (temperature, top-p, top-k, min-p, XTC,
+// repetition/presence/frequency penalties, logit bias).
 // All filtering happens on-device; only the chosen token id crosses to JS.
 // Seeded: each step derives a fresh key from (seed, step) so runs are
-// reproducible without sharing global RNG state.
+// reproducible without sharing global RNG state. (mlx-lm draws XTC's
+// uniform and the categorical from one global stream; we split the
+// per-step key into (xtc, categorical) subkeys instead — deterministic,
+// and the XTC-off path is unchanged.)
 
 import { Dtype } from "./mlx/ffi";
 import { MlxArray } from "./mlx/array";
@@ -63,6 +67,18 @@ export interface SamplerOptions {
   temperature?: number;
   topP?: number;
   topK?: number;
+  /** min-p: keep tokens whose probability ≥ minP · p(top token). mlx-lm's
+   *  `min_p`. Applied after top-p, before XTC/top-k (reference chain order). */
+  minP?: number;
+  /** Minimum number of tokens min-p may never filter out. Default 1. */
+  minTokensToKeep?: number;
+  /** XTC ("exclude top choices"): with this probability per step, drop every
+   *  token above the threshold EXCEPT the least likely of them. */
+  xtcProbability?: number;
+  /** Probability a token must exceed to be an XTC removal candidate. [0, 0.5]. */
+  xtcThreshold?: number;
+  /** Token ids XTC never removes (mlx-lm's server passes EOS + "\n"). */
+  xtcSpecialTokens?: number[];
   seed?: number;
   /** HLG tone-curve sampling. When enabled, replaces temperature's flat slope
    *  with the piecewise curve (temperature becomes the mid gain). Off/undefined
@@ -75,8 +91,19 @@ export interface SamplerOptions {
 }
 
 export interface LogitsProcessorOptions {
+  /** Additive per-token-id logit bias, applied before any penalty
+   *  (mlx-lm's `logit_bias`: {tokenId: bias}). */
+  logitBias?: Record<number, number>;
   repetitionPenalty?: number;
   repetitionContextSize?: number;
+  /** OpenAI-style presence penalty: subtracted ONCE from a token's logit if it
+   *  occurred at all in the recent window. May be negative. */
+  presencePenalty?: number;
+  presenceContextSize?: number;
+  /** OpenAI-style frequency penalty: subtracted once PER OCCURRENCE of the
+   *  token in the recent window. May be negative. */
+  frequencyPenalty?: number;
+  frequencyContextSize?: number;
 }
 
 /** logprobs [1, V] → sampled token array (uint32, shape [1]). */
@@ -98,7 +125,7 @@ function negInfLike(a: MlxArray): MlxArray {
 
 /** apply_top_p: keep the smallest set of tokens whose cumulative
  *  probability exceeds top_p; others → -inf. */
-function applyTopP(lp: MlxArray, topP: number): MlxArray {
+export function applyTopP(lp: MlxArray, topP: number): MlxArray {
   const probs = ops.exp(lp);
   const sortedIdx = ops.argsortAxis(lp, -1);
   const sortedProbs = ops.takeAlongAxis(probs, sortedIdx, -1);
@@ -120,7 +147,7 @@ function applyTopP(lp: MlxArray, topP: number): MlxArray {
 }
 
 /** apply_top_k: all but the k highest logprobs → -inf. */
-function applyTopK(lp: MlxArray, topK: number): MlxArray {
+export function applyTopK(lp: MlxArray, topK: number): MlxArray {
   const V = lp.shape[lp.shape.length - 1]!;
   const negLp = ops.neg(lp);
   const part = ops.argpartitionAxis(negLp, topK - 1, -1);
@@ -131,6 +158,88 @@ function applyTopK(lp: MlxArray, topK: number): MlxArray {
   return out;
 }
 
+/** apply_min_p: mask tokens whose probability is below minP · p(top token).
+ *  Faithful to the reference (mlx-lm 0.31.3 sample_utils.apply_min_p): the
+ *  threshold is computed in log space directly against the max logprob —
+ *  `scaled_min_p = max(lp) + log(minP)` — no full sort; min_tokens_to_keep
+ *  survivors are recovered with an argpartition on the top tail. */
+export function applyMinP(lp: MlxArray, minP: number, minTokensToKeep = 1): MlxArray {
+  if (!(minP >= 0 && minP <= 1))
+    throw new Error(`minP has to be a float in the [0, 1] interval, but is ${minP}`);
+  if (!Number.isInteger(minTokensToKeep) || minTokensToKeep < 1)
+    throw new Error(`minTokensToKeep has to be a positive integer, but is ${minTokensToKeep}`);
+
+  const owned: MlxArray[] = [];
+  const k = <T extends MlxArray>(a: T): T => { owned.push(a); return a; };
+
+  // Mask tokens that have a probability less than max(p) * min_p
+  const top = k(ops.maxAxis(lp, -1, true));
+  const scaledMinP = k(ops.add(top, k(ops.scalarLike(Math.log(minP), top))));
+  let remove = k(ops.less(lp, scaledMinP));
+
+  // Ensure at least min_tokens_to_keep survive the filter
+  if (minTokensToKeep > 1) {
+    const V = lp.shape[lp.shape.length - 1]!;
+    // mx.argpartition(lp, kth=-min_tokens_to_keep): negative kth ≡ V - k
+    const part = k(ops.argpartitionAxis(lp, V - minTokensToKeep, -1));
+    const topIdx = k(part.slice([0, V - minTokensToKeep], [1, V]));
+    const falses = k(ops.zeros([1, minTokensToKeep], Dtype.bool));
+    remove = k(ops.putAlongAxis(remove, topIdx, falses, -1));
+  }
+
+  const ninf = k(negInfLike(lp));
+  const out = ops.where(remove, ninf, lp);
+  for (const a of owned) a.dispose();
+  return out;
+}
+
+/** apply_xtc: with probability xtcProbability, remove every token whose
+ *  probability exceeds the minimum above-threshold probability — i.e. drop
+ *  the top choices, keeping the least likely token that still clears the
+ *  threshold (plus everything below it). Special tokens are never removed.
+ *  `key` seeds the per-step uniform draw (the reference uses the global RNG). */
+export function applyXtc(
+  logits: MlxArray,
+  xtcProbability: number,
+  xtcThreshold: number,
+  xtcSpecialTokens: number[],
+  key: MlxArray | null,
+): MlxArray {
+  if (!(xtcThreshold >= 0 && xtcThreshold <= 0.5))
+    throw new Error(`xtcThreshold has to be a float in the [0, 0.5] interval, but is ${xtcThreshold}`);
+  if (!(xtcProbability >= 0 && xtcProbability <= 1))
+    throw new Error(`xtcProbability has to be a float in the [0, 1] interval, but is ${xtcProbability}`);
+
+  const owned: MlxArray[] = [];
+  const k = <T extends MlxArray>(a: T): T => { owned.push(a); return a; };
+
+  const probs = k(ops.softmaxAxis(logits, -1, false));
+  // mask = probs > min(where(probs > threshold, probs, inf))
+  const thresh = k(ops.scalarLike(xtcThreshold, probs));
+  const above = k(ops.less(thresh, probs)); // probs > threshold
+  const inf = k(ops.scalarLike(Infinity, probs));
+  const aboveProbs = k(ops.where(above, probs, inf));
+  // min(x) = -max(-x) (no min reduction in the ops layer)
+  const minAbove = k(ops.neg(k(ops.maxAxis(k(ops.neg(aboveProbs)), -1, true))));
+  let mask = k(ops.less(minAbove, probs)); // probs > min above-threshold prob
+  if (xtcSpecialTokens.length > 0) {
+    const n = xtcSpecialTokens.length;
+    const idx = k(ops.fromInt32(xtcSpecialTokens, [1, n]));
+    const falses = k(ops.zeros([1, n], Dtype.bool));
+    mask = k(ops.putAlongAxis(mask, idx, falses, -1));
+  }
+
+  // where(uniform() > xtcProbability, logits, where(mask, -inf, logits))
+  const u = k(ops.randomUniform([1], Dtype.float32, 0, 1, key));
+  const pArr = k(ops.scalarLike(xtcProbability, u));
+  const skip = k(ops.less(pArr, u)); // u > xtcProbability → leave logits alone
+  const ninf = k(negInfLike(logits));
+  const culled = k(ops.where(mask, ninf, logits));
+  const out = ops.where(skip, logits, culled);
+  for (const a of owned) a.dispose();
+  return out;
+}
+
 /** HLG sampling parameters — a piecewise tone curve on the logprobs.
  *  See docs/design/hlg-sampling.md. Knobs are in nats (the logprob unit). */
 export interface HlgParams {
@@ -410,7 +519,11 @@ export function applyHlgShaper(lp: MlxArray, opts: HlgShaperParams = {}): MlxArr
 }
 
 export function makeSampler(opts: SamplerOptions = {}): Sampler {
-  const { temperature = 0, topP = 0, topK = 0, seed = 0, hlg, curve } = opts;
+  const {
+    temperature = 0, topP = 0, topK = 0, minP = 0, minTokensToKeep = 1,
+    xtcProbability = 0, xtcThreshold = 0, xtcSpecialTokens = [],
+    seed = 0, hlg, curve,
+  } = opts;
 
   // v2 curve sampler: the drawn log-prob transfer curve REPLACES temperature +
   // softmax. Stochastic and seeded by design (the curve IS the sampling shape),
@@ -461,13 +574,29 @@ export function makeSampler(opts: SamplerOptions = {}): Sampler {
     };
   }
 
+  // Filter chain order matches mlx-lm's make_sampler: top_p → min_p → xtc →
+  // top_k, then categorical(logprobs / temperature).
   return (lp, step) => {
     let cur = lp;
     const owned: MlxArray[] = [];
+    let key = stepKey(seed, step);
     if (topP > 0 && topP < 1) { cur = applyTopP(cur, topP); owned.push(cur); }
+    if (minP !== 0) { cur = applyMinP(cur, minP, minTokensToKeep); owned.push(cur); }
+    if (xtcProbability > 0) {
+      // Split the step key: subkey 0 → XTC's uniform, subkey 1 → categorical.
+      const split = ops.randomSplitNum(key, 2);
+      const row0 = split.slice([0, 0], [1, 2]);
+      const row1 = split.slice([1, 0], [2, 2]);
+      const xtcKey = ops.reshape(row0, [2]);
+      key.dispose();
+      key = ops.reshape(row1, [2]);
+      for (const a of [split, row0, row1]) a.dispose();
+      cur = applyXtc(cur, xtcProbability, xtcThreshold, xtcSpecialTokens, xtcKey);
+      owned.push(cur);
+      xtcKey.dispose();
+    }
     if (topK > 0) { cur = applyTopK(cur, topK); owned.push(cur); }
     const scaled = ops.mulScalar(cur, 1 / temperature);
-    const key = stepKey(seed, step);
     const tok = ops.randomCategorical(scaled, key);
     scaled.dispose();
     key.dispose();
@@ -476,18 +605,48 @@ export function makeSampler(opts: SamplerOptions = {}): Sampler {
   };
 }
 
+/** Start of the recent-token window, matching Python's `tokens[-context_size:]`
+ *  slice semantics exactly (context_size 0 ⇒ `[-0:]` ⇒ the WHOLE history). */
+function windowStart(n: number, contextSize: number): number {
+  return contextSize === 0 ? 0 : Math.max(0, n - contextSize);
+}
+
 export function makeLogitsProcessors(opts: LogitsProcessorOptions = {}): LogitsProcessor[] {
   const out: LogitsProcessor[] = [];
-  const { repetitionPenalty, repetitionContextSize = 20 } = opts;
+  const {
+    logitBias,
+    repetitionPenalty, repetitionContextSize = 20,
+    presencePenalty, presenceContextSize = 20,
+    frequencyPenalty, frequencyContextSize = 20,
+  } = opts;
+
+  // logit_bias first, as in mlx-lm's make_logits_processors. Reference:
+  // `logits.at[:, indices].add(values)` — dict keys are unique, so a
+  // gather → add → put round-trip is exactly equivalent to the scatter-add.
+  const biasIds = logitBias ? Object.keys(logitBias).map(Number) : [];
+  if (logitBias && biasIds.length > 0) {
+    const biasVals = Float32Array.from(biasIds.map((id) => logitBias[id]!));
+    out.push((_tokens, logits) => {
+      const K = biasIds.length;
+      const idx = ops.fromInt32(biasIds, [1, K]);
+      const vals32 = MlxArray.fromFloat32(biasVals, [1, K]);
+      const vals = vals32.astype(logits.dtype);
+      const selected = ops.takeAlongAxis(logits, idx, -1);
+      const biased = ops.add(selected, vals);
+      const updated = ops.putAlongAxis(logits, idx, biased, -1);
+      for (const a of [idx, vals32, vals, selected, biased]) a.dispose();
+      return updated;
+    });
+  }
 
-  if (repetitionPenalty && repetitionPenalty !== 0) {
+  if (repetitionPenalty !== undefined && repetitionPenalty !== 0) {
     if (repetitionPenalty < 0)
       throw new Error("repetitionPenalty must be non-negative");
     out.push((tokens, logits) => {
       if (!tokens) return logits;
       const n = tokens.shape[0]!;
       if (n === 0) return logits;
-      const start = Math.max(0, n - repetitionContextSize);
+      const start = windowStart(n, repetitionContextSize);
       const recent = tokens.slice([start], [n]);
       const idx = ops.reshape(recent, [1, n - start]);
       const selected = ops.takeAlongAxis(logits, idx, -1);
@@ -503,6 +662,58 @@ export function makeLogitsProcessors(opts: LogitsProcessorOptions = {}): LogitsP
       return updated;
     });
   }
+
+  // make_presence_penalty: `logits[:, tokens] -= penalty` — a plain
+  // fancy-index assignment, so duplicate occurrences all write the SAME
+  // value (original − penalty): the penalty applies once per distinct token.
+  if (presencePenalty !== undefined && presencePenalty !== 0) {
+    out.push((tokens, logits) => {
+      if (!tokens) return logits;
+      const n = tokens.shape[0]!;
+      if (n === 0) return logits;
+      const start = windowStart(n, presenceContextSize);
+      const recent = tokens.slice([start], [n]);
+      const idx = ops.reshape(recent, [1, n - start]);
+      const selected = ops.takeAlongAxis(logits, idx, -1);
+      const pen = ops.scalarLike(presencePenalty, selected);
+      const penalized = ops.sub(selected, pen);
+      const updated = ops.putAlongAxis(logits, idx, penalized, -1);
+      for (const a of [recent, idx, selected, pen, penalized]) a.dispose();
+      return updated;
+    });
+  }
+
+  // make_frequency_penalty: `logits.at[:, tokens].subtract(penalty)` — a
+  // scatter-subtract that ACCUMULATES over duplicates: each token loses
+  // penalty × (occurrence count in the window). We compute the count per
+  // position ([m,m] equality matrix summed over a row — m ≤ contextSize, tiny)
+  // so every duplicate writes the identical final value and a put_along_axis
+  // reproduces the scatter-accumulate exactly.
+  if (frequencyPenalty !== undefined && frequencyPenalty !== 0) {
+    out.push((tokens, logits) => {
+      if (!tokens) return logits;
+      const n = tokens.shape[0]!;
+      if (n === 0) return logits;
+      const start = windowStart(n, frequencyContextSize);
+      const m = n - start;
+      const recent = tokens.slice([start], [n]);
+      const col = ops.reshape(recent, [m, 1]);
+      const idx = ops.reshape(recent, [1, m]);
+      const eq = ops.equal(col, idx); // [m, m]
+      const counts = ops.sumAxis(eq, -1, false); // [m] occurrences of tokens[i]
+      const countsF = counts.astype(logits.dtype);
+      const countsRow = ops.reshape(countsF, [1, m]);
+      const selected = ops.takeAlongAxis(logits, idx, -1);
+      const pen = ops.scalarLike(frequencyPenalty, selected);
+      const penTotal = ops.mul(countsRow, pen); // penalty × count
+      const penalized = ops.sub(selected, penTotal);
+      const updated = ops.putAlongAxis(logits, idx, penalized, -1);
+      for (const a of [recent, col, idx, eq, counts, countsF, countsRow, selected, pen, penTotal, penalized])
+        a.dispose();
+      return updated;
+    });
+  }
+
   return out;
 }