mlx-bun 0.0.8 → 0.0.9
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +112 -20
- package/docs/design/batching-v2-plan.md +515 -0
- package/docs/design/bucketing-stage.md +196 -0
- package/docs/design/compat-cli-surface-design.md +11 -11
- package/docs/design/docs-reorg-plan.md +1 -1
- package/docs/design/dreaming-nightly-pipeline.md +291 -0
- package/docs/design/dspark-speculative-decoding.md +106 -0
- package/docs/design/generic-model-support.md +381 -0
- package/docs/design/memory-inference-path.md +422 -0
- package/docs/design/memory-synthesis.md +287 -0
- package/docs/design/minicpm5-decode-megakernel.md +484 -0
- package/docs/design/mlx-lm-tool-parity-plan.md +535 -0
- package/docs/design/orpo-dynamic-lambda.md +41 -0
- package/docs/design/orpo-training.md +35 -5
- package/docs/design/parallel-slots.md +12 -1
- package/docs/design/parity-tier-dag.md +13 -5
- package/docs/design/spec-decode-larger-targets.md +229 -0
- package/docs/design/tauri-desktop-app.md +2 -2
- package/docs/design/the-dreaming-handoff.md +141 -0
- package/docs/design/the-dreaming-master-plan.md +446 -0
- package/docs/design/write-pipeline-entity.md +468 -0
- package/docs/investigations/curve-bisector-routing.md +1 -1
- package/docs/investigations/decode-roofline-lookagain.md +262 -0
- package/docs/investigations/dspark-handoff.md +58 -0
- package/docs/investigations/kernel-perf-review-2026-07.md +92 -0
- package/docs/planning/PRODUCT_ROADMAP.md +2 -2
- package/docs/planning/curve-sampler-research-plan.md +331 -0
- package/docs/planning/memory-docs-and-dag-plan.md +185 -0
- package/docs/planning/pi-integration-review.md +178 -0
- package/docs/planning/release-notes-v0.0.9.md +122 -0
- package/docs/planning/web-ui-pass-plan.md +213 -0
- package/docs/planning/website-readme-pass-plan.md +224 -0
- package/docs/reference/cli.md +255 -0
- package/docs/reference/embedding.md +7 -2
- package/docs/reference/library-api.md +22 -3
- package/docs/reference/memory.md +5 -5
- package/docs/reference/models.md +137 -0
- package/docs/reference/orpo-quickstart.md +6 -0
- package/docs/reference/server-api.md +187 -9
- package/docs/reference/server-config.md +40 -9
- package/docs/reference/training.md +25 -11
- package/package.json +1 -1
- package/src/anthropic.ts +22 -0
- package/src/assets/curve-designer.html +1 -1
- package/src/chat-template.ts +25 -7
- package/src/cli.ts +1029 -86
- package/src/config.ts +4 -1
- package/src/download.ts +99 -8
- package/src/embed.ts +77 -0
- package/src/eval/perplexity.ts +164 -0
- package/src/eval/runner.ts +15 -6
- package/src/eval/tasks/bfcl.ts +1 -1
- package/src/generate.ts +141 -18
- package/src/harness-pi.ts +43 -17
- package/src/index.ts +2 -0
- package/src/lora.ts +2 -2
- package/src/memory/article.ts +558 -0
- package/src/memory/chunk-validation.ts +150 -0
- package/src/memory/chunk.ts +305 -12
- package/src/memory/cluster.ts +344 -17
- package/src/memory/coverage.ts +134 -0
- package/src/memory/crosslink.ts +0 -0
- package/src/memory/db.ts +184 -45
- package/src/memory/entity.ts +203 -0
- package/src/memory/gate.ts +187 -0
- package/src/memory/history.ts +152 -0
- package/src/memory/ingest.ts +324 -11
- package/src/memory/model.ts +321 -0
- package/src/memory/normalize.ts +671 -0
- package/src/memory/parse.ts +22 -0
- package/src/memory/pipeline.ts +242 -30
- package/src/memory/prompts.ts +42 -0
- package/src/memory/query.ts +334 -0
- package/src/memory/reconcile.ts +734 -0
- package/src/memory/redirect.ts +108 -0
- package/src/memory/reindex.ts +325 -0
- package/src/memory/resolve.ts +307 -0
- package/src/memory/route.ts +272 -0
- package/src/memory/schedule.ts +6 -4
- package/src/memory/stages.ts +774 -0
- package/src/memory/synthesize.ts +1629 -20
- package/src/memory/tools.ts +446 -96
- package/src/memory/vault.ts +462 -35
- package/src/memory/wikify.ts +1057 -16
- package/src/mlx/ffi.ts +3 -0
- package/src/mlx/ops.ts +46 -0
- package/src/model/compiled-decode.ts +43 -8
- package/src/model/factory.ts +29 -8
- package/src/model/flash-attention.ts +4 -2
- package/src/model/fused-decode-kernel.ts +14 -7
- package/src/model/fused-mlp-kernel.ts +169 -0
- package/src/model/fused-swiglu-kernel.ts +240 -0
- package/src/model/gemma4-base.ts +62 -6
- package/src/model/gemma4.ts +33 -2
- package/src/model/minicpm5.ts +106 -8
- package/src/model/qwen3.ts +276 -0
- package/src/model/steel-linear-kernel.ts +96 -0
- package/src/model/support.ts +38 -8
- package/src/model/universal/archs.ts +332 -0
- package/src/model/universal/dense.ts +576 -0
- package/src/model/universal/modules.ts +163 -0
- package/src/model/universal/rope.ts +366 -0
- package/src/pi-provider.ts +1 -1
- package/src/pi-terminal.ts +29 -8
- package/src/pi-web.ts +19 -2
- package/src/quantize/config-writer.ts +1 -1
- package/src/registry.ts +215 -7
- package/src/responses.ts +23 -0
- package/src/sampler.ts +220 -9
- package/src/serve/batch-scheduler.ts +272 -80
- package/src/serve/generation-gateway.ts +86 -18
- package/src/server.ts +612 -49
- package/src/spec/drafter.ts +81 -49
- package/src/spec/dspark/data-dflash.ts +156 -0
- package/src/spec/dspark/data.ts +201 -0
- package/src/spec/dspark/generate-dflash.ts +193 -0
- package/src/spec/dspark/generate.ts +290 -0
- package/src/spec/dspark/loss.ts +213 -0
- package/src/spec/dspark/module-dflash.ts +361 -0
- package/src/spec/dspark/module.ts +540 -0
- package/src/spec/dspark/sample.ts +112 -0
- package/src/spec/generate.ts +40 -19
- package/src/status-page.html +1 -1
- package/src/tokenizer.ts +7 -0
- package/src/train/flash-cce.ts +70 -15
- package/src/train/fuse.ts +243 -0
- package/src/train/index.ts +3 -0
- package/src/train/job.ts +8 -0
- package/src/train/loss.ts +289 -30
- package/src/train/prefix-shared.ts +33 -5
- package/src/train/segmented.ts +67 -11
- package/src/train/trainer.ts +92 -28
- package/src/web/skills/memory/SKILL.md +40 -17
package/src/server.ts
CHANGED
|
@@ -26,8 +26,8 @@ import { Gemma4Model } from "./model/gemma4";
|
|
|
26
26
|
import { DiffusionGemmaModel } from "./model/diffusion-gemma";
|
|
27
27
|
import { spliceImageTokens } from "./vision/diffusion-vision";
|
|
28
28
|
import { createModel, type RuntimeModel } from "./model/factory";
|
|
29
|
-
import { isMiniCPM5Config,
|
|
30
|
-
import { generate, type GenerateOptions } from "./generate";
|
|
29
|
+
import { isMiniCPM5Config, isSupportedModelRecord } from "./model/support";
|
|
30
|
+
import { generate, type GenerateOptions, type TokenLogprobs } from "./generate";
|
|
31
31
|
import type { HlgConfig } from "./sampler";
|
|
32
32
|
import { isMonotone, CURVE_UMIN, type CurveParams } from "./curve-sampler";
|
|
33
33
|
const CURVE_PAGE = curveDesignerHtml as unknown as string;
|
|
@@ -51,6 +51,7 @@ import {
|
|
|
51
51
|
type ResponsesRequest,
|
|
52
52
|
} from "./responses";
|
|
53
53
|
import { AdapterManager, listAvailableAdapters } from "./lora";
|
|
54
|
+
import { embedMany, isEmbeddingModel } from "./embed";
|
|
54
55
|
import { fit } from "./fit";
|
|
55
56
|
import { setMemoryLimit } from "./mlx/ffi";
|
|
56
57
|
import { VisionTower } from "./vision/embedder";
|
|
@@ -85,8 +86,10 @@ export interface ServerOptions {
|
|
|
85
86
|
* it). Exposed at /stats.server so other mlx-bun processes can
|
|
86
87
|
* warn before attaching to a server that may vanish. */
|
|
87
88
|
owner?: "serve" | "pi-session" | "embedded";
|
|
88
|
-
/** Interface to bind (Bun.serve hostname).
|
|
89
|
-
* (all interfaces)
|
|
89
|
+
/** Interface to bind (Bun.serve hostname). Unset ⇒ Bun's default
|
|
90
|
+
* (all interfaces) — embedded/library use. The CLI always passes one:
|
|
91
|
+
* "127.0.0.1" (loopback, mlx_lm.server parity) unless --host says
|
|
92
|
+
* otherwise. */
|
|
90
93
|
hostname?: string;
|
|
91
94
|
/** Server-wide default for the chat template's `enable_thinking`
|
|
92
95
|
* variable (MiniCPM5/CPM and other hybrid-reasoning models). A
|
|
@@ -101,6 +104,10 @@ export interface ServerOptions {
|
|
|
101
104
|
defaultTemperature?: number;
|
|
102
105
|
defaultTopP?: number;
|
|
103
106
|
defaultTopK?: number;
|
|
107
|
+
/** Completion cap when the request omits max_tokens (`--max-tokens`).
|
|
108
|
+
* mlx_lm.server's flag; its default there is 512 — ours stays 65,536 so
|
|
109
|
+
* thinking traces never truncate. `--max-tokens 512` = mlx-lm behavior. */
|
|
110
|
+
defaultMaxTokens?: number;
|
|
104
111
|
/** Max concurrent requests batched through the mlx-lm-parity engine
|
|
105
112
|
* (`--batch N`). Default 1 = today's serialized single-queue path. >1 opts
|
|
106
113
|
* the WHOLE server into the batched engine (continuous batching, B floats
|
|
@@ -111,6 +118,11 @@ export interface ServerOptions {
|
|
|
111
118
|
/** HLG tone-curve sampling default (set via --hlg-sampling on + sub-knobs).
|
|
112
119
|
* A per-request `hlg` object overrides it field-by-field. Off when unset. */
|
|
113
120
|
hlg?: HlgConfig;
|
|
121
|
+
/** Adapter id mounted at startup via `serve --adapter <dir>` (mlx-lm's
|
|
122
|
+
* `--adapter-path`). Used when a request sends no `adapter` field; a
|
|
123
|
+
* request's explicit `adapter` (including "none") always wins. Hot-swap
|
|
124
|
+
* via /v1/adapters is unchanged. */
|
|
125
|
+
defaultAdapter?: string;
|
|
114
126
|
}
|
|
115
127
|
|
|
116
128
|
export interface ServerContext {
|
|
@@ -262,6 +274,34 @@ interface ChatRequest {
|
|
|
262
274
|
top_k?: number;
|
|
263
275
|
seed?: number;
|
|
264
276
|
repetition_penalty?: number;
|
|
277
|
+
/** mlx-lm extension: recent-token window for repetition_penalty
|
|
278
|
+
* (default 20; 0 = whole history, Python `[-0:]` semantics). */
|
|
279
|
+
repetition_context_size?: number;
|
|
280
|
+
/** min-p sampling (mlx_lm.server's `min_p`): keep tokens whose probability
|
|
281
|
+
* is ≥ min_p · p(top token). 0 = off. */
|
|
282
|
+
min_p?: number;
|
|
283
|
+
/** XTC sampling (mlx_lm.server names): with probability `xtc_probability`
|
|
284
|
+
* per step, remove every token above `xtc_threshold` except the least
|
|
285
|
+
* likely of them. EOS + the newline token are always exempt (the server
|
|
286
|
+
* injects them as xtc special tokens, matching mlx_lm.server). */
|
|
287
|
+
xtc_probability?: number;
|
|
288
|
+
xtc_threshold?: number;
|
|
289
|
+
/** OpenAI logit_bias: {tokenId: additive bias}. JSON object keys arrive as
|
|
290
|
+
* strings; coerced to int keys / float values like mlx-lm (400 on failure). */
|
|
291
|
+
logit_bias?: Record<string, number>;
|
|
292
|
+
/** OpenAI presence/frequency penalties + mlx-lm's context-size extensions
|
|
293
|
+
* (window of recent tokens the penalty looks at; default 20). */
|
|
294
|
+
presence_penalty?: number;
|
|
295
|
+
presence_context_size?: number;
|
|
296
|
+
frequency_penalty?: number;
|
|
297
|
+
frequency_context_size?: number;
|
|
298
|
+
/** mlx_lm.server logprobs: `logprobs` is a BOOL (even on /v1/completions —
|
|
299
|
+
* not OpenAI's legacy int), `top_logprobs` an int in [0, 11] or the -1
|
|
300
|
+
* "unset" sentinel (server.py validates exactly that; OpenAI's cap is 20,
|
|
301
|
+
* mlx-lm's is 11 — we copy the reference). Non-stream responses carry
|
|
302
|
+
* mlx-lm's logprobs block; stream chunks never do (reference behavior). */
|
|
303
|
+
logprobs?: boolean;
|
|
304
|
+
top_logprobs?: number;
|
|
265
305
|
/** OpenAI stop sequences: plain string or array (spec allows up to 4).
|
|
266
306
|
* Matched on DECODED text, not token ids — see StopMatcher. */
|
|
267
307
|
stop?: string | string[];
|
|
@@ -290,6 +330,15 @@ interface ChatRequest {
|
|
|
290
330
|
};
|
|
291
331
|
}
|
|
292
332
|
|
|
333
|
+
/** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
|
|
334
|
+
* template). Sampling/penalty/stop fields are the same names as ChatRequest;
|
|
335
|
+
* `prompt` replaces `messages`. mlx_lm.server accepts only a STRING prompt
|
|
336
|
+
* (it calls `tokenizer.encode(request.prompt)` directly) — token-array
|
|
337
|
+
* prompts are rejected there too, so we match. No `echo` (mlx-lm has none). */
|
|
338
|
+
type TextCompletionRequest = Omit<ChatRequest, "messages" | "tools" | "tool_choice"> & {
|
|
339
|
+
prompt?: unknown;
|
|
340
|
+
};
|
|
341
|
+
|
|
293
342
|
/** Per-field default HLG knobs when enabling without specifying them. */
|
|
294
343
|
const HLG_DEFAULTS = { width: 4, shoulder: 4, toe: 6, pivotOffset: 6 } as const;
|
|
295
344
|
|
|
@@ -319,7 +368,23 @@ interface OpenAIToolCall {
|
|
|
319
368
|
function: { name: string; arguments: string };
|
|
320
369
|
}
|
|
321
370
|
|
|
322
|
-
type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
|
|
371
|
+
export type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
|
|
372
|
+
|
|
373
|
+
/** Pick the stream router for a model family. The token-id sentinel router is
|
|
374
|
+
* Gemma-4-ONLY: ids 48/49 (<|tool_call>/<tool_call|>) and 100/101
|
|
375
|
+
* (<|channel>/<channel|>) are special tokens of that tokenizer family
|
|
376
|
+
* (src/tool-call.ts). On every other tokenizer — MiniCPM5, Qwen3/3.5, and
|
|
377
|
+
* the Tier-0 generics (llama, qwen2, phi3, …) — those ids are ordinary
|
|
378
|
+
* low-id vocab, so the sentinel router would silently swallow output into a
|
|
379
|
+
* phantom tool/reasoning segment. Everyone else parses tool calls from
|
|
380
|
+
* decoded text (buffered-text; parseGeneratedToolCalls covers the
|
|
381
|
+
* OpenAI-JSON <tool_call>, Qwen <function=…>, and MiniCPM5 <function name=…>
|
|
382
|
+
* shapes) when tools are present, and streams plain otherwise. Models whose
|
|
383
|
+
* markup isn't covered fail soft: the markup stays in content. */
|
|
384
|
+
export function selectToolStreamMode(modelType: string, hasTools: boolean): ToolStreamMode {
|
|
385
|
+
if (modelType.startsWith("gemma4")) return "gemma-sentinel";
|
|
386
|
+
return hasTools ? "buffered-text" : "plain";
|
|
387
|
+
}
|
|
323
388
|
|
|
324
389
|
/** Routes generated tokens. Gemma uses family-specific sentinel token ids;
|
|
325
390
|
* MiniCPM5 and other text-template models use decoded-text parsing so
|
|
@@ -729,6 +794,129 @@ export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
|
|
|
729
794
|
});
|
|
730
795
|
}
|
|
731
796
|
|
|
797
|
+
/** Coerce a wire `logit_bias` ({"tokenId": bias}) to numeric keys/values —
|
|
798
|
+
* mlx_lm.server's `{int(k): float(v)}` coercion; throws its exact error
|
|
799
|
+
* message on anything non-numeric (surfaced as a 400). JSON object keys
|
|
800
|
+
* always arrive as strings, hence the coercion. */
|
|
801
|
+
export function parseLogitBias(
|
|
802
|
+
raw: Record<string, number> | undefined | null,
|
|
803
|
+
): Record<number, number> | undefined {
|
|
804
|
+
if (raw == null) return undefined;
|
|
805
|
+
if (typeof raw !== "object" || Array.isArray(raw))
|
|
806
|
+
throw new Error("logit_bias must be a dict of int to float");
|
|
807
|
+
const out: Record<number, number> = {};
|
|
808
|
+
for (const [k, v] of Object.entries(raw)) {
|
|
809
|
+
const id = Number(k);
|
|
810
|
+
const bias = Number(v);
|
|
811
|
+
if (!Number.isInteger(id) || !Number.isFinite(bias))
|
|
812
|
+
throw new Error("logit_bias must be a dict of int to float");
|
|
813
|
+
out[id] = bias;
|
|
814
|
+
}
|
|
815
|
+
return Object.keys(out).length ? out : undefined;
|
|
816
|
+
}
|
|
817
|
+
|
|
818
|
+
/** Default seed for a request that didn't pin one. `Date.now()` alone is NOT
|
|
819
|
+
* request-unique: under `--batch N` the batch lane serves ONLY default-seed
|
|
820
|
+
* requests (explicit-seed requests route serial, see GenerationGateway.willBatch),
|
|
821
|
+
* so two identical prompts arriving in the same millisecond would share a seed
|
|
822
|
+
* and — with per-row RNG keyed as stepKey(seed, generatedCount) — produce
|
|
823
|
+
* byte-identical completions, silently collapsing best-of-N diversity. Mix a
|
|
824
|
+
* per-process Weyl counter (golden-ratio increment, period 2^32) into the
|
|
825
|
+
* timestamp so every call yields a distinct uint32 within any given ms.
|
|
826
|
+
* Determinism contract unchanged: reproducibility is only promised for an
|
|
827
|
+
* EXPLICIT request seed (`req.seed ?? nextDefaultSeed()` — explicit wins,
|
|
828
|
+
* byte-identical to before); a default seed is fresh entropy per request. */
|
|
829
|
+
let seedWeyl = 0;
|
|
830
|
+
export function nextDefaultSeed(): number {
|
|
831
|
+
seedWeyl = (seedWeyl + 0x9e3779b9) >>> 0;
|
|
832
|
+
return ((Date.now() & 0xffffffff) ^ seedWeyl) >>> 0;
|
|
833
|
+
}
|
|
834
|
+
|
|
835
|
+
/** mlx_lm.server's logprobs request validation, copied exactly (server.py
|
|
836
|
+
* APIHandler.validate_model_parameters: `_validate("logprobs", bool)` and
|
|
837
|
+
* `_validate("top_logprobs", int, min_val=0, max_val=11, whitelist=[-1])`
|
|
838
|
+
* with defaults logprobs=False / top_logprobs=-1). Returns the reference's
|
|
839
|
+
* error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
|
|
840
|
+
* at 11, not OpenAI's 20 — we mirror the reference. */
|
|
841
|
+
export function validateLogprobsParams(body: {
|
|
842
|
+
logprobs?: unknown;
|
|
843
|
+
top_logprobs?: unknown;
|
|
844
|
+
}): string | null {
|
|
845
|
+
const lp = body.logprobs ?? false;
|
|
846
|
+
if (typeof lp !== "boolean") return "logprobs must be of type bool";
|
|
847
|
+
const tl = body.top_logprobs ?? -1;
|
|
848
|
+
if (typeof tl !== "number" || !Number.isInteger(tl))
|
|
849
|
+
return "top_logprobs must be of type int";
|
|
850
|
+
if (tl === -1) return null; // the "unset" whitelist sentinel
|
|
851
|
+
if (tl < 0) return "top_logprobs must be at least 0";
|
|
852
|
+
if (tl > 11) return "top_logprobs must be at most 11";
|
|
853
|
+
return null;
|
|
854
|
+
}
|
|
855
|
+
|
|
856
|
+
/** Collects per-token logprob info and shapes mlx_lm.server's response block
|
|
857
|
+
* (server.py generate_response L1317-1327). NOT OpenAI's shape — entries
|
|
858
|
+
* carry token *ids* (and raw vocab token strings in the top-k form), and the
|
|
859
|
+
* SAME block is attached under choices[0].logprobs for chat AND text
|
|
860
|
+
* completions:
|
|
861
|
+
* - top_logprobs > 0 → {content: [{id, token, logprob,
|
|
862
|
+
* top_logprobs: [{id, token, logprob}, …]}, …]} — each entry is the
|
|
863
|
+
* top-1 candidate merged with its own top-k list (`dict(i[0],
|
|
864
|
+
* top_logprobs=i)`); mlx-lm leaves argpartition order unspecified, we
|
|
865
|
+
* sort descending, so the entry is deterministically the argmax.
|
|
866
|
+
* - logprobs=true (and top_logprobs ≤ 0) → {content: [{id, logprob}, …]}
|
|
867
|
+
* with the SAMPLED token's logprob.
|
|
868
|
+
* When both are set, only the top_logprobs form is emitted (the reference's
|
|
869
|
+
* if/elif). Stream chunks never carry logprobs — mlx-lm's streaming
|
|
870
|
+
* generate_response calls pass no token_logprobs/top_tokens. */
|
|
871
|
+
class LogprobsCollector {
|
|
872
|
+
readonly #tokens: number[] = [];
|
|
873
|
+
readonly #tokenLogprobs: number[] = [];
|
|
874
|
+
readonly #topTokens: { id: number; token: string; logprob: number }[][] = [];
|
|
875
|
+
|
|
876
|
+
constructor(
|
|
877
|
+
private readonly wantLogprobs: boolean,
|
|
878
|
+
private readonly topK: number,
|
|
879
|
+
private readonly idToToken: (id: number) => string,
|
|
880
|
+
) {}
|
|
881
|
+
|
|
882
|
+
get active(): boolean {
|
|
883
|
+
return this.wantLogprobs || this.topK > 0;
|
|
884
|
+
}
|
|
885
|
+
|
|
886
|
+
push(token: number, info?: TokenLogprobs): void {
|
|
887
|
+
if (!this.active) return;
|
|
888
|
+
this.#tokens.push(token);
|
|
889
|
+
if (this.wantLogprobs) this.#tokenLogprobs.push(info?.logprob ?? NaN);
|
|
890
|
+
if (this.topK > 0)
|
|
891
|
+
this.#topTokens.push(
|
|
892
|
+
(info?.top ?? []).map((t) => ({
|
|
893
|
+
id: t.id,
|
|
894
|
+
token: this.idToToken(t.id),
|
|
895
|
+
logprob: t.logprob,
|
|
896
|
+
})),
|
|
897
|
+
);
|
|
898
|
+
}
|
|
899
|
+
|
|
900
|
+
/** choices[0].logprobs value, or null when nothing to attach (mirrors the
|
|
901
|
+
* reference: the key is omitted entirely for zero collected tokens). */
|
|
902
|
+
payload(): { content: Record<string, unknown>[] } | null {
|
|
903
|
+
if (this.#topTokens.length)
|
|
904
|
+
return {
|
|
905
|
+
content: this.#topTokens.map((t) =>
|
|
906
|
+
t.length ? { ...t[0]!, top_logprobs: t } : {},
|
|
907
|
+
),
|
|
908
|
+
};
|
|
909
|
+
if (this.#tokenLogprobs.length)
|
|
910
|
+
return {
|
|
911
|
+
content: this.#tokens.map((id, i) => ({
|
|
912
|
+
id,
|
|
913
|
+
logprob: this.#tokenLogprobs[i]!,
|
|
914
|
+
})),
|
|
915
|
+
};
|
|
916
|
+
return null;
|
|
917
|
+
}
|
|
918
|
+
}
|
|
919
|
+
|
|
732
920
|
/** Incremental detokenizer: emits the longest stable decoded prefix. */
|
|
733
921
|
class StreamDecoder {
|
|
734
922
|
#ids: number[] = [];
|
|
@@ -788,13 +976,6 @@ export function createServer(
|
|
|
788
976
|
// lane (see GenerationGateway.willBatch).
|
|
789
977
|
const batch = Math.max(1, Math.floor(serverOptions.batch ?? 1));
|
|
790
978
|
|
|
791
|
-
let queue: Promise<unknown> = Promise.resolve();
|
|
792
|
-
const enqueue = <T>(fn: () => Promise<T>): Promise<T> => {
|
|
793
|
-
const run = queue.then(fn, fn);
|
|
794
|
-
queue = run.catch(() => {});
|
|
795
|
-
return run;
|
|
796
|
-
};
|
|
797
|
-
|
|
798
979
|
const promptCache = new PromptCache(serverOptions.promptCacheBytes ?? 2e9);
|
|
799
980
|
// Responses-API store for previous_response_id resumption (Phase 11):
|
|
800
981
|
// TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
|
|
@@ -802,8 +983,11 @@ export function createServer(
|
|
|
802
983
|
// so its KV prefill is already cached.
|
|
803
984
|
const responseStore = new ResponseStore();
|
|
804
985
|
|
|
805
|
-
/** Run one generation with prompt-cache reuse. Must be called
|
|
806
|
-
* the
|
|
986
|
+
/** Run one generation with prompt-cache reuse. Must be called under the
|
|
987
|
+
* gateway's exclusive lock — the gateway invokes it on the serial lane;
|
|
988
|
+
* the curve /generate endpoint wraps it in gateway.runExclusive. (The old
|
|
989
|
+
* `enqueue` promise queue is gone; the gateway's AsyncMutex is THE lock.)
|
|
990
|
+
* onToken returning `false` halts generation early (stop
|
|
807
991
|
* sequence fired); the cache snapshot stays valid and is still kept.
|
|
808
992
|
* Vision requests bypass the prompt cache: image tokens are
|
|
809
993
|
* identical placeholder ids, so prefix matching across different
|
|
@@ -811,7 +995,7 @@ export function createServer(
|
|
|
811
995
|
const runGeneration = async (
|
|
812
996
|
promptIds: number[],
|
|
813
997
|
options: GenerateOptions,
|
|
814
|
-
onToken: (token: number) => void | boolean | Promise<void | boolean>,
|
|
998
|
+
onToken: (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>,
|
|
815
999
|
vision?: { embeddings: import("./mlx/array").MlxArray; imageMask: import("./mlx/array").MlxArray },
|
|
816
1000
|
) => {
|
|
817
1001
|
// Cache entries are adapter-specific: KV computed under one adapter
|
|
@@ -826,7 +1010,7 @@ export function createServer(
|
|
|
826
1010
|
...(vision ? { promptEmbeddings: vision.embeddings, imageMask: vision.imageMask } : {}),
|
|
827
1011
|
});
|
|
828
1012
|
for await (const t of gen) {
|
|
829
|
-
if ((await onToken(t.token)) === false) break;
|
|
1013
|
+
if ((await onToken(t.token, t.logprobs)) === false) break;
|
|
830
1014
|
}
|
|
831
1015
|
const s = gen.stats!; // set on completion AND on early break
|
|
832
1016
|
if (vision) {
|
|
@@ -914,6 +1098,17 @@ export function createServer(
|
|
|
914
1098
|
`Omit --kv-quant to batch in bf16. (docs/design/parallel-slots.md)`,
|
|
915
1099
|
);
|
|
916
1100
|
|
|
1101
|
+
// XTC never removes EOS or the newline token: mlx_lm.server passes
|
|
1102
|
+
// [tokenizer.eos_token_id, tokenizer.encode("\n")] as xtc_special_tokens.
|
|
1103
|
+
// We pass the flat equivalent — ALL configured EOS ids plus the bare newline
|
|
1104
|
+
// id(s), encoded WITHOUT special tokens (mlx-lm's encode("\n") can drag a
|
|
1105
|
+
// BOS along on some tokenizers; protecting BOS from XTC is a no-op, so the
|
|
1106
|
+
// flat/no-BOS form is behaviorally identical and cleaner).
|
|
1107
|
+
const xtcSpecialTokens = [
|
|
1108
|
+
...ctx.model.config.eosTokenIds,
|
|
1109
|
+
...ctx.tokenizer.encode("\n", false),
|
|
1110
|
+
];
|
|
1111
|
+
|
|
917
1112
|
// Effective enable_thinking for a request, with the same precedence the chat
|
|
918
1113
|
// template uses (extracted so sampling and template rendering can't disagree):
|
|
919
1114
|
// explicit chat_template_kwargs.enable_thinking → reasoning_effort ("none" =
|
|
@@ -947,12 +1142,25 @@ export function createServer(
|
|
|
947
1142
|
// truncates the visible reply. Default very generously (the model's
|
|
948
1143
|
// context is far larger); only an explicit max_tokens narrows it. The
|
|
949
1144
|
// model still stops at its eos_token well before this in normal replies.
|
|
950
|
-
maxTokens: req.max_completion_tokens ?? req.max_tokens ?? 65_536,
|
|
1145
|
+
maxTokens: req.max_completion_tokens ?? req.max_tokens ?? serverOptions.defaultMaxTokens ?? 65_536,
|
|
951
1146
|
temperature: req.temperature ?? serverOptions.defaultTemperature ?? defaultTemp,
|
|
952
1147
|
topP: req.top_p ?? serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? 0,
|
|
953
1148
|
topK: req.top_k ?? serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? 0,
|
|
954
|
-
seed: req.seed ?? (
|
|
1149
|
+
seed: req.seed ?? nextDefaultSeed(),
|
|
955
1150
|
repetitionPenalty: req.repetition_penalty ?? ctx.genDefaults.repetitionPenalty,
|
|
1151
|
+
// mlx_lm.server sampling extensions (defaults mirror server.py: min_p /
|
|
1152
|
+
// xtc off, context windows 20). parseLogitBias throws on non-numeric
|
|
1153
|
+
// input — callers surface it as a 400 (mlx-lm's coercion error).
|
|
1154
|
+
minP: req.min_p ?? 0,
|
|
1155
|
+
xtcProbability: req.xtc_probability ?? 0,
|
|
1156
|
+
xtcThreshold: req.xtc_threshold ?? 0,
|
|
1157
|
+
...(req.xtc_probability ? { xtcSpecialTokens } : {}),
|
|
1158
|
+
logitBias: parseLogitBias(req.logit_bias),
|
|
1159
|
+
repetitionContextSize: req.repetition_context_size ?? 20,
|
|
1160
|
+
presencePenalty: req.presence_penalty ?? 0,
|
|
1161
|
+
presenceContextSize: req.presence_context_size ?? 20,
|
|
1162
|
+
frequencyPenalty: req.frequency_penalty ?? 0,
|
|
1163
|
+
frequencyContextSize: req.frequency_context_size ?? 20,
|
|
956
1164
|
hlg: resolveHlg(req.hlg, serverOptions.hlg),
|
|
957
1165
|
// generate() yields token ids; stop sequences match on decoded text
|
|
958
1166
|
// (they can span token boundaries), so the StopMatcher sits at the
|
|
@@ -980,15 +1188,11 @@ export function createServer(
|
|
|
980
1188
|
return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
|
|
981
1189
|
};
|
|
982
1190
|
|
|
983
|
-
const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
|
|
984
|
-
//
|
|
985
|
-
//
|
|
986
|
-
//
|
|
987
|
-
|
|
988
|
-
if (isMiniCPM5Config(ctx.model.config) || isQwen35Config(ctx.model.config))
|
|
989
|
-
return tools?.length ? "buffered-text" : "plain";
|
|
990
|
-
return "gemma-sentinel";
|
|
991
|
-
};
|
|
1191
|
+
const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
|
|
1192
|
+
// Gemma-4 keeps the token-level sentinel path; every other family
|
|
1193
|
+
// (MiniCPM5, Qwen3/3.5, Tier-0 generics) emits tool calls as DECODED TEXT
|
|
1194
|
+
// — see selectToolStreamMode for why sentinel ids must stay family-gated.
|
|
1195
|
+
selectToolStreamMode(ctx.model.config.modelType, !!tools?.length);
|
|
992
1196
|
|
|
993
1197
|
const toolRouter = (tools: ToolDefinition[] | null): ToolAwareStream =>
|
|
994
1198
|
new ToolAwareStream(ctx.tokenizer, toolStreamMode(tools), tools);
|
|
@@ -1067,9 +1271,14 @@ export function createServer(
|
|
|
1067
1271
|
// surface on their own. (scan() is INSERT-OR-REPLACE + prunes deleted,
|
|
1068
1272
|
// so it's idempotent and cheap for a local cache.)
|
|
1069
1273
|
await reg.scan();
|
|
1274
|
+
const { visionCapable } = await import("./registry");
|
|
1275
|
+
const { supportTier } = await import("./model/support");
|
|
1070
1276
|
const rows = [];
|
|
1071
|
-
|
|
1072
|
-
|
|
1277
|
+
// listCanonical: one row per repo (refs/main) — duplicate snapshots
|
|
1278
|
+
// from upstream re-pushes stay visible only in `ls --all-revisions`.
|
|
1279
|
+
for (const m of reg.listCanonical()) {
|
|
1280
|
+
const tier = supportTier(m.modelType, m.repoId);
|
|
1281
|
+
const supported = tier !== null;
|
|
1073
1282
|
let assessment = null;
|
|
1074
1283
|
try {
|
|
1075
1284
|
const config = await loadModelConfig(m.path);
|
|
@@ -1083,7 +1292,7 @@ export function createServer(
|
|
|
1083
1292
|
rows.push({
|
|
1084
1293
|
repo_id: m.repoId, model_type: m.modelType,
|
|
1085
1294
|
size_bytes: m.sizeBytes, quant_bits: m.quantBits,
|
|
1086
|
-
vision: m
|
|
1295
|
+
vision: visionCapable(m), supported, support_tier: tier,
|
|
1087
1296
|
serving: m.repoId === ctx.modelId,
|
|
1088
1297
|
assessment,
|
|
1089
1298
|
});
|
|
@@ -1098,6 +1307,39 @@ export function createServer(
|
|
|
1098
1307
|
return Response.json({ downloads: downloadsSnapshot() });
|
|
1099
1308
|
}
|
|
1100
1309
|
|
|
1310
|
+
// Memory synthesis progress (P8-T5). The same DAG the nightly launchd job
|
|
1311
|
+
// runs (`mlx-bun memory synthesize`), streamed as Server-Sent Events so the
|
|
1312
|
+
// status page can show live stage/log/done progress. `?dry=1` plans the DAG
|
|
1313
|
+
// without any model call or vault write — the safe wiring-verification path
|
|
1314
|
+
// (the FULL-corpus run is USER-ACTION, P6-T5). GET so EventSource can drive it.
|
|
1315
|
+
if (url.pathname === "/v1/memory/synthesize" && request.method === "GET") {
|
|
1316
|
+
const dryRun = url.searchParams.get("dry") === "1";
|
|
1317
|
+
const { runSynthesis } = await import("./memory/pipeline");
|
|
1318
|
+
const enc = new TextEncoder();
|
|
1319
|
+
const stream = new ReadableStream({
|
|
1320
|
+
async start(controller) {
|
|
1321
|
+
const send = (e: unknown) =>
|
|
1322
|
+
controller.enqueue(enc.encode(`data: ${JSON.stringify(e)}\n\n`));
|
|
1323
|
+
try {
|
|
1324
|
+
const summary = await runSynthesis({ dryRun }, (ev) => send(ev));
|
|
1325
|
+
send({ type: "summary", ...summary });
|
|
1326
|
+
controller.enqueue(enc.encode("data: [DONE]\n\n"));
|
|
1327
|
+
} catch (e) {
|
|
1328
|
+
send({ type: "error", message: (e as Error).message });
|
|
1329
|
+
} finally {
|
|
1330
|
+
controller.close();
|
|
1331
|
+
}
|
|
1332
|
+
},
|
|
1333
|
+
});
|
|
1334
|
+
return new Response(stream, {
|
|
1335
|
+
headers: {
|
|
1336
|
+
"content-type": "text/event-stream",
|
|
1337
|
+
"cache-control": "no-cache",
|
|
1338
|
+
connection: "keep-alive",
|
|
1339
|
+
},
|
|
1340
|
+
});
|
|
1341
|
+
}
|
|
1342
|
+
|
|
1101
1343
|
if (url.pathname === "/fit" && request.method === "GET") {
|
|
1102
1344
|
// Fit assessment for the status page: this-machine report at the
|
|
1103
1345
|
// admission ceiling + the Apple SKU matrix at a fixed 32k.
|
|
@@ -1160,8 +1402,9 @@ export function createServer(
|
|
|
1160
1402
|
return Response.json({
|
|
1161
1403
|
name: "mlx-bun", version: pkgVersion, model: ctx.modelId,
|
|
1162
1404
|
endpoints: [
|
|
1163
|
-
"POST /v1/chat/completions", "POST /v1/
|
|
1164
|
-
"
|
|
1405
|
+
"POST /v1/chat/completions", "POST /v1/completions", "POST /v1/messages",
|
|
1406
|
+
"POST /v1/responses", "POST /v1/embeddings", "GET /v1/models",
|
|
1407
|
+
"GET/POST/DELETE /v1/adapters", "GET /health",
|
|
1165
1408
|
"GET /stats", "GET /fit", "GET /library", "GET /downloads",
|
|
1166
1409
|
],
|
|
1167
1410
|
});
|
|
@@ -1229,16 +1472,95 @@ export function createServer(
|
|
|
1229
1472
|
});
|
|
1230
1473
|
}
|
|
1231
1474
|
|
|
1232
|
-
|
|
1475
|
+
// mlx_lm.server parity: GET /health → the exact body it writes
|
|
1476
|
+
// ('{"status": "ok"}', note the space) so byte-for-byte health checks pass.
|
|
1477
|
+
if (url.pathname === "/health" && request.method === "GET") {
|
|
1478
|
+
return new Response('{"status": "ok"}', {
|
|
1479
|
+
headers: { "content-type": "application/json" },
|
|
1480
|
+
});
|
|
1481
|
+
}
|
|
1482
|
+
|
|
1483
|
+
// GET /v1/models — the served model FIRST (with our capability extras),
|
|
1484
|
+
// then every other servable model the registry knows (mlx-lm scans the
|
|
1485
|
+
// HF cache here; our registry is that scan, filtered to supported
|
|
1486
|
+
// architectures). GET /v1/models/<id> filters to that id — same list
|
|
1487
|
+
// shape, matching mlx_lm.server's handle_models_request.
|
|
1488
|
+
if (
|
|
1489
|
+
(url.pathname === "/v1/models" || url.pathname.startsWith("/v1/models/")) &&
|
|
1490
|
+
request.method === "GET"
|
|
1491
|
+
) {
|
|
1492
|
+
const filterId = url.pathname.length > "/v1/models/".length - 1
|
|
1493
|
+
? decodeURIComponent(url.pathname.slice("/v1/models/".length))
|
|
1494
|
+
: null;
|
|
1495
|
+
const created = Math.floor(startedAt / 1000);
|
|
1496
|
+
const data: Array<Record<string, unknown>> = [{
|
|
1497
|
+
id: ctx.modelId, object: "model", created, owned_by: "mlx-bun",
|
|
1498
|
+
context_window: ctx.model.config.text.maxPositionEmbeddings,
|
|
1499
|
+
// Capability flags for clients (CLI/external pi) that build a
|
|
1500
|
+
// provider from discovery — `reasoning` gates the thinking toggle,
|
|
1501
|
+
// `vision` the image input declaration (tower loaded or lazily
|
|
1502
|
+
// loadable; same signal the web embed uses).
|
|
1503
|
+
reasoning: ctx.template.supportsThinking,
|
|
1504
|
+
vision: !!(ctx.vision || ctx.loadVision),
|
|
1505
|
+
}];
|
|
1506
|
+
try {
|
|
1507
|
+
const { Registry } = await import("./registry");
|
|
1508
|
+
const reg = new Registry();
|
|
1509
|
+
try {
|
|
1510
|
+
if (reg.list().length === 0) await reg.scan();
|
|
1511
|
+
// Canonical rows only — a repo with N snapshots is ONE model id.
|
|
1512
|
+
for (const m of reg.listCanonical()) {
|
|
1513
|
+
if (m.repoId === ctx.modelId) continue;
|
|
1514
|
+
if (!isSupportedModelRecord(m.modelType, m.repoId)) continue;
|
|
1515
|
+
data.push({ id: m.repoId, object: "model", created });
|
|
1516
|
+
}
|
|
1517
|
+
} finally {
|
|
1518
|
+
reg.close();
|
|
1519
|
+
}
|
|
1520
|
+
} catch { /* registry unavailable → served model only */ }
|
|
1521
|
+
return Response.json({
|
|
1522
|
+
object: "list",
|
|
1523
|
+
data: filterId ? data.filter((m) => m.id === filterId) : data,
|
|
1524
|
+
});
|
|
1525
|
+
}
|
|
1526
|
+
|
|
1527
|
+
// OpenAI embeddings API. Works when the SERVED model is an embedding model
|
|
1528
|
+
// (plain Qwen3 / Qwen3-Embedding) — consistent with the single-model server
|
|
1529
|
+
// design: `mlx-bun serve <embedding-model>` to use this. Optional non-standard
|
|
1530
|
+
// `instruction` applies Qwen3-Embedding's query format. Embedding is a pure
|
|
1531
|
+
// forward (no decode loop / gateway), so it runs inline.
|
|
1532
|
+
if (url.pathname === "/v1/embeddings" && request.method === "POST") {
|
|
1533
|
+
if (!isEmbeddingModel(ctx.model))
|
|
1534
|
+
return Response.json({
|
|
1535
|
+
error: {
|
|
1536
|
+
message: `served model "${ctx.modelId}" is not an embedding model; ` +
|
|
1537
|
+
`serve an embedding model (e.g. Qwen3-Embedding) to use /v1/embeddings`,
|
|
1538
|
+
type: "invalid_request_error",
|
|
1539
|
+
},
|
|
1540
|
+
}, { status: 400 });
|
|
1541
|
+
let body: { input?: string | string[]; instruction?: string };
|
|
1542
|
+
try {
|
|
1543
|
+
body = (await request.json()) as typeof body;
|
|
1544
|
+
} catch {
|
|
1545
|
+
return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
|
|
1546
|
+
}
|
|
1547
|
+
const inputs = Array.isArray(body.input) ? body.input : body.input != null ? [body.input] : [];
|
|
1548
|
+
if (inputs.length === 0 || !inputs.every((s) => typeof s === "string"))
|
|
1549
|
+
return Response.json({
|
|
1550
|
+
error: { message: "`input` must be a string or array of strings", type: "invalid_request_error" },
|
|
1551
|
+
}, { status: 400 });
|
|
1552
|
+
const instruction = typeof body.instruction === "string" ? body.instruction : undefined;
|
|
1553
|
+
const results = embedMany(ctx.model, ctx.tokenizer, inputs, instruction);
|
|
1554
|
+
let totalTokens = 0;
|
|
1555
|
+
const data = results.map((r, index) => {
|
|
1556
|
+
totalTokens += r.tokens;
|
|
1557
|
+
return { object: "embedding", index, embedding: Array.from(r.vector) };
|
|
1558
|
+
});
|
|
1233
1559
|
return Response.json({
|
|
1234
1560
|
object: "list",
|
|
1235
|
-
data
|
|
1236
|
-
|
|
1237
|
-
|
|
1238
|
-
// Capability flags for clients (CLI/external pi) that build a
|
|
1239
|
-
// provider from discovery — `reasoning` gates the thinking toggle.
|
|
1240
|
-
reasoning: ctx.template.supportsThinking,
|
|
1241
|
-
}],
|
|
1561
|
+
data,
|
|
1562
|
+
model: ctx.modelId,
|
|
1563
|
+
usage: { prompt_tokens: totalTokens, total_tokens: totalTokens },
|
|
1242
1564
|
});
|
|
1243
1565
|
}
|
|
1244
1566
|
|
|
@@ -1287,7 +1609,10 @@ export function createServer(
|
|
|
1287
1609
|
if (!body.id || !body.path)
|
|
1288
1610
|
return Response.json({ error: { message: "id and path required" } }, { status: 400 });
|
|
1289
1611
|
try {
|
|
1290
|
-
|
|
1612
|
+
// Under the gateway lock: mount/unmount mutate the shared adapter
|
|
1613
|
+
// registry (unmount disposes arrays a running generation could still
|
|
1614
|
+
// hold) — one mutual-exclusion domain with generation (D3).
|
|
1615
|
+
const info = await gateway.runExclusive(() => ctx.adapters.mount(body.id!, body.path!));
|
|
1291
1616
|
return Response.json({
|
|
1292
1617
|
id: info.id, mounted_layers: info.mountedLayers,
|
|
1293
1618
|
rank: info.rank, scale: info.scale,
|
|
@@ -1298,7 +1623,7 @@ export function createServer(
|
|
|
1298
1623
|
}
|
|
1299
1624
|
if (url.pathname.startsWith("/v1/adapters/") && request.method === "DELETE") {
|
|
1300
1625
|
const id = decodeURIComponent(url.pathname.slice("/v1/adapters/".length));
|
|
1301
|
-
const removed = await
|
|
1626
|
+
const removed = await gateway.runExclusive(async () => ctx.adapters.unmount(id));
|
|
1302
1627
|
return removed > 0
|
|
1303
1628
|
? Response.json({ id, removed_layers: removed })
|
|
1304
1629
|
: Response.json({ error: { message: `adapter ${id} not mounted` } }, { status: 404 });
|
|
@@ -1317,7 +1642,10 @@ export function createServer(
|
|
|
1317
1642
|
if (sids[0] === sids[1] && sids[0] === ctx.tokenizer.bosTokenId) sids = sids.slice(1);
|
|
1318
1643
|
try {
|
|
1319
1644
|
const NB = 80;
|
|
1320
|
-
|
|
1645
|
+
// gateway.runExclusive: this raw forward must not run concurrently
|
|
1646
|
+
// with batched decode steps or a serial generation (GPU + shared
|
|
1647
|
+
// model state are single-owner — D3, one lock).
|
|
1648
|
+
const result = await gateway.runExclusive(async () => {
|
|
1321
1649
|
const cache = ctx.model.makeCache();
|
|
1322
1650
|
try {
|
|
1323
1651
|
const logits = ctx.model.forward(sids, cache); // [1, L, V]
|
|
@@ -1372,7 +1700,9 @@ export function createServer(
|
|
|
1372
1700
|
const genOpts: GenerateOptions = useCurve
|
|
1373
1701
|
? { curve, seed: baseSeed + i, maxTokens, ...kvScheme }
|
|
1374
1702
|
: { temperature: recipe.temperature, topP: recipe.topP, topK: recipe.topK, seed: baseSeed + i, maxTokens, ...kvScheme };
|
|
1375
|
-
|
|
1703
|
+
// runExclusive: runGeneration touches the GPU + prompt cache, so
|
|
1704
|
+
// it needs the gateway lock (this endpoint bypasses gateway.run).
|
|
1705
|
+
await gateway.runExclusive(() => runGeneration(ids, genOpts, (t) => { toks.push(t); }));
|
|
1376
1706
|
const text = ctx.tokenizer.decode(toks, true).trim();
|
|
1377
1707
|
samples.push({ text, junk: curveJunk(text) });
|
|
1378
1708
|
}
|
|
@@ -1390,6 +1720,10 @@ export function createServer(
|
|
|
1390
1720
|
const handleChat = async (body: ChatRequest): Promise<Response> => {
|
|
1391
1721
|
if (!Array.isArray(body.messages) || body.messages.length === 0)
|
|
1392
1722
|
return Response.json({ error: { message: "messages required" } }, { status: 400 });
|
|
1723
|
+
// mlx-lm validates logprobs params up front (ValueError → 400)
|
|
1724
|
+
const lpParamError = validateLogprobsParams(body);
|
|
1725
|
+
if (lpParamError)
|
|
1726
|
+
return Response.json({ error: { message: lpParamError } }, { status: 400 });
|
|
1393
1727
|
|
|
1394
1728
|
const id = `chatcmpl-${crypto.randomUUID()}`;
|
|
1395
1729
|
const created = Math.floor(Date.now() / 1000);
|
|
@@ -1459,7 +1793,16 @@ export function createServer(
|
|
|
1459
1793
|
{ status: 400 },
|
|
1460
1794
|
);
|
|
1461
1795
|
}
|
|
1462
|
-
|
|
1796
|
+
let options: ReturnType<typeof toOptions>;
|
|
1797
|
+
try {
|
|
1798
|
+
options = toOptions(body);
|
|
1799
|
+
} catch (e) {
|
|
1800
|
+
// bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
|
|
1801
|
+
vision?.embeddings.dispose();
|
|
1802
|
+
vision?.imageMask.dispose();
|
|
1803
|
+
diffusionPixels?.dispose();
|
|
1804
|
+
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
1805
|
+
}
|
|
1463
1806
|
if (diffusionPixels) options.visionPixels = diffusionPixels;
|
|
1464
1807
|
// Admission: reject what cannot finish within the memory budget
|
|
1465
1808
|
// (the GPU OOM it would otherwise hit is uncatchable and kills
|
|
@@ -1485,18 +1828,41 @@ export function createServer(
|
|
|
1485
1828
|
);
|
|
1486
1829
|
}
|
|
1487
1830
|
try {
|
|
1488
|
-
|
|
1831
|
+
// A request's explicit `adapter` (incl. "none") wins over the
|
|
1832
|
+
// startup default from `serve --adapter <dir>`.
|
|
1833
|
+
const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
|
|
1489
1834
|
if (adapterIds.length) options.adapters = adapterIds;
|
|
1490
1835
|
} catch (e) {
|
|
1491
1836
|
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
1492
1837
|
}
|
|
1493
1838
|
|
|
1494
|
-
//
|
|
1839
|
+
// logprobs capture: non-stream only — stream chunks never carry
|
|
1840
|
+
// logprobs (mirroring mlx_lm.server, whose streaming generate_response
|
|
1841
|
+
// calls pass no token_logprobs/top_tokens), so streaming requests skip
|
|
1842
|
+
// the capture cost entirely and stay batchable.
|
|
1843
|
+
const wantLogprobs = body.logprobs === true;
|
|
1844
|
+
const topLogprobs =
|
|
1845
|
+
typeof body.top_logprobs === "number" && body.top_logprobs > 0
|
|
1846
|
+
? body.top_logprobs : 0;
|
|
1847
|
+
const captureLogprobs = !body.stream && (wantLogprobs || topLogprobs > 0);
|
|
1848
|
+
if (captureLogprobs) {
|
|
1849
|
+
options.logprobs = wantLogprobs;
|
|
1850
|
+
options.topLogprobs = topLogprobs;
|
|
1851
|
+
}
|
|
1852
|
+
|
|
1853
|
+
// What lane this request takes (vision / adapters / logits processors /
|
|
1495
1854
|
// a user-fixed seed → serial; everything else batches when --batch N).
|
|
1496
1855
|
const shape = {
|
|
1497
1856
|
hasVision: !!vision,
|
|
1498
1857
|
hasAdapters: !!options.adapters?.length,
|
|
1499
1858
|
hasRepetitionPenalty: !!options.repetitionPenalty,
|
|
1859
|
+
// min_p / XTC / logit_bias / presence+frequency penalties: serial-only
|
|
1860
|
+
// in v1 (see GenerationGateway.RequestShape.hasLogitsExtras).
|
|
1861
|
+
hasLogitsExtras: !!(
|
|
1862
|
+
options.minP || options.xtcProbability || options.logitBias ||
|
|
1863
|
+
options.presencePenalty || options.frequencyPenalty
|
|
1864
|
+
),
|
|
1865
|
+
wantsLogprobs: captureLogprobs,
|
|
1500
1866
|
userSeed: body.seed !== undefined,
|
|
1501
1867
|
kvQuant: !!(options.kvConfig?.length || options.kvBits),
|
|
1502
1868
|
};
|
|
@@ -1611,9 +1977,15 @@ export function createServer(
|
|
|
1611
1977
|
const router = toolRouter(tools);
|
|
1612
1978
|
const stopper = new StopMatcher(options.stopSequences);
|
|
1613
1979
|
const thinking = new ThinkingTagSplitter(ctx.template.thinkingFormat === "think-tag", startInThinking);
|
|
1980
|
+
// mlx-lm collects logprobs across EVERY generated token (reasoning
|
|
1981
|
+
// and tool tokens included), not just visible content — same here.
|
|
1982
|
+
const lpc = captureLogprobs
|
|
1983
|
+
? new LogprobsCollector(wantLogprobs, topLogprobs, (tid) => ctx.tokenizer.idToToken(tid))
|
|
1984
|
+
: null;
|
|
1614
1985
|
let content = "";
|
|
1615
1986
|
let reasoning = "";
|
|
1616
|
-
const s = await gateway.run(promptIds, options, (token) => {
|
|
1987
|
+
const s = await gateway.run(promptIds, options, (token, lpInfo) => {
|
|
1988
|
+
lpc?.push(token, lpInfo);
|
|
1617
1989
|
const rawContent = router.push(token);
|
|
1618
1990
|
reasoning += router.takeReasoning(); // gemma-channel thinking
|
|
1619
1991
|
const parts = thinking.push(stopper.push(rawContent));
|
|
@@ -1640,6 +2012,7 @@ export function createServer(
|
|
|
1640
2012
|
? "tool_calls"
|
|
1641
2013
|
: stopper.stopped ? "stop"
|
|
1642
2014
|
: s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
|
|
2015
|
+
const logprobsBlock = lpc?.payload() ?? null;
|
|
1643
2016
|
return Response.json({
|
|
1644
2017
|
id, object: "chat.completion", created, model: ctx.modelId,
|
|
1645
2018
|
choices: [{
|
|
@@ -1650,6 +2023,7 @@ export function createServer(
|
|
|
1650
2023
|
...(reasoning ? { reasoning } : {}),
|
|
1651
2024
|
...(toolCalls.length ? { tool_calls: toolCalls } : {}),
|
|
1652
2025
|
},
|
|
2026
|
+
...(logprobsBlock ? { logprobs: logprobsBlock } : {}),
|
|
1653
2027
|
finish_reason: finish,
|
|
1654
2028
|
}],
|
|
1655
2029
|
usage: {
|
|
@@ -1675,6 +2049,195 @@ export function createServer(
|
|
|
1675
2049
|
return handleChat(body);
|
|
1676
2050
|
}
|
|
1677
2051
|
|
|
2052
|
+
// Raw text completion (mlx_lm.server's /v1/completions, request_type
|
|
2053
|
+
// "text"): NO chat template — the prompt string is tokenized directly
|
|
2054
|
+
// (tokenizer.encode with the tokenizer's own special-token handling,
|
|
2055
|
+
// exactly mlx-lm's `tokenizer.encode(request.prompt)`). Rides the same
|
|
2056
|
+
// GenerationGateway + admission + adapter path as chat; no tool router
|
|
2057
|
+
// or thinking splitter (raw text in, raw text out).
|
|
2058
|
+
if (url.pathname === "/v1/completions" && request.method === "POST") {
|
|
2059
|
+
let body: TextCompletionRequest;
|
|
2060
|
+
try {
|
|
2061
|
+
body = (await request.json()) as TextCompletionRequest;
|
|
2062
|
+
} catch {
|
|
2063
|
+
return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
|
|
2064
|
+
}
|
|
2065
|
+
if (typeof body.prompt !== "string" || body.prompt.length === 0)
|
|
2066
|
+
return Response.json(
|
|
2067
|
+
{
|
|
2068
|
+
error: {
|
|
2069
|
+
message: "prompt (a non-empty string) is required " +
|
|
2070
|
+
"(token-array prompts are not accepted, matching mlx_lm.server)",
|
|
2071
|
+
},
|
|
2072
|
+
},
|
|
2073
|
+
{ status: 400 },
|
|
2074
|
+
);
|
|
2075
|
+
// mlx-lm validates logprobs params up front (ValueError → 400)
|
|
2076
|
+
const lpParamError = validateLogprobsParams(body);
|
|
2077
|
+
if (lpParamError)
|
|
2078
|
+
return Response.json({ error: { message: lpParamError } }, { status: 400 });
|
|
2079
|
+
const id = `cmpl-${crypto.randomUUID()}`;
|
|
2080
|
+
const created = Math.floor(Date.now() / 1000);
|
|
2081
|
+
const promptIds = ctx.tokenizer.encode(body.prompt);
|
|
2082
|
+
let options: ReturnType<typeof toOptions>;
|
|
2083
|
+
try {
|
|
2084
|
+
options = toOptions(body as unknown as ChatRequest);
|
|
2085
|
+
} catch (e) {
|
|
2086
|
+
// bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
|
|
2087
|
+
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
2088
|
+
}
|
|
2089
|
+
// mlx_lm.server's default max_tokens is 512 (its --max-tokens CLI
|
|
2090
|
+
// default). The chat lane's very generous default is wrong for raw
|
|
2091
|
+
// completion: with no template an EOS may never come.
|
|
2092
|
+
options.maxTokens = body.max_completion_tokens ?? body.max_tokens ?? serverOptions.defaultMaxTokens ?? 512;
|
|
2093
|
+
const requiredCtx = promptIds.length + options.maxTokens;
|
|
2094
|
+
if (requiredCtx > admission.maxSafeContext)
|
|
2095
|
+
return Response.json(
|
|
2096
|
+
{
|
|
2097
|
+
error: {
|
|
2098
|
+
message:
|
|
2099
|
+
`request needs ${requiredCtx} tokens of context ` +
|
|
2100
|
+
`(prompt ${promptIds.length} + max_tokens ${options.maxTokens}) but the ` +
|
|
2101
|
+
`memory budget caps safe context at ${admission.maxSafeContext} — ` +
|
|
2102
|
+
`shorten the prompt or lower max_tokens`,
|
|
2103
|
+
type: "memory_admission",
|
|
2104
|
+
code: "context_over_budget",
|
|
2105
|
+
},
|
|
2106
|
+
},
|
|
2107
|
+
{ status: 400 },
|
|
2108
|
+
);
|
|
2109
|
+
try {
|
|
2110
|
+
const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
|
|
2111
|
+
if (adapterIds.length) options.adapters = adapterIds;
|
|
2112
|
+
} catch (e) {
|
|
2113
|
+
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
2114
|
+
}
|
|
2115
|
+
// logprobs capture — same mlx-lm block as chat (generate_response is
|
|
2116
|
+
// shared in the reference too); non-stream only, stream chunks never
|
|
2117
|
+
// carry logprobs.
|
|
2118
|
+
const wantLogprobs = body.logprobs === true;
|
|
2119
|
+
const topLogprobs =
|
|
2120
|
+
typeof body.top_logprobs === "number" && body.top_logprobs > 0
|
|
2121
|
+
? body.top_logprobs : 0;
|
|
2122
|
+
const captureLogprobs = !body.stream && (wantLogprobs || topLogprobs > 0);
|
|
2123
|
+
if (captureLogprobs) {
|
|
2124
|
+
options.logprobs = wantLogprobs;
|
|
2125
|
+
options.topLogprobs = topLogprobs;
|
|
2126
|
+
}
|
|
2127
|
+
const shape = {
|
|
2128
|
+
hasVision: false,
|
|
2129
|
+
hasAdapters: !!options.adapters?.length,
|
|
2130
|
+
hasRepetitionPenalty: !!options.repetitionPenalty,
|
|
2131
|
+
hasLogitsExtras: !!(
|
|
2132
|
+
options.minP || options.xtcProbability || options.logitBias ||
|
|
2133
|
+
options.presencePenalty || options.frequencyPenalty
|
|
2134
|
+
),
|
|
2135
|
+
wantsLogprobs: captureLogprobs,
|
|
2136
|
+
userSeed: body.seed !== undefined,
|
|
2137
|
+
kvQuant: !!(options.kvConfig?.length || options.kvBits),
|
|
2138
|
+
};
|
|
2139
|
+
const finishReason = (stopped: boolean, generated: number): "stop" | "length" =>
|
|
2140
|
+
stopped ? "stop" : generated >= (options.maxTokens ?? 512) ? "length" : "stop";
|
|
2141
|
+
|
|
2142
|
+
if (body.stream) {
|
|
2143
|
+
const stream = new ReadableStream<Uint8Array>({
|
|
2144
|
+
async start(controller) {
|
|
2145
|
+
const enc = new TextEncoder();
|
|
2146
|
+
const send = (obj: unknown) =>
|
|
2147
|
+
controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
|
|
2148
|
+
const chunk = (text: string, finish: string | null) => ({
|
|
2149
|
+
id, object: "text_completion", created, model: ctx.modelId,
|
|
2150
|
+
choices: [{ index: 0, text, finish_reason: finish }],
|
|
2151
|
+
});
|
|
2152
|
+
try {
|
|
2153
|
+
const decoder = new StreamDecoder(ctx.tokenizer);
|
|
2154
|
+
const stopper = new StopMatcher(options.stopSequences);
|
|
2155
|
+
// ≥25 ms macrotask hop, same reason as the chat lane: keep the
|
|
2156
|
+
// serial decode loop from starving the socket (SSE bursts).
|
|
2157
|
+
let lastFlush = performance.now();
|
|
2158
|
+
const s = await gateway.run(promptIds, options, (token) => {
|
|
2159
|
+
const text = stopper.push(decoder.push(token));
|
|
2160
|
+
if (text) send(chunk(text, null));
|
|
2161
|
+
if (stopper.stopped) return false; // halt generation
|
|
2162
|
+
if (text) {
|
|
2163
|
+
const now = performance.now();
|
|
2164
|
+
if (now - lastFlush >= 25) {
|
|
2165
|
+
lastFlush = now;
|
|
2166
|
+
return new Promise<void>((r) => setImmediate(r));
|
|
2167
|
+
}
|
|
2168
|
+
}
|
|
2169
|
+
}, undefined, shape);
|
|
2170
|
+
if (!stopper.stopped) {
|
|
2171
|
+
let tail = stopper.push(decoder.flush());
|
|
2172
|
+
if (!stopper.stopped) tail += stopper.flush();
|
|
2173
|
+
if (tail) send(chunk(tail, null));
|
|
2174
|
+
}
|
|
2175
|
+
// final chunk: finish_reason + usage (mlx-lm gates usage behind
|
|
2176
|
+
// stream_options.include_usage; we always attach it, matching
|
|
2177
|
+
// our chat lane — an additive superset OpenAI clients ignore)
|
|
2178
|
+
send({
|
|
2179
|
+
...chunk("", finishReason(stopper.stopped, s.generatedTokens)),
|
|
2180
|
+
usage: {
|
|
2181
|
+
prompt_tokens: s.promptTokens,
|
|
2182
|
+
completion_tokens: s.generatedTokens,
|
|
2183
|
+
total_tokens: s.promptTokens + s.generatedTokens,
|
|
2184
|
+
prompt_tokens_details: { cached_tokens: s.cachedTokens },
|
|
2185
|
+
},
|
|
2186
|
+
});
|
|
2187
|
+
controller.enqueue(enc.encode("data: [DONE]\n\n"));
|
|
2188
|
+
} catch (e) {
|
|
2189
|
+
send({ error: { message: (e as Error).message } });
|
|
2190
|
+
} finally {
|
|
2191
|
+
controller.close();
|
|
2192
|
+
}
|
|
2193
|
+
},
|
|
2194
|
+
});
|
|
2195
|
+
return new Response(stream, {
|
|
2196
|
+
headers: {
|
|
2197
|
+
"content-type": "text/event-stream",
|
|
2198
|
+
"cache-control": "no-cache",
|
|
2199
|
+
connection: "keep-alive",
|
|
2200
|
+
},
|
|
2201
|
+
});
|
|
2202
|
+
}
|
|
2203
|
+
|
|
2204
|
+
try {
|
|
2205
|
+
const decoder = new StreamDecoder(ctx.tokenizer);
|
|
2206
|
+
const stopper = new StopMatcher(options.stopSequences);
|
|
2207
|
+
const lpc = captureLogprobs
|
|
2208
|
+
? new LogprobsCollector(wantLogprobs, topLogprobs, (tid) => ctx.tokenizer.idToToken(tid))
|
|
2209
|
+
: null;
|
|
2210
|
+
let text = "";
|
|
2211
|
+
const s = await gateway.run(promptIds, options, (token, lpInfo) => {
|
|
2212
|
+
lpc?.push(token, lpInfo);
|
|
2213
|
+
text += stopper.push(decoder.push(token));
|
|
2214
|
+
if (stopper.stopped) return false; // halt generation
|
|
2215
|
+
}, undefined, shape);
|
|
2216
|
+
if (!stopper.stopped) {
|
|
2217
|
+
let tail = stopper.push(decoder.flush());
|
|
2218
|
+
if (!stopper.stopped) tail += stopper.flush();
|
|
2219
|
+
text += tail;
|
|
2220
|
+
}
|
|
2221
|
+
const logprobsBlock = lpc?.payload() ?? null;
|
|
2222
|
+
return Response.json({
|
|
2223
|
+
id, object: "text_completion", created, model: ctx.modelId,
|
|
2224
|
+
choices: [{
|
|
2225
|
+
index: 0, text,
|
|
2226
|
+
...(logprobsBlock ? { logprobs: logprobsBlock } : {}),
|
|
2227
|
+
finish_reason: finishReason(stopper.stopped, s.generatedTokens),
|
|
2228
|
+
}],
|
|
2229
|
+
usage: {
|
|
2230
|
+
prompt_tokens: s.promptTokens,
|
|
2231
|
+
completion_tokens: s.generatedTokens,
|
|
2232
|
+
total_tokens: s.promptTokens + s.generatedTokens,
|
|
2233
|
+
prompt_tokens_details: { cached_tokens: s.cachedTokens },
|
|
2234
|
+
},
|
|
2235
|
+
});
|
|
2236
|
+
} catch (e) {
|
|
2237
|
+
return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
|
|
2238
|
+
}
|
|
2239
|
+
}
|
|
2240
|
+
|
|
1678
2241
|
// Anthropic Messages API (Phase 11) — on by default, mirroring
|
|
1679
2242
|
// optiq serve (--anthropic defaults True; the drop-in claim
|
|
1680
2243
|
// depends on it). Oracle: optiq/anthropic_shim.py, ported in
|
|
@@ -1932,7 +2495,7 @@ export function createServer(
|
|
|
1932
2495
|
const snapMatch = body.model_id.match(/(models--[^/]+)\/snapshots\//);
|
|
1933
2496
|
let org = "local", name: string;
|
|
1934
2497
|
if (snapMatch) {
|
|
1935
|
-
const parts = snapMatch[1]
|
|
2498
|
+
const parts = snapMatch[1]!.split("--"); // ["models", org, ...name]
|
|
1936
2499
|
org = parts[1] ?? "local";
|
|
1937
2500
|
name = parts.slice(2).join("--");
|
|
1938
2501
|
} else if (body.model_id.includes("/") && !body.model_id.startsWith("/") && !body.model_id.startsWith("~")) {
|