@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/sampling/step.ts

@@ -0,0 +1,343 @@
+// 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 * as ops from "@mlx-bun/mlx/ops";
+import { runtimeFlag } from "../runtime/config";
+import { normalizedArgmax } from "../kernels/sampling/normalized-argmax";
+import { disposeStepExtras } from "./extras";
+import { applyTopKRows,applyTopPRows } from "./window-filters";
+import { logitsHistoryLimit,makeLogitsProcessors } from "./processors";
+import { makeSampler,stepKey,toLogprobs } from "./sampler";
+import { DeviceStepSampler,DeviceStepSamplerConfig,NumberStepSampler,NumberStepSamplerConfig,StepExtras,StepSample,StepSamplerOptions } from "./types";
+
+/** Whether scores can use stateless greedy sampling across independent rows. */
+export function isPlainGreedy(
+  opts: StepSamplerOptions,
+  hasProcessors = makeLogitsProcessors(opts).length > 0,
+): boolean {
+  return (opts.temperature ?? 0) === 0 && !opts.curve && !hasProcessors && !opts.grammar;
+}
+
+export function makeStepSampler(
+  options: StepSamplerOptions,
+  config: DeviceStepSamplerConfig,
+): DeviceStepSampler;
+export function makeStepSampler(
+  options: StepSamplerOptions,
+  config: NumberStepSamplerConfig,
+): NumberStepSampler;
+export function makeStepSampler(
+  options: StepSamplerOptions,
+  config: DeviceStepSamplerConfig | NumberStepSamplerConfig,
+): DeviceStepSampler | NumberStepSampler {
+  const sampler = config.sampler ?? makeSampler(options);
+  const processors = makeLogitsProcessors(options);
+  const grammar = options.grammar ?? null;
+  const captureSelected = config.captureSelectedLogprob === true;
+  const captureTop = Math.max(0, config.captureTopLogprobs ?? 0);
+  const capture = captureSelected || captureTop > 0;
+  const greedyWithoutMetadata = (options.temperature ?? 0) === 0 && !options.curve && !capture && !config.sampler;
+  const historyLimit = logitsHistoryLimit(options);
+  let history: MlxArray | null = null;
+
+  const seedHistory = (tokens: readonly number[]): void => {
+    if (historyLimit === 0) return;
+    if (tokens.length === 0) {
+      history?.dispose();
+      history = null;
+      return;
+    }
+    const recent = tokens.slice(Math.max(0, tokens.length - historyLimit));
+    const next = ops.fromInt32(recent, [recent.length]);
+    const previous = history;
+    history = next;
+    previous?.dispose();
+  };
+
+  // Takes ownership of next. Trim before concatenation so the appended array
+  // never exceeds the largest active finite window.
+  const appendHistory = (next: MlxArray): void => {
+    const previous = history;
+    if (!previous || next.size >= historyLimit) {
+      history = next;
+      previous?.dispose();
+      return;
+    }
+    const keep = Math.min(previous.size, historyLimit - next.size);
+    let tail: MlxArray | null = null;
+    try {
+      if (keep < previous.size) tail = previous.slice([previous.size - keep], [previous.size]);
+      history = ops.concatAxis([tail ?? previous, next], 0);
+      previous.dispose();
+    } finally {
+      tail?.dispose();
+      next.dispose();
+    }
+  };
+
+  const commitDevice = (token: MlxArray): void => {
+    if (historyLimit === 0) return;
+    appendHistory(ops.reshape(token, [1]));
+  };
+
+  const commitNumbers = (tokens: readonly number[]): void => {
+    if (tokens.length === 0 || historyLimit === 0) return;
+    const recent = tokens.slice(Math.max(0, tokens.length - historyLimit));
+    appendHistory(ops.fromInt32(recent, [recent.length]));
+  };
+
+  // Borrows `scores` [1, V]; returns them processed over `tokens` (the history the processors read), owned unless
+  // no processor replaced them.
+  const applyProcessors = (tokens: MlxArray | null, scores: MlxArray): MlxArray => {
+    let current = scores;
+    try {
+      for (const processor of processors) {
+        const next = processor(tokens, current);
+        if (current !== scores && next !== current) current.dispose();
+        current = next;
+      }
+      return current;
+    } catch (error) {
+      if (current !== scores) current.dispose();
+      throw error;
+    }
+  };
+
+  const sampleDevice = (input: MlxArray, step: number): StepSample<MlxArray> => {
+    const vocab = input.shape[input.shape.length - 1]!;
+    let current = input.shape.length === 2 ? input : ops.reshape(input, [1, vocab]);
+    let ownsCurrent = current !== input;
+    let logprobs: MlxArray | null = null;
+    let token: MlxArray | null = null;
+    let extras: StepExtras | null = null;
+    try {
+      if (processors.length) {
+        const next = applyProcessors(history, current);
+        if (ownsCurrent && next !== current) current.dispose();
+        current = next;
+        ownsCurrent = current !== input;
+      }
+      if (grammar && !grammar.isTerminated) {
+        const next = grammar.applyMask(current);
+        if (ownsCurrent && next !== current) current.dispose();
+        current = next;
+        ownsCurrent = current !== input;
+      }
+      if (greedyWithoutMetadata)
+        return { token: normalizedArgmax(current), extras: null };
+      logprobs = toLogprobs(current);
+      if (ownsCurrent) current.dispose();
+      current = input;
+      ownsCurrent = false;
+      token = sampler(logprobs, step);
+
+      if (capture) {
+        extras = { sel: null, topIdx: null, topVals: null };
+        if (captureSelected) {
+          const idx = ops.reshape(token, [1, 1]);
+          try {
+            extras.sel = ops.takeAlongAxis(logprobs, idx, -1);
+          } finally {
+            idx.dispose();
+          }
+        }
+        if (captureTop > 0) {
+          const k = Math.min(captureTop, vocab);
+          let neg: MlxArray | null = null;
+          let partition: MlxArray | null = null;
+          let view: MlxArray | null = null;
+          try {
+            neg = ops.neg(logprobs);
+            partition = ops.argpartitionAxis(neg, k - 1, -1);
+            view = partition.slice([0, 0], [1, k]);
+            extras.topIdx = ops.contiguous(view);
+            extras.topVals = ops.takeAlongAxis(logprobs, extras.topIdx, -1);
+          } finally {
+            neg?.dispose();
+            partition?.dispose();
+            view?.dispose();
+          }
+        }
+      }
+
+      logprobs.dispose();
+      logprobs = null;
+      const result = { token, extras };
+      token = null;
+      extras = null;
+      return result;
+    } finally {
+      if (ownsCurrent) current.dispose();
+      logprobs?.dispose();
+      token?.dispose();
+      disposeStepExtras(extras);
+    }
+  };
+
+  if (config.initialHistory) seedHistory(config.initialHistory);
+
+  // Per-position sampling for a verify window: with no grammar or capture, position p's draw depends only on its
+  // scores, its step key and its processor history, which is this request's history extended by the window's inputs
+  // up to p (the drafts the walk accepted before reaching p). Every position of the window can then sit in ONE device
+  // graph and read back together. Row p reproduces sample(scores[p], steps[p]) exactly, processors included: they run
+  // per position on its own [1, V] row over that history, as sample() runs them.
+  // The plain filter chain (top-p -> top-k -> temperature -> categorical) can run
+  // its filters once over the whole window. Other samplers keep per-row calls.
+  const windowFilters = !config.sampler && !options.curve && options.hlg?.enabled !== true &&
+    (options.minP ?? 0) === 0 && (options.xtcProbability ?? 0) === 0 && (options.temperature ?? 0) > 0 &&
+    runtimeFlag("MLX_BUN_SAMPLER_WINDOW_FILTERS", true)
+    ? { temperature: options.temperature!, topP: options.topP ?? 0, topK: options.topK ?? 0, seed: options.seed ?? 0 }
+    : null;
+  const positional = !grammar && !capture
+    ? {
+      sample(scores: MlxArray, steps: readonly number[], inputs?: MlxArray): MlxArray {
+        if (scores.shape.length !== 2 || scores.shape[0] !== steps.length)
+          throw new Error(`positional sampling expects [${steps.length}, V] scores, got [${scores.shape.join(", ")}]`);
+        if (historyLimit > 0 && inputs?.size !== steps.length)
+          throw new Error(`positional sampling with processor history needs the window's ${steps.length} inputs`);
+        const vocab = scores.shape[1]!;
+        const tokens: MlxArray[] = [];
+        const processed: MlxArray[] = [];
+        let packed: MlxArray | null = null, stacked: MlxArray | null = null;
+        try {
+          // Each position's scores after its processors, over the history sample() would hold there.
+          const rowScores = (position: number): MlxArray => {
+            const row = scores.slice([position, 0], [position + 1, vocab]);
+            if (!processors.length) return row;
+            let extended: MlxArray | null = null;
+            try {
+              if (historyLimit > 0 && position > 0) {
+                using drafts = inputs!.slice([1], [position + 1]);
+                extended = history ? ops.concatAxis([history, drafts], 0) : ops.contiguous(drafts);
+              }
+              const next = applyProcessors(extended ?? history, row);
+              if (next !== row) row.dispose();
+              return next;
+            } catch (error) { row.dispose(); throw error; }
+            finally { extended?.dispose(); }
+          };
+          if (windowFilters && steps.length > 1) {
+            // One filter chain over [W, V] (bit-identical per row to the [1, V]
+            // chain, packages/inference/tests/sampling/window.test.ts), then each row's own
+            // keyed draw on its [1, V] slice: that shape keeps MLX's inverse-CDF
+            // categorical, which the draft's coupled draw relies on.
+            if (processors.length) {
+              for (let position = 0; position < steps.length; position++) processed.push(rowScores(position));
+              stacked = ops.concatAxis(processed, 0);
+            }
+            using logprobs = toLogprobs(stacked ?? scores);
+            let current = logprobs, owned: MlxArray | null = null;
+            try {
+              if (windowFilters.topP > 0 && windowFilters.topP < 1) { owned = applyTopPRows(current, windowFilters.topP); current = owned; }
+              if (windowFilters.topK > 0) {
+                const next = applyTopKRows(current, windowFilters.topK);
+                owned?.dispose(); owned = next; current = next;
+              }
+              for (const [position, step] of steps.entries()) {
+                using row = current.slice([position, 0], [position + 1, vocab]);
+                using scaled = ops.mulScalar(row, 1 / windowFilters.temperature);
+                using key = stepKey(windowFilters.seed, step);
+                tokens.push(ops.randomCategorical(scaled, key));
+              }
+            } finally { owned?.dispose(); }
+          } else
+          for (const [position, step] of steps.entries()) {
+            using row = rowScores(position);
+            if (greedyWithoutMetadata) { tokens.push(normalizedArgmax(row)); continue; }
+            using logprobs = toLogprobs(row);
+            tokens.push(sampler(logprobs, step));
+          }
+          if (tokens.length === 1) return tokens.pop()!;
+          packed = ops.concatAxis(tokens, 0);
+          const result = packed;
+          packed = null;
+          return result;
+        } finally {
+          packed?.dispose();
+          stacked?.dispose();
+          for (const array of [...tokens, ...processed]) array.dispose();
+        }
+      },
+      commit(tokens: readonly number[]): void {
+        if (config.historyUpdate === "after-sample") commitNumbers(tokens);
+      },
+    }
+    : undefined;
+
+  const unread = !grammar && !capture
+    ? {
+      sample(scores: MlxArray, step: number): MlxArray {
+        const { token } = sampleDevice(scores, step);
+        if (config.historyUpdate === "after-sample") commitDevice(token);
+        return token;
+      },
+    }
+    : undefined;
+
+  const common = {
+    unread,
+    independent: isPlainGreedy(options, processors.length > 0) && !capture && !config.sampler
+      ? independentGreedySampling : undefined,
+    positional,
+    isPlainGreedy: isPlainGreedy(options, processors.length > 0),
+    capturesLogprobs: capture,
+    needsHistory: historyLimit > 0,
+    seedHistory,
+    commitDevice,
+    commitNumbers,
+    dispose: () => {
+      history?.dispose();
+      history = null;
+    },
+  };
+
+  if (config.tokenRepresentation === "device") {
+    return {
+      ...common,
+      sample(logits, step) {
+        const result = sampleDevice(logits, step);
+        if (config.historyUpdate === "after-sample") commitDevice(result.token);
+        return result;
+      },
+    };
+  }
+
+  return {
+    ...common,
+    async sample(logits, step) {
+      if (grammar && !grammar.isTerminated) await grammar.ready();
+      const result = sampleDevice(logits, step);
+      let tokenArray: MlxArray | null = result.token;
+      try {
+        // Read the sampler's integer result directly. Adding an astype here
+        // would append work behind the next dispatched decode step.
+        const token = ops.itemUint32(tokenArray);
+        tokenArray.dispose();
+        tokenArray = null;
+        if (config.acceptGrammar && grammar && !config.eosTokenIds?.includes(token))
+          grammar.accept(token);
+        if (config.historyUpdate === "after-sample") commitNumbers([token]);
+        return { token, extras: result.extras };
+      } catch (error) {
+        disposeStepExtras(result.extras);
+        throw error;
+      } finally {
+        tokenArray?.dispose();
+      }
+    },
+  };
+}
+
+/** Stateless sampling across rows or verification positions. Preserve the
+ * normalized-score argmax, including rounding-created ties. */
+export const independentGreedySampling = Object.freeze({
+  sample: normalizedArgmax,
+});

package/src/sampling/types.ts

@@ -0,0 +1,161 @@
+// 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 { SampledToken,SamplingSession } from "../contracts/portable/sampling";
+import { type CurveParams } from "./curve";
+
+/** Resolved HLG sampling config (the user-facing knobs). The mid gain is NOT
+ *  here — it folds from temperature (m = 1/T) when the sampler is built. */
+export interface HlgConfig {
+  enabled: boolean;
+  width: number;
+  shoulder: number;
+  toe: number;
+  pivot?: "top" | "entropy" | "median";
+  pivotOffset: number;
+  /** Use the LITERAL HLG OETF (gamma below 1/12, log above) instead of the
+   *  parametric toe/mid/shoulder curve. Ignores gain/width/shoulder/toe/pivot. */
+  oetf?: boolean;
+  /** Use the LITERAL HLG EOTF (inverse OETF → OOTF gamma) — the decode/display
+   *  direction (suppresses the tail). Ignores the parametric knobs. */
+  eotf?: boolean;
+  /** EOTF only: nominal peak display luminance L_W (cd/m²) → OOTF gamma.
+   *  Default 1000 (γ = 1.2). Lower = flatter (more diverse), higher = sharper. */
+  lw?: number;
+  /** Use the FULL HLG chain: OETF signal shape → OOTF scaling (applyHlgPipeline).
+   *  γ from `lw`. Ignores the parametric knobs. */
+  pipeline?: boolean;
+  /** Pipeline only: α = max brightness — the top logit the OOTF scales everything
+   *  under ("the most confidence we want"). Default 5. For the shaper, reused as
+   *  `out_scale` (default 12). */
+  maxBrightness?: number;
+  /** Use the user-specified HLGShaper (windowed anchor → piecewise OETF with a
+   *  cubic suppress-toe → OOTF → ×out_scale). `lw`→L_W, `maxBrightness`→out_scale. */
+  shaper?: boolean;
+  /** Shaper only: W = nats of headroom below the top logit that span the curve
+   *  (`x = clamp((ℓ−ℓmax)/W + 1, 0, 1)`). Default 10. */
+  window?: number;
+  /** Shaper only: s_m — mid sharpness (how hard above-median tokens are elevated). */
+  sM?: number;
+  /** Shaper only: A — shoulder compression (how hard top-end confidence is tamed). */
+  shoulderA?: number;
+  /** Shaper only: target top-to-reference logit gap (out_scale auto-derives from this
+   *  so it's decoupled from W). Default 15. */
+  targetGap?: number;
+  /** Shaper only: reference token nats below the top for the auto out_scale. Default 4. */
+  refGap?: number;
+  /** Shaper only: curve geometry (the pivots / toe / mid power). Default to the
+   *  HLGShaper constants (xM 0.55, yM 0.5, xFloor 0.2, yFloor 0.18, p 2.0) when unset. */
+  xM?: number; yM?: number; xFloor?: number; yFloor?: number; p?: number;
+  /** Optional explicit mid gain m. Default: folds from temperature (m = 1/T),
+   *  so temperature stays the contrast knob. An explicit value decouples mid
+   *  contrast from temperature — needed to probe mid-boost while holding the
+   *  model's recommended temperature fixed. */
+  gain?: number;
+}
+
+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
+   *  ⇒ the plain temperature path, unchanged. `02d723a:docs/archive/hlg-sampling.md`. */
+  hlg?: HlgConfig;
+  /** v2 log-prob transfer-curve sampler. When set, the drawn monotone curve
+   *  REPLACES temperature+softmax entirely (stochastic, seeded) — see
+   *  packages/inference/src/sampling/curve.ts. Identity curve ≡ temperature 1. */
+  curve?: CurveParams;
+}
+
+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]). */
+export type Sampler = (logprobs: MlxArray, step: number) => MlxArray;
+
+/** (deviceTokens [n] | null, logits [1, V]) → logits [1, V]. */
+export type LogitsProcessor = (tokens: MlxArray | null, logits: MlxArray) => MlxArray;
+
+/** Optional arrays captured from the same post-processor logprobs used to
+ * sample a token. The caller controls readback timing and owns disposal. */
+export interface StepExtras {
+  sel: MlxArray | null;
+  topIdx: MlxArray | null;
+  topVals: MlxArray | null;
+}
+
+export type StepSample<T> = SampledToken<T, StepExtras>;
+
+interface StepGrammar {
+  readonly isTerminated: boolean;
+  ready(): Promise<void>;
+  applyMask(logits: MlxArray): MlxArray;
+  accept(token: number): void;
+}
+
+export interface StepSamplerOptions extends SamplerOptions, LogitsProcessorOptions {
+  grammar?: StepGrammar | null;
+}
+
+interface StepSamplerConfigBase {
+  /** The grammar owner already awaited ready(), or this sampler must do it. */
+  grammarWait: "external" | "before-sample";
+  /** Update processor history in sample(), or let the decode loop commit later. */
+  historyUpdate: "after-sample" | "manual";
+  initialHistory?: readonly number[];
+  captureSelectedLogprob?: boolean;
+  captureTopLogprobs?: number;
+  /** Eager-number mode can advance grammar after sampling. */
+  acceptGrammar?: boolean;
+  eosTokenIds?: readonly number[];
+  /** Reuse a sampler already shared with a speculative draft source. */
+  sampler?: Sampler;
+}
+
+export interface DeviceStepSamplerConfig extends StepSamplerConfigBase {
+  tokenRepresentation: "device";
+  grammarWait: "external";
+}
+
+export interface NumberStepSamplerConfig extends StepSamplerConfigBase {
+  tokenRepresentation: "number";
+  grammarWait: "before-sample";
+}
+
+export type DeviceStepSampler = SamplingSession<MlxArray, MlxArray, StepSample<MlxArray>>;
+export type NumberStepSampler = SamplingSession<MlxArray, MlxArray, Promise<StepSample<number>>>;

package/src/sampling/window-filters.ts

@@ -0,0 +1,46 @@
+// Row-batched top-p / top-k filtering for a verify window.
+//
+// The accept walk samples every position of a speculative window. Each position
+// runs the same filter chain over the full vocabulary: exp, argsort, gather,
+// cumsum, inverse permutation, compare (top-p), then argpartition and a scatter
+// (top-k). Done per position that is W independent chains of full-vocabulary
+// kernels. Every one of those ops is row-wise along the last axis, so the chain
+// over [W, V] computes each row with the same arithmetic as the [1, V] chain.
+// Only the FILTERS batch: each row is still drawn separately with its own step
+// key on a [1, V] slice, because MLX's categorical takes its inverse-CDF path
+// only for that shape, and the draft is coupled to exactly that draw.
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+function negInfLike(a: MlxArray): MlxArray {
+  return ops.scalarLike(-Infinity, a);
+}
+
+/** apply_top_p over rows: logprobs [W, V] -> logprobs with the tail at -inf. */
+export function applyTopPRows(lp: MlxArray, topP: number): MlxArray {
+  const V = lp.shape[1]!;
+  using probs = ops.exp(lp);
+  using sortedIdx = ops.argsortAxis(lp, -1);
+  using sortedProbs = ops.takeAlongAxis(probs, sortedIdx, -1);
+  using cum = ops.cumsum(sortedProbs, -1);
+  using zerosIdx = ops.zeros(sortedIdx.shape, sortedIdx.dtype);
+  using ar = ops.arange(0, V, 1, sortedIdx.dtype);
+  using arRow = ops.reshape(ar, [1, V]);
+  using arRows = ops.add(zerosIdx, arRow); // integer broadcast: exact
+  using inverse = ops.putAlongAxis(zerosIdx, sortedIdx, arRows, -1);
+  using cumOrig = ops.takeAlongAxis(cum, inverse, -1);
+  using threshold = ops.scalarLike(1 - topP, cumOrig);
+  using keep = ops.less(threshold, cumOrig);
+  using ninf = negInfLike(lp);
+  return ops.where(keep, lp, ninf);
+}
+
+/** apply_top_k over rows: all but each row's k highest logprobs -> -inf. */
+export function applyTopKRows(lp: MlxArray, topK: number): MlxArray {
+  const [W, V] = lp.shape as [number, number];
+  using negLp = ops.neg(lp);
+  using part = ops.argpartitionAxis(negLp, topK - 1, -1);
+  using maskIdx = part.slice([0, topK], [W, V]);
+  using ninf = negInfLike(lp);
+  return ops.putAlongAxis(lp, maskIdx, ninf, -1);
+}

package/src/scoring/decisions.ts

@@ -0,0 +1,96 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../contracts/mlx/cache";
+import type { MlxHeadForward, MlxSequenceForward } from "../contracts/mlx/graph";
+import type { MlxPromptInput } from "../contracts/mlx/media";
+import { logitsFromHiddenPadM, sequenceLogProbabilities } from "./response-logps";
+
+/** The operations a closed-answer decision consumes. The graph and its weights
+ * are borrowed; each call owns fresh state. No sampling or generation is required. */
+export interface DecisionScoringGraph {
+  makeCache(): Cache[];
+  forward: MlxSequenceForward;
+  logits: MlxHeadForward;
+}
+
+/** Teacher-forced conditional likelihoods for complete supplied answers.
+ * Each forward reads [prompt; answer[:-1]]. Only hidden rows that predict answer
+ * tokens reach the head: prompt targets are masked out. Average log p(token | prefix)
+ * over each answer (negative mean NLL), then softmax across candidate averages. No vocabulary restriction
+ * changes a token's likelihood. Include a terminator in each answer to score
+ * complete, mutually exclusive responses (especially when one prefixes another).
+ *
+ * Inputs and graph are borrowed; every candidate gets fresh caches, released with
+ * its intermediates on success, error or cancellation. The caller serializes
+ * shared graph access. This is inference, with no training or reference model.
+ * The average includes the terminator and excludes every prompt token. */
+export async function scoreDecisionAnswers(graph: DecisionScoringGraph, prompt: readonly number[],
+  candidates: readonly (readonly number[])[], options: { signal?: AbortSignal; input?: MlxPromptInput;
+    /** Run one candidate's owned work at a scheduler boundary. Direct library
+     * callers may omit this when they already own exclusive graph access. */
+    execute?: <T>(work: () => Promise<T>) => Promise<T>;
+  } = {}): Promise<{
+    probabilities: number[]; meanLogLikelihoods: number[];
+  }> {
+  options.signal?.throwIfAborted();
+  if (!prompt.length) throw new RangeError("a decision prompt must contain tokens");
+  if (!candidates.length || candidates.some(answer => !answer.length || answer.some(id => !Number.isSafeInteger(id) || id < 0)))
+    throw new RangeError("a decision requires nonempty candidate token sequences");
+  const meanLogLikelihoods: number[] = [];
+  for (const answer of candidates) {
+    options.signal?.throwIfAborted();
+    const score = await (options.execute ?? (work => work()))(async () => {
+      options.signal?.throwIfAborted();
+      const caches = graph.makeCache();
+      try {
+        const tokens = [...prompt, ...answer.slice(0, -1)];
+        using ids = ops.fromInt32(tokens, [1, tokens.length]);
+        using hidden = options.input ? mediaResponseHidden(options.input, prompt, answer, caches, options.signal)
+          : graph.forward(ids, caches, tokens.length);
+        const [batch, length, width] = hidden.shape;
+        if (batch !== 1 || length !== (options.input ? answer.length : tokens.length) || !width)
+          throw new Error("decision forward returned an invalid hidden shape");
+        // Prediction at prompt.length - 1 targets answer[0], not prompt's last token.
+        using response = hidden.slice([0, options.input ? 0 : prompt.length - 1, 0], [1, length, width]);
+        using logits = logitsFromHiddenPadM({ logitsFromHidden: graph.logits }, response);
+        const vocab = logits.shape.at(-1)!;
+        if (answer.some(id => id >= vocab)) throw new RangeError("decision token id is outside the graph vocabulary");
+        // The DPO reduction sees only supervised rows; equivalent to a zero mask
+        // on all earlier prompt positions, without projecting them to vocabulary.
+        using logp = sequenceLogProbabilities(logits, [[prompt.at(-1)!, ...answer]], [[0, ...answer.map(() => 1)]]);
+        const score = logp.toFloat32()[0]! / answer.length;
+        options.signal?.throwIfAborted();
+        if (!Number.isFinite(score)) throw new Error("decision answer has a nonfinite log likelihood");
+        return score;
+      } finally {
+        const errors: unknown[] = [];
+        for (const cache of caches) { try { cache.dispose(); } catch (error) { errors.push(error); } }
+        if (errors.length) throw new AggregateError(errors, "decision cache cleanup failed");
+      }
+    });
+    meanLogLikelihoods.push(score);
+    await new Promise<void>(resolve => setImmediate(resolve));
+  }
+  options.signal?.throwIfAborted();
+  const maximum = Math.max(...meanLogLikelihoods), weights = meanLogLikelihoods.map(value => Math.exp(value - maximum));
+  const total = weights.reduce((sum, value) => sum + value, 0);
+  return { meanLogLikelihoods, probabilities: weights.map(value => value / total) };
+}
+
+/** Media ports expose the prompt tip and position-aware continuation. Feed
+ * supplied answer tokens through those operations, collecting their prediction
+ * rows; no sampling, generated ids, or generation state is involved. */
+function mediaResponseHidden(input: MlxPromptInput, prompt: readonly number[], answer: readonly number[],
+  caches: Cache[], signal?: AbortSignal): MlxArray {
+  const rows: MlxArray[] = [];
+  try {
+    using ids = ops.fromInt32([...prompt], [1, prompt.length]);
+    rows.push(input.prefill(ids, caches, 0, "tail", prompt.length));
+    for (let i = 0; i < answer.length - 1; i++) {
+      signal?.throwIfAborted();
+      using token = ops.fromInt32([answer[i]!], [1, 1]);
+      rows.push(input.prefill(token, caches, prompt.length + i, "decode", 1));
+    }
+    return ops.concatAxis(rows, 1);
+  } finally { for (const row of rows) row.dispose(); }
+}