mlx-bun 0.0.12 → 0.2.0

package/src/sampler.ts

@@ -112,6 +112,75 @@ 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 interface StepSample<T> {
+  token: T;
+  extras: StepExtras | null;
+}
+
+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 interface DeviceStepSampler {
+  readonly isPlainGreedy: boolean;
+  readonly needsHistory: boolean;
+  sample(logits: MlxArray, step: number): StepSample<MlxArray>;
+  seedHistory(tokens: readonly number[]): void;
+  commitDevice(token: MlxArray): void;
+  commitNumbers(tokens: readonly number[]): void;
+  dispose(): void;
+}
+
+export interface NumberStepSampler {
+  readonly isPlainGreedy: boolean;
+  readonly needsHistory: boolean;
+  sample(logits: MlxArray, step: number): Promise<StepSample<number>>;
+  seedHistory(tokens: readonly number[]): void;
+  commitDevice(token: MlxArray): void;
+  commitNumbers(tokens: readonly number[]): void;
+  dispose(): void;
+}
+
 const GOLDEN = 0x9e3779b97f4a7c15n;
 
 function stepKey(seed: number, step: number): MlxArray {
@@ -717,6 +786,209 @@ export function makeLogitsProcessors(opts: LogitsProcessorOptions = {}): LogitsP
   return out;
 }
 
+/** Whether a whole [B,V] row can use the scheduler's vectorized argmax path. */
+export function isPlainGreedy(
+  opts: StepSamplerOptions,
+  hasProcessors = makeLogitsProcessors(opts).length > 0,
+): boolean {
+  return (opts.temperature ?? 0) === 0 && !opts.curve && !hasProcessors && !opts.grammar;
+}
+
+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();
+}
+
+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;
+  let history: MlxArray | null = null;
+
+  const seedHistory = (tokens: readonly number[]): void => {
+    if (processors.length === 0) return;
+    const next = ops.fromInt32([...tokens], [tokens.length]);
+    const previous = history;
+    history = next;
+    previous?.dispose();
+  };
+
+  const commitDevice = (token: MlxArray): void => {
+    if (processors.length === 0) return;
+    const one = ops.reshape(token, [1]);
+    if (!history) {
+      history = one;
+      return;
+    }
+    const previous = history;
+    try {
+      history = ops.concatAxis([previous, one], 0);
+      previous.dispose();
+    } finally {
+      one.dispose();
+    }
+  };
+
+  const commitNumbers = (tokens: readonly number[]): void => {
+    if (tokens.length === 0 || processors.length === 0) return;
+    const next = ops.fromInt32([...tokens], [tokens.length]);
+    if (!history) {
+      history = next;
+      return;
+    }
+    const previous = history;
+    try {
+      history = ops.concatAxis([previous, next], 0);
+      previous.dispose();
+    } finally {
+      next.dispose();
+    }
+  };
+
+  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 {
+      for (const processor of processors) {
+        const next = processor(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;
+      }
+      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);
+
+  const common = {
+    isPlainGreedy: isPlainGreedy(options, processors.length > 0),
+    needsHistory: processors.length > 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();
+      }
+    },
+  };
+}
+
 /** logits [1, V] → logprobs [1, V] (logits - logsumexp). */
 export function toLogprobs(logits: MlxArray): MlxArray {
   const lse = ops.logsumexpAxis(logits, -1, true);

package/src/serve/admin-routes.ts

@@ -0,0 +1,184 @@
+import type { JobStore } from "../jobs";
+
+export interface AdminRouteDeps {
+  ensureJobs(): Promise<JobStore>;
+  invalidateLibrary(): void;
+}
+
+export type AdminRoute =
+  | { kind: "hf-token-get" }
+  | { kind: "hf-token-set" }
+  | { kind: "tool-approvals-get" }
+  | { kind: "tool-approvals-delete" }
+  | { kind: "push"; artifactKind: "quantize" | "finetune" | "dataset" }
+  | { kind: "gc-plan" }
+  | { kind: "gc-execute" }
+  | { kind: "jobs-list" }
+  | { kind: "job-get"; jobId: string; stream: boolean };
+
+export function matchAdminRoute(method: string, pathname: string): AdminRoute | null {
+  switch (`${method} ${pathname}`) {
+    case "GET /api/settings/hf-token": return { kind: "hf-token-get" };
+    case "POST /api/settings/hf-token": return { kind: "hf-token-set" };
+    case "GET /api/settings/tool-approvals": return { kind: "tool-approvals-get" };
+    case "DELETE /api/settings/tool-approvals": return { kind: "tool-approvals-delete" };
+    case "GET /api/gc/plan": return { kind: "gc-plan" };
+    case "POST /api/gc/execute": return { kind: "gc-execute" };
+    case "GET /api/jobs": return { kind: "jobs-list" };
+  }
+
+  if (method === "POST") {
+    const push = pathname.match(/^\/api\/(quantize|finetune|dataset)\/push$/);
+    if (push) {
+      return {
+        kind: "push",
+        artifactKind: push[1] as "quantize" | "finetune" | "dataset",
+      };
+    }
+  }
+
+  if (method === "GET") {
+    const job = pathname.match(/^\/api\/jobs\/([^/]+?)(\/stream)?$/);
+    if (job) {
+      return {
+        kind: "job-get",
+        jobId: job[1]!,
+        stream: !!job[2],
+      };
+    }
+  }
+
+  return null;
+}
+
+export async function handleAdminRoute(
+  url: URL,
+  request: Request,
+  deps: AdminRouteDeps,
+): Promise<Response | null> {
+  const route = matchAdminRoute(request.method, url.pathname);
+  if (!route) return null;
+
+  switch (route.kind) {
+    case "hf-token-get": {
+      const { hasHfToken } = await import("../hf-push");
+      return Response.json({ ok: true, hasToken: hasHfToken() });
+    }
+    case "hf-token-set": {
+      const body = (await request.json().catch(() => ({}))) as { token?: string };
+      if (!body.token)
+        return Response.json({ ok: false, error: "token required" }, { status: 400 });
+      const { saveHfToken } = await import("../hf-push");
+      saveHfToken(body.token);
+      return Response.json({ ok: true });
+    }
+    case "tool-approvals-get": {
+      const { listAlwaysAllowedTools } = await import("../tool-approvals");
+      return Response.json({ ok: true, alwaysAllow: listAlwaysAllowedTools() });
+    }
+    case "tool-approvals-delete": {
+      const body = (await request.json().catch(() => ({}))) as { tool?: string };
+      if (!body.tool)
+        return Response.json({ ok: false, error: "tool required" }, { status: 400 });
+      const { revokeToolAlwaysAllowed } = await import("../tool-approvals");
+      const file = revokeToolAlwaysAllowed(body.tool);
+      return Response.json({ ok: true, alwaysAllow: Object.keys(file.allows).sort() });
+    }
+    case "push": {
+      const body = (await request.json().catch(() => ({}))) as {
+        job_id?: string;
+        repo_id?: string;
+        private?: boolean;
+        source_path?: string;
+      };
+      if (!body.repo_id)
+        return Response.json({ ok: false, error: "repo_id required" }, { status: 400 });
+      const { getHfToken, uploadFolder } = await import("../hf-push");
+      const token = getHfToken();
+      if (!token) {
+        return Response.json(
+          { ok: false, error: "no HF token saved — add one in Settings → Hugging Face" },
+          { status: 400 },
+        );
+      }
+      const store = await deps.ensureJobs();
+      let dir = body.source_path;
+      if (!dir && body.job_id) dir = store.get(body.job_id)?.output_path ?? undefined;
+      if (!dir) {
+        return Response.json(
+          { ok: false, error: "no source dir (pass job_id or source_path)" },
+          { status: 400 },
+        );
+      }
+      try {
+        const result = await uploadFolder(dir, body.repo_id, {
+          repoType: route.artifactKind === "dataset" ? "dataset" : "model",
+          private: !!body.private,
+          token,
+        });
+        return Response.json({ ok: true, url: result.url });
+      } catch (error) {
+        return Response.json({ ok: false, error: (error as Error).message }, { status: 400 });
+      }
+    }
+    case "gc-plan": {
+      const { planGc } = await import("../registry");
+      const plans = planGc().filter(
+        (plan) => plan.pruneSnapshots.length || plan.skippedSnapshots.length || plan.deadBlobs.length,
+      );
+      const superseded = plans.map((plan) => ({
+        repo_id: plan.repoId,
+        prune_snapshots: plan.pruneSnapshots.length,
+        skipped_snapshots: plan.skippedSnapshots.length,
+        dead_blobs: plan.deadBlobs.length,
+        reclaim_bytes: plan.reclaimBytes,
+      }));
+      const reclaim_bytes = plans.reduce((total, plan) => total + plan.reclaimBytes, 0);
+      return Response.json({ ok: true, superseded, reclaim_bytes });
+    }
+    case "gc-execute": {
+      const body = (await request.json().catch(() => ({}))) as { yes?: boolean };
+      if (body.yes !== true) {
+        return Response.json(
+          { ok: false, error: "pass {\"yes\": true} to confirm deletion" },
+          { status: 400 },
+        );
+      }
+      const { planGc, executeGc, Registry } = await import("../registry");
+      const plans = planGc().filter(
+        (plan) => plan.pruneSnapshots.length || plan.skippedSnapshots.length || plan.deadBlobs.length,
+      );
+      const result = executeGc(plans);
+      const registry = new Registry();
+      try {
+        await registry.scan();
+      } finally {
+        registry.close();
+      }
+      deps.invalidateLibrary();
+      return Response.json({
+        ok: true,
+        snapshots: result.snapshots,
+        blobs: result.blobs,
+        reclaimed_bytes: result.reclaimedBytes,
+      });
+    }
+    case "jobs-list": {
+      const store = await deps.ensureJobs();
+      const limit = Number(url.searchParams.get("limit") ?? "50");
+      const kind = url.searchParams.get("kind") ?? undefined;
+      return Response.json({ ok: true, jobs: store.recent(limit, kind) });
+    }
+    case "job-get": {
+      const store = await deps.ensureJobs();
+      if (route.stream) {
+        const { streamJobResponse } = await import("../jobs");
+        return streamJobResponse(store, route.jobId);
+      }
+      const job = store.get(route.jobId);
+      if (!job)
+        return Response.json({ ok: false, error: "job not found" }, { status: 404 });
+      return Response.json({ ok: true, job });
+    }
+  }
+}

package/src/serve/aux-routes.ts

@@ -0,0 +1,132 @@
+export type AuxiliaryRoute =
+  | "memory-status"
+  | "memory-list"
+  | "memory-search"
+  | "memory-article"
+  | "memory-links"
+  | "memory-history"
+  | "memory-diff"
+  | "memory-init"
+  | "memory-synthesize"
+  | "hub-local"
+  | "hub-search"
+  | "hub-download"
+  | "hub-serve"
+  | "sessions-search"
+  | "sessions-export";
+
+/** Match request-only routes that do not depend on the loaded model or
+ * createServer's live generation state. */
+export function matchAuxiliaryRoute(method: string, pathname: string): AuxiliaryRoute | null {
+  switch (`${method} ${pathname}`) {
+    case "GET /api/memory/status": return "memory-status";
+    case "GET /api/memory/list": return "memory-list";
+    case "GET /api/memory/search": return "memory-search";
+    case "GET /api/memory/article": return "memory-article";
+    case "GET /api/memory/links": return "memory-links";
+    case "GET /api/memory/history": return "memory-history";
+    case "GET /api/memory/diff": return "memory-diff";
+    case "POST /api/memory/init": return "memory-init";
+    case "GET /v1/memory/synthesize": return "memory-synthesize";
+    case "GET /api/hub/local": return "hub-local";
+    case "GET /api/hub/search": return "hub-search";
+    case "POST /api/hub/download": return "hub-download";
+    case "POST /api/hub/serve": return "hub-serve";
+    case "GET /api/sessions/search": return "sessions-search";
+    case "GET /api/sessions/export": return "sessions-export";
+    default: return null;
+  }
+}
+
+export async function handleAuxiliaryRoute(
+  url: URL,
+  request: Request,
+): Promise<Response | null> {
+  switch (matchAuxiliaryRoute(request.method, url.pathname)) {
+    case "memory-status": {
+      const { handleMemoryStatus } = await import("../memory/rest");
+      return handleMemoryStatus();
+    }
+    case "memory-list": {
+      const { handleMemoryList } = await import("../memory/rest");
+      return handleMemoryList();
+    }
+    case "memory-search": {
+      const { handleMemorySearch } = await import("../memory/rest");
+      return handleMemorySearch(url);
+    }
+    case "memory-article": {
+      const { handleMemoryArticle } = await import("../memory/rest");
+      return handleMemoryArticle(url);
+    }
+    case "memory-links": {
+      const { handleMemoryLinks } = await import("../memory/rest");
+      return handleMemoryLinks(url);
+    }
+    case "memory-history": {
+      const { handleMemoryHistory } = await import("../memory/rest");
+      return handleMemoryHistory(url);
+    }
+    case "memory-diff": {
+      const { handleMemoryDiff } = await import("../memory/rest");
+      return handleMemoryDiff(url);
+    }
+    case "memory-init": {
+      const { handleMemoryInit } = await import("../memory/rest");
+      return handleMemoryInit(request);
+    }
+    case "memory-synthesize": {
+      const dryRun = url.searchParams.get("dry") === "1";
+      const { runSynthesis } = await import("../memory/pipeline");
+      const encoder = new TextEncoder();
+      const stream = new ReadableStream({
+        async start(controller) {
+          const send = (event: unknown) =>
+            controller.enqueue(encoder.encode(`data: ${JSON.stringify(event)}\n\n`));
+          try {
+            const summary = await runSynthesis({ dryRun }, (event) => send(event));
+            send({ type: "summary", ...summary });
+            controller.enqueue(encoder.encode("data: [DONE]\n\n"));
+          } catch (error) {
+            send({ type: "error", message: (error as Error).message });
+          } finally {
+            controller.close();
+          }
+        },
+      });
+      return new Response(stream, {
+        headers: {
+          "content-type": "text/event-stream",
+          "cache-control": "no-cache",
+          connection: "keep-alive",
+        },
+      });
+    }
+    case "hub-local": {
+      const { handleHubLocal } = await import("../hub-rest");
+      return handleHubLocal();
+    }
+    case "hub-search": {
+      const { handleHubSearch } = await import("../hub-rest");
+      return handleHubSearch(url);
+    }
+    case "hub-download": {
+      const { handleHubDownload } = await import("../hub-rest");
+      return handleHubDownload(request);
+    }
+    case "hub-serve": {
+      const { handleHubServe } = await import("../hub-rest");
+      return handleHubServe(request);
+    }
+    case "sessions-search": {
+      const { handleSessionsSearch } = await import("./session-search");
+      return handleSessionsSearch(url);
+    }
+    case "sessions-export": {
+      const { handleSessionsExport } = await import("./session-search");
+      return handleSessionsExport(url);
+    }
+    default:
+      return null;
+  }
+}