@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/runtime/resources.ts

@@ -0,0 +1,53 @@
+import type { DisposableResource, PreparedStateChange, ResourceOwner } from "../contracts/portable/resources";
+
+export function ownResource<T>(value: T, release: (value: T) => void): ResourceOwner<T> {
+  let held: { value: T } | undefined = { value };
+  const borrow = () => {
+    if (!held) throw new Error("resource ownership has ended");
+    return held.value;
+  };
+  return {
+    borrow,
+    transfer() { const result = borrow(); held = undefined; return result; },
+    close() {
+      if (!held) return;
+      const result = held.value;
+      held = undefined;
+      release(result);
+    },
+  };
+}
+
+/** Attempt every release once. One failed destructor must not leak siblings. */
+export function disposeResources(resources: Iterable<DisposableResource>): void {
+  const errors: unknown[] = [];
+  for (const resource of new Set(resources)) {
+    try { resource.dispose(); } catch (error) { errors.push(error); }
+  }
+  if (errors.length === 1) throw errors[0];
+  if (errors.length) throw new AggregateError(errors, "resource cleanup failed");
+}
+
+/** Used at failure boundaries to retain the execution error as well as cleanup. */
+export function cleanupFailure(error: unknown, cleanup: () => void): never {
+  try { cleanup(); }
+  catch (cleanupError) { throw new AggregateError([error, cleanupError], "execution and cleanup failed"); }
+  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()); }
+  finally { owner.close(); }
+}

package/src/runtime/row-operations.ts

@@ -0,0 +1,42 @@
+type RowTable<Operation> = Readonly<Partial<Record<number, Operation>> & { readonly default?: Operation }>;
+type RowOperations<Operation> = RowTable<Operation> |
+  readonly { readonly from: number; readonly to: number; readonly operation: Operation }[];
+
+/** Index authored exact rows, an explicit default, or authored contiguous ranges. */
+export function rowOperationFor<Operation>(operations: readonly {readonly from:number;readonly to:number;readonly operation:Operation}[], rows:number):Operation;
+export function rowOperationFor<Operation>(operations: RowTable<Operation>, rows:number):Operation;
+export function rowOperationFor<Operation>(operations: RowOperations<Operation>, rows:number):Operation;
+export function rowOperationFor<Operation>(operations: RowOperations<Operation>, rows: number): Operation {
+  let operation: Operation | undefined;
+  if (Number.isSafeInteger(rows) && rows > 0) {
+    operation = Array.isArray(operations)
+      ? operations.find(range => rows >= range.from && rows <= range.to)?.operation
+      : (Object.hasOwn(operations, rows) ? (operations as RowTable<Operation>)[rows] : undefined)
+        ?? (Object.hasOwn(operations, "default") ? (operations as RowTable<Operation>).default : undefined);
+  }
+  if (operation === undefined) throw new Error(`no declared operation for ${rows} rows`);
+  return operation;
+}
+
+/** Compatibility data helper. Its default is an ordinary inspectable own property. */
+export function rowTableWithDefault<Operation>(entries: RowTable<Operation>, defaultOperation: Operation): RowTable<Operation> {
+  return Object.freeze({ ...entries, default: defaultOperation });
+}
+
+/** Map only authored entries, ranges and an explicitly present default. */
+export function mapOperations<Input,Output>(operations: readonly {readonly from:number;readonly to:number;readonly operation:Input}[], bind:(operation:Input)=>Output): readonly {readonly from:number;readonly to:number;readonly operation:Output}[];
+export function mapOperations<Input,Output>(operations: RowTable<Input>, bind:(operation:Input)=>Output):RowTable<Output>;
+export function mapOperations<Input,Output>(operations: RowOperations<Input>, bind:(operation:Input)=>Output):RowOperations<Output>;
+export function mapOperations<Input, Output>(operations: RowOperations<Input>,
+  bind: (operation: Input) => Output): RowOperations<Output> {
+  if (Array.isArray(operations))
+    return Object.freeze(operations.map(range => Object.freeze({ from: range.from, to: range.to, operation: bind(range.operation) })));
+  const source = operations as RowTable<Input>, entries: Partial<Record<number, Output>> & { default?: Output } = {};
+  for (const [key, operation] of Object.entries(source)) {
+    if (operation === undefined) continue;
+    const rows = Number(key);
+    if (key === "default") entries.default = bind(operation);
+    else if (Number.isSafeInteger(rows) && rows > 0 && String(rows) === key) entries[rows] = bind(operation);
+  }
+  return Object.freeze(entries);
+}

package/src/runtime/trace.ts

@@ -0,0 +1,161 @@
+import { runtimeValue } from "./config";
+
+export const P2R_TRACE_PREFIX = "[p2r] ";
+export const P2R_TRACE_VERSION = 1;
+
+export type P2RTraceOutcome = "success" | "error" | "abort";
+
+export type P2RTracePhase =
+  | "request.body_parse"
+  | "request.prompt_prepare"
+  | "completion.total"
+  | "completion.placement"
+  | "engine.request_wait"
+  | "engine.admission_wait"
+  | "engine.mixed_forward"
+  | "cache.lookup_restore"
+  | "cache.prefetch"
+  | "prefill.total"
+  | "prefill.batch_setup"
+  | "prefill.chunk"
+  | "prefill.row_sync"
+  | "prefill.forward"
+  | "prefill.evaluate"
+  | "prefill.checkpoint"
+  | "prefill.project"
+  | "prefill.complete"
+  | "prefill.companion"
+  | "prefill.kv_maintenance"
+  | "token_zero.total"
+  | "token_zero.forward"
+  | "token_zero.head"
+  | "token_zero.sample"
+  | "response.token_route"
+  | "response.first_write"
+  | "response.final_write";
+
+export type P2RTraceAttributes = Readonly<Record<
+  string,
+  string | number | boolean | null
+>>;
+
+export interface P2RTraceEvent {
+  phase: P2RTracePhase;
+  startMs: number;
+  durationMs: number;
+  attributes?: P2RTraceAttributes;
+}
+
+export interface P2RTraceRecord {
+  version: typeof P2R_TRACE_VERSION;
+  traceId: string;
+  requestId: string;
+  route: string;
+  clock: "monotonic-ms";
+  outcome: P2RTraceOutcome;
+  /** Process-local monotonic origin; aligns concurrent requests in one server. */
+  startedAtMs?: number;
+  totalMs: number;
+  events: P2RTraceEvent[];
+  attributes?: P2RTraceAttributes;
+}
+
+type Clock = () => number;
+type Emit = (record: P2RTraceRecord) => void;
+
+export interface PromptResponseTraceOptions {
+  traceId: string;
+  requestId: string;
+  route: string;
+  clock?: Clock;
+  emit?: Emit;
+}
+
+/** Request-local diagnostic trace. Callers pass it explicitly because the
+ * continuous scheduler interleaves rows and cannot safely use ambient async
+ * context. No instance is created unless MLX_BUN_P2R_TRACE=1. */
+export class PromptResponseTrace {
+  readonly #startedAt: number;
+  readonly #events: P2RTraceEvent[] = [];
+  readonly #clock: Clock;
+  readonly #emit: Emit;
+  #finished = false;
+
+  readonly traceId: string;
+  readonly requestId: string;
+  readonly route: string;
+
+  constructor(options: PromptResponseTraceOptions) {
+    this.traceId = options.traceId;
+    this.requestId = options.requestId;
+    this.route = options.route;
+    this.#clock = options.clock ?? performance.now.bind(performance);
+    this.#emit = options.emit ?? ((record) => {
+      console.error(`${P2R_TRACE_PREFIX}${JSON.stringify(record)}`);
+    });
+    this.#startedAt = this.#clock();
+  }
+
+  get finished(): boolean {
+    return this.#finished;
+  }
+
+  /** Start one span. The returned closer is idempotent so finally blocks and
+   * early exits cannot double-record a phase. */
+  begin(phase: P2RTracePhase, attributes?: P2RTraceAttributes): () => void {
+    const startedAt = this.#clock();
+    let closed = false;
+    return () => {
+      if (closed || this.#finished) return;
+      closed = true;
+      const endedAt = this.#clock();
+      this.#events.push({
+        phase,
+        startMs: startedAt - this.#startedAt,
+        durationMs: endedAt - startedAt,
+        ...(attributes ? { attributes } : {}),
+      });
+    };
+  }
+
+  mark(phase: P2RTracePhase, attributes?: P2RTraceAttributes): void {
+    if (this.#finished) return;
+    this.#events.push({
+      phase,
+      startMs: this.#clock() - this.#startedAt,
+      durationMs: 0,
+      ...(attributes ? { attributes } : {}),
+    });
+  }
+
+  finish(
+    outcome: P2RTraceOutcome,
+    attributes?: P2RTraceAttributes,
+  ): P2RTraceRecord | null {
+    if (this.#finished) return null;
+    this.#finished = true;
+    const totalMs = this.#clock() - this.#startedAt;
+    const record: P2RTraceRecord = {
+      version: P2R_TRACE_VERSION,
+      traceId: this.traceId,
+      requestId: this.requestId,
+      route: this.route,
+      clock: "monotonic-ms",
+      outcome,
+      startedAtMs: this.#startedAt,
+      totalMs,
+      events: [...this.#events].sort((a, b) => a.startMs - b.startMs),
+      ...(attributes ? { attributes } : {}),
+    };
+    this.#emit(record);
+    return record;
+  }
+}
+
+export function createPromptResponseTrace(
+  options: PromptResponseTraceOptions,
+): PromptResponseTrace | undefined {
+  return runtimeValue("MLX_BUN_P2R_TRACE") === "1"
+    ? new PromptResponseTrace(options)
+    : undefined;
+}

package/src/sampling/curve.ts

@@ -0,0 +1,122 @@
+// curve.ts — the v2 log-prob transfer-curve sampler. Replaces the whole
+// temperature+softmax stage with a drawn MONOTONE CUBIC curve in log-probability
+// space:
+//
+//   p_out_i = exp(curve(log p_in_i)) / Σ_j exp(curve(log p_in_j))
+//
+// `lp` reaching the sampler is already log p_in (toLogprobs = logits − logsumexp),
+// so the curve maps it directly and an identity curve reproduces temperature 1.
+// The curve is a Fritsch–Carlson monotone cubic Hermite (PCHIP) through N movable
+// control points in nat log-prob space (DaVinci-curves style — the endpoints start
+// at the corners but are draggable too). Smooth, never overshoots, monotone. A
+// control point dragged below the diagonal suppresses those tokens; the gate is
+// emergent from the curve shape, not a separate threshold.
+//
+// Reading the curve: at a given p_in, ABOVE the diagonal raises those tokens,
+// BELOW lowers them (renorm makes the literal number an intent); the SLOPE across
+// a region sets how evenly the tokens inside it spread (flat<1 even, steep>1 separated).
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+export interface CurvePoint { x_pct: number; y_pct: number; }
+export interface CurveParams {
+  space?: "logprob";
+  points: CurvePoint[]; // ≥2 control points, all movable; first/last are the endpoints
+  monotonic?: boolean;
+}
+
+/** Bottom-left default anchor / plot floor (p = 1e-6). */
+export const CURVE_UMIN = Math.log(1e-6);
+const L = (pt: CurvePoint): [number, number] => [Math.log(Math.max(pt.x_pct, 1e-9) / 100), Math.log(Math.max(pt.y_pct, 1e-9) / 100)];
+
+/** Control points in nat log-prob space, sorted ascending in x. */
+function knotArrays(p: CurveParams): { X: number[]; Y: number[] } {
+  const k = p.points.map(L).sort((a, b) => a[0] - b[0]);
+  return { X: k.map((q) => q[0]), Y: k.map((q) => q[1]) };
+}
+
+/** Consecutive-point secants, top→bottom (region slopes; what the editor readout shows). */
+export function curveSecants(p: CurveParams): number[] {
+  const { X, Y } = knotArrays(p);
+  const s: number[] = [];
+  for (let i = X.length - 1; i > 0; i--) s.push((Y[i]! - Y[i - 1]!) / (X[i]! - X[i - 1]!));
+  return s;
+}
+/** True iff every consecutive-point secant is ≥ 0 (monotone non-decreasing). */
+export function isMonotone(p: CurveParams): boolean {
+  return p.points.length >= 2 && curveSecants(p).every((s) => s >= -1e-12);
+}
+
+interface Seg { A: number; B: number; C: number; D: number } // cubic in u: A+Bu+Cu²+Du³
+export interface Spline { X: number[]; Y: number[]; M: number[]; segs: Seg[] }
+
+function cubicInU(x0: number, h: number, y0: number, y1: number, m0: number, m1: number): Seg {
+  const c0 = y0, c1 = h * m0;
+  const c2 = -3 * y0 - 2 * h * m0 + 3 * y1 - h * m1;
+  const c3 = 2 * y0 + h * m0 - 2 * y1 + h * m1;
+  const s = 1 / h, s2 = s * s, s3 = s2 * s;
+  let A = c0, B = 0, C = 0, D = 0;
+  B += c1 * s; A += -c1 * s * x0;
+  C += c2 * s2; B += -2 * x0 * c2 * s2; A += x0 * x0 * c2 * s2;
+  D += c3 * s3; C += -3 * x0 * c3 * s3; B += 3 * x0 * x0 * c3 * s3; A += -x0 * x0 * x0 * c3 * s3;
+  return { A, B, C, D };
+}
+
+/** Monotone cubic (Fritsch–Carlson) over the N control points. */
+export function buildSpline(p: CurveParams): Spline {
+  const { X, Y } = knotArrays(p);
+  const n = X.length;
+  const Dsec: number[] = [];
+  for (let i = 0; i < n - 1; i++) Dsec.push((Y[i + 1]! - Y[i]!) / (X[i + 1]! - X[i]!));
+  const M = new Array<number>(n);
+  M[0] = Dsec[0]!; M[n - 1] = Dsec[n - 2]!;
+  for (let i = 1; i < n - 1; i++) {
+    if (Dsec[i - 1]! * Dsec[i]! <= 0) { M[i] = 0; continue; }
+    const h1 = X[i]! - X[i - 1]!, h2 = X[i + 1]! - X[i]!;
+    const w1 = 2 * h2 + h1, w2 = h2 + 2 * h1;
+    M[i] = (w1 + w2) / (w1 / Dsec[i - 1]! + w2 / Dsec[i]!);
+  }
+  const segs: Seg[] = [];
+  for (let i = 0; i < n - 1; i++) segs.push(cubicInU(X[i]!, X[i + 1]! - X[i]!, Y[i]!, Y[i + 1]!, M[i]!, M[i + 1]!));
+  return { X, Y, M, segs };
+}
+
+/** Scalar eval (editor preview / tests). Extrapolates linearly past both endpoints. */
+export function evalSpline(sp: Spline, u: number): number {
+  const n = sp.X.length;
+  if (u < sp.X[0]!) return sp.Y[0]! + Math.max(sp.M[0]!, 1) * (u - sp.X[0]!); // deep tail: slope ≥1, stays gated
+  if (u > sp.X[n - 1]!) return sp.Y[n - 1]! + sp.M[n - 1]! * (u - sp.X[n - 1]!);
+  let i = sp.segs.length - 1;
+  while (i > 0 && u < sp.X[i]!) i--;
+  const s = sp.segs[i]!;
+  return s.A + s.B * u + s.C * u * u + s.D * u * u * u;
+}
+export function curveAt(p: CurveParams, u: number): number { return evalSpline(buildSpline(p), u); }
+
+/** Apply the monotone-cubic transfer to normalized logprobs `lp` [.., V] (on-device). */
+export function applyCurve(lp: MlxArray, p: CurveParams): MlxArray {
+  const sp = buildSpline(p);
+  const n = sp.X.length;
+  const tmp: MlxArray[] = [];
+  const k = <T extends MlxArray>(a: T): T => { tmp.push(a); return a; };
+  const C = (v: number): MlxArray => k(ops.scalarLike(v, lp));
+  const lp2 = k(ops.mul(lp, lp));
+  const lp3 = k(ops.mul(lp2, lp));
+  const seg = (s: Seg): MlxArray => k(ops.add(
+    k(ops.add(C(s.A), k(ops.mulScalar(lp, s.B)))),
+    k(ops.add(k(ops.mulScalar(lp2, s.C)), k(ops.mulScalar(lp3, s.D)))),
+  ));
+  // below the bottom point: slope ≥ 1 (deep tail stays gated)
+  const lowSlope = Math.max(sp.M[0]!, 1);
+  let out: MlxArray = k(ops.add(C(sp.Y[0]! - lowSlope * sp.X[0]!), k(ops.mulScalar(lp, lowSlope))));
+  for (let i = 0; i < sp.segs.length; i++) {
+    const next = ops.where(k(ops.greaterEqual(lp, C(sp.X[i]!))), seg(sp.segs[i]!), out);
+    tmp.push(next); out = next;
+  }
+  // above the top point: linear extrapolation with the top tangent (if it was dragged left of 100%)
+  const aboveExt = k(ops.add(C(sp.Y[n - 1]! - sp.M[n - 1]! * sp.X[n - 1]!), k(ops.mulScalar(lp, sp.M[n - 1]!))));
+  const result = ops.where(k(ops.greaterEqual(lp, C(sp.X[n - 1]!))), aboveExt, out);
+  for (const a of tmp) a.dispose();
+  return result;
+}

package/src/sampling/draft.ts

@@ -0,0 +1,55 @@
+// Sampling support for the DSpark/DFlash drafter at temperature > 0: the
+// drafter samples each position from the server sampler's exact
+// top-p/top-k/temperature distribution (sampler.ts), using a seeded key stream
+// that does not touch global RNG state. HLG/curve samplers are out of scope;
+// the caller rejects them.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { applyTopP, applyTopK, toLogprobs } from "./index";
+import { stepKey } from "./sampler";
+
+export interface DSparkSampleConfig {
+  /** temperature > 0 enables sampling; 0 ⇒ greedy (handled by the caller). */
+  temperature: number;
+  topP?: number;
+  topK?: number;
+  seed?: number;
+}
+
+/** Seeded stream of mlx random keys — reproducible draws without sharing global
+ *  RNG state: draw n is `stepKey(seed >>> 0, n)`. */
+export class KeyStream {
+  #seed: number;
+  #ctr = 0;
+  constructor(seed = 0) { this.#seed = seed >>> 0; }
+  next(): MlxArray {
+    return stepKey(this.#seed, this.#ctr++);
+  }
+}
+
+/**
+ * Processed categorical logits for one position: logits [1,V] → masked +
+ * temperature-scaled logits [1,V] (the input to randomCategorical;
+ * softmax(·) is the sampling distribution). Mirrors makeSampler's top-p/top-k/
+ * temperature path exactly. temperature must be > 0.
+ */
+export function processLogits(logits: MlxArray, cfg: DSparkSampleConfig): MlxArray {
+  if (cfg.temperature <= 0) throw new Error("processLogits requires temperature > 0");
+  const lp = toLogprobs(logits); // [1,V]
+  let cur = lp;
+  const owned: MlxArray[] = [lp];
+  if (cfg.topP && cfg.topP > 0 && cfg.topP < 1) { cur = applyTopP(cur, cfg.topP); owned.push(cur); }
+  if (cfg.topK && cfg.topK > 0) { cur = applyTopK(cur, cfg.topK); owned.push(cur); }
+  const scaled = ops.mulScalar(cur, 1 / cfg.temperature);
+  for (const a of owned) a.dispose();
+  return scaled;
+}
+
+/** Sample one token id from processed logits [1,V] (∝ softmax). */
+export function sampleToken(scaled: MlxArray, key: MlxArray): number {
+  const t = ops.randomCategorical(scaled, key); // [1] uint32
+  const id = ops.itemUint32(t);
+  t.dispose();
+  return id;
+}

package/src/sampling/extras.ts

@@ -0,0 +1,45 @@
+// Sampling — port of mlx-lm's sample_utils.make_sampler /
+// 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. (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 { MlxArray } from "@mlx-bun/mlx/array";
+import type { TokenLogprobs } from "../contracts/portable/generation";
+import { StepExtras } from "./types";
+
+export function stepExtrasArrays(extras: StepExtras | null): MlxArray[] {
+  return extras
+    ? [extras.sel, extras.topIdx, extras.topVals].filter(
+        (array): array is MlxArray => array !== null,
+      )
+    : [];
+}
+
+export function disposeStepExtras(extras: StepExtras | null): void {
+  if (!extras) return;
+  extras.sel?.dispose();
+  extras.topIdx?.dispose();
+  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); }
+}

package/src/sampling/filters.ts

@@ -0,0 +1,134 @@
+// Sampling — port of mlx-lm's sample_utils.make_sampler /
+// 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. (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 { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+
+function negInfLike(a: MlxArray): MlxArray {
+  return ops.scalarLike(-Infinity, a);
+}
+
+/** apply_top_p: keep the smallest set of tokens whose cumulative
+ *  probability exceeds top_p; others → -inf. */
+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);
+  const cum = ops.cumsum(sortedProbs, -1);
+  // scatter arange back through sortedIdx to invert the permutation
+  const V = lp.shape[lp.shape.length - 1]!;
+  const zerosIdx = ops.zeros(sortedIdx.shape, sortedIdx.dtype);
+  const ar = ops.arange(0, V, 1, sortedIdx.dtype);
+  const arB = ops.reshape(ar, sortedIdx.shape);
+  const inverse = ops.putAlongAxis(zerosIdx, sortedIdx, arB, -1);
+  const cumOrig = ops.takeAlongAxis(cum, inverse, -1);
+  const threshold = ops.scalarLike(1 - topP, cumOrig);
+  const keep = ops.less(threshold, cumOrig); // cum > 1 - topP
+  const ninf = negInfLike(lp);
+  const out = ops.where(keep, lp, ninf);
+  for (const a of [probs, sortedIdx, sortedProbs, cum, zerosIdx, ar, arB, inverse, cumOrig, threshold, keep, ninf])
+    a.dispose();
+  return out;
+}
+
+/** apply_top_k: all but the k highest logprobs → -inf. */
+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);
+  const maskIdx = part.slice([0, topK], [1, V]);
+  const ninf = negInfLike(lp);
+  const out = ops.putAlongAxis(lp, maskIdx, ninf, -1);
+  for (const a of [negLp, part, maskIdx, ninf]) a.dispose();
+  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;
+}