mlx-bun 0.0.10 → 0.0.11
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 +66 -21
- package/THIRD_PARTY_LICENSES.md +1 -0
- package/docs/archive/mlx-bun-lab-report.html +275 -0
- package/docs/archive/wwdc-mlx-bun.html +848 -0
- package/docs/design/audio-input-plan.md +356 -0
- package/docs/design/batching-perf-path.md +58 -35
- package/docs/design/batching-v2-plan.md +23 -9
- package/docs/design/bucketing-stage.md +2 -1
- package/docs/design/compat-cli-surface-design.md +4 -3
- package/docs/design/decode-speed-program.md +172 -0
- package/docs/design/docs-reorg-plan.md +1 -1
- package/docs/design/dspark-serving-program.md +371 -0
- package/docs/design/dspark-speculative-decoding.md +167 -5
- package/docs/design/faithful-l1-consolidation.md +310 -0
- package/docs/design/generic-model-support.md +8 -5
- package/docs/design/grammar-spec-batching-integration.md +301 -0
- package/docs/design/minicpm5-decode-megakernel.md +16 -11
- package/docs/design/mlx-lm-tool-parity-plan.md +11 -1
- package/docs/design/omlx-adoption-map.md +14 -6
- package/docs/design/orpo-training.md +34 -30
- package/docs/design/paged-kv-cache.md +214 -0
- package/docs/design/parallel-slots.md +82 -49
- package/docs/design/parity-tier-dag.md +43 -33
- package/docs/design/repo-cleanup-plan.md +189 -0
- package/docs/design/runtime-isolation.md +19 -3
- package/docs/design/segmented-backward-training.md +3 -2
- package/docs/design/ssd-kv-cold-tier.md +130 -0
- package/docs/design/structured-output.md +484 -0
- package/docs/design/the-dreaming-master-plan.md +2 -3
- package/docs/design/turboquant-kv.md +235 -0
- package/docs/design/unified-engine-frontier-plan.md +774 -0
- package/docs/design/web-chat-beat-matrix.md +315 -0
- package/docs/design/web-chat-redesign.md +1112 -0
- package/docs/investigations/dspark-handoff.md +182 -16
- package/docs/investigations/lab-build-journal.md +2 -2
- package/docs/investigations/openwiki-evaluation.md +249 -0
- package/docs/planning/IDEAS.md +14 -0
- package/docs/planning/PRODUCT_ROADMAP.md +2 -3
- package/docs/planning/chat-ui-vision-memory-first.md +199 -0
- package/docs/planning/chat-ui-vision.md +50 -0
- package/docs/planning/memory-docs-and-dag-plan.md +15 -6
- package/docs/planning/release-notes-v0.0.11.md +239 -0
- package/docs/reference/cli.md +35 -16
- package/docs/reference/distribution.md +13 -6
- package/docs/reference/embedding.md +35 -14
- package/docs/reference/features-matrix.md +97 -0
- package/docs/reference/library-api.md +30 -9
- package/docs/reference/memory.md +26 -6
- package/docs/reference/models.md +11 -1
- package/docs/reference/orpo-quickstart.md +3 -4
- package/docs/reference/server-api.md +634 -22
- package/docs/reference/server-config.md +250 -121
- package/docs/reference/training.md +14 -15
- package/package.json +5 -3
- package/src/anthropic.ts +13 -0
- package/src/audio/conformer.ts +647 -0
- package/src/audio/decode.ts +133 -0
- package/src/audio/features.ts +202 -0
- package/src/audio/transcode.ts +67 -0
- package/src/cli.ts +356 -105
- package/src/config.ts +49 -2
- package/src/download.ts +40 -0
- package/src/eval/README.md +7 -6
- package/src/eval/kl.ts +59 -3
- package/src/eval/quality-db.ts +2 -2
- package/src/eval/runner.ts +3 -3
- package/src/expert-offload.ts +15 -2
- package/src/flags.ts +13 -0
- package/src/generate.ts +330 -24
- package/src/grammar.ts +585 -0
- package/src/hub-rest.ts +242 -0
- package/src/kv-store.ts +442 -88
- package/src/lora.ts +8 -0
- package/src/md-modules.d.ts +42 -0
- package/src/media-fetch.ts +284 -0
- package/src/memory/rest.ts +245 -0
- package/src/memory/vault.ts +74 -0
- package/src/mlx/array.ts +106 -19
- package/src/mlx/autograd.ts +2 -2
- package/src/mlx/ffi.ts +27 -3
- package/src/mlx/metal-kernel.ts +7 -4
- package/src/mlx/ops.ts +153 -23
- package/src/mlx/turboquant-ops.ts +831 -0
- package/src/mlx/turboquant-tables.ts +237 -0
- package/src/model/batched-mask.ts +62 -1
- package/src/model/batched-quant.ts +206 -0
- package/src/model/batched-rotating-quant.ts +346 -0
- package/src/model/batched-rotating.ts +46 -1
- package/src/model/compiled-decode.ts +4 -15
- package/src/model/factory.ts +10 -4
- package/src/model/gemma4-base.ts +363 -39
- package/src/model/gemma4.ts +122 -33
- package/src/model/generated/gemma4-12b.ts +9 -13
- package/src/model/generated/gemma4-26b.ts +10 -5
- package/src/model/generated/gemma4-e4b.ts +14 -9
- package/src/model/minicpm5.ts +58 -102
- package/src/model/paged-kv.ts +253 -0
- package/src/model/qwen3-delta.ts +76 -15
- package/src/model/qwen3-moe.ts +404 -0
- package/src/model/qwen3.ts +9 -8
- package/src/model/qwen3_5.ts +88 -23
- package/src/model/support.ts +12 -1
- package/src/model/universal/dense.ts +40 -5
- package/src/model/universal/rope.ts +36 -21
- package/src/pi-web.ts +1299 -82
- package/src/preflight.ts +27 -8
- package/src/prompt-cache.ts +293 -33
- package/src/quantize/quantizer.ts +16 -1
- package/src/registry.ts +57 -2
- package/src/serve/batch-scheduler.ts +965 -72
- package/src/serve/generation-gateway.ts +214 -11
- package/src/serve/isolate.ts +417 -0
- package/src/serve/lane-registry.ts +44 -0
- package/src/serve/session-search.ts +326 -0
- package/src/server.ts +1380 -135
- package/src/spec/assistant-source.ts +150 -0
- package/src/spec/deepspec-source.ts +143 -0
- package/src/spec/dflash-source.ts +124 -0
- package/src/spec/dspark/ab-stats.ts +136 -0
- package/src/spec/dspark/calibration.ts +112 -0
- package/src/spec/dspark/deepspec-module.ts +801 -0
- package/src/spec/dspark/generate-dflash.ts +72 -20
- package/src/spec/dspark/loader.ts +23 -0
- package/src/spec/dspark/module-dflash.ts +278 -36
- package/src/spec/dspark/quantize-drafter.ts +91 -0
- package/src/spec/generate.ts +59 -25
- package/src/spec/ngram-source.ts +132 -0
- package/src/spec/serve-loop.ts +536 -0
- package/src/spec/source.ts +93 -0
- package/src/spec/two-model.ts +171 -0
- package/src/ssd-cache.ts +107 -35
- package/src/tokenizer.ts +288 -2
- package/src/tool-approvals.ts +131 -0
- package/src/tool-call.ts +408 -20
- package/src/train/flash-cce.ts +2 -2
- package/src/train/job.ts +0 -2
- package/src/train/trainer.ts +1 -4
- package/src/vision/prompt.ts +217 -48
- package/src/web/app.html +1184 -1789
- package/src/web/app.js +5667 -0
- package/src/web/icon.svg +13 -0
- package/src/web/manifest.webmanifest +24 -0
- package/src/web/src/adapters-panel.ts +288 -0
- package/src/web/src/api.ts +57 -0
- package/src/web/src/assistant.ts +485 -0
- package/src/web/src/chat.ts +1050 -0
- package/src/web/src/composer.ts +959 -0
- package/src/web/src/dataset.ts +132 -0
- package/src/web/src/finetune.ts +267 -0
- package/src/web/src/hub.ts +428 -0
- package/src/web/src/main.ts +98 -0
- package/src/web/src/markdown.ts +0 -0
- package/src/web/src/memory-panel.ts +572 -0
- package/src/web/src/model-picker.ts +159 -0
- package/src/web/src/palette.ts +346 -0
- package/src/web/src/protocol.ts +62 -0
- package/src/web/src/quantize.ts +167 -0
- package/src/web/src/rag.ts +274 -0
- package/src/web/src/sessions.ts +469 -0
- package/src/web/src/shell.ts +804 -0
- package/src/web/src/status.ts +252 -0
- package/src/web/src/ui-catalog.ts +92 -0
- package/src/web/sw.js +71 -0
- package/src/web/vendor/README.md +61 -0
- package/src/web/vendor/hljs-LICENSE +29 -0
- package/src/web/vendor/hljs-theme.css +59 -0
- package/src/web/vendor/hljs.js +26 -0
- package/src/model/fused-decode-kernel.ts +0 -312
- package/src/model/fused-geglu-kernel.ts +0 -198
- package/src/model/fused-mlp-kernel.ts +0 -169
- package/src/model/fused-swiglu-kernel.ts +0 -240
- package/src/model/steel-linear-kernel.ts +0 -96
- package/src/status-page.html +0 -497
package/src/server.ts
CHANGED
|
@@ -16,22 +16,56 @@ import type { Server } from "bun";
|
|
|
16
16
|
// the runtime value a string — hence the double cast.
|
|
17
17
|
import appHtml from "./web/app.html" with { type: "text" };
|
|
18
18
|
import curveDesignerHtml from "./assets/curve-designer.html" with { type: "text" };
|
|
19
|
+
// Vendored, no-CDN static assets referenced by app.html (convention: any
|
|
20
|
+
// self-contained JS/CSS too big to inline gets `with { type: "text" }`
|
|
21
|
+
// imported here and served under /assets/<name>; see src/web/vendor/README
|
|
22
|
+
// for how hljs.js was built). Add new vendored assets the same way.
|
|
23
|
+
import hljsJs from "./web/vendor/hljs.js" with { type: "text" };
|
|
24
|
+
import hljsCss from "./web/vendor/hljs-theme.css" with { type: "text" };
|
|
25
|
+
// The frontend bundle (plan §7/§9 Phase 2 module split): GENERATED from
|
|
26
|
+
// src/web/src/*.ts by `bun scripts/build-web.ts` — see that file's header
|
|
27
|
+
// and tests/web-build.test.ts (the freshness gate). Same
|
|
28
|
+
// with { type: "text" } + /assets/<name> pattern as the vendored assets
|
|
29
|
+
// above; app.html's <script defer src="/assets/app.js"> loads it.
|
|
30
|
+
import appJs from "./web/app.js" with { type: "text" };
|
|
31
|
+
// PWA installability (plan §9 Phase 3, beat-matrix Axis 10): a manifest +
|
|
32
|
+
// a single inline SVG icon (no binary PNGs — the hygiene gate forbids
|
|
33
|
+
// tracked binary files; some browsers won't show an SVG app icon, which is
|
|
34
|
+
// an accepted tradeoff, see docs/reference/features-matrix.md) + a shell-
|
|
35
|
+
// only service worker. Same with { type: "text" } + /assets/<name>-shaped
|
|
36
|
+
// pattern as the vendored assets above; see src/web/sw.js's header for why
|
|
37
|
+
// it deliberately does NOT cache API/WS traffic.
|
|
38
|
+
import manifestWebmanifest from "./web/manifest.webmanifest" with { type: "text" };
|
|
39
|
+
import iconSvg from "./web/icon.svg" with { type: "text" };
|
|
40
|
+
import swJs from "./web/sw.js" with { type: "text" };
|
|
19
41
|
import pkgJson from "../package.json" with { type: "json" };
|
|
20
42
|
import { readFileSync } from "node:fs";
|
|
21
43
|
const APP_PAGE = appHtml as unknown as string;
|
|
44
|
+
const HLJS_JS = hljsJs as unknown as string;
|
|
45
|
+
const HLJS_CSS = hljsCss as unknown as string;
|
|
46
|
+
const APP_JS = appJs as unknown as string;
|
|
47
|
+
const MANIFEST_WEBMANIFEST = manifestWebmanifest as unknown as string;
|
|
48
|
+
const ICON_SVG = iconSvg as unknown as string;
|
|
49
|
+
const SW_JS = swJs as unknown as string;
|
|
22
50
|
const pkgVersion = (pkgJson as { version: string }).version;
|
|
23
|
-
import { loadModelConfig, type KvQuantSpec, type ModelConfig } from "./config";
|
|
51
|
+
import { loadModelConfig, type KvQuantSpec, type ModelConfig, type TurboQuantScheme } from "./config";
|
|
24
52
|
import { Weights } from "./weights";
|
|
25
53
|
import { Gemma4Model } from "./model/gemma4";
|
|
26
54
|
import { DiffusionGemmaModel } from "./model/diffusion-gemma";
|
|
27
55
|
import { spliceImageTokens } from "./vision/diffusion-vision";
|
|
28
56
|
import { createModel, type RuntimeModel } from "./model/factory";
|
|
29
|
-
import { isMiniCPM5Config
|
|
57
|
+
import { isMiniCPM5Config } from "./model/support";
|
|
30
58
|
import { generate, type GenerateOptions, type TokenLogprobs } from "./generate";
|
|
59
|
+
import { cloneKvCaches, SpillQueue } from "./kv-store";
|
|
60
|
+
import { TURBOQUANT_HEAD_DIMS } from "./mlx/turboquant-tables";
|
|
61
|
+
import {
|
|
62
|
+
compileGrammarRequest, grammarEnabled, type GrammarRequest,
|
|
63
|
+
} from "./grammar";
|
|
31
64
|
import type { HlgConfig } from "./sampler";
|
|
32
65
|
import { isMonotone, CURVE_UMIN, type CurveParams } from "./curve-sampler";
|
|
33
66
|
const CURVE_PAGE = curveDesignerHtml as unknown as string;
|
|
34
67
|
import { GenerationGateway } from "./serve/generation-gateway";
|
|
68
|
+
import { recordLane, type Lane } from "./serve/lane-registry";
|
|
35
69
|
import {
|
|
36
70
|
ChatTemplate, type ChatMessage, type ToolDefinition,
|
|
37
71
|
} from "./chat-template";
|
|
@@ -40,7 +74,7 @@ import {
|
|
|
40
74
|
CHANNEL_END, CHANNEL_START, parseGeneratedToolCalls, parseToolCalls,
|
|
41
75
|
TOOL_CALL_END, TOOL_CALL_START,
|
|
42
76
|
} from "./tool-call";
|
|
43
|
-
import { PromptCache } from "./prompt-cache";
|
|
77
|
+
import { PromptCache, cacheBytes } from "./prompt-cache";
|
|
44
78
|
import { SsdCacheStore } from "./ssd-cache";
|
|
45
79
|
import { configFingerprint } from "./model/fingerprint";
|
|
46
80
|
import {
|
|
@@ -59,14 +93,23 @@ import { setMemoryLimit } from "./mlx/ffi";
|
|
|
59
93
|
import { VisionTower } from "./vision/embedder";
|
|
60
94
|
import { SiglipVisionTower, parseSiglipConfig } from "./vision/siglip";
|
|
61
95
|
import {
|
|
62
|
-
buildVisionPrompt, extractImages,
|
|
63
|
-
type VisionTokenIds, type VisionEncoder,
|
|
96
|
+
buildMultimodalPrompt, buildVisionPrompt, extractAudio, extractImages,
|
|
97
|
+
type AudioTokenIds, type VisionTokenIds, type VisionEncoder,
|
|
64
98
|
} from "./vision/prompt";
|
|
99
|
+
import { AudioTower, parseAudioConfig } from "./audio/conformer";
|
|
100
|
+
import { ensureWav } from "./audio/transcode";
|
|
101
|
+
import { sidecarShipsAudioTower } from "./registry";
|
|
65
102
|
import { makePiWsHandler, type PiWsData } from "./pi-web";
|
|
66
103
|
|
|
67
104
|
export interface ServerOptions {
|
|
68
|
-
/** Byte cap for the prompt (KV) cache. Default
|
|
105
|
+
/** Byte cap for the prompt (KV) cache. Default 8 GB. */
|
|
69
106
|
promptCacheBytes?: number;
|
|
107
|
+
/** Aggregate KV-byte budget across concurrently-admitted batch rows
|
|
108
|
+
* (`--kv-budget`, batching-perf-path P3). Joiners whose projected KV
|
|
109
|
+
* (prompt + max_tokens, window-capped) would exceed it QUEUE until rows
|
|
110
|
+
* evict; a request over the budget alone is rejected. Unset = no
|
|
111
|
+
* aggregate cap (per-request admission via memoryBudget still applies). */
|
|
112
|
+
kvBudgetBytes?: number;
|
|
70
113
|
/** KV quantization override. When unset: apply ctx.kvConfig (mixed
|
|
71
114
|
* per-layer) for serial serving, but bf16 under `--batch N` (the batched
|
|
72
115
|
* engine is bf16-only — a mode switch, see `batch` below). "config" forces
|
|
@@ -74,6 +117,19 @@ export interface ServerOptions {
|
|
|
74
117
|
* the serial path); "off" forces bf16; a number forces uniform bits
|
|
75
118
|
* (group size 64, start 0) ignoring the config file. */
|
|
76
119
|
kvQuant?: "off" | "config" | number;
|
|
120
|
+
/** TurboQuant scheme (docs/design/turboquant-kv.md): a separate axis from
|
|
121
|
+
* kvQuant above, mutually exclusive with it (`--kv-quant turbo[:k<bits>v
|
|
122
|
+
* <bits>]` sets this instead of kvQuant). Solo-only in v1 — see
|
|
123
|
+
* GenerationGateway's explicit refusal. */
|
|
124
|
+
turboQuant?: TurboQuantScheme;
|
|
125
|
+
/** OPTIONAL paged KV cache (`--paged-kv`, docs/design/paged-kv-cache.md):
|
|
126
|
+
* vLLM-style block-pool storage for full-attention layers,
|
|
127
|
+
* gather-to-contiguous before the stock SDPA. Default off (unset = the
|
|
128
|
+
* plain KVCache path, byte-identical). v1 scope: serial batch=1,
|
|
129
|
+
* Gemma4-family, bf16 — startup refuses `--batch N>1`, `--kv-quant`,
|
|
130
|
+
* turbo, and `--draft-model` combinations; paged requests bypass the
|
|
131
|
+
* prompt cache and run uncompiled. Gated bit-exact vs the plain path. */
|
|
132
|
+
pagedKv?: { blockSize?: number };
|
|
77
133
|
/** Memory budget for the serving process (admission control — Phase 5).
|
|
78
134
|
* Requests whose prompt + max_tokens exceed the budget's max safe
|
|
79
135
|
* context are rejected with 400 instead of crashing the GPU: the OOM
|
|
@@ -88,6 +144,11 @@ export interface ServerOptions {
|
|
|
88
144
|
* it). Exposed at /stats.server so other mlx-bun processes can
|
|
89
145
|
* warn before attaching to a server that may vanish. */
|
|
90
146
|
owner?: "serve" | "pi-session" | "embedded";
|
|
147
|
+
/** Listen on a UNIX DOMAIN SOCKET instead of a TCP port — the engine-
|
|
148
|
+
* child mode of the isolation architecture (runtime-isolation.md): the
|
|
149
|
+
* parent process reverse-proxies HTTP to this socket. Stale socket
|
|
150
|
+
* files are unlinked before bind. When set, `hostname`/port are ignored. */
|
|
151
|
+
unixSocket?: string;
|
|
91
152
|
/** Interface to bind (Bun.serve hostname). Unset ⇒ Bun's default
|
|
92
153
|
* (all interfaces) — embedded/library use. The CLI always passes one:
|
|
93
154
|
* "127.0.0.1" (loopback, mlx_lm.server parity) unless --host says
|
|
@@ -111,9 +172,9 @@ export interface ServerOptions {
|
|
|
111
172
|
* thinking traces never truncate. `--max-tokens 512` = mlx-lm behavior. */
|
|
112
173
|
defaultMaxTokens?: number;
|
|
113
174
|
/** Max concurrent requests batched through the mlx-lm-parity engine
|
|
114
|
-
* (`--batch N`). Default
|
|
115
|
-
*
|
|
116
|
-
*
|
|
175
|
+
* (`--batch N`). Default 8 (continuous batching, B floats 1..N,
|
|
176
|
+
* bit-parity with mlx-lm B=N); `--batch 1` pins the serialized
|
|
177
|
+
* single-queue path — a mode switch, not a load-dependent
|
|
117
178
|
* fallback. See docs/design/parallel-slots.md. NOTE: the batched executor
|
|
118
179
|
* is mid-build; until it lands, >1 warns and runs serially. */
|
|
119
180
|
batch?: number;
|
|
@@ -132,6 +193,11 @@ export interface ServerOptions {
|
|
|
132
193
|
ssdCacheDir?: string;
|
|
133
194
|
/** Byte cap for the SSD tier (default 32 GiB). */
|
|
134
195
|
ssdCacheMaxBytes?: number;
|
|
196
|
+
/** Idle-demotion threshold in seconds (`--ssd-demote-idle`): prompt-cache
|
|
197
|
+
* entries unused this long spill to the SSD tier and free their GPU
|
|
198
|
+
* memory (Layer 0 — RAM drains between bursts, prefixes stay reachable
|
|
199
|
+
* via zero-copy restore). Default 300 when the tier is on; 0 disables. */
|
|
200
|
+
ssdDemoteIdleSec?: number;
|
|
135
201
|
/** Verify every tensor hash on restore (`--ssd-cache-verify`) — reads all
|
|
136
202
|
* bytes eagerly, defeating lazy fault-in; integrity paranoia only. */
|
|
137
203
|
ssdCacheVerify?: boolean;
|
|
@@ -151,6 +217,17 @@ export interface ServerContext {
|
|
|
151
217
|
* `vision`. */
|
|
152
218
|
loadVision: (() => VisionEncoder) | null;
|
|
153
219
|
visionTokenIds: VisionTokenIds;
|
|
220
|
+
/** Lazily-loaded Conformer audio tower — null until the first audio
|
|
221
|
+
* request (see `getAudioTower`). Same sidecar file as vision
|
|
222
|
+
* (optiq_vision.safetensors), separate tower; text-only sessions never
|
|
223
|
+
* pay for it. */
|
|
224
|
+
audio: AudioTower | null;
|
|
225
|
+
/** Loads the audio tower on demand; null when the model can't do audio
|
|
226
|
+
* (no `audio_config` in config.json or no sidecar). No flags — audio
|
|
227
|
+
* auto-enables exactly like vision. */
|
|
228
|
+
loadAudio: (() => AudioTower) | null;
|
|
229
|
+
/** null when the model has no `audio_config` (audio-incapable). */
|
|
230
|
+
audioTokenIds: AudioTokenIds | null;
|
|
154
231
|
adapters: AdapterManager;
|
|
155
232
|
/** Per-layer KV quantization from the repo's kv_config.json (null if
|
|
156
233
|
* absent). Applied by default — optiq serve's headline behavior;
|
|
@@ -160,6 +237,14 @@ export interface ServerContext {
|
|
|
160
237
|
* optiq serve injects these as server defaults (gen_config.py);
|
|
161
238
|
* explicit request fields always win. */
|
|
162
239
|
genDefaults: GenSamplingDefaults;
|
|
240
|
+
/** Speculative decoding (`serve --draft-model`, mlx_lm.server parity).
|
|
241
|
+
* Server-level: when set, EVERY request routes to the serial lane
|
|
242
|
+
* (upstream: is_batchable = draft is None) and spec-eligible ones decode
|
|
243
|
+
* through src/spec/serve-loop.ts. null = no draft configured. */
|
|
244
|
+
draft?: {
|
|
245
|
+
provider: import("./spec/source").DraftProvider;
|
|
246
|
+
numDraftTokens: number;
|
|
247
|
+
} | null;
|
|
163
248
|
}
|
|
164
249
|
|
|
165
250
|
export interface GenSamplingDefaults {
|
|
@@ -169,6 +254,30 @@ export interface GenSamplingDefaults {
|
|
|
169
254
|
repetitionPenalty?: number;
|
|
170
255
|
}
|
|
171
256
|
|
|
257
|
+
export type DraftKind = "dspark" | "deepspec" | "assistant" | "two-model" | "ngram";
|
|
258
|
+
|
|
259
|
+
/** Detect the draft artifact's kind so the right provider is loaded. All
|
|
260
|
+
* providers share ONE serve loop (src/spec/serve-loop.ts). Exported for the
|
|
261
|
+
* bench harness (scripts/bench-feature-matrix.ts) — one detection, no drift.
|
|
262
|
+
* "ngram" is never detected — it has no artifact (model-free prompt lookup,
|
|
263
|
+
* src/spec/ngram-source.ts) and mounts via an explicit `--draft-kind ngram`. */
|
|
264
|
+
export async function detectDraftKind(dir: string): Promise<DraftKind> {
|
|
265
|
+
if (await Bun.file(`${dir}/dspark.json`).exists()) return "dspark"; // our trained module
|
|
266
|
+
try {
|
|
267
|
+
const cfg = (await Bun.file(`${dir}/config.json`).json()) as {
|
|
268
|
+
model_type?: string;
|
|
269
|
+
architectures?: string[];
|
|
270
|
+
};
|
|
271
|
+
// DeepSeek's released DSpark drafters (DeepSpec reference): no
|
|
272
|
+
// dspark.json, plain HF config stamped Gemma4DSparkModel.
|
|
273
|
+
if (cfg.architectures?.[0] === "Gemma4DSparkModel") return "deepspec";
|
|
274
|
+
if (String(cfg.model_type ?? "").includes("assistant")) return "assistant";
|
|
275
|
+
} catch {
|
|
276
|
+
// no/unreadable config → fall through to a full second model
|
|
277
|
+
}
|
|
278
|
+
return "two-model";
|
|
279
|
+
}
|
|
280
|
+
|
|
172
281
|
async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDefaults> {
|
|
173
282
|
const file = Bun.file(`${modelDir}/generation_config.json`);
|
|
174
283
|
if (!(await file.exists())) return {};
|
|
@@ -188,7 +297,24 @@ async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDef
|
|
|
188
297
|
|
|
189
298
|
export async function loadContext(
|
|
190
299
|
modelDir: string, modelId?: string,
|
|
191
|
-
opts: {
|
|
300
|
+
opts: {
|
|
301
|
+
memoryBudgetBytes?: number;
|
|
302
|
+
/** Snapshot dir of a draft model for speculative decoding
|
|
303
|
+
* (`--draft-model`). Loaded alongside the target; the pair must share
|
|
304
|
+
* a tokenizer family (hard check below — upstream's silent-garbage
|
|
305
|
+
* mode on mismatched tokenizers isn't worth inheriting). */
|
|
306
|
+
draftModelDir?: string;
|
|
307
|
+
/** Drafts per round (`--num-draft-tokens`, mlx_lm.server default 3). */
|
|
308
|
+
numDraftTokens?: number;
|
|
309
|
+
/** Draft-provider kind override (`--draft-kind`); auto-detected from the
|
|
310
|
+
* draft artifact when absent (dspark.json → dspark; *_assistant config →
|
|
311
|
+
* assistant; otherwise two-model). "ngram" mounts WITHOUT a draft dir. */
|
|
312
|
+
draftKind?: DraftKind;
|
|
313
|
+
/** Prompt-lookup window bounds (`--ngram-max` / `--ngram-min`), ngram
|
|
314
|
+
* kind only. Defaults 3 / 1 (Saxena's reference values). */
|
|
315
|
+
ngramMax?: number;
|
|
316
|
+
ngramMin?: number;
|
|
317
|
+
} = {},
|
|
192
318
|
): Promise<ServerContext> {
|
|
193
319
|
const config = await loadModelConfig(modelDir);
|
|
194
320
|
const weights = await Weights.open(modelDir);
|
|
@@ -219,7 +345,116 @@ export async function loadContext(
|
|
|
219
345
|
// the tokenizer eos; union it in. No-op when already present (Gemma, 27B).
|
|
220
346
|
if (tokenizer.eosTokenId != null && !config.eosTokenIds.includes(tokenizer.eosTokenId))
|
|
221
347
|
config.eosTokenIds = [...config.eosTokenIds, tokenizer.eosTokenId];
|
|
348
|
+
|
|
349
|
+
// Speculative decoding: load the draft (mlx_lm.server --draft-model). The
|
|
350
|
+
// draft artifact's KIND selects the provider — all three share ONE serve
|
|
351
|
+
// loop (src/spec/serve-loop.ts): dspark.json → DSpark (KV-injected), a
|
|
352
|
+
// *_assistant config → the optiq KV-borrowing Gemma drafter, otherwise a
|
|
353
|
+
// full second model (mlx-lm parity). `--draft-kind` overrides the detect.
|
|
354
|
+
let draft: ServerContext["draft"] = null;
|
|
355
|
+
if (opts.draftKind === "ngram") {
|
|
356
|
+
// Model-free prompt lookup: no artifact, no dir, no probe/budget concerns
|
|
357
|
+
// (weightsBytes 0, open() never throws). Default γ=10 per the reference
|
|
358
|
+
// implementation — drafting is free, so wide blocks cost only verify-window
|
|
359
|
+
// width when wrong.
|
|
360
|
+
if (opts.draftModelDir)
|
|
361
|
+
throw new Error(
|
|
362
|
+
"--draft-kind ngram is model-free — drop --draft-model (it would be ignored)",
|
|
363
|
+
);
|
|
364
|
+
const { NgramProvider } = await import("./spec/ngram-source");
|
|
365
|
+
draft = {
|
|
366
|
+
provider: new NgramProvider({ max: opts.ngramMax, min: opts.ngramMin }),
|
|
367
|
+
numDraftTokens: Math.max(1, opts.numDraftTokens ?? 10),
|
|
368
|
+
};
|
|
369
|
+
} else if (opts.draftModelDir) {
|
|
370
|
+
const dir = opts.draftModelDir;
|
|
371
|
+
const kind = opts.draftKind ?? (await detectDraftKind(dir));
|
|
372
|
+
let provider: import("./spec/source").DraftProvider;
|
|
373
|
+
let numDraftTokens = Math.max(1, opts.numDraftTokens ?? 3);
|
|
374
|
+
if (kind === "dspark") {
|
|
375
|
+
const { DflashProvider } = await import("./spec/dflash-source");
|
|
376
|
+
const p = await DflashProvider.load(dir);
|
|
377
|
+
provider = p;
|
|
378
|
+
// Pin to the trained block width — the serve loop must never ask for
|
|
379
|
+
// more positions than the DSpark block was trained for (n ≤ cfg.gamma).
|
|
380
|
+
numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
|
|
381
|
+
} else if (kind === "deepspec") {
|
|
382
|
+
const { DeepspecProvider } = await import("./spec/deepspec-source");
|
|
383
|
+
const p = await DeepspecProvider.load(dir);
|
|
384
|
+
provider = p;
|
|
385
|
+
// Same pin, from their config's block_size (e.g. 7 for the released
|
|
386
|
+
// dspark_gemma4_12b_block7).
|
|
387
|
+
numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
|
|
388
|
+
} else if (kind === "assistant") {
|
|
389
|
+
const { AssistantProvider } = await import("./spec/assistant-source");
|
|
390
|
+
provider = await AssistantProvider.load(dir);
|
|
391
|
+
} else {
|
|
392
|
+
const { TwoModelProvider } = await import("./spec/two-model");
|
|
393
|
+
provider = await TwoModelProvider.load(dir, config.text.vocabSize);
|
|
394
|
+
// Tokenizer-family hard check — two-model ONLY (it ships its own
|
|
395
|
+
// tokenizer). Exact-token-match acceptance is meaningful only when both
|
|
396
|
+
// models tokenize identically; a probe that ENCODES differently means
|
|
397
|
+
// different families — refuse instead of silently accepting ~0% of
|
|
398
|
+
// drafts. The KV-borrowing drafters share the target's tokenization by
|
|
399
|
+
// construction, so this check does not apply to them.
|
|
400
|
+
const draftTok = await loadTokenizer(dir);
|
|
401
|
+
const probe = "The 3 quick brown foxes jumped över the lazy dog?! 🦊";
|
|
402
|
+
if (JSON.stringify(tokenizer.encode(probe)) !== JSON.stringify(draftTok.encode(probe))) {
|
|
403
|
+
provider.dispose();
|
|
404
|
+
throw new Error(
|
|
405
|
+
`--draft-model tokenizer differs from the target's (probe string encodes ` +
|
|
406
|
+
`differently) — speculation needs the same tokenizer family`,
|
|
407
|
+
);
|
|
408
|
+
}
|
|
409
|
+
}
|
|
410
|
+
// Fail-fast pairing validation (2026-07-07 review): the KV-borrowing
|
|
411
|
+
// sources validate the (target, drafter) pairing in open() — non-Gemma4
|
|
412
|
+
// target, DeepSpec target-layer-count mismatch — which used to surface
|
|
413
|
+
// as a 500 from inside specServeRun on EVERY text request. Probe-open
|
|
414
|
+
// once with throwaway caches here so a mismatch refuses at load; open()
|
|
415
|
+
// allocates no per-request tensors before prefill/draft, so this is
|
|
416
|
+
// free. The probe sampler is never called during open().
|
|
417
|
+
{
|
|
418
|
+
const probeCaches = model.makeCache();
|
|
419
|
+
try {
|
|
420
|
+
provider
|
|
421
|
+
.open({
|
|
422
|
+
sampler: () => { throw new Error("probe sampler never samples"); },
|
|
423
|
+
target: { model, caches: probeCaches },
|
|
424
|
+
})
|
|
425
|
+
.dispose();
|
|
426
|
+
} catch (err) {
|
|
427
|
+
provider.dispose();
|
|
428
|
+
throw new Error(
|
|
429
|
+
`--draft-model is incompatible with this target: ${err instanceof Error ? err.message : String(err)}`,
|
|
430
|
+
);
|
|
431
|
+
} finally {
|
|
432
|
+
for (const c of probeCaches) c.dispose();
|
|
433
|
+
}
|
|
434
|
+
}
|
|
435
|
+
if (opts.memoryBudgetBytes) {
|
|
436
|
+
const targetBytes = [...weights.shards.files.values()].reduce((a, f) => a + f.mmap.size, 0);
|
|
437
|
+
// Draft weights shrink the target's envelope. Draft KV is not modeled
|
|
438
|
+
// (small relative to its weights at serve contexts); admission stays
|
|
439
|
+
// approximately conservative via the combined-weights term.
|
|
440
|
+
const report = fit(config, targetBytes + provider.weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
|
|
441
|
+
if (report.maxSafeContext < 1) {
|
|
442
|
+
provider.dispose();
|
|
443
|
+
throw new Error(
|
|
444
|
+
`target + draft do not fit the memory budget (draft adds ` +
|
|
445
|
+
`${(provider.weightsBytes / 1e9).toFixed(2)} GB)`,
|
|
446
|
+
);
|
|
447
|
+
}
|
|
448
|
+
}
|
|
449
|
+
draft = { provider, numDraftTokens };
|
|
450
|
+
} else if (opts.draftKind) {
|
|
451
|
+
// Every other kind names an artifact to load — refuse instead of silently
|
|
452
|
+
// serving without speculation.
|
|
453
|
+
throw new Error(`--draft-kind ${opts.draftKind} requires --draft-model`);
|
|
454
|
+
}
|
|
455
|
+
|
|
222
456
|
return {
|
|
457
|
+
draft,
|
|
223
458
|
model,
|
|
224
459
|
adapters: new AdapterManager(model),
|
|
225
460
|
kvConfig: config.kvQuant,
|
|
@@ -239,6 +474,17 @@ export async function loadContext(
|
|
|
239
474
|
boiTokenId: (config.raw.boi_token_id as number) ?? 255999,
|
|
240
475
|
eoiTokenId: (config.raw.eoi_token_id as number) ?? 258882,
|
|
241
476
|
},
|
|
477
|
+
// Audio mirrors vision: lazy tower from the same sidecar, loaded on the
|
|
478
|
+
// first audio request only (docs/design/audio-input-plan.md A4).
|
|
479
|
+
audio: null,
|
|
480
|
+
loadAudio: makeAudioLoader(modelDir, model, config),
|
|
481
|
+
audioTokenIds: config.raw.audio_config
|
|
482
|
+
? {
|
|
483
|
+
audioTokenId: (config.raw.audio_token_id as number) ?? 258881,
|
|
484
|
+
boaTokenId: (config.raw.boa_token_id as number) ?? 256000,
|
|
485
|
+
eoaTokenId: (config.raw.eoa_token_id as number) ?? 258883,
|
|
486
|
+
}
|
|
487
|
+
: null,
|
|
242
488
|
};
|
|
243
489
|
}
|
|
244
490
|
|
|
@@ -276,6 +522,46 @@ function getVisionTower(ctx: ServerContext): VisionEncoder | null {
|
|
|
276
522
|
}
|
|
277
523
|
}
|
|
278
524
|
|
|
525
|
+
/** Build the on-demand audio-tower loader (gemma-4 Conformer, A4 of
|
|
526
|
+
* docs/design/audio-input-plan.md). Auto-enables — no flags — when
|
|
527
|
+
* config.json carries an `audio_config` AND the optiq_vision.safetensors
|
|
528
|
+
* sidecar exists (the audio tensors ship in the same sidecar as vision).
|
|
529
|
+
* Returns null when the model can't do audio (no audio_config: 26B-A4B,
|
|
530
|
+
* DiffusionGemma, bf16 assistants — architectural, not a porting gap). */
|
|
531
|
+
function makeAudioLoader(
|
|
532
|
+
modelDir: string, model: RuntimeModel, config: ModelConfig,
|
|
533
|
+
): (() => AudioTower) | null {
|
|
534
|
+
if (!(config.hasVisionSidecar && config.raw.audio_config && model instanceof Gemma4Model))
|
|
535
|
+
return null;
|
|
536
|
+
// The sidecar header must actually name the Conformer tensors: the local
|
|
537
|
+
// 12B snapshot pairs audio_config with a STUB sidecar (embed_audio only),
|
|
538
|
+
// and a non-null loader here is what advertises `audio: true` on every
|
|
539
|
+
// capability surface (ws handshake, /v1/models). Header-only read.
|
|
540
|
+
if (!sidecarShipsAudioTower(`${modelDir}/optiq_vision.safetensors`)) return null;
|
|
541
|
+
const audioCfg = parseAudioConfig(config.raw.audio_config as Record<string, any>);
|
|
542
|
+
return () => AudioTower.load(modelDir, audioCfg, model.embedScale);
|
|
543
|
+
}
|
|
544
|
+
|
|
545
|
+
/** Lazily load + cache the audio tower on first use. Unlike vision's
|
|
546
|
+
* warn-and-continue (text-only degrade is correct for requests WITHOUT
|
|
547
|
+
* images), this is only ever consulted for requests WITH audio — the
|
|
548
|
+
* caller turns null into an explicit 400, never a silent text-only
|
|
549
|
+
* degrade. A failed load is not retried every request. (The stub-sidecar
|
|
550
|
+
* case — the local 12B state — never gets here: makeAudioLoader checks the
|
|
551
|
+
* sidecar header and returns a null loader.) */
|
|
552
|
+
function getAudioTower(ctx: ServerContext): AudioTower | null {
|
|
553
|
+
if (ctx.audio) return ctx.audio;
|
|
554
|
+
if (!ctx.loadAudio) return null;
|
|
555
|
+
try {
|
|
556
|
+
ctx.audio = ctx.loadAudio();
|
|
557
|
+
return ctx.audio;
|
|
558
|
+
} catch (e) {
|
|
559
|
+
console.warn(`audio sidecar not loadable (${(e as Error).message})`);
|
|
560
|
+
ctx.loadAudio = null;
|
|
561
|
+
return null;
|
|
562
|
+
}
|
|
563
|
+
}
|
|
564
|
+
|
|
279
565
|
interface ChatRequest {
|
|
280
566
|
messages: ChatMessage[];
|
|
281
567
|
stream?: boolean;
|
|
@@ -340,6 +626,20 @@ interface ChatRequest {
|
|
|
340
626
|
toe?: number;
|
|
341
627
|
pivot_offset?: number;
|
|
342
628
|
};
|
|
629
|
+
/** OpenAI structured output: {type:"json_object"} | {type:"json_schema",
|
|
630
|
+
* json_schema:{name,schema,strict?}} | {type:"text"}. Enforced at the
|
|
631
|
+
* sampler via xgrammar token-bitmasks (src/grammar.ts). L2-class (oMLX
|
|
632
|
+
* oracle). On compile failure, degrades to a system-prompt injection +
|
|
633
|
+
* Warning header (oMLX parity), never 500. */
|
|
634
|
+
response_format?: unknown;
|
|
635
|
+
/** vLLM/oMLX grammar aliases (all compiled via xgrammar): raw EBNF/LARK
|
|
636
|
+
* grammar, regex, enum choice, and a bare JSON-schema object. Precedence
|
|
637
|
+
* (guided_grammar > json_schema > json_object > structured_outputs >
|
|
638
|
+
* guided_regex > guided_choice) mirrors oMLX _effective_guided_grammar. */
|
|
639
|
+
guided_grammar?: string;
|
|
640
|
+
guided_regex?: string;
|
|
641
|
+
guided_choice?: string[];
|
|
642
|
+
structured_outputs?: unknown;
|
|
343
643
|
}
|
|
344
644
|
|
|
345
645
|
/** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
|
|
@@ -419,9 +719,9 @@ export class ToolAwareStream {
|
|
|
419
719
|
* chain-of-thought as `<|channel>thought\n…<channel|>` using special tokens
|
|
420
720
|
* 100/101 that the content decoder strips, so reasoning is captured here at
|
|
421
721
|
* the token level. A SEPARATE decoder keeps the reasoning byte-stream's
|
|
422
|
-
* incremental state independent of content's. The `thought
|
|
423
|
-
*
|
|
424
|
-
*
|
|
722
|
+
* incremental state independent of content's. The `thought` channel-name
|
|
723
|
+
* word is stripped before the reasoning text (mlx-lm's think-start marker
|
|
724
|
+
* is `<|channel>thought`; the "\n" after it is reasoning content). */
|
|
425
725
|
readonly #channelDecoder: StreamDecoder;
|
|
426
726
|
#inChannel = false;
|
|
427
727
|
#channelNamePending = "";
|
|
@@ -442,15 +742,20 @@ export class ToolAwareStream {
|
|
|
442
742
|
this.#channelDecoder = new StreamDecoder(tokenizer, true);
|
|
443
743
|
}
|
|
444
744
|
|
|
445
|
-
/** Feed decoded channel text, stripping the leading `thought
|
|
446
|
-
* (which may arrive across tokens). Returns the reasoning delta.
|
|
745
|
+
/** Feed decoded channel text, stripping the leading `thought` channel-name
|
|
746
|
+
* word (which may arrive across tokens). Returns the reasoning delta.
|
|
747
|
+
* mlx-lm parity: its think-start marker is exactly `<|channel>thought`
|
|
748
|
+
* (tokenizer_utils.py `_infer_thinking`) and only the MARKER tokens get
|
|
749
|
+
* their text blanked (server.py `_process_control_tokens`) — the "\n"
|
|
750
|
+
* after the name is an ordinary generated token, so it is the FIRST byte
|
|
751
|
+
* of the reasoning stream and must be kept, not swallowed with the name. */
|
|
447
752
|
#feedChannel(text: string): string {
|
|
448
753
|
if (this.#channelNameDone) return text;
|
|
449
754
|
this.#channelNamePending += text;
|
|
450
755
|
const nl = this.#channelNamePending.indexOf("\n");
|
|
451
|
-
if (nl === -1) return ""; // still inside the channel-name
|
|
756
|
+
if (nl === -1) return ""; // still inside the channel-name word
|
|
452
757
|
this.#channelNameDone = true;
|
|
453
|
-
const rest = this.#channelNamePending.slice(nl
|
|
758
|
+
const rest = this.#channelNamePending.slice(nl); // keep the "\n" (mlx-lm does)
|
|
454
759
|
this.#channelNamePending = "";
|
|
455
760
|
return rest;
|
|
456
761
|
}
|
|
@@ -515,7 +820,8 @@ export class ToolAwareStream {
|
|
|
515
820
|
// Reasoning channel: tokens between <|channel> and <channel|> are thought,
|
|
516
821
|
// captured to #reasoning (drained via takeReasoning), never content. An
|
|
517
822
|
// empty block (<|channel>thought\n<channel|>, emitted by larger Gemmas even
|
|
518
|
-
// with thinking off) yields
|
|
823
|
+
// with thinking off) yields only the "\n" as reasoning (mlx-lm parity) and
|
|
824
|
+
// leaks nothing into content.
|
|
519
825
|
if (this.#inChannel) {
|
|
520
826
|
if (token === CHANNEL_END) {
|
|
521
827
|
this.#inChannel = false;
|
|
@@ -756,10 +1062,15 @@ function contentPartsToText(parts: Array<Record<string, unknown>>): string {
|
|
|
756
1062
|
.join("");
|
|
757
1063
|
}
|
|
758
1064
|
|
|
759
|
-
/** True if any content part is
|
|
760
|
-
* array form for extractImages
|
|
761
|
-
function
|
|
762
|
-
return parts.some(
|
|
1065
|
+
/** True if any content part is media — image OR audio — so the multimodal
|
|
1066
|
+
* path must keep the array form for extractImages/extractAudio. */
|
|
1067
|
+
function hasMediaPart(parts: Array<Record<string, unknown>>): boolean {
|
|
1068
|
+
return parts.some(
|
|
1069
|
+
(p) =>
|
|
1070
|
+
p &&
|
|
1071
|
+
(p.type === "image" || p.type === "image_url" ||
|
|
1072
|
+
p.type === "audio" || p.type === "input_audio" || p.type === "audio_url"),
|
|
1073
|
+
);
|
|
763
1074
|
}
|
|
764
1075
|
|
|
765
1076
|
/** OpenAI sends assistant tool_call arguments as JSON strings; the
|
|
@@ -780,13 +1091,14 @@ function hasImagePart(parts: Array<Record<string, unknown>>): boolean {
|
|
|
780
1091
|
* OpenAI multimodal client) sends user content as `[{type:"text",text}]`,
|
|
781
1092
|
* but non-vision chat templates expect `content` to be a STRING and render
|
|
782
1093
|
* nothing for an array — so the user's turn silently vanishes and the model
|
|
783
|
-
* replies "I don't see any message." Arrays that carry
|
|
784
|
-
* left intact for the
|
|
785
|
-
*
|
|
1094
|
+
* replies "I don't see any message." Arrays that carry a media part
|
|
1095
|
+
* (image or audio) are left intact for the multimodal path
|
|
1096
|
+
* (extractImages/extractAudio); only text-only arrays are collapsed
|
|
1097
|
+
* here, which is a no-op for the vision templates too. */
|
|
786
1098
|
export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
|
|
787
1099
|
return messages.map((raw) => {
|
|
788
1100
|
let m = raw.role === "developer" ? { ...raw, role: "system" } : raw;
|
|
789
|
-
if (Array.isArray(m.content) && !
|
|
1101
|
+
if (Array.isArray(m.content) && !hasMediaPart(m.content)) {
|
|
790
1102
|
m = { ...m, content: contentPartsToText(m.content) };
|
|
791
1103
|
}
|
|
792
1104
|
if (!m.tool_calls) return m;
|
|
@@ -850,6 +1162,33 @@ export function nextDefaultSeed(): number {
|
|
|
850
1162
|
* with defaults logprobs=False / top_logprobs=-1). Returns the reference's
|
|
851
1163
|
* error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
|
|
852
1164
|
* at 11, not OpenAI's 20 — we mirror the reference. */
|
|
1165
|
+
/** Build a degrade-path system-prompt instruction for JSON output, mirroring
|
|
1166
|
+
* oMLX's api.tool_calling.build_json_system_prompt (used when xgrammar
|
|
1167
|
+
* compile fails — the response_format degrades to prompt injection rather
|
|
1168
|
+
* than a 500, oMLX parity). Returns null for {type:"text"} / unset. */
|
|
1169
|
+
export function degradeJsonSystemPrompt(body: ChatRequest): string | null {
|
|
1170
|
+
const rf = body.response_format as { type?: string; json_schema?: { name?: string; description?: string; schema?: unknown } } | undefined;
|
|
1171
|
+
if (!rf || typeof rf !== "object") return null;
|
|
1172
|
+
const type = rf.type ?? "text";
|
|
1173
|
+
if (type === "text") return null;
|
|
1174
|
+
if (type === "json_object") {
|
|
1175
|
+
return "You must respond with valid JSON only. " +
|
|
1176
|
+
"Do not include any explanation or text outside the JSON object.";
|
|
1177
|
+
}
|
|
1178
|
+
if (type === "json_schema") {
|
|
1179
|
+
const spec = rf.json_schema ?? {};
|
|
1180
|
+
const schema = spec.schema ?? body.structured_outputs ?? {};
|
|
1181
|
+
const name = spec.name ?? "response";
|
|
1182
|
+
const description = spec.description ?? "";
|
|
1183
|
+
let prompt = `You must respond with valid JSON matching the '${name}' schema.`;
|
|
1184
|
+
if (description) prompt += ` ${description}`;
|
|
1185
|
+
prompt += `\n\nJSON Schema:\n\`\`\`json\n${JSON.stringify(schema, null, 2)}\n\`\`\`\n\n` +
|
|
1186
|
+
"Respond with only the JSON object, no additional text or explanation.";
|
|
1187
|
+
return prompt;
|
|
1188
|
+
}
|
|
1189
|
+
return null;
|
|
1190
|
+
}
|
|
1191
|
+
|
|
853
1192
|
export function validateLogprobsParams(body: {
|
|
854
1193
|
logprobs?: unknown;
|
|
855
1194
|
top_logprobs?: unknown;
|
|
@@ -929,24 +1268,79 @@ class LogprobsCollector {
|
|
|
929
1268
|
}
|
|
930
1269
|
}
|
|
931
1270
|
|
|
932
|
-
/** Incremental detokenizer: emits the longest stable decoded prefix.
|
|
933
|
-
|
|
1271
|
+
/** Incremental detokenizer: emits the longest stable decoded prefix.
|
|
1272
|
+
*
|
|
1273
|
+
* Byte parity with mlx-lm's streaming detokenizers (the drop-in contract is
|
|
1274
|
+
* rendered BYTES, not just token ids). For BPE/ByteLevel tokenizers, two
|
|
1275
|
+
* mlx-lm 0.31.3 BPEStreamingDetokenizer behaviors our full-sequence decode
|
|
1276
|
+
* lacks (tokenizer_utils.py:195-226):
|
|
1277
|
+
*
|
|
1278
|
+
* 1. `trimsLeadingSpace` — mlx-lm drops ONE leading " " at the start of the
|
|
1279
|
+
* generated sequence (`_maybe_trim_space`); trim it here. SPM decode
|
|
1280
|
+
* already matches (see LoadedTokenizer.trimsLeadingSpace).
|
|
1281
|
+
* 2. `bareSpaceTokenId` — add_token WITHHOLDS a single-char byte-32 token
|
|
1282
|
+
* ("Ġ") in `_unflushed` ("For single spaces wait until the next token"),
|
|
1283
|
+
* flushing it together with the NEXT token — and mlx_lm.server NEVER
|
|
1284
|
+
* calls detokenizer.finalize() (zero hits in server.py 0.31.3), so a
|
|
1285
|
+
* generation ENDING on bare-space token(s) silently drops their spaces
|
|
1286
|
+
* from the served bytes. push() withholds those spaces; flush() drops a
|
|
1287
|
+
* trailing bare-space run (the held text dies with the request, exactly
|
|
1288
|
+
* like mlx-lm serve). Re-check both if upstream ever adds a finalize()
|
|
1289
|
+
* call. LATENT HAZARD (deliberately not emulated): models with
|
|
1290
|
+
* clean_up_tokenization_spaces=true get an ADDITIONAL mid-stream rule
|
|
1291
|
+
* (`_space_matches`: held space dropped before "." "," "'s" …) — both
|
|
1292
|
+
* current BPE targets have it false (MiniCPM5), so it never fires here.
|
|
1293
|
+
*
|
|
1294
|
+
* Exported for unit tests (serve-detok mlx-lm byte parity). */
|
|
1295
|
+
export class StreamDecoder {
|
|
934
1296
|
#ids: number[] = [];
|
|
935
1297
|
#emitted = "";
|
|
1298
|
+
#warnedRevision = false;
|
|
1299
|
+
readonly #trimLeadingSpace: boolean;
|
|
1300
|
+
readonly #bareSpaceId: number | undefined;
|
|
936
1301
|
|
|
937
1302
|
constructor(
|
|
938
1303
|
readonly tokenizer: LoadedTokenizer,
|
|
939
1304
|
readonly skipSpecialTokens = true,
|
|
940
|
-
) {
|
|
1305
|
+
) {
|
|
1306
|
+
this.#trimLeadingSpace = tokenizer.trimsLeadingSpace === true;
|
|
1307
|
+
this.#bareSpaceId = tokenizer.bareSpaceTokenId;
|
|
1308
|
+
}
|
|
1309
|
+
|
|
1310
|
+
#decode(ids: number[]): string {
|
|
1311
|
+
const full = this.tokenizer.decode(ids, this.skipSpecialTokens);
|
|
1312
|
+
return this.#trimLeadingSpace && full.startsWith(" ") ? full.slice(1) : full;
|
|
1313
|
+
}
|
|
941
1314
|
|
|
942
1315
|
push(token: number): string {
|
|
943
1316
|
this.#ids.push(token);
|
|
944
|
-
|
|
1317
|
+
// Bare-space hold-back (mlx-lm add_token keeps "Ġ" in _unflushed): don't
|
|
1318
|
+
// advance #emitted; the held space(s) flush as part of the next
|
|
1319
|
+
// non-bare-space token's delta — consecutive bare spaces accumulate.
|
|
1320
|
+
if (token === this.#bareSpaceId) return "";
|
|
1321
|
+
const full = this.#decode(this.#ids);
|
|
945
1322
|
// hold back a trailing replacement char (partial multi-byte sequence)
|
|
946
1323
|
const stable = full.endsWith("�") ? full.slice(0, -1) : full;
|
|
947
1324
|
if (!stable.startsWith(this.#emitted)) {
|
|
948
|
-
// decoder revised
|
|
949
|
-
|
|
1325
|
+
// The decoder revised already-streamed text (cleanup rules — never
|
|
1326
|
+
// fires for the shipped tokenizers, which all have
|
|
1327
|
+
// clean_up_tokenization_spaces=false; see the LATENT HAZARD note in
|
|
1328
|
+
// the class doc). An SSE client cannot un-receive bytes, so the old
|
|
1329
|
+
// "re-emit from scratch" answer DUPLICATED the whole stream
|
|
1330
|
+
// (2026-07-07 review). Truncate-safe resync instead: emit only the
|
|
1331
|
+
// length-extension and keep the emitted watermark monotone — byte
|
|
1332
|
+
// drift is confined to the revised span, and the once-per-stream
|
|
1333
|
+
// warning makes the tokenizer that needs real _space_matches
|
|
1334
|
+
// emulation loud instead of silently wrong.
|
|
1335
|
+
if (!this.#warnedRevision) {
|
|
1336
|
+
this.#warnedRevision = true;
|
|
1337
|
+
console.warn(
|
|
1338
|
+
"[detok] decoder revised already-streamed text (clean_up_tokenization_spaces?) — " +
|
|
1339
|
+
"resyncing without re-emit; add mlx-lm's cleanup-rule emulation for this tokenizer",
|
|
1340
|
+
);
|
|
1341
|
+
}
|
|
1342
|
+
if (stable.length <= this.#emitted.length) return "";
|
|
1343
|
+
const out = stable.slice(this.#emitted.length);
|
|
950
1344
|
this.#emitted = stable;
|
|
951
1345
|
return out;
|
|
952
1346
|
}
|
|
@@ -956,7 +1350,15 @@ class StreamDecoder {
|
|
|
956
1350
|
}
|
|
957
1351
|
|
|
958
1352
|
flush(): string {
|
|
959
|
-
|
|
1353
|
+
// mlx_lm.server never finalize()s: a trailing bare-space run stays
|
|
1354
|
+
// withheld forever, so its text is dropped from the served bytes.
|
|
1355
|
+
let ids = this.#ids;
|
|
1356
|
+
if (this.#bareSpaceId !== undefined) {
|
|
1357
|
+
let n = ids.length;
|
|
1358
|
+
while (n > 0 && ids[n - 1] === this.#bareSpaceId) n--;
|
|
1359
|
+
if (n < ids.length) ids = ids.slice(0, n);
|
|
1360
|
+
}
|
|
1361
|
+
const full = this.#decode(ids);
|
|
960
1362
|
const delta = full.slice(this.#emitted.length);
|
|
961
1363
|
this.#emitted = full;
|
|
962
1364
|
return delta;
|
|
@@ -986,41 +1388,139 @@ export function createServer(
|
|
|
986
1388
|
// cache by attention type. Non-batchable requests (vision / adapters /
|
|
987
1389
|
// repetition penalty / user seed / explicit kv-quant) drain to the serial
|
|
988
1390
|
// lane (see GenerationGateway.willBatch).
|
|
989
|
-
|
|
990
|
-
|
|
991
|
-
//
|
|
992
|
-
//
|
|
993
|
-
//
|
|
994
|
-
//
|
|
995
|
-
//
|
|
996
|
-
|
|
997
|
-
|
|
998
|
-
//
|
|
1391
|
+
// DEFAULT 8 (flipped 2026-07-05, Josh's call, after GATE-B1-SPEED): a
|
|
1392
|
+
// lone request through the batch lane IS the serial engine (adopted
|
|
1393
|
+
// serial-class caches, compiled decode, prompt cache + SSD restore;
|
|
1394
|
+
// 0.992-0.996 paired decode ratios, byte-identical output), so the cap
|
|
1395
|
+
// only changes behavior when concurrent requests actually arrive — the
|
|
1396
|
+
// agentic sub-agent workload. --batch 1 pins strict serial for
|
|
1397
|
+
// arrival-independent numerics. 8 = optiq's Mac-safe concurrency.
|
|
1398
|
+
const batch = Math.max(1, Math.floor(serverOptions.batch ?? 8));
|
|
1399
|
+
|
|
1400
|
+
// KV-quant scheme, resolved once. UNSET now means bf16 (flipped
|
|
1401
|
+
// 2026-07-05 with the naked-=-L1 default): quantized KV measured 5–20%
|
|
1402
|
+
// SLOWER decode than bf16 at ≤16k on every model — on mlx-lm too (its
|
|
1403
|
+
// kv8 oracle trails its own bf16) — so it pays only in memory headroom
|
|
1404
|
+
// and must be an explicit opt-in (--kv-quant config|4|8, or --l2
|
|
1405
|
+
// whose presets pass it explicitly). The CLI always passes kvQuant now;
|
|
1406
|
+
// this fallback is the library-user default and matches the CLI's.
|
|
999
1407
|
const configScheme = ctx.kvConfig?.length ? { kvConfig: ctx.kvConfig } : {};
|
|
1000
|
-
|
|
1001
|
-
|
|
1408
|
+
// Mutually exclusive by contract (GenerateOptions.turboQuant doc): a
|
|
1409
|
+
// programmatic caller setting both gets turboQuant — say so, like the
|
|
1410
|
+
// other risky-combination warnings below.
|
|
1411
|
+
if (serverOptions.turboQuant && serverOptions.kvQuant && serverOptions.kvQuant !== "off")
|
|
1412
|
+
console.warn(
|
|
1413
|
+
`[kv-quant] both turboQuant and --kv-quant ${serverOptions.kvQuant} are set; ` +
|
|
1414
|
+
`turboQuant wins (they are mutually exclusive).`,
|
|
1415
|
+
);
|
|
1416
|
+
const kvScheme: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart" | "turboQuant"> =
|
|
1417
|
+
serverOptions.turboQuant ? { turboQuant: serverOptions.turboQuant, quantizedKvStart: 0 }
|
|
1418
|
+
: serverOptions.kvQuant === "off" ? {}
|
|
1002
1419
|
: serverOptions.kvQuant === "config" ? configScheme
|
|
1003
1420
|
: typeof serverOptions.kvQuant === "number"
|
|
1004
1421
|
? { kvBits: serverOptions.kvQuant, quantizedKvStart: 0 }
|
|
1005
|
-
:
|
|
1006
|
-
|
|
1007
|
-
|
|
1422
|
+
: {}; // unset → bf16 (L1 default; quantized KV is opt-in)
|
|
1423
|
+
// Phase 3.1: kvConfig whose layers are all full-attention BATCHES (the
|
|
1424
|
+
// scheduler applies the mixed scheme per row); uniform kvBits and configs
|
|
1425
|
+
// touching rotating layers still route those requests serial. The warning
|
|
1426
|
+
// only fires for the still-serial compositions.
|
|
1427
|
+
if (batch > 1 && kvScheme.kvBits)
|
|
1008
1428
|
console.warn(
|
|
1009
|
-
`[batch] --batch ${batch} with
|
|
1010
|
-
`
|
|
1011
|
-
`
|
|
1429
|
+
`[batch] --batch ${batch} with uniform --kv-quant ${kvScheme.kvBits}: uniform ` +
|
|
1430
|
+
`quantized KV is serial-only (it touches rotating layers) — those requests ` +
|
|
1431
|
+
`won't batch. Use --kv-quant config for batched mixed-precision KV, or omit ` +
|
|
1432
|
+
`--kv-quant to batch in bf16. (docs/design/unified-engine-frontier-plan.md)`,
|
|
1012
1433
|
);
|
|
1434
|
+
// TurboQuant is solo-only in v1 (novel cache class, not batchable by
|
|
1435
|
+
// construction — see GenerationGateway#modelCachesBatchable/willBatch).
|
|
1436
|
+
if (batch > 1 && kvScheme.turboQuant)
|
|
1437
|
+
console.warn(
|
|
1438
|
+
`[batch] --batch ${batch} with --kv-quant turbo: TurboQuant is serial-only in v1 ` +
|
|
1439
|
+
`— those requests won't batch. Omit --kv-quant to batch in bf16. ` +
|
|
1440
|
+
`(docs/design/turboquant-kv.md)`,
|
|
1441
|
+
);
|
|
1442
|
+
// Quantized KV (any axis) excludes the spec lane: drafted requests fall to
|
|
1443
|
+
// the normal serial path WITH the KV scheme applied rather than losing it
|
|
1444
|
+
// silently (mirrors the affine kvBits/kvConfig exclusion; the spec loop is
|
|
1445
|
+
// bf16-KV-only in v1).
|
|
1446
|
+
if (ctx.draft && (kvScheme.turboQuant || kvScheme.kvBits || kvScheme.kvConfig?.length))
|
|
1447
|
+
console.warn(
|
|
1448
|
+
`[spec] --draft-model with quantized KV (--kv-quant ${serverOptions.turboQuant ? "turbo" : String(serverOptions.kvQuant)}): ` +
|
|
1449
|
+
`the speculative lane is bf16-KV-only in v1 — requests keep the KV scheme and ` +
|
|
1450
|
+
`decode serially WITHOUT speculation. Omit --kv-quant to speculate. ` +
|
|
1451
|
+
`(docs/design/dspark-serving-program.md Phase 4)`,
|
|
1452
|
+
);
|
|
1453
|
+
// TurboQuant head-dim fail-fast (2026-07-07 review): the cache class only
|
|
1454
|
+
// supports {64,128,256,512} (sign-vector + Lloyd-Max table coverage) and
|
|
1455
|
+
// used to validate LAZILY on the first append — an unsupported model
|
|
1456
|
+
// (e.g. a 72/80/96 head dim) accepted --kv-quant turbo at startup and
|
|
1457
|
+
// then 500'd EVERY request from inside prefill. The config knows the
|
|
1458
|
+
// full-attention head dim (the only kind that converts; sliding layers
|
|
1459
|
+
// stay bf16) — refuse at createServer instead.
|
|
1460
|
+
if (kvScheme.turboQuant) {
|
|
1461
|
+
const dim = ctx.model.config.text.globalHeadDim;
|
|
1462
|
+
if (!(TURBOQUANT_HEAD_DIMS as readonly number[]).includes(dim))
|
|
1463
|
+
throw new Error(
|
|
1464
|
+
`--kv-quant turbo: this model's full-attention head_dim is ${dim}; ` +
|
|
1465
|
+
`TurboQuant supports {${TURBOQUANT_HEAD_DIMS.join(",")}} ` +
|
|
1466
|
+
`(docs/design/turboquant-kv.md) — use --kv-quant config|4|8 or omit it`,
|
|
1467
|
+
);
|
|
1468
|
+
}
|
|
1469
|
+
// Paged KV v1 (docs/design/paged-kv-cache.md): explicit refusals, not
|
|
1470
|
+
// silent downgrades — the incompatible combos would otherwise degrade
|
|
1471
|
+
// quietly (batch: paged caches can't merge; kv-quant: the swap would
|
|
1472
|
+
// drop the scheme, the exact composition bug the kvQuant gate above
|
|
1473
|
+
// exists to prevent; draft: the spec lane assumes serial-class caches).
|
|
1474
|
+
if (serverOptions.pagedKv) {
|
|
1475
|
+
if (batch > 1)
|
|
1476
|
+
throw new Error(
|
|
1477
|
+
`--paged-kv is serial-only in v1 — add --batch 1 (got --batch ${batch}). ` +
|
|
1478
|
+
`Batched paging is the follow-up PR (docs/design/paged-kv-cache.md).`,
|
|
1479
|
+
);
|
|
1480
|
+
if (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant)
|
|
1481
|
+
throw new Error(
|
|
1482
|
+
`--paged-kv is bf16-only in v1 — omit --kv-quant (quantized paged ` +
|
|
1483
|
+
`blocks are a documented follow-up, docs/design/paged-kv-cache.md).`,
|
|
1484
|
+
);
|
|
1485
|
+
if (ctx.draft)
|
|
1486
|
+
throw new Error(
|
|
1487
|
+
`--paged-kv cannot combine with --draft-model in v1 ` +
|
|
1488
|
+
`(docs/design/paged-kv-cache.md non-goals).`,
|
|
1489
|
+
);
|
|
1490
|
+
if (!ctx.model.config.modelType.startsWith("gemma4"))
|
|
1491
|
+
throw new Error(
|
|
1492
|
+
`--paged-kv v1 supports Gemma4-family models only ` +
|
|
1493
|
+
`(this model: ${ctx.model.config.modelType}) — docs/design/paged-kv-cache.md.`,
|
|
1494
|
+
);
|
|
1495
|
+
const bs = serverOptions.pagedKv.blockSize;
|
|
1496
|
+
if (bs !== undefined && (!Number.isInteger(bs) || bs <= 0))
|
|
1497
|
+
throw new Error(`--paged-kv-block-size must be a positive integer (got ${bs})`);
|
|
1498
|
+
if (serverOptions.ssdCacheDir)
|
|
1499
|
+
console.warn(
|
|
1500
|
+
"[paged-kv] --ssd-cache has no effect: paged requests bypass the prompt " +
|
|
1501
|
+
"cache (v1 non-goal), so nothing reaches the SSD tier.",
|
|
1502
|
+
);
|
|
1503
|
+
}
|
|
1013
1504
|
|
|
1014
1505
|
// SSD cold tier (docs/design/ssd-kv-cold-tier.md): prefix KV survives RAM
|
|
1015
1506
|
// eviction and restarts. Compatibility key = configFingerprint (graph
|
|
1016
1507
|
// shape) + the EFFECTIVE kv scheme (flags pick bf16 vs config vs uniform
|
|
1017
1508
|
// on the same model — restored caches must match what serving produces) +
|
|
1018
1509
|
// tokenizer hash (ids must keep meaning the same text).
|
|
1510
|
+
// RAM prompt-cache cap: 8 GB default (Josh's call, 2026-07-06). The old
|
|
1511
|
+
// flat 2e9 was an anti-OOM reflex sized for a 1B model in 24 GB — a
|
|
1512
|
+
// single full-context 12B entry is ~0.6 GB, so it silently shrank to
|
|
1513
|
+
// "three entries" on big models. 8 GB holds a dozen 12B contexts; idle
|
|
1514
|
+
// entries still demote to SSD and LRU eviction bounds pressure.
|
|
1515
|
+
// --prompt-cache <GB> overrides; 0 disables.
|
|
1516
|
+
const promptCacheCap = serverOptions.promptCacheBytes ?? 8e9;
|
|
1019
1517
|
let ssdStore: SsdCacheStore | null = null;
|
|
1020
1518
|
if (serverOptions.ssdCacheDir) {
|
|
1021
|
-
if (
|
|
1519
|
+
if (promptCacheCap <= 0)
|
|
1022
1520
|
throw new Error("--ssd-cache requires the RAM prompt cache (--prompt-cache 0 disables it)");
|
|
1023
|
-
const schemeKey = kvScheme.
|
|
1521
|
+
const schemeKey = kvScheme.turboQuant
|
|
1522
|
+
? `turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`
|
|
1523
|
+
: kvScheme.kvBits
|
|
1024
1524
|
? `kv${kvScheme.kvBits}`
|
|
1025
1525
|
: kvScheme.kvConfig?.length ? "config" : "bf16";
|
|
1026
1526
|
const tokJson = readFileSync(`${ctx.model.config.modelDir}/tokenizer.json`);
|
|
@@ -1039,11 +1539,51 @@ export function createServer(
|
|
|
1039
1539
|
);
|
|
1040
1540
|
}
|
|
1041
1541
|
|
|
1542
|
+
// Layer 0 (unified-engine plan): the cold tier is bound INTO the prompt
|
|
1543
|
+
// cache — take() itself tiers over it, so the batch scheduler and every
|
|
1544
|
+
// other consumer get SSD restores through the same API; eviction AND
|
|
1545
|
+
// idle demotion spill into it.
|
|
1546
|
+
const coldTier = ssdStore
|
|
1547
|
+
? {
|
|
1548
|
+
find: (prompt: number[], ns: string) => {
|
|
1549
|
+
const h = ssdStore!.find(prompt, ns);
|
|
1550
|
+
return h ? { prefixLen: h.prefixLen, handle: h.entry } : null;
|
|
1551
|
+
},
|
|
1552
|
+
restore: (handle: unknown) => {
|
|
1553
|
+
const loaded = ssdStore!.restore(handle as import("./ssd-cache").SsdIndexEntry, ctx.model);
|
|
1554
|
+
if (!loaded) return null;
|
|
1555
|
+
// Restore is a STREAMED COPY (2026-07-07 A7-restore): the caches
|
|
1556
|
+
// own their bytes and no mapping outlives loadKvCache — nothing
|
|
1557
|
+
// to pin, nothing to unmap. retain stays in the entry contract
|
|
1558
|
+
// as a no-op so callers' dispose ordering is unchanged.
|
|
1559
|
+
return { tokens: loaded.tokens, caches: loaded.caches, retain: () => {} };
|
|
1560
|
+
},
|
|
1561
|
+
store: (tokens: number[], caches: import("./model/gemma4").Cache[], ns: string) => {
|
|
1562
|
+
ssdStore!.store(tokens, caches, ns);
|
|
1563
|
+
},
|
|
1564
|
+
}
|
|
1565
|
+
: null;
|
|
1042
1566
|
const promptCache = new PromptCache(
|
|
1043
|
-
|
|
1044
|
-
// RAM eviction
|
|
1045
|
-
//
|
|
1046
|
-
|
|
1567
|
+
promptCacheCap,
|
|
1568
|
+
// RAM eviction AND idle demotion spill to the cold tier NON-BLOCKING
|
|
1569
|
+
// (2026-07-06, same contract as the write-behind snapshot below): the
|
|
1570
|
+
// cache hands us OWNED zero-copy clones + copied tokens (made under
|
|
1571
|
+
// the generation lock — microseconds; entries are immutable so the
|
|
1572
|
+
// clones stay consistent), and the flush goes through the BOUNDED
|
|
1573
|
+
// SpillQueue -> storeAsync (idle-gated per tensor, see below) ->
|
|
1574
|
+
// clone disposal on every settle/drop path (that dispose is what
|
|
1575
|
+
// actually frees the demoted GPU memory; the queue's byte cap keeps
|
|
1576
|
+
// starved-gate retention bounded — 2026-07-07 review fix). spillQueue
|
|
1577
|
+
// is declared with the snapshot scheduler below; safe to close over
|
|
1578
|
+
// here because spills only fire at put()/demoteIdle time, long after
|
|
1579
|
+
// init.
|
|
1580
|
+
ssdStore
|
|
1581
|
+
? {
|
|
1582
|
+
spillOwned: (entry) =>
|
|
1583
|
+
spillQueue!.enqueue({ tokens: entry.tokens, caches: entry.caches, ns: entry.ns }),
|
|
1584
|
+
}
|
|
1585
|
+
: null,
|
|
1586
|
+
coldTier,
|
|
1047
1587
|
);
|
|
1048
1588
|
// Responses-API store for previous_response_id resumption (Phase 11):
|
|
1049
1589
|
// TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
|
|
@@ -1057,68 +1597,143 @@ export function createServer(
|
|
|
1057
1597
|
* `enqueue` promise queue is gone; the gateway's AsyncMutex is THE lock.)
|
|
1058
1598
|
* onToken returning `false` halts generation early (stop
|
|
1059
1599
|
* sequence fired); the cache snapshot stays valid and is still kept.
|
|
1060
|
-
*
|
|
1061
|
-
* identical placeholder ids, so prefix matching
|
|
1062
|
-
* images would false-hit. */
|
|
1600
|
+
* Media (vision/audio embeddings-prefill) requests bypass the prompt
|
|
1601
|
+
* cache: soft tokens are identical placeholder ids, so prefix matching
|
|
1602
|
+
* across different images/clips would false-hit. */
|
|
1063
1603
|
const runGeneration = async (
|
|
1064
1604
|
promptIds: number[],
|
|
1065
1605
|
options: GenerateOptions,
|
|
1066
1606
|
onToken: (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>,
|
|
1067
|
-
vision?: {
|
|
1607
|
+
vision?: {
|
|
1608
|
+
embeddings: import("./mlx/array").MlxArray;
|
|
1609
|
+
imageMask?: import("./mlx/array").MlxArray;
|
|
1610
|
+
multimodalMask?: import("./mlx/array").MlxArray;
|
|
1611
|
+
},
|
|
1068
1612
|
) => {
|
|
1613
|
+
// Speculative decoding (serve --draft-model): spec-ELIGIBLE requests
|
|
1614
|
+
// decode through the verify loop; the rest fall through to the normal
|
|
1615
|
+
// serial path (never wrong results, just no speedup — logged once per
|
|
1616
|
+
// combination class would be noise, so silent). Eligibility v1:
|
|
1617
|
+
// text-only, base weights, no logprobs capture, bf16 KV — ALL quantized
|
|
1618
|
+
// KV axes excluded (affine kvBits/kvConfig AND turboQuant; the spec loop
|
|
1619
|
+
// builds fresh bf16 caches and never calls maybeQuantizeKv, so routing a
|
|
1620
|
+
// turbo request here would silently drop the operator's KV scheme —
|
|
1621
|
+
// 2026-07-07 review). Grammar COMPOSES (Phase C constrained verify walk
|
|
1622
|
+
// — see the serve-loop.ts header). Prompt-cache reuse is bypassed on
|
|
1623
|
+
// the spec path v1.
|
|
1624
|
+
if (
|
|
1625
|
+
ctx.draft &&
|
|
1626
|
+
!vision &&
|
|
1627
|
+
!options.adapters?.length &&
|
|
1628
|
+
!options.logprobs &&
|
|
1629
|
+
!options.kvBits &&
|
|
1630
|
+
!options.kvConfig &&
|
|
1631
|
+
!options.turboQuant &&
|
|
1632
|
+
// Belt: --paged-kv + --draft-model is refused at createServer; this
|
|
1633
|
+
// keeps the spec lane paged-free even for programmatic callers.
|
|
1634
|
+
!options.pagedKv
|
|
1635
|
+
) {
|
|
1636
|
+
const { specServeRun } = await import("./spec/serve-loop");
|
|
1637
|
+
return specServeRun(
|
|
1638
|
+
ctx.model, ctx.draft.provider, ctx.draft.numDraftTokens,
|
|
1639
|
+
promptIds, options, onToken,
|
|
1640
|
+
);
|
|
1641
|
+
}
|
|
1069
1642
|
// Cache entries are adapter-specific: KV computed under one adapter
|
|
1070
1643
|
// must never seed another's (or the base's) prefill.
|
|
1071
1644
|
const cacheNs = options.adapters?.join("+") ?? "";
|
|
1072
|
-
|
|
1073
|
-
//
|
|
1074
|
-
//
|
|
1075
|
-
//
|
|
1076
|
-
//
|
|
1077
|
-
|
|
1078
|
-
|
|
1079
|
-
|
|
1080
|
-
|
|
1081
|
-
|
|
1082
|
-
|
|
1083
|
-
|
|
1084
|
-
|
|
1085
|
-
// this file can't seed the prompt — fall back to the RAM hit.
|
|
1086
|
-
for (const c of loaded.caches) c.dispose();
|
|
1087
|
-
loaded.mmap.unmap();
|
|
1088
|
-
} else {
|
|
1089
|
-
if (trimNeeded > 0) for (const c of loaded.caches) c.trim(trimNeeded);
|
|
1090
|
-
if (entry) promptCache.put(entry.tokens, entry.caches, cacheNs, entry.retain); // return the RAM hit untouched
|
|
1091
|
-
entry = { tokens: loaded.tokens.slice(0, ssdHit.prefixLen), caches: loaded.caches,
|
|
1092
|
-
ns: cacheNs, retain: () => loaded.mmap.unmap() };
|
|
1093
|
-
}
|
|
1094
|
-
}
|
|
1095
|
-
}
|
|
1096
|
-
}
|
|
1645
|
+
// Paged-KV request scope (docs/design/paged-kv-cache.md): media
|
|
1646
|
+
// prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
|
|
1647
|
+
// they run the PLAIN cache path even under --paged-kv (scope the flag
|
|
1648
|
+
// per request, never 400). Effective value computed ONCE so the
|
|
1649
|
+
// prompt-cache bypass below and the generate() options can't disagree.
|
|
1650
|
+
const pagedKv = vision || options.adapters?.length ? undefined : options.pagedKv;
|
|
1651
|
+
// Paged requests bypass the prompt cache entirely (v1 non-goal:
|
|
1652
|
+
// PagedKVCache has no cloneKvCaches/restore path — the vision
|
|
1653
|
+
// precedent). Fresh caches per request, disposed on completion.
|
|
1654
|
+
const skipPromptCache = Boolean(vision || pagedKv);
|
|
1655
|
+
// Both tiers in one call (Layer 0): take() prefers a strictly-longer
|
|
1656
|
+
// SSD prefix, restores it zero-copy, and trims — see PromptCache.take.
|
|
1657
|
+
const entry = skipPromptCache ? null : promptCache.take(promptIds, cacheNs);
|
|
1097
1658
|
const caches = entry?.caches ?? ctx.model.makeCache();
|
|
1659
|
+
// Prompt-boundary snapshot (the multi-turn agent fix, 2026-07-04): the
|
|
1660
|
+
// prompt+gen entry put() below is UNTRIMMABLE at context > sliding
|
|
1661
|
+
// window (wrapped rings) and under quantized KV (mid-group), so any
|
|
1662
|
+
// decode→encode roundtrip drift in the reply the client sends back
|
|
1663
|
+
// turns the next turn into a total miss (measured: 12B turn-2 TTFT
|
|
1664
|
+
// 8.9 s instead of ~0.2 s). A prompt-ONLY entry is always an exact
|
|
1665
|
+
// prefix of the next turn's rendering regardless of reply drift.
|
|
1666
|
+
// Zero-copy (cloneKvCaches = slice views); only for substantial cold
|
|
1667
|
+
// prefills, where the re-prefill it saves is worth an extra entry.
|
|
1668
|
+
// The oracle invariant (mlx-lm insert_segments): a trim-free STRICT
|
|
1669
|
+
// prefix of the prompt exists for EVERY substantial request — cap the
|
|
1670
|
+
// boundary at len-1 so even a stableLen == len prompt (e4b: the
|
|
1671
|
+
// template tail survives the probe render) snapshots prompt[:-1]. An
|
|
1672
|
+
// exact repeat then matches with trimNeeded == 0, bypassing
|
|
1673
|
+
// isTrimmable() entirely — the only reuse path a wrapped ring has.
|
|
1674
|
+
const boundary = Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1);
|
|
1675
|
+
// Re-snapshot on EVERY substantial request whose stable boundary extends
|
|
1676
|
+
// past the cached prefix; the clone is zero-copy views, so re-putting
|
|
1677
|
+
// is ~free.
|
|
1678
|
+
const snapshotBoundary =
|
|
1679
|
+
!skipPromptCache && boundary >= 256 && boundary > (entry?.tokens.length ?? 0);
|
|
1098
1680
|
try {
|
|
1099
1681
|
const gen = generate(ctx.model, promptIds, {
|
|
1100
1682
|
...options,
|
|
1683
|
+
pagedKv, // request-scoped (undefined strips the server-wide flag)
|
|
1101
1684
|
cache: caches,
|
|
1102
|
-
...(
|
|
1685
|
+
...(snapshotBoundary
|
|
1686
|
+
? {
|
|
1687
|
+
// snapshotAt MUST travel with the hook: generate() splits the
|
|
1688
|
+
// prefill at exactly this many tokens and fires the hook while
|
|
1689
|
+
// the caches hold exactly that prefix — putting boundary
|
|
1690
|
+
// tokens against caches at any other offset is silent KV
|
|
1691
|
+
// corruption.
|
|
1692
|
+
snapshotAt: boundary,
|
|
1693
|
+
onPrefillDone: () => {
|
|
1694
|
+
try {
|
|
1695
|
+
promptCache.put(promptIds.slice(0, boundary), cloneKvCaches(caches), cacheNs);
|
|
1696
|
+
} catch (e) {
|
|
1697
|
+
console.warn(`prompt-boundary snapshot skipped: ${(e as Error).message}`);
|
|
1698
|
+
}
|
|
1699
|
+
},
|
|
1700
|
+
}
|
|
1701
|
+
: {}),
|
|
1702
|
+
...(vision
|
|
1703
|
+
? {
|
|
1704
|
+
promptEmbeddings: vision.embeddings,
|
|
1705
|
+
...(vision.imageMask ? { imageMask: vision.imageMask } : {}),
|
|
1706
|
+
...(vision.multimodalMask ? { multimodalMask: vision.multimodalMask } : {}),
|
|
1707
|
+
}
|
|
1708
|
+
: {}),
|
|
1103
1709
|
});
|
|
1104
1710
|
for await (const t of gen) {
|
|
1105
1711
|
if ((await onToken(t.token, t.logprobs)) === false) break;
|
|
1106
1712
|
}
|
|
1107
1713
|
const s = gen.stats!; // set on completion AND on early break
|
|
1108
|
-
if (
|
|
1714
|
+
if (skipPromptCache) {
|
|
1715
|
+
// Vision and paged-KV requests own their caches for exactly one
|
|
1716
|
+
// generation (paged: v1 non-goal — no PromptCache integration).
|
|
1109
1717
|
for (const c of caches) c.dispose();
|
|
1110
1718
|
} else {
|
|
1719
|
+
// put() fires onPut → the debounced write-behind SSD snapshot
|
|
1720
|
+
// (wired below), covering the batch lane's puts too.
|
|
1111
1721
|
promptCache.put(s.cacheTokens, caches, cacheNs, entry?.retain);
|
|
1112
|
-
scheduleSsdSnapshot(s.cacheTokens, cacheNs);
|
|
1113
1722
|
}
|
|
1114
1723
|
return s;
|
|
1115
1724
|
} catch (e) {
|
|
1116
1725
|
for (const c of caches) c.dispose();
|
|
1117
1726
|
entry?.retain?.();
|
|
1727
|
+
// A throw BEFORE generate()/the gateway took ownership (promptCache
|
|
1728
|
+
// take / makeCache) would leak the grammar's WASM matcher; dispose()
|
|
1729
|
+
// is idempotent, so this is safe when the throw came from inside run
|
|
1730
|
+
// (whose finally already disposed it).
|
|
1731
|
+
options.grammar?.dispose();
|
|
1118
1732
|
throw e;
|
|
1119
1733
|
} finally {
|
|
1120
1734
|
vision?.embeddings.dispose();
|
|
1121
|
-
vision?.imageMask
|
|
1735
|
+
vision?.imageMask?.dispose();
|
|
1736
|
+
vision?.multimodalMask?.dispose();
|
|
1122
1737
|
options.visionPixels?.dispose();
|
|
1123
1738
|
}
|
|
1124
1739
|
};
|
|
@@ -1132,24 +1747,134 @@ export function createServer(
|
|
|
1132
1747
|
// extraction must not race a generation mutating the entry); if the entry
|
|
1133
1748
|
// was evicted (already spilled) or extended (a newer schedule pending)
|
|
1134
1749
|
// meanwhile, findExact misses and nothing is written.
|
|
1750
|
+
// Keyed by ns AND entry length: a sub-second final [prompt+gen] put must
|
|
1751
|
+
// not cancel the pending boundary-snapshot write (they are DIFFERENT
|
|
1752
|
+
// entries; for wrapped-ring models the boundary file is the only one a
|
|
1753
|
+
// restart can use). Same-length reschedules still coalesce; stale keys
|
|
1754
|
+
// self-clean at fire time (findExact misses once the RAM tier superseded
|
|
1755
|
+
// the entry, so nothing extra is written).
|
|
1756
|
+
//
|
|
1757
|
+
// NON-BLOCKING (2026-07-06, the write-behind persistence contract): the
|
|
1758
|
+
// gateway lock is held only for a zero-copy SNAPSHOT (findExact +
|
|
1759
|
+
// cloneKvCaches — microseconds; entries are immutable so the clones are
|
|
1760
|
+
// consistent forever). The flush itself runs OFF the lock via
|
|
1761
|
+
// storeAsync, and writes chain serially so two multi-hundred-MB flushes
|
|
1762
|
+
// never overlap.
|
|
1763
|
+
//
|
|
1764
|
+
// IDLE-GATED (2026-07-07, the decode@ctx fix): "off the lock" was not
|
|
1765
|
+
// enough — every per-tensor flush step is a blocking GPU sync on the
|
|
1766
|
+
// SAME stream decode uses (ops.contiguous enqueues a kernel, rawBytesView
|
|
1767
|
+
// evals) plus a synchronous multi-MB writeSync, and the setImmediate
|
|
1768
|
+
// pacing interleaved those slices exactly between decode tokens. A ~16k
|
|
1769
|
+
// entry's flush overlapping the bench's ctx repeats depressed decode@ctx
|
|
1770
|
+
// ~9% on e4b (mlx-lm runs no equivalent background work). Now every step
|
|
1771
|
+
// — including the first — awaits gateway.onIdle() (ssdFlushGate), so the
|
|
1772
|
+
// flush only progresses while NOTHING is generating and pauses when a
|
|
1773
|
+
// request arrives mid-flush. Tradeoffs, accepted: durability waits for a
|
|
1774
|
+
// quiet moment (single-user serving quiesces constantly; sustained
|
|
1775
|
+
// hammering defers the flush AND the spill clones' GPU-memory release —
|
|
1776
|
+
// bounded by the chain), and one in-flight tensor step (~10-15 MB) can
|
|
1777
|
+
// still land ahead of a just-arrived request. MLX_BUN_SSD_WRITEBEHIND=0
|
|
1778
|
+
// disables write-behind snapshots entirely — the paired-A/B lever + kill
|
|
1779
|
+
// switch (restart survival then degrades to spill-on-evict only).
|
|
1780
|
+
const writeBehindOn = process.env.MLX_BUN_SSD_WRITEBEHIND !== "0";
|
|
1781
|
+
const ssdFlushGate = (): Promise<void> => gateway.onIdle();
|
|
1135
1782
|
const ssdPending = new Map<string, ReturnType<typeof setTimeout>>();
|
|
1783
|
+
// Bounded write-behind queue (2026-07-07 review fix — see SpillQueue in
|
|
1784
|
+
// kv-store.ts): pending clones pin their entries' GPU buffers while the
|
|
1785
|
+
// idle gate starves under sustained traffic, so QUEUED bytes are capped —
|
|
1786
|
+
// over cap the oldest queued spill drops (clone disposed immediately, a
|
|
1787
|
+
// future cache miss, never a wrong result) instead of accumulating past
|
|
1788
|
+
// the prompt-cache cap. Default 2 GB (a quarter of the 8 GB RAM-cache
|
|
1789
|
+
// default); MLX_BUN_SSD_SPILL_QUEUE_GB overrides. There is deliberately
|
|
1790
|
+
// NO shutdown flush: exiting under traffic loses queued write-behind
|
|
1791
|
+
// entries (restart survival degrades to whatever flushed) — accepted for
|
|
1792
|
+
// a best-effort cache tier; drops/pending are visible in /stats.
|
|
1793
|
+
// `|| 2` would coerce an explicit "0" back to 2 GB — parse so 0 works
|
|
1794
|
+
// (cap 0 = keep only the newest + in-flight clone pinned; the soft cap
|
|
1795
|
+
// never drops the item just enqueued).
|
|
1796
|
+
const spillQueueGbRaw = Number(process.env.MLX_BUN_SSD_SPILL_QUEUE_GB);
|
|
1797
|
+
const spillQueueCapBytes =
|
|
1798
|
+
(Number.isFinite(spillQueueGbRaw) && spillQueueGbRaw >= 0 ? spillQueueGbRaw : 2) * 1024 ** 3;
|
|
1799
|
+
const spillQueue = ssdStore
|
|
1800
|
+
? new SpillQueue(
|
|
1801
|
+
spillQueueCapBytes,
|
|
1802
|
+
cacheBytes,
|
|
1803
|
+
(item) => ssdStore!.storeAsync(item.tokens, item.caches, item.ns, ssdFlushGate),
|
|
1804
|
+
(caches) => { for (const c of caches) c.dispose(); },
|
|
1805
|
+
)
|
|
1806
|
+
: null;
|
|
1136
1807
|
const scheduleSsdSnapshot = (tokens: number[], ns: string): void => {
|
|
1137
|
-
if (!ssdStore || tokens.length === 0) return;
|
|
1138
|
-
const
|
|
1808
|
+
if (!ssdStore || !writeBehindOn || tokens.length === 0) return;
|
|
1809
|
+
const key = `${ns}|${tokens.length}`;
|
|
1810
|
+
const prev = ssdPending.get(key);
|
|
1139
1811
|
if (prev) clearTimeout(prev);
|
|
1140
|
-
|
|
1141
|
-
ssdPending.delete(
|
|
1812
|
+
const fire = (): void => {
|
|
1813
|
+
ssdPending.delete(key);
|
|
1814
|
+
// Same guard as demoteIdle below (2026-07-07 review fix): the
|
|
1815
|
+
// runExclusive registers a serial waiter, which HOLDS BATCH ADMISSION
|
|
1816
|
+
// until every running row drains — a background durability timer must
|
|
1817
|
+
// never do that. Busy → re-arm and try again in 5 s (entries are
|
|
1818
|
+
// immutable; a late snapshot is just as valid).
|
|
1819
|
+
if (gateway.activeRows > 0 || gateway.pendingRows > 0) {
|
|
1820
|
+
const t = setTimeout(fire, 5000);
|
|
1821
|
+
t.unref?.();
|
|
1822
|
+
ssdPending.set(key, t);
|
|
1823
|
+
return;
|
|
1824
|
+
}
|
|
1142
1825
|
void gateway.runExclusive(async () => {
|
|
1143
1826
|
const e = promptCache.findExact(tokens, ns);
|
|
1144
|
-
if (e)
|
|
1827
|
+
if (!e) return null;
|
|
1828
|
+
return { tokens: e.tokens, caches: cloneKvCaches(e.caches) };
|
|
1829
|
+
}).then((snap) => {
|
|
1830
|
+
if (!snap) return;
|
|
1831
|
+
spillQueue!.enqueue({ tokens: snap.tokens, caches: snap.caches, ns });
|
|
1145
1832
|
}).catch(() => { /* cold tier is best-effort */ });
|
|
1146
|
-
}
|
|
1833
|
+
};
|
|
1834
|
+
ssdPending.set(key, setTimeout(fire, 1000));
|
|
1147
1835
|
};
|
|
1836
|
+
// Every put() — serial lane AND batch scheduler — schedules the snapshot
|
|
1837
|
+
// (Layer 0: batch-lane entries survive restarts too, not just evictions).
|
|
1838
|
+
promptCache.onPut = (tokens, ns) => scheduleSsdSnapshot(tokens, ns);
|
|
1839
|
+
|
|
1840
|
+
// Idle demotion (Layer 0): entries unused for --ssd-demote-idle seconds
|
|
1841
|
+
// spill to SSD and free their GPU memory — the RAM tier drains between
|
|
1842
|
+
// bursts while every prefix stays reachable (take() restores via
|
|
1843
|
+
// zero-copy mmap). The exclusive section below covers only the sweep's
|
|
1844
|
+
// zero-copy clone + dispose (spillOwned above); the SSD write itself
|
|
1845
|
+
// runs off-lock on ssdWriteChain. Swept only when the engine is TRULY idle: the
|
|
1846
|
+
// runExclusive below registers as a serial waiter, which would DRAIN a
|
|
1847
|
+
// running batch, so the activity check guards it (a momentary exclusive
|
|
1848
|
+
// grab of an idle engine is free). 0 disables.
|
|
1849
|
+
const demoteIdleMs = (serverOptions.ssdDemoteIdleSec ?? (ssdStore ? 300 : 0)) * 1000;
|
|
1850
|
+
let demoteTimer: ReturnType<typeof setInterval> | null = null;
|
|
1851
|
+
if (ssdStore && demoteIdleMs > 0) {
|
|
1852
|
+
demoteTimer = setInterval(() => {
|
|
1853
|
+
if (gateway.activeRows > 0 || gateway.pendingRows > 0) return;
|
|
1854
|
+
void gateway.runExclusive(async () => {
|
|
1855
|
+
const n = promptCache.demoteIdle(demoteIdleMs);
|
|
1856
|
+
if (n > 0) console.log(`[ssd-cache] demoted ${n} idle entr${n === 1 ? "y" : "ies"} to disk`);
|
|
1857
|
+
}).catch(() => {});
|
|
1858
|
+
}, Math.max(30_000, demoteIdleMs / 4));
|
|
1859
|
+
demoteTimer.unref?.();
|
|
1860
|
+
}
|
|
1148
1861
|
|
|
1149
1862
|
// The lane picker: routes each request to the serial path (runGeneration,
|
|
1150
1863
|
// above) or the continuous-batching scheduler, keeping the two off the GPU
|
|
1151
1864
|
// (and shared loraState) at the same time. See src/serve/generation-gateway.ts.
|
|
1152
|
-
const gateway = new GenerationGateway(ctx.model, batch, runGeneration
|
|
1865
|
+
const gateway = new GenerationGateway(ctx.model, batch, runGeneration, {
|
|
1866
|
+
kvBudgetBytes: serverOptions.kvBudgetBytes,
|
|
1867
|
+
// Phase 3.1: the server-wide scheme travels to the gateway once; the
|
|
1868
|
+
// gateway batches kv-quant requests when the scheme is the batchable
|
|
1869
|
+
// kvConfig/full-attention composition (see #kvBatchable) and the
|
|
1870
|
+
// scheduler applies it — otherwise those requests route serial as before.
|
|
1871
|
+
kvScheme,
|
|
1872
|
+
// Phase 3.2: batch-lane prompt-cache reuse — joiners take() the longest
|
|
1873
|
+
// usable prefix (multi-turn chat TTFT under --batch N); never-merged
|
|
1874
|
+
// lone rows put() back on finish. Vision/adapter requests never batch,
|
|
1875
|
+
// so the serial lane's bypass rules are preserved by routing.
|
|
1876
|
+
promptCache,
|
|
1877
|
+
});
|
|
1153
1878
|
|
|
1154
1879
|
// Admission ceiling, resolved once (Phase 5 memoryBudget enforcement).
|
|
1155
1880
|
// fit() solves max safe context from weights + KV growth + prefill
|
|
@@ -1265,9 +1990,30 @@ export function createServer(
|
|
|
1265
1990
|
stopSequences: (typeof req.stop === "string" ? [req.stop] : req.stop ?? [])
|
|
1266
1991
|
.filter((s) => typeof s === "string" && s.length > 0),
|
|
1267
1992
|
...kvScheme,
|
|
1993
|
+
// Paged KV is a server-wide mode like the kv scheme above (v1:
|
|
1994
|
+
// serial bf16 gemma4 — the createServer refusals hold the invariants).
|
|
1995
|
+
...(serverOptions.pagedKv ? { pagedKv: serverOptions.pagedKv } : {}),
|
|
1268
1996
|
};
|
|
1269
1997
|
};
|
|
1270
1998
|
|
|
1999
|
+
/** Compile a grammar controller from the request's structured-output fields
|
|
2000
|
+
* (response_format / guided_* / structured_outputs). Returns null when no
|
|
2001
|
+
* constraint is requested OR when compile failed (degrade path). On degrade,
|
|
2002
|
+
* `degradeHint` carries a human description for the system-prompt injection
|
|
2003
|
+
* + Warning header (oMLX parity — never 500). Honors MLX_BUN_GRAMMAR=0. */
|
|
2004
|
+
const compileGrammarForRequest = async (
|
|
2005
|
+
req: ChatRequest,
|
|
2006
|
+
): Promise<{ controller: import("./grammar").GrammarController | null; degradeHint: string | null }> => {
|
|
2007
|
+
if (!grammarEnabled()) return { controller: null, degradeHint: null };
|
|
2008
|
+
const r = await compileGrammarRequest(
|
|
2009
|
+
req as GrammarRequest,
|
|
2010
|
+
ctx.tokenizer,
|
|
2011
|
+
ctx.model.config.text.vocabSize,
|
|
2012
|
+
);
|
|
2013
|
+
if (!r) return { controller: null, degradeHint: null };
|
|
2014
|
+
return { controller: r.controller, degradeHint: r.degradeHint };
|
|
2015
|
+
};
|
|
2016
|
+
|
|
1271
2017
|
const templateOptionsFor = (req: ChatRequest, tools: ToolDefinition[] | null) => {
|
|
1272
2018
|
// enableThinking resolution (and its precedence) lives in
|
|
1273
2019
|
// resolveEnableThinking so template rendering and sampling stay in sync.
|
|
@@ -1278,13 +2024,60 @@ export function createServer(
|
|
|
1278
2024
|
req: ChatRequest,
|
|
1279
2025
|
tools: ToolDefinition[] | null,
|
|
1280
2026
|
): { ids: number[]; startInThinking: boolean } => {
|
|
1281
|
-
const
|
|
2027
|
+
const opts = templateOptionsFor(req, tools);
|
|
2028
|
+
const rendered = ctx.template.render(normalizeMessages(req.messages), opts);
|
|
1282
2029
|
const ids = ctx.tokenizer.encode(rendered);
|
|
1283
2030
|
// template includes <bos>; tokenizer post-processor also prepends one
|
|
1284
2031
|
const trimmed = ids[0] === ids[1] && ids[0] === ctx.tokenizer.bosTokenId ? ids.slice(1) : ids;
|
|
1285
2032
|
return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
|
|
1286
2033
|
};
|
|
1287
2034
|
|
|
2035
|
+
/** STABLE cache boundary: the prompt's tail is a generation primer (e.g.
|
|
2036
|
+
* 12B's `<|channel>thought` tokens) that the NEXT turn's re-render does
|
|
2037
|
+
* not contain — a cache entry ending there can never prefix-match again
|
|
2038
|
+
* once the rings wrap. Probe by rendering this conversation as if a
|
|
2039
|
+
* reply existed; the common prefix is what every future turn preserves.
|
|
2040
|
+
* The probe is EXPENSIVE (a second full render + encode — ~150 ms at
|
|
2041
|
+
* 16k with our tokenizer; 30-50% of the 12B ctx-repeat TTFT,
|
|
2042
|
+
* 2026-07-06b), but the PRIMER LENGTH it measures is constant per
|
|
2043
|
+
* template mode: the divergence sits after a special-token turn
|
|
2044
|
+
* delimiter, and special tokens break BPE merges, so message content
|
|
2045
|
+
* can't shift it. So the full probe runs ONCE per mode and later
|
|
2046
|
+
* requests pay a subtraction. (A wrong boundary is a QUALITY knob, not
|
|
2047
|
+
* a correctness one — any prefix ≤ len-1 is a valid snapshot point.) */
|
|
2048
|
+
const primerLenByMode = new Map<string, number>();
|
|
2049
|
+
const stableLenFor = (
|
|
2050
|
+
req: ChatRequest,
|
|
2051
|
+
tools: ToolDefinition[] | null,
|
|
2052
|
+
trimmed: number[],
|
|
2053
|
+
): number => {
|
|
2054
|
+
const opts = templateOptionsFor(req, tools);
|
|
2055
|
+
const mode = `${opts.enableThinking}|${!!tools?.length}`;
|
|
2056
|
+
const primer = primerLenByMode.get(mode);
|
|
2057
|
+
if (primer !== undefined) return Math.max(0, trimmed.length - primer);
|
|
2058
|
+
let stableLen = trimmed.length;
|
|
2059
|
+
try {
|
|
2060
|
+
const probe = ctx.tokenizer.encode(
|
|
2061
|
+
ctx.template.render(
|
|
2062
|
+
[...normalizeMessages(req.messages), { role: "assistant", content: "x" }],
|
|
2063
|
+
opts,
|
|
2064
|
+
),
|
|
2065
|
+
);
|
|
2066
|
+
const probeTrimmed =
|
|
2067
|
+
probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
|
|
2068
|
+
let i = 0;
|
|
2069
|
+
const n = Math.min(trimmed.length, probeTrimmed.length);
|
|
2070
|
+
while (i < n && trimmed[i] === probeTrimmed[i]) i++;
|
|
2071
|
+
stableLen = i;
|
|
2072
|
+
// Memoize the mode's primer length (sanity-capped: a "primer" longer
|
|
2073
|
+
// than 64 tokens means the probe diverged for content reasons —
|
|
2074
|
+
// don't generalize that).
|
|
2075
|
+
const p = trimmed.length - stableLen;
|
|
2076
|
+
if (p >= 0 && p <= 64) primerLenByMode.set(mode, p);
|
|
2077
|
+
} catch { /* probe is best-effort; full boundary is the fallback */ }
|
|
2078
|
+
return stableLen;
|
|
2079
|
+
};
|
|
2080
|
+
|
|
1288
2081
|
const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
|
|
1289
2082
|
// Gemma-4 keeps the token-level sentinel path; every other family
|
|
1290
2083
|
// (MiniCPM5, Qwen3/3.5, Tier-0 generics) emits tool calls as DECODED TEXT
|
|
@@ -1294,9 +2087,18 @@ export function createServer(
|
|
|
1294
2087
|
const toolRouter = (tools: ToolDefinition[] | null): ToolAwareStream =>
|
|
1295
2088
|
new ToolAwareStream(ctx.tokenizer, toolStreamMode(tools), tools);
|
|
1296
2089
|
|
|
2090
|
+
if (serverOptions.unixSocket) {
|
|
2091
|
+
try { require("node:fs").unlinkSync(serverOptions.unixSocket); } catch {}
|
|
2092
|
+
}
|
|
2093
|
+
// Bun's types make unix/port mutually exclusive variants that a spread
|
|
2094
|
+
// union can't prove — runtime accepts either; cast once at the call.
|
|
1297
2095
|
serverRef = Bun.serve({
|
|
1298
|
-
|
|
1299
|
-
|
|
2096
|
+
...(serverOptions.unixSocket
|
|
2097
|
+
? ({ unix: serverOptions.unixSocket } as unknown as Record<string, never>)
|
|
2098
|
+
: {
|
|
2099
|
+
port,
|
|
2100
|
+
...(serverOptions.hostname ? { hostname: serverOptions.hostname } : {}),
|
|
2101
|
+
}),
|
|
1300
2102
|
idleTimeout: 0,
|
|
1301
2103
|
// Web chat rides pi's AgentSession events over a WebSocket; the embedded
|
|
1302
2104
|
// pi provider points back at THIS server's own loopback /v1 (port
|
|
@@ -1305,9 +2107,21 @@ export function createServer(
|
|
|
1305
2107
|
port: () => serverRef.port ?? port,
|
|
1306
2108
|
modelId: ctx.modelId,
|
|
1307
2109
|
contextWindow: ctx.model.config.text.maxPositionEmbeddings,
|
|
1308
|
-
// capability
|
|
2110
|
+
// capability flags — true if a tower is loaded or loadable (lazy)
|
|
1309
2111
|
vision: !!(ctx.vision || ctx.loadVision),
|
|
2112
|
+
audio: !!(ctx.audio || ctx.loadAudio),
|
|
1310
2113
|
thinking: ctx.template.supportsThinking,
|
|
2114
|
+
// Model-author sampling defaults (generation_config.json, server-CLI
|
|
2115
|
+
// overrides applied) for the sampling popover's per-model recommended
|
|
2116
|
+
// values (web-ui-pass-plan.md #14 groundwork) — same resolution
|
|
2117
|
+
// /v1/models' gen_defaults uses, so the two surfaces agree. The
|
|
2118
|
+
// think-mode-vs-not distinction lives in toOptions (per-request); this
|
|
2119
|
+
// is the single "model's own defaults" snapshot sent once at ready.
|
|
2120
|
+
genDefaults: {
|
|
2121
|
+
temperature: serverOptions.defaultTemperature ?? ctx.genDefaults.temperature ?? null,
|
|
2122
|
+
topP: serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? null,
|
|
2123
|
+
topK: serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? null,
|
|
2124
|
+
},
|
|
1311
2125
|
}),
|
|
1312
2126
|
async fetch(request, server) {
|
|
1313
2127
|
const url = new URL(request.url);
|
|
@@ -1318,6 +2132,81 @@ export function createServer(
|
|
|
1318
2132
|
return new Response("expected websocket", { status: 426 });
|
|
1319
2133
|
}
|
|
1320
2134
|
|
|
2135
|
+
// Memory REST wrappers (web-chat-redesign.md §2.3/§9 Phase 2): thin
|
|
2136
|
+
// loopback JSON routes over src/memory/vault.ts for the web chat's
|
|
2137
|
+
// Memory panel. Handlers live in src/memory/rest.ts (pure functions,
|
|
2138
|
+
// no `ctx` dependency) so they're unit-testable without a loaded
|
|
2139
|
+
// model; dispatch here just matches path+method. The agent-tool
|
|
2140
|
+
// surface (src/memory/tools.ts) is untouched by this block.
|
|
2141
|
+
if (url.pathname === "/api/memory/status" && request.method === "GET") {
|
|
2142
|
+
const { handleMemoryStatus } = await import("./memory/rest");
|
|
2143
|
+
return handleMemoryStatus();
|
|
2144
|
+
}
|
|
2145
|
+
if (url.pathname === "/api/memory/list" && request.method === "GET") {
|
|
2146
|
+
const { handleMemoryList } = await import("./memory/rest");
|
|
2147
|
+
return handleMemoryList();
|
|
2148
|
+
}
|
|
2149
|
+
if (url.pathname === "/api/memory/search" && request.method === "GET") {
|
|
2150
|
+
const { handleMemorySearch } = await import("./memory/rest");
|
|
2151
|
+
return handleMemorySearch(url);
|
|
2152
|
+
}
|
|
2153
|
+
if (url.pathname === "/api/memory/article" && request.method === "GET") {
|
|
2154
|
+
const { handleMemoryArticle } = await import("./memory/rest");
|
|
2155
|
+
return handleMemoryArticle(url);
|
|
2156
|
+
}
|
|
2157
|
+
if (url.pathname === "/api/memory/links" && request.method === "GET") {
|
|
2158
|
+
const { handleMemoryLinks } = await import("./memory/rest");
|
|
2159
|
+
return handleMemoryLinks(url);
|
|
2160
|
+
}
|
|
2161
|
+
if (url.pathname === "/api/memory/history" && request.method === "GET") {
|
|
2162
|
+
const { handleMemoryHistory } = await import("./memory/rest");
|
|
2163
|
+
return handleMemoryHistory(url);
|
|
2164
|
+
}
|
|
2165
|
+
if (url.pathname === "/api/memory/diff" && request.method === "GET") {
|
|
2166
|
+
const { handleMemoryDiff } = await import("./memory/rest");
|
|
2167
|
+
return handleMemoryDiff(url);
|
|
2168
|
+
}
|
|
2169
|
+
if (url.pathname === "/api/memory/init" && request.method === "POST") {
|
|
2170
|
+
const { handleMemoryInit } = await import("./memory/rest");
|
|
2171
|
+
return handleMemoryInit(request);
|
|
2172
|
+
}
|
|
2173
|
+
|
|
2174
|
+
// Model Hub (web-chat-redesign.md §9 Phase 3, beat-matrix Axis 3):
|
|
2175
|
+
// thin loopback JSON routes over the registry/download/fit machinery
|
|
2176
|
+
// for the web chat's Hub panel. Handlers live in src/hub-rest.ts
|
|
2177
|
+
// (pure functions, no `ctx` dependency); dispatch here just matches
|
|
2178
|
+
// path+method, same convention as the memory wrappers above.
|
|
2179
|
+
if (url.pathname === "/api/hub/local" && request.method === "GET") {
|
|
2180
|
+
const { handleHubLocal } = await import("./hub-rest");
|
|
2181
|
+
return handleHubLocal();
|
|
2182
|
+
}
|
|
2183
|
+
if (url.pathname === "/api/hub/search" && request.method === "GET") {
|
|
2184
|
+
const { handleHubSearch } = await import("./hub-rest");
|
|
2185
|
+
return handleHubSearch(url);
|
|
2186
|
+
}
|
|
2187
|
+
if (url.pathname === "/api/hub/download" && request.method === "POST") {
|
|
2188
|
+
const { handleHubDownload } = await import("./hub-rest");
|
|
2189
|
+
return handleHubDownload(request);
|
|
2190
|
+
}
|
|
2191
|
+
if (url.pathname === "/api/hub/serve" && request.method === "POST") {
|
|
2192
|
+
const { handleHubServe } = await import("./hub-rest");
|
|
2193
|
+
return handleHubServe(request);
|
|
2194
|
+
}
|
|
2195
|
+
|
|
2196
|
+
// Session full-text search + export (docs/design/web-chat-redesign.md
|
|
2197
|
+
// §9 Phase 3, beat-matrix Axis 10/11 — the full-text-search BEAT and
|
|
2198
|
+
// Markdown/JSON chat export). Handlers live in src/serve/session-search.ts
|
|
2199
|
+
// (pure, read-only, no `ctx` dependency), dispatched here by
|
|
2200
|
+
// path+method, same convention as the memory/hub blocks above.
|
|
2201
|
+
if (url.pathname === "/api/sessions/search" && request.method === "GET") {
|
|
2202
|
+
const { handleSessionsSearch } = await import("./serve/session-search");
|
|
2203
|
+
return handleSessionsSearch(url);
|
|
2204
|
+
}
|
|
2205
|
+
if (url.pathname === "/api/sessions/export" && request.method === "GET") {
|
|
2206
|
+
const { handleSessionsExport } = await import("./serve/session-search");
|
|
2207
|
+
return handleSessionsExport(url);
|
|
2208
|
+
}
|
|
2209
|
+
|
|
1321
2210
|
// The unified SPA is served at "/"; legacy deep links redirect into
|
|
1322
2211
|
// the hash router so old bookmarks still land on the right section.
|
|
1323
2212
|
if (url.pathname === "/" && request.method === "GET") {
|
|
@@ -1325,6 +2214,50 @@ export function createServer(
|
|
|
1325
2214
|
headers: { "content-type": "text/html; charset=utf-8" },
|
|
1326
2215
|
});
|
|
1327
2216
|
}
|
|
2217
|
+
// Vendored syntax-highlighting assets (src/web/vendor/*, no CDN — see
|
|
2218
|
+
// that dir's README). Long cache: content is versioned by the vendor
|
|
2219
|
+
// rebuild step, not by URL, but a hard reload always wins during dev.
|
|
2220
|
+
if (url.pathname === "/assets/hljs.js" && request.method === "GET") {
|
|
2221
|
+
return new Response(HLJS_JS, {
|
|
2222
|
+
headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "public, max-age=3600" },
|
|
2223
|
+
});
|
|
2224
|
+
}
|
|
2225
|
+
if (url.pathname === "/assets/hljs.css" && request.method === "GET") {
|
|
2226
|
+
return new Response(HLJS_CSS, {
|
|
2227
|
+
headers: { "content-type": "text/css; charset=utf-8", "cache-control": "public, max-age=3600" },
|
|
2228
|
+
});
|
|
2229
|
+
}
|
|
2230
|
+
// The frontend bundle (see the import comment above) — cache header
|
|
2231
|
+
// matches the other /assets/* entries; content is versioned by the
|
|
2232
|
+
// build step, not the URL, and a hard reload always wins during dev.
|
|
2233
|
+
if (url.pathname === "/assets/app.js" && request.method === "GET") {
|
|
2234
|
+
return new Response(APP_JS, {
|
|
2235
|
+
headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "public, max-age=3600" },
|
|
2236
|
+
});
|
|
2237
|
+
}
|
|
2238
|
+
// PWA manifest + icon (plan §9 Phase 3): linked from app.html's
|
|
2239
|
+
// <link rel="manifest">. Same versioned-by-content cache posture as
|
|
2240
|
+
// the other /assets/* entries above.
|
|
2241
|
+
if (url.pathname === "/manifest.webmanifest" && request.method === "GET") {
|
|
2242
|
+
return new Response(MANIFEST_WEBMANIFEST, {
|
|
2243
|
+
headers: { "content-type": "application/manifest+json; charset=utf-8", "cache-control": "public, max-age=3600" },
|
|
2244
|
+
});
|
|
2245
|
+
}
|
|
2246
|
+
if (url.pathname === "/assets/icon.svg" && request.method === "GET") {
|
|
2247
|
+
return new Response(ICON_SVG, {
|
|
2248
|
+
headers: { "content-type": "image/svg+xml; charset=utf-8", "cache-control": "public, max-age=3600" },
|
|
2249
|
+
});
|
|
2250
|
+
}
|
|
2251
|
+
// Service worker MUST be served at the root scope (/sw.js, not
|
|
2252
|
+
// /assets/sw.js) — a worker's default scope is the directory it's
|
|
2253
|
+
// served from, and the shell it manages includes "/" itself.
|
|
2254
|
+
// no-store: the browser's own update check needs a fresh byte
|
|
2255
|
+
// comparison every time to notice a new worker, not a cached 304.
|
|
2256
|
+
if (url.pathname === "/sw.js" && request.method === "GET") {
|
|
2257
|
+
return new Response(SW_JS, {
|
|
2258
|
+
headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "no-store" },
|
|
2259
|
+
});
|
|
2260
|
+
}
|
|
1328
2261
|
// v2 HLG Curve Designer — served same-origin so /generate + /signal need no CORS.
|
|
1329
2262
|
// Read fresh from disk in dev (edits show on reload, no restart); fall back to the
|
|
1330
2263
|
// embedded copy when running as the compiled single binary.
|
|
@@ -1368,7 +2301,7 @@ export function createServer(
|
|
|
1368
2301
|
// surface on their own. (scan() is INSERT-OR-REPLACE + prunes deleted,
|
|
1369
2302
|
// so it's idempotent and cheap for a local cache.)
|
|
1370
2303
|
await reg.scan();
|
|
1371
|
-
const { visionCapable } = await import("./registry");
|
|
2304
|
+
const { visionCapable, audioCapable } = await import("./registry");
|
|
1372
2305
|
const { supportTier } = await import("./model/support");
|
|
1373
2306
|
const rows = [];
|
|
1374
2307
|
// listCanonical: one row per repo (refs/main) — duplicate snapshots
|
|
@@ -1389,7 +2322,8 @@ export function createServer(
|
|
|
1389
2322
|
rows.push({
|
|
1390
2323
|
repo_id: m.repoId, model_type: m.modelType,
|
|
1391
2324
|
size_bytes: m.sizeBytes, quant_bits: m.quantBits,
|
|
1392
|
-
vision: visionCapable(m),
|
|
2325
|
+
vision: visionCapable(m), audio: audioCapable(m),
|
|
2326
|
+
supported, support_tier: tier,
|
|
1393
2327
|
serving: m.repoId === ctx.modelId,
|
|
1394
2328
|
assessment,
|
|
1395
2329
|
});
|
|
@@ -1516,7 +2450,13 @@ export function createServer(
|
|
|
1516
2450
|
const layerTypes = ctx.model.config.text.layerTypes;
|
|
1517
2451
|
const kvLayers: Record<string, number> = {};
|
|
1518
2452
|
let kvMode = "bf16";
|
|
1519
|
-
if (kvScheme.
|
|
2453
|
+
if (kvScheme.turboQuant) {
|
|
2454
|
+
// v1: full-attention layers only (sliding-window stays bf16 —
|
|
2455
|
+
// docs/design/turboquant-kv.md non-goal).
|
|
2456
|
+
const fullAttn = layerTypes.filter((l) => l !== "sliding_attention").length;
|
|
2457
|
+
kvMode = `turbo k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`;
|
|
2458
|
+
kvLayers[`turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`] = fullAttn;
|
|
2459
|
+
} else if (kvScheme.kvBits) {
|
|
1520
2460
|
kvMode = `uniform-kv${kvScheme.kvBits}`;
|
|
1521
2461
|
kvLayers[`kv${kvScheme.kvBits}`] = layerTypes.length;
|
|
1522
2462
|
} else if (kvScheme.kvConfig) {
|
|
@@ -1548,6 +2488,13 @@ export function createServer(
|
|
|
1548
2488
|
restores: ssdStore.stats.restores,
|
|
1549
2489
|
spills: ssdStore.stats.spills,
|
|
1550
2490
|
restore_ms_last: Math.round(ssdStore.stats.restoreMsLast),
|
|
2491
|
+
demotions: promptCache.demotions,
|
|
2492
|
+
// Bounded write-behind queue (2026-07-07): clones waiting for
|
|
2493
|
+
// the idle-gated flush, and spills the byte cap dropped
|
|
2494
|
+
// (each = one future cache miss, never a wrong result).
|
|
2495
|
+
pending_spills: spillQueue?.pendingCount ?? 0,
|
|
2496
|
+
pending_spill_bytes: spillQueue?.pendingBytes ?? 0,
|
|
2497
|
+
dropped_spills: spillQueue?.droppedCount ?? 0,
|
|
1551
2498
|
},
|
|
1552
2499
|
} : {}),
|
|
1553
2500
|
response_store: {
|
|
@@ -1576,12 +2523,31 @@ export function createServer(
|
|
|
1576
2523
|
configured: batch,
|
|
1577
2524
|
batched: gateway.batchingEnabled,
|
|
1578
2525
|
active_rows: gateway.activeRows,
|
|
2526
|
+
pending_rows: gateway.pendingRows,
|
|
2527
|
+
submitted_rows: gateway.submittedRows,
|
|
2528
|
+
kv_bytes: gateway.kvBytes.projected,
|
|
2529
|
+
kv_budget_bytes: gateway.kvBytes.budget,
|
|
1579
2530
|
},
|
|
1580
2531
|
});
|
|
1581
2532
|
}
|
|
1582
2533
|
|
|
1583
2534
|
// mlx_lm.server parity: GET /health → the exact body it writes
|
|
1584
2535
|
// ('{"status": "ok"}', note the space) so byte-for-byte health checks pass.
|
|
2536
|
+
// Engine-mode admin (unix-socket children only — never exposed on
|
|
2537
|
+
// TCP): drain = quiesce the gateway + demote the whole prompt cache
|
|
2538
|
+
// to the SSD tier. The pool calls this before evicting a model
|
|
2539
|
+
// child, so its state survives the eviction losslessly.
|
|
2540
|
+
if (
|
|
2541
|
+
serverOptions.unixSocket &&
|
|
2542
|
+
url.pathname === "/admin/drain" &&
|
|
2543
|
+
request.method === "POST"
|
|
2544
|
+
) {
|
|
2545
|
+
await gateway.runExclusive(async () => {
|
|
2546
|
+
promptCache.demoteIdle(0);
|
|
2547
|
+
});
|
|
2548
|
+
return Response.json({ drained: true, demotions: promptCache.demotions });
|
|
2549
|
+
}
|
|
2550
|
+
|
|
1585
2551
|
if (url.pathname === "/health" && request.method === "GET") {
|
|
1586
2552
|
return new Response('{"status": "ok"}', {
|
|
1587
2553
|
headers: { "content-type": "application/json" },
|
|
@@ -1601,26 +2567,43 @@ export function createServer(
|
|
|
1601
2567
|
? decodeURIComponent(url.pathname.slice("/v1/models/".length))
|
|
1602
2568
|
: null;
|
|
1603
2569
|
const created = Math.floor(startedAt / 1000);
|
|
2570
|
+
// Model-author sampling defaults (generation_config.json), resolved once
|
|
2571
|
+
// at load (web-ui-pass-plan.md #14 groundwork): lets a client (the web
|
|
2572
|
+
// chat's sampling popover, external pi) show the SERVED model's actual
|
|
2573
|
+
// recommended values instead of one hardcoded shape.
|
|
2574
|
+
const genDefaultsOut = {
|
|
2575
|
+
temperature: ctx.genDefaults.temperature ?? null,
|
|
2576
|
+
top_p: ctx.genDefaults.topP ?? null,
|
|
2577
|
+
top_k: ctx.genDefaults.topK ?? null,
|
|
2578
|
+
};
|
|
1604
2579
|
const data: Array<Record<string, unknown>> = [{
|
|
1605
2580
|
id: ctx.modelId, object: "model", created, owned_by: "mlx-bun",
|
|
1606
2581
|
context_window: ctx.model.config.text.maxPositionEmbeddings,
|
|
1607
2582
|
// Capability flags for clients (CLI/external pi) that build a
|
|
1608
2583
|
// provider from discovery — `reasoning` gates the thinking toggle,
|
|
1609
|
-
// `vision` the image input declaration
|
|
1610
|
-
//
|
|
2584
|
+
// `vision` the image input declaration, `audio` whether
|
|
2585
|
+
// `input_audio` content parts are accepted (tower loaded or lazily
|
|
2586
|
+
// loadable; same signals the web embed uses).
|
|
1611
2587
|
reasoning: ctx.template.supportsThinking,
|
|
1612
2588
|
vision: !!(ctx.vision || ctx.loadVision),
|
|
2589
|
+
audio: !!(ctx.audio || ctx.loadAudio),
|
|
2590
|
+
gen_defaults: genDefaultsOut,
|
|
1613
2591
|
}];
|
|
1614
2592
|
try {
|
|
1615
|
-
const { Registry } = await import("./registry");
|
|
2593
|
+
const { Registry, visionCapable } = await import("./registry");
|
|
2594
|
+
const { supportTier } = await import("./model/support");
|
|
1616
2595
|
const reg = new Registry();
|
|
1617
2596
|
try {
|
|
1618
2597
|
if (reg.list().length === 0) await reg.scan();
|
|
1619
2598
|
// Canonical rows only — a repo with N snapshots is ONE model id.
|
|
1620
2599
|
for (const m of reg.listCanonical()) {
|
|
1621
2600
|
if (m.repoId === ctx.modelId) continue;
|
|
1622
|
-
|
|
1623
|
-
|
|
2601
|
+
const tier = supportTier(m.modelType, m.repoId);
|
|
2602
|
+
if (tier === null) continue;
|
|
2603
|
+
data.push({
|
|
2604
|
+
id: m.repoId, object: "model", created,
|
|
2605
|
+
vision: visionCapable(m), tier,
|
|
2606
|
+
});
|
|
1624
2607
|
}
|
|
1625
2608
|
} finally {
|
|
1626
2609
|
reg.close();
|
|
@@ -1676,13 +2659,16 @@ export function createServer(
|
|
|
1676
2659
|
// unmount. Mount and unmount go through the generation queue so
|
|
1677
2660
|
// they never race an in-flight forward pass.
|
|
1678
2661
|
if (url.pathname === "/v1/adapters/available" && request.method === "GET") {
|
|
1679
|
-
// On-disk adapters that can be mounted (the chat selector's source)
|
|
1680
|
-
//
|
|
1681
|
-
//
|
|
1682
|
-
//
|
|
1683
|
-
//
|
|
1684
|
-
// adapter
|
|
1685
|
-
//
|
|
2662
|
+
// On-disk adapters that can be mounted (the chat selector's source).
|
|
2663
|
+
// Every on-disk adapter is returned — the web chat routing chip (§5.2)
|
|
2664
|
+
// needs incompatible entries visible-but-grayed, not silently dropped,
|
|
2665
|
+
// so a user understands "10 adapters exist, 2 apply here" rather than
|
|
2666
|
+
// seeing a shorter list with no explanation. `compatible` = the
|
|
2667
|
+
// adapter's recorded base repo id matches the served model (compared
|
|
2668
|
+
// on the bare name, lenient about the org); an adapter with no
|
|
2669
|
+
// recorded base is treated as compatible (mount validates it for
|
|
2670
|
+
// real). `mounted` flags ones already loaded so the UI auto-loads on
|
|
2671
|
+
// select only if needed.
|
|
1686
2672
|
const { homedir } = await import("node:os");
|
|
1687
2673
|
const stores = [
|
|
1688
2674
|
`${homedir()}/.cache/mlx-bun-finetunes`,
|
|
@@ -1692,10 +2678,10 @@ export function createServer(
|
|
|
1692
2678
|
const bareName = (s: string) => s.split("/").pop()!.toLowerCase();
|
|
1693
2679
|
const servedName = bareName(ctx.modelId);
|
|
1694
2680
|
const adapters = (await listAvailableAdapters(stores))
|
|
1695
|
-
.filter((a) => a.baseModel == null || bareName(a.baseModel) === servedName)
|
|
1696
2681
|
.map((a) => ({
|
|
1697
2682
|
id: a.id, path: a.path, rank: a.rank, scale: a.scale,
|
|
1698
2683
|
base_model: a.baseModel, mounted: mounted.has(a.id),
|
|
2684
|
+
compatible: a.baseModel == null || bareName(a.baseModel) === servedName,
|
|
1699
2685
|
}));
|
|
1700
2686
|
return Response.json({ adapters });
|
|
1701
2687
|
}
|
|
@@ -1704,6 +2690,7 @@ export function createServer(
|
|
|
1704
2690
|
adapters: ctx.adapters.list().map((a) => ({
|
|
1705
2691
|
id: a.id, path: a.path, rank: a.rank, scale: a.scale,
|
|
1706
2692
|
size_bytes: a.sizeBytes, mounted_layers: a.mountedLayers,
|
|
2693
|
+
ram_bytes: a.ramBytes,
|
|
1707
2694
|
})),
|
|
1708
2695
|
});
|
|
1709
2696
|
}
|
|
@@ -1723,7 +2710,7 @@ export function createServer(
|
|
|
1723
2710
|
const info = await gateway.runExclusive(() => ctx.adapters.mount(body.id!, body.path!));
|
|
1724
2711
|
return Response.json({
|
|
1725
2712
|
id: info.id, mounted_layers: info.mountedLayers,
|
|
1726
|
-
rank: info.rank, scale: info.scale,
|
|
2713
|
+
rank: info.rank, scale: info.scale, ram_bytes: info.ramBytes,
|
|
1727
2714
|
});
|
|
1728
2715
|
} catch (e) {
|
|
1729
2716
|
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
@@ -1841,29 +2828,140 @@ export function createServer(
|
|
|
1841
2828
|
(m) => Array.isArray(m.content) &&
|
|
1842
2829
|
m.content.some((p: any) => p.type === "image_url" || p.type === "image"),
|
|
1843
2830
|
);
|
|
2831
|
+
// Same shapes extractAudio accepts: OpenAI-canonical input_audio plus
|
|
2832
|
+
// optiq's audio / audio_url aliases (audio-input-plan.md §3.2).
|
|
2833
|
+
const hasAudio = body.messages.some(
|
|
2834
|
+
(m) => Array.isArray(m.content) &&
|
|
2835
|
+
m.content.some(
|
|
2836
|
+
(p: any) => p.type === "input_audio" || p.type === "audio" || p.type === "audio_url",
|
|
2837
|
+
),
|
|
2838
|
+
);
|
|
1844
2839
|
let promptIds: number[];
|
|
2840
|
+
let stableLen: number | null = null;
|
|
1845
2841
|
// Whether the prompt primed an open <think> (Qwen3.5/MiniCPM5 thinking
|
|
1846
2842
|
// on) so the model's output starts mid-reasoning — seeds the splitter.
|
|
1847
2843
|
// Vision is Gemma4-only (no switchable thinking channel), so it's false.
|
|
1848
2844
|
let startInThinking = false;
|
|
1849
2845
|
let vision: Parameters<typeof runGeneration>[3];
|
|
1850
2846
|
let diffusionPixels: import("./mlx/array").MlxArray | null = null;
|
|
2847
|
+
// Grammar-constrained decoding (src/grammar.ts). Compile BEFORE prompt
|
|
2848
|
+
// rendering: on the degrade path (compile failed but a constraint was
|
|
2849
|
+
// requested) inject a system message instructing valid JSON so the
|
|
2850
|
+
// model still best-efforts schema-conformant output (oMLX parity —
|
|
2851
|
+
// _compile_grammar_for_request returning None + build_json_system_prompt).
|
|
2852
|
+
let grammarCtrl: import("./grammar").GrammarController | null = null;
|
|
2853
|
+
let grammarWarning: string | null = null;
|
|
2854
|
+
const grammarReq =
|
|
2855
|
+
body.response_format != null || !!body.guided_grammar ||
|
|
2856
|
+
!!body.guided_regex || !!body.guided_choice?.length ||
|
|
2857
|
+
body.structured_outputs != null;
|
|
2858
|
+
if (grammarReq) {
|
|
2859
|
+
const g = await compileGrammarForRequest(body);
|
|
2860
|
+
grammarCtrl = g.controller;
|
|
2861
|
+
if (!g.controller && g.degradeHint) {
|
|
2862
|
+
grammarWarning = `grammar not enforced: ${g.degradeHint} — falling back to prompt injection`;
|
|
2863
|
+
body = {
|
|
2864
|
+
...body,
|
|
2865
|
+
messages: [
|
|
2866
|
+
{ role: "system", content: degradeJsonSystemPrompt(body) },
|
|
2867
|
+
...body.messages,
|
|
2868
|
+
],
|
|
2869
|
+
};
|
|
2870
|
+
}
|
|
2871
|
+
}
|
|
2872
|
+
// Free what a rejected request has already allocated (grammar WASM
|
|
2873
|
+
// matcher/compiled-grammar, vision embeddings, diffusion pixels).
|
|
2874
|
+
// Every early return between here and gateway.run() must call this —
|
|
2875
|
+
// generate()/the gateway only take disposal ownership once running
|
|
2876
|
+
// (leaked one matcher per SSRF-blocked-media + response_format
|
|
2877
|
+
// request before the 2026-07-07 sweep).
|
|
2878
|
+
const disposeRejected = () => {
|
|
2879
|
+
grammarCtrl?.dispose();
|
|
2880
|
+
vision?.embeddings.dispose();
|
|
2881
|
+
vision?.imageMask?.dispose();
|
|
2882
|
+
vision?.multimodalMask?.dispose();
|
|
2883
|
+
diffusionPixels?.dispose();
|
|
2884
|
+
};
|
|
1851
2885
|
try {
|
|
1852
|
-
if (
|
|
2886
|
+
if (hasAudio) {
|
|
2887
|
+
// Audio (and MIXED image+audio) input — A4 of
|
|
2888
|
+
// docs/design/audio-input-plan.md. One buildMultimodalPrompt call
|
|
2889
|
+
// splices both media kinds in document order. A request WITH
|
|
2890
|
+
// audio on a model whose tower is unavailable is an explicit 400
|
|
2891
|
+
// — never a silent text-only degrade (that leniency is only for
|
|
2892
|
+
// requests without the media, getVisionTower's contract).
|
|
2893
|
+
const audioTower = getAudioTower(ctx);
|
|
2894
|
+
if (!audioTower || !ctx.audioTokenIds) {
|
|
2895
|
+
disposeRejected();
|
|
2896
|
+
return Response.json(
|
|
2897
|
+
{
|
|
2898
|
+
error: {
|
|
2899
|
+
message:
|
|
2900
|
+
`model ${ctx.modelId} has no audio tower — audio input needs ` +
|
|
2901
|
+
`a model whose config.json carries audio_config and whose ` +
|
|
2902
|
+
`sidecar ships the audio tensors (e.g. gemma-4 e4b OptiQ)`,
|
|
2903
|
+
},
|
|
2904
|
+
},
|
|
2905
|
+
{ status: 400 },
|
|
2906
|
+
);
|
|
2907
|
+
}
|
|
2908
|
+
let visionSide:
|
|
2909
|
+
| { tower: VisionEncoder; tokenIds: VisionTokenIds }
|
|
2910
|
+
| undefined;
|
|
2911
|
+
if (hasImages) {
|
|
2912
|
+
const tower = getVisionTower(ctx);
|
|
2913
|
+
if (!tower) {
|
|
2914
|
+
disposeRejected();
|
|
2915
|
+
return Response.json(
|
|
2916
|
+
{ error: { message: "model has no vision sidecar" } }, { status: 400 },
|
|
2917
|
+
);
|
|
2918
|
+
}
|
|
2919
|
+
visionSide = { tower, tokenIds: ctx.visionTokenIds };
|
|
2920
|
+
}
|
|
2921
|
+
const { messages: withAudioParts, images } =
|
|
2922
|
+
await extractImages(normalizeMessages(body.messages));
|
|
2923
|
+
const { messages, audio } = await extractAudio(withAudioParts);
|
|
2924
|
+
// Non-WAV containers (mp3/m4a/flac/ogg/aiff/…) transcode to
|
|
2925
|
+
// 16 kHz WAV via CoreAudio; RIFF bytes pass through untouched.
|
|
2926
|
+
// Failures throw into the prompt-build 400 below.
|
|
2927
|
+
const wavs = await Promise.all(audio.map(ensureWav));
|
|
2928
|
+
// The towers are only ever non-null for Gemma4 (loader gates).
|
|
2929
|
+
const mp = await buildMultimodalPrompt(
|
|
2930
|
+
ctx.model as Gemma4Model,
|
|
2931
|
+
{
|
|
2932
|
+
...(visionSide ? { vision: visionSide } : {}),
|
|
2933
|
+
audio: { tower: audioTower, tokenIds: ctx.audioTokenIds },
|
|
2934
|
+
},
|
|
2935
|
+
ctx.tokenizer, ctx.template, messages, images, wavs, tools,
|
|
2936
|
+
);
|
|
2937
|
+
promptIds = mp.ids;
|
|
2938
|
+
// bidirMask is null whenever audio is present (§3.3 Q1: mixed
|
|
2939
|
+
// prompts run fully causal); the union mask does the per-layer
|
|
2940
|
+
// id zeroing either way.
|
|
2941
|
+
vision = {
|
|
2942
|
+
embeddings: mp.embeddings,
|
|
2943
|
+
...(mp.bidirMask ? { imageMask: mp.bidirMask } : {}),
|
|
2944
|
+
multimodalMask: mp.multimodalMask,
|
|
2945
|
+
};
|
|
2946
|
+
} else if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
|
|
1853
2947
|
// DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
|
|
1854
2948
|
// encoder vision merge feed the denoising engine (NOT the AR
|
|
1855
2949
|
// forwardEmbeddings path). v1 supports a single image.
|
|
1856
2950
|
const dm = ctx.model;
|
|
1857
|
-
if (!dm.visionTower)
|
|
2951
|
+
if (!dm.visionTower) {
|
|
2952
|
+
disposeRejected();
|
|
1858
2953
|
return Response.json(
|
|
1859
2954
|
{ error: { message: "this checkpoint has no vision tower" } }, { status: 400 },
|
|
1860
2955
|
);
|
|
2956
|
+
}
|
|
1861
2957
|
const { messages, images } = await extractImages(normalizeMessages(body.messages));
|
|
1862
|
-
if (images.length !== 1)
|
|
2958
|
+
if (images.length !== 1) {
|
|
2959
|
+
disposeRejected();
|
|
1863
2960
|
return Response.json(
|
|
1864
2961
|
{ error: { message: "DiffusionGemma image input supports exactly one image" } },
|
|
1865
2962
|
{ status: 400 },
|
|
1866
2963
|
);
|
|
2964
|
+
}
|
|
1867
2965
|
const rendered = ctx.template.render(messages, { tools, addGenerationPrompt: true });
|
|
1868
2966
|
const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
|
|
1869
2967
|
const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
|
|
@@ -1877,10 +2975,12 @@ export function createServer(
|
|
|
1877
2975
|
// Loads (and caches) the tower on first image request — text-only
|
|
1878
2976
|
// sessions never pay for it.
|
|
1879
2977
|
const tower = getVisionTower(ctx);
|
|
1880
|
-
if (!tower)
|
|
2978
|
+
if (!tower) {
|
|
2979
|
+
disposeRejected();
|
|
1881
2980
|
return Response.json(
|
|
1882
2981
|
{ error: { message: "model has no vision sidecar" } }, { status: 400 },
|
|
1883
2982
|
);
|
|
2983
|
+
}
|
|
1884
2984
|
const { messages, images } = await extractImages(normalizeMessages(body.messages));
|
|
1885
2985
|
// The tower is only ever non-null for Gemma4 (sidecar gate in
|
|
1886
2986
|
// makeVisionLoader), so the model narrow is safe here.
|
|
@@ -1894,8 +2994,10 @@ export function createServer(
|
|
|
1894
2994
|
const built = promptIdsFor(body, tools);
|
|
1895
2995
|
promptIds = built.ids;
|
|
1896
2996
|
startInThinking = built.startInThinking;
|
|
2997
|
+
stableLen = -1; // marker: chat path, probe lazily below
|
|
1897
2998
|
}
|
|
1898
2999
|
} catch (e) {
|
|
3000
|
+
disposeRejected();
|
|
1899
3001
|
return Response.json(
|
|
1900
3002
|
{ error: { message: `prompt build failed: ${(e as Error).message}` } },
|
|
1901
3003
|
{ status: 400 },
|
|
@@ -1904,22 +3006,33 @@ export function createServer(
|
|
|
1904
3006
|
let options: ReturnType<typeof toOptions>;
|
|
1905
3007
|
try {
|
|
1906
3008
|
options = toOptions(body);
|
|
3009
|
+
// Stable cache boundary for the prompt-boundary snapshot (chat
|
|
3010
|
+
// prompts end in a generation primer the next turn won't contain).
|
|
3011
|
+
// The probe costs a second full render+encode (~150 ms at 16k), so
|
|
3012
|
+
// run it ONLY when a snapshot can actually be taken: the cache
|
|
3013
|
+
// must not already hold the strict prefix (warm repeats peek at
|
|
3014
|
+
// len-1 → skip; the snapshot gate needs boundary > cached prefix).
|
|
3015
|
+
if (stableLen !== null) {
|
|
3016
|
+
const cachePeek = promptCache.peekPrefixLen(
|
|
3017
|
+
promptIds, options.adapters?.join("+") ?? "");
|
|
3018
|
+
if (cachePeek < promptIds.length - 1)
|
|
3019
|
+
options.snapshotAt = stableLenFor(body, tools, promptIds);
|
|
3020
|
+
}
|
|
1907
3021
|
} catch (e) {
|
|
1908
3022
|
// bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
|
|
1909
|
-
|
|
1910
|
-
vision?.imageMask.dispose();
|
|
1911
|
-
diffusionPixels?.dispose();
|
|
3023
|
+
disposeRejected();
|
|
1912
3024
|
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
1913
3025
|
}
|
|
1914
3026
|
if (diffusionPixels) options.visionPixels = diffusionPixels;
|
|
3027
|
+
// Attach the compiled grammar controller (null when no constraint /
|
|
3028
|
+
// degrade — generate() runs the unmasked fast pipelined loop).
|
|
3029
|
+
if (grammarCtrl) options.grammar = grammarCtrl;
|
|
1915
3030
|
// Admission: reject what cannot finish within the memory budget
|
|
1916
3031
|
// (the GPU OOM it would otherwise hit is uncatchable and kills
|
|
1917
3032
|
// the process — Phase 6 finding).
|
|
1918
3033
|
const requiredCtx = promptIds.length + (options.maxTokens ?? 1024);
|
|
1919
3034
|
if (requiredCtx > admission.maxSafeContext) {
|
|
1920
|
-
|
|
1921
|
-
vision?.imageMask.dispose();
|
|
1922
|
-
diffusionPixels?.dispose();
|
|
3035
|
+
disposeRejected();
|
|
1923
3036
|
return Response.json(
|
|
1924
3037
|
{
|
|
1925
3038
|
error: {
|
|
@@ -1941,6 +3054,7 @@ export function createServer(
|
|
|
1941
3054
|
const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
|
|
1942
3055
|
if (adapterIds.length) options.adapters = adapterIds;
|
|
1943
3056
|
} catch (e) {
|
|
3057
|
+
disposeRejected();
|
|
1944
3058
|
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
1945
3059
|
}
|
|
1946
3060
|
|
|
@@ -1958,14 +3072,18 @@ export function createServer(
|
|
|
1958
3072
|
options.topLogprobs = topLogprobs;
|
|
1959
3073
|
}
|
|
1960
3074
|
|
|
1961
|
-
// What lane this request takes (vision
|
|
1962
|
-
//
|
|
3075
|
+
// What lane this request takes (vision/audio media / adapters /
|
|
3076
|
+
// logprobs / seed / explicit kv-quant / a mounted draft → serial;
|
|
3077
|
+
// sampler extras, repetition penalty, and grammar all BATCH — see
|
|
3078
|
+
// willBatch). `vision` is the embeddings-prefill payload for BOTH
|
|
3079
|
+
// media kinds, so hasVision routes audio serial too — batching is a
|
|
3080
|
+
// mode, not a fallback, and batched audio has no oracle (§3.2).
|
|
1963
3081
|
const shape = {
|
|
1964
3082
|
hasVision: !!vision,
|
|
1965
3083
|
hasAdapters: !!options.adapters?.length,
|
|
1966
3084
|
hasRepetitionPenalty: !!options.repetitionPenalty,
|
|
1967
|
-
//
|
|
1968
|
-
//
|
|
3085
|
+
// Informational since 2026-07-02: per-row logits processors batch;
|
|
3086
|
+
// willBatch no longer gates on these fields.
|
|
1969
3087
|
hasLogitsExtras: !!(
|
|
1970
3088
|
options.minP || options.xtcProbability || options.logitBias ||
|
|
1971
3089
|
options.presencePenalty || options.frequencyPenalty
|
|
@@ -1973,10 +3091,26 @@ export function createServer(
|
|
|
1973
3091
|
wantsLogprobs: captureLogprobs,
|
|
1974
3092
|
userSeed: body.seed !== undefined,
|
|
1975
3093
|
kvQuant: !!(options.kvConfig?.length || options.kvBits),
|
|
3094
|
+
turboQuant: !!options.turboQuant,
|
|
3095
|
+
// grammarCtrl is null on the degrade path (prompt injection) —
|
|
3096
|
+
// those stay batchable. A real controller batches via per-row
|
|
3097
|
+
// matchers (MLX_BUN_GRAMMAR_BATCH=0 forces it serial).
|
|
3098
|
+
hasGrammar: !!grammarCtrl,
|
|
3099
|
+
hasDraft: !!ctx.draft,
|
|
1976
3100
|
};
|
|
1977
3101
|
const batched = gateway.willBatch(shape);
|
|
1978
3102
|
if (process.env.MLX_BUN_LANE_DEBUG === "1")
|
|
1979
3103
|
console.error(`[lane] batched=${batched} shape=${JSON.stringify(shape)} t=${Date.now() % 100000}`);
|
|
3104
|
+
// Per-turn lane (docs/design/web-chat-redesign.md §2.3 caveat / risk #5):
|
|
3105
|
+
// reported on usage AND recorded in the in-process lane registry (keyed
|
|
3106
|
+
// by this response's `id`) so pi-web's WS bridge — which the pi SDK's
|
|
3107
|
+
// usage parsing can't carry a custom field through — can correlate it
|
|
3108
|
+
// via AssistantMessage.responseId. batched wins over hasDraft (mutually
|
|
3109
|
+
// exclusive by construction: willBatch already excludes hasDraft), so
|
|
3110
|
+
// the only ambiguity is serial vs serial+spec, resolved once s.spec is
|
|
3111
|
+
// known post-generation.
|
|
3112
|
+
const lane: Lane = batched ? "batched" : shape.hasDraft ? "serial+spec" : "serial";
|
|
3113
|
+
recordLane(id, lane);
|
|
1980
3114
|
|
|
1981
3115
|
if (body.stream) {
|
|
1982
3116
|
const stream = new ReadableStream<Uint8Array>({
|
|
@@ -2054,6 +3188,11 @@ export function createServer(
|
|
|
2054
3188
|
? "tool_calls"
|
|
2055
3189
|
: stopper.stopped ? "stop"
|
|
2056
3190
|
: s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
|
|
3191
|
+
// Refine serial vs serial+spec now that s.spec is known (a mounted
|
|
3192
|
+
// draft can still decode zero speculative tokens on a very short
|
|
3193
|
+
// reply) and re-record for pi-web's WS correlation.
|
|
3194
|
+
const finalLane: Lane = batched ? "batched" : s.spec ? "serial+spec" : lane;
|
|
3195
|
+
if (finalLane !== lane) recordLane(id, finalLane);
|
|
2057
3196
|
send({
|
|
2058
3197
|
...chunk({}, finish),
|
|
2059
3198
|
usage: {
|
|
@@ -2061,6 +3200,8 @@ export function createServer(
|
|
|
2061
3200
|
completion_tokens: s.generatedTokens,
|
|
2062
3201
|
total_tokens: s.promptTokens + s.generatedTokens,
|
|
2063
3202
|
prompt_tokens_details: { cached_tokens: s.cachedTokens },
|
|
3203
|
+
lane: finalLane,
|
|
3204
|
+
...(s.spec ? { speculation: s.spec } : {}),
|
|
2064
3205
|
},
|
|
2065
3206
|
});
|
|
2066
3207
|
// bare sentinel per the OpenAI spec — JSON.stringify would
|
|
@@ -2078,6 +3219,7 @@ export function createServer(
|
|
|
2078
3219
|
"content-type": "text/event-stream",
|
|
2079
3220
|
"cache-control": "no-cache",
|
|
2080
3221
|
connection: "keep-alive",
|
|
3222
|
+
...(grammarWarning ? { Warning: grammarWarning } : {}),
|
|
2081
3223
|
},
|
|
2082
3224
|
});
|
|
2083
3225
|
}
|
|
@@ -2123,6 +3265,10 @@ export function createServer(
|
|
|
2123
3265
|
: stopper.stopped ? "stop"
|
|
2124
3266
|
: s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
|
|
2125
3267
|
const logprobsBlock = lpc?.payload() ?? null;
|
|
3268
|
+
// Refine serial vs serial+spec now that s.spec is known; re-record
|
|
3269
|
+
// for pi-web's WS correlation (see the streaming branch above).
|
|
3270
|
+
const finalLane: Lane = batched ? "batched" : s.spec ? "serial+spec" : lane;
|
|
3271
|
+
if (finalLane !== lane) recordLane(id, finalLane);
|
|
2126
3272
|
return Response.json({
|
|
2127
3273
|
id, object: "chat.completion", created, model: ctx.modelId,
|
|
2128
3274
|
choices: [{
|
|
@@ -2141,8 +3287,10 @@ export function createServer(
|
|
|
2141
3287
|
completion_tokens: s.generatedTokens,
|
|
2142
3288
|
total_tokens: s.promptTokens + s.generatedTokens,
|
|
2143
3289
|
prompt_tokens_details: { cached_tokens: s.cachedTokens },
|
|
3290
|
+
lane: finalLane,
|
|
3291
|
+
...(s.spec ? { speculation: s.spec } : {}),
|
|
2144
3292
|
},
|
|
2145
|
-
});
|
|
3293
|
+
}, grammarWarning ? { headers: { Warning: grammarWarning } } : undefined);
|
|
2146
3294
|
}
|
|
2147
3295
|
} catch (e) {
|
|
2148
3296
|
return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
|
|
@@ -2196,12 +3344,34 @@ export function createServer(
|
|
|
2196
3344
|
// bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
|
|
2197
3345
|
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
2198
3346
|
}
|
|
3347
|
+
// Grammar-constrained decoding on raw completions too (response_format /
|
|
3348
|
+
// guided_* / structured_outputs). The degrade path for /v1/completions
|
|
3349
|
+
// has no chat template to inject a system message into, so a degrade
|
|
3350
|
+
// only emits the Warning header (no prompt injection) — documented gap
|
|
3351
|
+
// vs the chat lane, which mirrors oMLX's text-completions behavior.
|
|
3352
|
+
let textGrammarWarning: string | null = null;
|
|
3353
|
+
const textGrammarReq =
|
|
3354
|
+
body.response_format != null || !!body.guided_grammar ||
|
|
3355
|
+
!!body.guided_regex || !!body.guided_choice?.length ||
|
|
3356
|
+
body.structured_outputs != null;
|
|
3357
|
+
if (textGrammarReq) {
|
|
3358
|
+
const g = await compileGrammarForRequest(body as unknown as ChatRequest);
|
|
3359
|
+
if (g.controller) options.grammar = g.controller;
|
|
3360
|
+
else if (g.degradeHint)
|
|
3361
|
+
textGrammarWarning = `grammar not enforced: ${g.degradeHint} — no prompt injection on /v1/completions`;
|
|
3362
|
+
}
|
|
3363
|
+
// Mirrors the chat lane: a real controller (not the degrade path)
|
|
3364
|
+
// shapes the request for per-row grammar batching.
|
|
3365
|
+
const textGrammarCtrl = options.grammar ?? null;
|
|
2199
3366
|
// mlx_lm.server's default max_tokens is 512 (its --max-tokens CLI
|
|
2200
3367
|
// default). The chat lane's very generous default is wrong for raw
|
|
2201
3368
|
// completion: with no template an EOS may never come.
|
|
2202
3369
|
options.maxTokens = body.max_completion_tokens ?? body.max_tokens ?? serverOptions.defaultMaxTokens ?? 512;
|
|
2203
3370
|
const requiredCtx = promptIds.length + options.maxTokens;
|
|
2204
|
-
if (requiredCtx > admission.maxSafeContext)
|
|
3371
|
+
if (requiredCtx > admission.maxSafeContext) {
|
|
3372
|
+
// Pre-run reject: the gateway only takes grammar disposal ownership
|
|
3373
|
+
// once run — free the WASM matcher here (chat lane: disposeRejected).
|
|
3374
|
+
textGrammarCtrl?.dispose();
|
|
2205
3375
|
return Response.json(
|
|
2206
3376
|
{
|
|
2207
3377
|
error: {
|
|
@@ -2216,10 +3386,12 @@ export function createServer(
|
|
|
2216
3386
|
},
|
|
2217
3387
|
{ status: 400 },
|
|
2218
3388
|
);
|
|
3389
|
+
}
|
|
2219
3390
|
try {
|
|
2220
3391
|
const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
|
|
2221
3392
|
if (adapterIds.length) options.adapters = adapterIds;
|
|
2222
3393
|
} catch (e) {
|
|
3394
|
+
textGrammarCtrl?.dispose();
|
|
2223
3395
|
return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
|
|
2224
3396
|
}
|
|
2225
3397
|
// logprobs capture — same mlx-lm block as chat (generate_response is
|
|
@@ -2245,6 +3417,11 @@ export function createServer(
|
|
|
2245
3417
|
wantsLogprobs: captureLogprobs,
|
|
2246
3418
|
userSeed: body.seed !== undefined,
|
|
2247
3419
|
kvQuant: !!(options.kvConfig?.length || options.kvBits),
|
|
3420
|
+
turboQuant: !!options.turboQuant,
|
|
3421
|
+
// /v1/completions grammar (textGrammarCtrl). Same null-on-degrade
|
|
3422
|
+
// contract as the chat lane.
|
|
3423
|
+
hasGrammar: !!textGrammarCtrl,
|
|
3424
|
+
hasDraft: !!ctx.draft,
|
|
2248
3425
|
};
|
|
2249
3426
|
const finishReason = (stopped: boolean, generated: number): "stop" | "length" =>
|
|
2250
3427
|
stopped ? "stop" : generated >= (options.maxTokens ?? 512) ? "length" : "stop";
|
|
@@ -2292,6 +3469,7 @@ export function createServer(
|
|
|
2292
3469
|
completion_tokens: s.generatedTokens,
|
|
2293
3470
|
total_tokens: s.promptTokens + s.generatedTokens,
|
|
2294
3471
|
prompt_tokens_details: { cached_tokens: s.cachedTokens },
|
|
3472
|
+
...(s.spec ? { speculation: s.spec } : {}),
|
|
2295
3473
|
},
|
|
2296
3474
|
});
|
|
2297
3475
|
controller.enqueue(enc.encode("data: [DONE]\n\n"));
|
|
@@ -2307,6 +3485,7 @@ export function createServer(
|
|
|
2307
3485
|
"content-type": "text/event-stream",
|
|
2308
3486
|
"cache-control": "no-cache",
|
|
2309
3487
|
connection: "keep-alive",
|
|
3488
|
+
...(textGrammarWarning ? { Warning: textGrammarWarning } : {}),
|
|
2310
3489
|
},
|
|
2311
3490
|
});
|
|
2312
3491
|
}
|
|
@@ -2341,8 +3520,9 @@ export function createServer(
|
|
|
2341
3520
|
completion_tokens: s.generatedTokens,
|
|
2342
3521
|
total_tokens: s.promptTokens + s.generatedTokens,
|
|
2343
3522
|
prompt_tokens_details: { cached_tokens: s.cachedTokens },
|
|
3523
|
+
...(s.spec ? { speculation: s.spec } : {}),
|
|
2344
3524
|
},
|
|
2345
|
-
});
|
|
3525
|
+
}, textGrammarWarning ? { headers: { Warning: textGrammarWarning } } : undefined);
|
|
2346
3526
|
} catch (e) {
|
|
2347
3527
|
return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
|
|
2348
3528
|
}
|
|
@@ -2708,6 +3888,27 @@ export function createServer(
|
|
|
2708
3888
|
saveHfToken(body.token);
|
|
2709
3889
|
return Response.json({ ok: true });
|
|
2710
3890
|
}
|
|
3891
|
+
|
|
3892
|
+
// --- Durable tool-approvals settings (plan §5.4/§6.5/§9 Phase 2) ---
|
|
3893
|
+
// The web chat's own "always allow this tool" list
|
|
3894
|
+
// (~/.mlx-bun/tool-approvals.json, src/tool-approvals.ts). The
|
|
3895
|
+
// approval card itself grants an always-allow via the `approval` WS
|
|
3896
|
+
// frame (alwaysAllow:true) — these REST routes are read/forget only,
|
|
3897
|
+
// for the settings panel's list view (matches the hf-token GET/POST
|
|
3898
|
+
// pattern's separation of "the card that grants" vs "settings that
|
|
3899
|
+
// manage" — there is no POST here on purpose, granting always goes
|
|
3900
|
+
// through the approval card, never a bare settings form).
|
|
3901
|
+
if (url.pathname === "/api/settings/tool-approvals" && request.method === "GET") {
|
|
3902
|
+
const { listAlwaysAllowedTools } = await import("./tool-approvals");
|
|
3903
|
+
return Response.json({ ok: true, alwaysAllow: listAlwaysAllowedTools() });
|
|
3904
|
+
}
|
|
3905
|
+
if (url.pathname === "/api/settings/tool-approvals" && request.method === "DELETE") {
|
|
3906
|
+
const body = (await request.json().catch(() => ({}))) as { tool?: string };
|
|
3907
|
+
if (!body.tool) return Response.json({ ok: false, error: "tool required" }, { status: 400 });
|
|
3908
|
+
const { revokeToolAlwaysAllowed } = await import("./tool-approvals");
|
|
3909
|
+
const file = revokeToolAlwaysAllowed(body.tool);
|
|
3910
|
+
return Response.json({ ok: true, alwaysAllow: Object.keys(file.allows).sort() });
|
|
3911
|
+
}
|
|
2711
3912
|
{
|
|
2712
3913
|
const m = url.pathname.match(/^\/api\/(quantize|finetune|dataset)\/push$/);
|
|
2713
3914
|
if (m && request.method === "POST") {
|
|
@@ -2736,6 +3937,50 @@ export function createServer(
|
|
|
2736
3937
|
}
|
|
2737
3938
|
}
|
|
2738
3939
|
|
|
3940
|
+
// gc plan/execute (web-ui-pass-plan.md #9): thin wrappers over the CLI's
|
|
3941
|
+
// own `mlx-bun gc` planner/executor (src/registry.ts planGc/executeGc).
|
|
3942
|
+
// Loopback-served admin route — no auth beyond the server's own bind
|
|
3943
|
+
// (matches every other /api/* route here). GET plans (read-only, cheap:
|
|
3944
|
+
// config.json + safetensors index reads, no tensor bytes); POST executes
|
|
3945
|
+
// and requires an explicit {yes:true} body (mirrors the CLI's --yes gate)
|
|
3946
|
+
// so a stray call can't delete anything by accident.
|
|
3947
|
+
if (url.pathname === "/api/gc/plan" && request.method === "GET") {
|
|
3948
|
+
const { planGc } = await import("./registry");
|
|
3949
|
+
const plans = planGc().filter(
|
|
3950
|
+
(p) => p.pruneSnapshots.length || p.skippedSnapshots.length || p.deadBlobs.length,
|
|
3951
|
+
);
|
|
3952
|
+
const superseded = plans.map((p) => ({
|
|
3953
|
+
repo_id: p.repoId,
|
|
3954
|
+
prune_snapshots: p.pruneSnapshots.length,
|
|
3955
|
+
skipped_snapshots: p.skippedSnapshots.length,
|
|
3956
|
+
dead_blobs: p.deadBlobs.length,
|
|
3957
|
+
reclaim_bytes: p.reclaimBytes,
|
|
3958
|
+
}));
|
|
3959
|
+
const reclaim_bytes = plans.reduce((a, p) => a + p.reclaimBytes, 0);
|
|
3960
|
+
return Response.json({ ok: true, superseded, reclaim_bytes });
|
|
3961
|
+
}
|
|
3962
|
+
if (url.pathname === "/api/gc/execute" && request.method === "POST") {
|
|
3963
|
+
const body = (await request.json().catch(() => ({}))) as { yes?: boolean };
|
|
3964
|
+
if (body.yes !== true)
|
|
3965
|
+
return Response.json({ ok: false, error: "pass {\"yes\": true} to confirm deletion" }, { status: 400 });
|
|
3966
|
+
const { planGc, executeGc, Registry } = await import("./registry");
|
|
3967
|
+
const plans = planGc().filter(
|
|
3968
|
+
(p) => p.pruneSnapshots.length || p.skippedSnapshots.length || p.deadBlobs.length,
|
|
3969
|
+
);
|
|
3970
|
+
const res = executeGc(plans);
|
|
3971
|
+
const reg = new Registry();
|
|
3972
|
+
try {
|
|
3973
|
+
await reg.scan(); // reap deleted snapshots from the registry
|
|
3974
|
+
} finally {
|
|
3975
|
+
reg.close();
|
|
3976
|
+
}
|
|
3977
|
+
libraryCache = null; // force /library to reflect the deletion
|
|
3978
|
+
return Response.json({
|
|
3979
|
+
ok: true,
|
|
3980
|
+
snapshots: res.snapshots, blobs: res.blobs, reclaimed_bytes: res.reclaimedBytes,
|
|
3981
|
+
});
|
|
3982
|
+
}
|
|
3983
|
+
|
|
2739
3984
|
if (url.pathname === "/api/jobs" && request.method === "GET") {
|
|
2740
3985
|
const store = await ensureJobs();
|
|
2741
3986
|
const limit = Number(url.searchParams.get("limit") ?? "50");
|