mlx-bun 0.4.0 → 0.5.0
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 +57 -32
- package/docs/README.md +55 -42
- package/docs/archive/investigations/benchmark-ledger-through-v0.4.0.md +3617 -0
- package/docs/archive/planning/release-notes-v0.5.0.md +37 -0
- package/docs/design/batching.md +220 -7
- package/docs/design/decode-speed-program.md +121 -66
- package/docs/design/generic-model-support.md +116 -0
- package/docs/design/kv-cache.md +380 -211
- package/docs/design/speculative-decoding.md +165 -26
- package/docs/design/turboquant.md +20 -7
- package/docs/design/unified-engine-frontier-plan.md +120 -13
- package/docs/planning/PRODUCT_ROADMAP.md +5 -6
- package/docs/planning/public-readiness-handoff.md +244 -0
- package/docs/planning/public-readiness.md +149 -0
- package/docs/reference/benchmarks.md +3463 -3489
- package/docs/reference/cli.md +62 -1
- package/docs/reference/distribution.md +29 -4
- package/docs/reference/environment.md +104 -8
- package/docs/reference/library-api.md +18 -4
- package/docs/reference/memory.md +1 -1
- package/docs/reference/models.md +17 -5
- package/docs/reference/server-api.md +181 -11
- package/docs/reference/server-config.md +232 -105
- package/package.json +2 -2
- package/src/audio/audiotoolbox.ts +112 -0
- package/src/audio/conformer.ts +8 -0
- package/src/audio/mic-capture.ts +98 -0
- package/src/audio/silero-vad.ts +379 -0
- package/src/audio/whisper-decode.ts +774 -0
- package/src/audio/whisper-fast.ts +732 -0
- package/src/audio/whisper-format.ts +73 -0
- package/src/audio/whisper-mel.ts +185 -0
- package/src/audio/whisper-timing.ts +274 -0
- package/src/audio/whisper-tokenizer.ts +239 -0
- package/src/audio/whisper-transcribe.ts +410 -0
- package/src/backends/mlx/autoregressive.ts +22 -1
- package/src/backends/mlx/batch-group.ts +111 -27
- package/src/backends/mlx/cache-layout.ts +3 -3
- package/src/backends/mlx/continuation-persistence.ts +5 -4
- package/src/backends/mlx/encoder-cache.ts +25 -0
- package/src/backends/mlx/fill-append.ts +21 -8
- package/src/backends/mlx/fill-group.ts +497 -0
- package/src/backends/mlx/gateway-binding.ts +56 -12
- package/src/backends/mlx/kv-maintenance.ts +30 -2
- package/src/backends/mlx/mixed-iteration.ts +62 -0
- package/src/backends/mlx/model-prompt.ts +46 -32
- package/src/backends/mlx/model-serving.ts +4 -3
- package/src/backends/mlx/paged-cache-codec.ts +27 -0
- package/src/backends/mlx/prefill-cohort.ts +15 -4
- package/src/backends/mlx/prefill-policy.ts +26 -0
- package/src/backends/mlx/prefill-rows.ts +57 -13
- package/src/backends/mlx/prompt-input.ts +29 -0
- package/src/backends/mlx/qwen-prompt-input.ts +40 -0
- package/src/backends/mlx/request-state-policy.ts +18 -5
- package/src/backends/mlx/rollback.ts +6 -3
- package/src/backends/mlx/serial-executor.ts +16 -7
- package/src/backends/mlx/speculative-group.ts +105 -26
- package/src/backends/mlx/speculative.ts +24 -7
- package/src/backends/mlx/target-layout-capability.ts +1 -1
- package/src/backends/mlx/transitioning-kv-rows.ts +8 -0
- package/src/cli.ts +322 -0
- package/src/contracts/execution.ts +10 -0
- package/src/contracts/object-cache.ts +8 -0
- package/src/contracts/pi-provider.ts +1 -1
- package/src/contracts/pi.ts +3 -1
- package/src/contracts/prefix-cache.ts +13 -3
- package/src/contracts/scheduling.ts +9 -0
- package/src/engine/execution-plan.ts +23 -12
- package/src/engine/execution-tasks.ts +44 -0
- package/src/engine/namespaced-cache.ts +13 -0
- package/src/engine/scheduler.ts +22 -2
- package/src/fill/fill-session.ts +14 -3
- package/src/generate.ts +28 -22
- package/src/grammar.ts +13 -0
- package/src/inference/prefill.ts +6 -0
- package/src/kv-scheme.ts +22 -0
- package/src/kv-store.ts +115 -79
- package/src/lab/README.md +5 -1
- package/src/lab/paged-kv/paged-attention.ts +117 -0
- package/src/lab/paged-kv/paged-kv-rows.ts +33 -5
- package/src/lab/paged-kv/paged-kv.ts +108 -42
- package/src/md-modules.d.ts +5 -0
- package/src/mlx/array.ts +37 -0
- package/src/mlx/ffi.ts +19 -0
- package/src/mlx/normalized-argmax.ts +116 -0
- package/src/mlx/ops.ts +97 -10
- package/src/model/batched-turboquant-kv.ts +6 -5
- package/src/model/delayed-quantized-kv.ts +3 -3
- package/src/model/delayed-rotating-quantized-kv.ts +1 -1
- package/src/model/delayed-turboquant-kv.ts +6 -4
- package/src/model/factory.ts +26 -0
- package/src/model/fingerprint.ts +12 -0
- package/src/model/full-prefill-row.ts +1 -1
- package/src/model/gemma4-base.ts +25 -6
- package/src/model/gemma4.ts +82 -2
- package/src/model/kv-attention-view.ts +4 -1
- package/src/model/profile.ts +29 -4
- package/src/model/quantized-append-attention.ts +31 -0
- package/src/model/qwen3-mrope.ts +7 -9
- package/src/model/qwen3_5.ts +99 -37
- package/src/model/support.ts +8 -0
- package/src/model/token-groups.ts +66 -0
- package/src/model/trellis-linear.ts +2 -2
- package/src/model/turboquant-codec.ts +6 -0
- package/src/model/whisper.ts +435 -0
- package/src/native/mic_capture.swift +92 -0
- package/src/native-pack.ts +6 -5
- package/src/pi-web.ts +9 -1
- package/src/prompt-cache.ts +265 -43
- package/src/runtime-config.ts +1 -5
- package/src/safetensors.ts +9 -2
- package/src/sampler.ts +5 -4
- package/src/serve/audio-routes.ts +323 -0
- package/src/serve/chat-request.ts +2 -0
- package/src/serve/chat-stage.ts +17 -12
- package/src/serve/completion-executor.ts +9 -1
- package/src/serve/discovery-routes.ts +17 -1
- package/src/serve/generated-token-history.ts +57 -0
- package/src/serve/generation-gateway.ts +52 -11
- package/src/serve/model-binding.ts +7 -1
- package/src/serve/preparation.ts +3 -3
- package/src/serve/prompt-response-trace.ts +14 -0
- package/src/serve/request-plan.ts +2 -0
- package/src/serve/request-prep.ts +38 -21
- package/src/serve/transcription-server.ts +64 -0
- package/src/serve/transcription-service.ts +502 -0
- package/src/server.ts +132 -164
- package/src/spec/assistant-rows.ts +1 -0
- package/src/spec/deepspec-rows.ts +1 -1
- package/src/spec/dflash-source.ts +1 -1
- package/src/spec/glm52-mtp-source.ts +1 -0
- package/src/spec/ngram-source.ts +27 -5
- package/src/spec/projected-draft-rows.ts +2 -2
- package/src/spec/qwen-mtp-rows.ts +11 -0
- package/src/spec/qwen-mtp-source.ts +3 -0
- package/src/spec/serve-loop.ts +5 -3
- package/src/spec/source.ts +14 -1
- package/src/spec/two-model-rows.ts +1 -0
- package/src/spec/two-model.ts +2 -2
- package/src/ssd-cache.ts +76 -17
- package/src/ssd-durability.ts +20 -22
- package/src/storage/host-buffer.ts +24 -0
- package/src/storage/kv-writer.ts +83 -0
- package/src/storage/kv-writer.worker.js +264 -0
- package/src/storage/retention-policy.ts +37 -0
- package/src/tiered-prompt-cache.ts +45 -0
- package/src/vision/embedder.ts +8 -0
- package/src/vision/prompt.ts +101 -65
- package/src/vision/qwen3vl-prompt.ts +101 -64
- package/src/vision/qwen3vl-tower.ts +11 -0
- package/src/vision/siglip.ts +8 -0
- package/src/web/app.html +12 -0
- package/src/web/app.js +160 -0
- package/src/web/src/chat.ts +4 -1
- package/src/web/src/composer.ts +3 -0
- package/src/web/src/voice.ts +141 -0
- package/src/web/sw.js +1 -1
- /package/docs/{planning → archive/planning}/release-notes-v0.4.0.md +0 -0
|
@@ -0,0 +1,323 @@
|
|
|
1
|
+
// OpenAI-compatible speech-to-text routes:
|
|
2
|
+
// POST /v1/audio/transcriptions (task transcribe)
|
|
3
|
+
// POST /v1/audio/translations (task translate → English)
|
|
4
|
+
// POST /admin/transcription/unload (explicit page-out of the Whisper weights)
|
|
5
|
+
// POST /v1/audio/sessions · POST /v1/audio/sessions/:id/audio ·
|
|
6
|
+
// POST /v1/audio/sessions/:id/finish · DELETE /v1/audio/sessions/:id
|
|
7
|
+
// (streaming dictation: chunks arrive during capture, each finished 30 s
|
|
8
|
+
// window transcribes immediately, finish returns the whole transcript)
|
|
9
|
+
// Bodies: multipart/form-data (OpenAI: file, model, language, prompt,
|
|
10
|
+
// response_format, temperature, stream, timestamp_granularities[]) or JSON
|
|
11
|
+
// ({ file: <base64 | data: URL>, ...same fields }). Non-standard fields:
|
|
12
|
+
// beam_size, condition_on_previous_text, no_speech_threshold, vocabulary
|
|
13
|
+
// (array of hint terms fitted into the prompt budget; usage reported),
|
|
14
|
+
// faithful (oracle graph instead of the fast path), audio_ctx (Lab), vad /
|
|
15
|
+
// vad_threshold / vad_min_speech_ms / vad_trim (Silero gate: silence returns
|
|
16
|
+
// an empty transcript without running Whisper).
|
|
17
|
+
// Streaming (`stream=true`): SSE `transcript.text.delta` per segment,
|
|
18
|
+
// then `transcript.text.done` carrying the final body.
|
|
19
|
+
|
|
20
|
+
import { formatTranscription, TRANSCRIPTION_FORMATS, toVerboseJson, type TranscriptionFormat } from "../audio/whisper-format";
|
|
21
|
+
import type { WhisperSegment } from "../audio/whisper-transcribe";
|
|
22
|
+
import { TranscriptionError, TranscriptionService, type TranscriptionOutcome, type TranscriptionParams } from "./transcription-service";
|
|
23
|
+
|
|
24
|
+
export type AudioRoute =
|
|
25
|
+
| { kind: "transcriptions"; task: "transcribe" | "translate" }
|
|
26
|
+
| { kind: "unload" }
|
|
27
|
+
| { kind: "session-create" }
|
|
28
|
+
| { kind: "session-audio"; id: string }
|
|
29
|
+
| { kind: "session-finish"; id: string }
|
|
30
|
+
| { kind: "session-delete"; id: string };
|
|
31
|
+
|
|
32
|
+
export function matchAudioRoute(method: string, pathname: string): AudioRoute | null {
|
|
33
|
+
switch (`${method} ${pathname}`) {
|
|
34
|
+
case "POST /v1/audio/transcriptions": return { kind: "transcriptions", task: "transcribe" };
|
|
35
|
+
case "POST /v1/audio/translations": return { kind: "transcriptions", task: "translate" };
|
|
36
|
+
case "POST /admin/transcription/unload": return { kind: "unload" };
|
|
37
|
+
case "POST /v1/audio/sessions": return { kind: "session-create" };
|
|
38
|
+
}
|
|
39
|
+
if (pathname.startsWith("/v1/audio/sessions/")) {
|
|
40
|
+
const rest = pathname.slice("/v1/audio/sessions/".length).split("/");
|
|
41
|
+
const id = decodeURIComponent(rest[0] ?? "");
|
|
42
|
+
if (!id) return null;
|
|
43
|
+
if (method === "POST" && rest[1] === "audio" && rest.length === 2) return { kind: "session-audio", id };
|
|
44
|
+
if (method === "POST" && rest[1] === "finish" && rest.length === 2) return { kind: "session-finish", id };
|
|
45
|
+
if (method === "DELETE" && rest.length === 1) return { kind: "session-delete", id };
|
|
46
|
+
}
|
|
47
|
+
return null;
|
|
48
|
+
}
|
|
49
|
+
|
|
50
|
+
/** Everything the route needs from the host: a service (or a reason there is none). */
|
|
51
|
+
export interface AudioRouteHost {
|
|
52
|
+
service(): Promise<TranscriptionService | null>;
|
|
53
|
+
/** Message for the 503 when no Whisper model is available. */
|
|
54
|
+
unavailableMessage?: string;
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
export interface ParsedAudioRequest {
|
|
58
|
+
audio: Uint8Array;
|
|
59
|
+
format: TranscriptionFormat;
|
|
60
|
+
stream: boolean;
|
|
61
|
+
params: TranscriptionParams;
|
|
62
|
+
}
|
|
63
|
+
|
|
64
|
+
const sessionFormats = new Map<string, TranscriptionFormat>();
|
|
65
|
+
|
|
66
|
+
const segmentJson = (s: WhisperSegment) => ({ id: s.id, seek: s.seek, start: s.start, end: s.end, text: s.text, tokens: s.tokens, temperature: s.temperature, avg_logprob: s.avgLogprob, compression_ratio: s.compressionRatio, no_speech_prob: s.noSpeechProb, ...(s.words ? { words: s.words } : {}) });
|
|
67
|
+
|
|
68
|
+
function errorJson(message: string, status: number, type = "invalid_request_error"): Response {
|
|
69
|
+
return Response.json({ error: { message, type } }, { status });
|
|
70
|
+
}
|
|
71
|
+
|
|
72
|
+
const B64_MAX = 200 * 1024 * 1024;
|
|
73
|
+
|
|
74
|
+
function decodeBase64Audio(v: unknown): Uint8Array | null {
|
|
75
|
+
if (typeof v !== "string" || v.length === 0 || v.length > B64_MAX * 1.4) return null;
|
|
76
|
+
const data = v.startsWith("data:") ? v.slice(v.indexOf(",") + 1) : v;
|
|
77
|
+
try {
|
|
78
|
+
return new Uint8Array(Buffer.from(data, "base64"));
|
|
79
|
+
} catch {
|
|
80
|
+
return null;
|
|
81
|
+
}
|
|
82
|
+
}
|
|
83
|
+
|
|
84
|
+
function num(v: unknown): number | null {
|
|
85
|
+
if (v === null || v === undefined || v === "") return null;
|
|
86
|
+
const n = typeof v === "number" ? v : Number(v);
|
|
87
|
+
return Number.isFinite(n) ? n : null;
|
|
88
|
+
}
|
|
89
|
+
|
|
90
|
+
function bool(v: unknown): boolean | null {
|
|
91
|
+
if (typeof v === "boolean") return v;
|
|
92
|
+
if (v === "true" || v === "1") return true;
|
|
93
|
+
if (v === "false" || v === "0") return false;
|
|
94
|
+
return null;
|
|
95
|
+
}
|
|
96
|
+
|
|
97
|
+
/** Decoding params shared by one-shot requests and sessions (no audio). */
|
|
98
|
+
function parseParams(fields: Record<string, unknown>, task: "transcribe" | "translate"): TranscriptionParams | Response {
|
|
99
|
+
const language = typeof fields.language === "string" && fields.language.trim() ? fields.language.trim() : null;
|
|
100
|
+
const temperature = num(fields.temperature);
|
|
101
|
+
if (temperature !== null && (temperature < 0 || temperature > 1)) return errorJson("temperature must be in [0, 1]", 400);
|
|
102
|
+
const beamSize = num(fields.beam_size);
|
|
103
|
+
if (beamSize !== null && (!Number.isInteger(beamSize) || beamSize < 1 || beamSize > 16))
|
|
104
|
+
return errorJson("beam_size must be an integer in [1, 16]", 400);
|
|
105
|
+
if (temperature !== null && temperature > 0 && beamSize !== null && beamSize > 1)
|
|
106
|
+
return errorJson("beam_size requires temperature 0", 400);
|
|
107
|
+
const vocabulary = Array.isArray(fields.vocabulary)
|
|
108
|
+
? (fields.vocabulary as unknown[]).map(String)
|
|
109
|
+
: typeof fields.vocabulary === "string" && fields.vocabulary.trim()
|
|
110
|
+
? fields.vocabulary.split(",").map((s) => s.trim()).filter(Boolean)
|
|
111
|
+
: undefined;
|
|
112
|
+
if (vocabulary && (vocabulary.length > 8192 || vocabulary.some((t) => t.length > 16384)))
|
|
113
|
+
return errorJson("vocabulary: at most 8192 terms of at most 16 KiB each", 400);
|
|
114
|
+
const noSpeech = num(fields.no_speech_threshold);
|
|
115
|
+
const vadOn = bool(fields.vad) ?? false;
|
|
116
|
+
const vadThreshold = num(fields.vad_threshold);
|
|
117
|
+
if (vadThreshold !== null && (vadThreshold <= 0 || vadThreshold >= 1)) return errorJson("vad_threshold must be in (0, 1)", 400);
|
|
118
|
+
const vadMinSpeech = num(fields.vad_min_speech_ms);
|
|
119
|
+
const audioCtx = num(fields.audio_ctx);
|
|
120
|
+
if (audioCtx !== null && (!Number.isInteger(audioCtx) || audioCtx < 64 || audioCtx > 1500))
|
|
121
|
+
return errorJson("audio_ctx must be an integer in [64, 1500]", 400);
|
|
122
|
+
const granularities = Array.isArray(fields.timestamp_granularities) ? (fields.timestamp_granularities as string[]) : [];
|
|
123
|
+
return {
|
|
124
|
+
task, language,
|
|
125
|
+
prompt: typeof fields.prompt === "string" && fields.prompt.trim() ? fields.prompt : null,
|
|
126
|
+
vocabulary, temperature, beamSize,
|
|
127
|
+
conditionOnPreviousText: bool(fields.condition_on_previous_text) ?? true,
|
|
128
|
+
...(noSpeech !== null ? { noSpeechThreshold: noSpeech } : {}),
|
|
129
|
+
withoutTimestamps: bool(fields.without_timestamps) ?? false,
|
|
130
|
+
vad: vadOn ? { threshold: vadThreshold ?? undefined, minSpeechMs: vadMinSpeech ?? undefined, trim: bool(fields.vad_trim) ?? false } : null,
|
|
131
|
+
wordTimestamps: granularities.includes("word"),
|
|
132
|
+
faithful: bool(fields.faithful) ?? false,
|
|
133
|
+
audioCtx,
|
|
134
|
+
};
|
|
135
|
+
}
|
|
136
|
+
|
|
137
|
+
/** Parse a multipart or JSON transcription request into audio bytes + params. */
|
|
138
|
+
export async function parseAudioRequest(request: Request, task: "transcribe" | "translate"): Promise<ParsedAudioRequest | Response> {
|
|
139
|
+
const ct = request.headers.get("content-type") ?? "";
|
|
140
|
+
let fields: Record<string, unknown> = {};
|
|
141
|
+
let audio: Uint8Array | null = null;
|
|
142
|
+
if (ct.includes("multipart/form-data")) {
|
|
143
|
+
let form: FormData;
|
|
144
|
+
try { form = await request.formData(); }
|
|
145
|
+
catch (e) { return errorJson(`invalid multipart body: ${(e as Error).message}`, 400); }
|
|
146
|
+
const file = form.get("file");
|
|
147
|
+
if (!(file instanceof Blob)) return errorJson("`file` (audio upload) is required", 400);
|
|
148
|
+
audio = new Uint8Array(await file.arrayBuffer());
|
|
149
|
+
for (const [k, v] of form.entries()) {
|
|
150
|
+
if (k === "file") continue;
|
|
151
|
+
if (k.endsWith("[]")) {
|
|
152
|
+
const key = k.slice(0, -2);
|
|
153
|
+
(fields[key] ??= [] as string[]);
|
|
154
|
+
(fields[key] as string[]).push(String(v));
|
|
155
|
+
} else fields[k] = typeof v === "string" ? v : String(v);
|
|
156
|
+
}
|
|
157
|
+
} else if (ct.includes("application/json")) {
|
|
158
|
+
try { fields = (await request.json()) as Record<string, unknown>; }
|
|
159
|
+
catch { return errorJson("invalid JSON body", 400); }
|
|
160
|
+
const raw = fields.file ?? (fields.input_audio as { data?: unknown } | undefined)?.data ?? fields.audio;
|
|
161
|
+
audio = decodeBase64Audio(raw);
|
|
162
|
+
if (!audio) return errorJson("`file` must be base64 audio (optionally a data: URL)", 400);
|
|
163
|
+
} else {
|
|
164
|
+
return errorJson("content-type must be multipart/form-data or application/json", 415);
|
|
165
|
+
}
|
|
166
|
+
if (audio.length === 0) return errorJson("empty audio upload", 400);
|
|
167
|
+
const format = (fields.response_format as string | undefined) ?? "json";
|
|
168
|
+
if (!TRANSCRIPTION_FORMATS.includes(format as TranscriptionFormat))
|
|
169
|
+
return errorJson(`response_format must be one of ${TRANSCRIPTION_FORMATS.join(", ")}`, 400);
|
|
170
|
+
const params = parseParams(fields, task);
|
|
171
|
+
if (params instanceof Response) return params;
|
|
172
|
+
const stream = bool(fields.stream) ?? false;
|
|
173
|
+
if (params.wordTimestamps && format !== "verbose_json" && !stream)
|
|
174
|
+
return errorJson("timestamp_granularities[]=word requires response_format=verbose_json", 400);
|
|
175
|
+
return { audio, format: format as TranscriptionFormat, stream, params };
|
|
176
|
+
}
|
|
177
|
+
|
|
178
|
+
/** Body for a finished transcription in the requested format. */
|
|
179
|
+
export function transcriptionResponse(outcome: TranscriptionOutcome, format: TranscriptionFormat, task: string): Response {
|
|
180
|
+
const extra = {
|
|
181
|
+
model: outcome.modelId,
|
|
182
|
+
language: outcome.result.language,
|
|
183
|
+
duration: outcome.durationSeconds,
|
|
184
|
+
timings: outcome.timings,
|
|
185
|
+
...(outcome.vocabulary ? { vocabulary: outcome.vocabulary } : {}),
|
|
186
|
+
...(outcome.vad ? { vad: outcome.vad } : {}),
|
|
187
|
+
};
|
|
188
|
+
const headers: Record<string, string> = {
|
|
189
|
+
"x-mlx-bun-language": outcome.result.language,
|
|
190
|
+
"x-mlx-bun-timings": JSON.stringify(outcome.timings),
|
|
191
|
+
};
|
|
192
|
+
if (format === "json")
|
|
193
|
+
return Response.json({ text: outcome.result.text.trim(), mlx_bun: extra }, { headers });
|
|
194
|
+
if (format === "verbose_json")
|
|
195
|
+
return Response.json({ ...toVerboseJson(outcome.result, outcome.durationSeconds, task), mlx_bun: extra }, { headers });
|
|
196
|
+
const body = formatTranscription(outcome.result, format, outcome.durationSeconds, task);
|
|
197
|
+
const type = format === "srt" ? "application/x-subrip" : format === "vtt" ? "text/vtt" : "text/plain";
|
|
198
|
+
return new Response(body, { headers: { ...headers, "content-type": `${type}; charset=utf-8` } });
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
export async function handleAudioRoute(url: URL, request: Request, host: AudioRouteHost): Promise<Response | null> {
|
|
202
|
+
const route = matchAudioRoute(request.method, url.pathname);
|
|
203
|
+
if (!route) return null;
|
|
204
|
+
const service = await host.service();
|
|
205
|
+
if (!service)
|
|
206
|
+
return errorJson(
|
|
207
|
+
host.unavailableMessage ??
|
|
208
|
+
"no Whisper model is available — download one (mlx-bun get mlx-community/whisper-large-v3-turbo) or pass --whisper-model",
|
|
209
|
+
503, "model_unavailable",
|
|
210
|
+
);
|
|
211
|
+
if (route.kind === "unload") {
|
|
212
|
+
const unloaded = service.unload();
|
|
213
|
+
return Response.json({ unloaded, ...service.stats });
|
|
214
|
+
}
|
|
215
|
+
if (route.kind === "session-create") {
|
|
216
|
+
let fields: Record<string, unknown> = {};
|
|
217
|
+
if ((request.headers.get("content-type") ?? "").includes("application/json")) {
|
|
218
|
+
try { fields = (await request.json()) as Record<string, unknown>; }
|
|
219
|
+
catch { return errorJson("invalid JSON body", 400); }
|
|
220
|
+
}
|
|
221
|
+
const params = parseParams(fields, fields.task === "translate" ? "translate" : "transcribe");
|
|
222
|
+
if (params instanceof Response) return params;
|
|
223
|
+
const format = (fields.response_format as string | undefined) ?? "json";
|
|
224
|
+
if (!["json", "verbose_json"].includes(format)) return errorJson("sessions return json or verbose_json", 400);
|
|
225
|
+
try {
|
|
226
|
+
const session = await service.createSession(params);
|
|
227
|
+
sessionFormats.set(session.id, format as TranscriptionFormat);
|
|
228
|
+
return Response.json({ id: session.id, model: service.modelId, vocabulary: session.vocabulary ?? null });
|
|
229
|
+
} catch (e) {
|
|
230
|
+
return transcriptionFailure(e);
|
|
231
|
+
}
|
|
232
|
+
}
|
|
233
|
+
if (route.kind === "session-audio" || route.kind === "session-finish" || route.kind === "session-delete") {
|
|
234
|
+
const session = service.session(route.id);
|
|
235
|
+
if (!session) return errorJson("unknown transcription session", 404);
|
|
236
|
+
if (route.kind === "session-delete") {
|
|
237
|
+
session.close();
|
|
238
|
+
sessionFormats.delete(route.id);
|
|
239
|
+
return new Response(null, { status: 204 });
|
|
240
|
+
}
|
|
241
|
+
if (route.kind === "session-audio") {
|
|
242
|
+
const ct = request.headers.get("content-type") ?? "";
|
|
243
|
+
const bytes = new Uint8Array(await request.arrayBuffer());
|
|
244
|
+
if (bytes.length === 0) return errorJson("empty audio chunk", 400);
|
|
245
|
+
let pcm: Float32Array;
|
|
246
|
+
try {
|
|
247
|
+
if (ct.startsWith("audio/pcm") || ct.startsWith("audio/l32") || ct.includes("float32")) {
|
|
248
|
+
if (bytes.length % 4 !== 0) return errorJson("float32 PCM chunk length must be a multiple of 4", 400);
|
|
249
|
+
pcm = new Float32Array(bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength));
|
|
250
|
+
const rate = Number(/rate=(\d+)/.exec(ct)?.[1] ?? 16000);
|
|
251
|
+
if (rate !== 16000) return errorJson("float32 PCM chunks must be 16 kHz mono", 400);
|
|
252
|
+
} else pcm = await (service.constructor as typeof TranscriptionService).decodeAudio(bytes);
|
|
253
|
+
} catch (e) {
|
|
254
|
+
return transcriptionFailure(e);
|
|
255
|
+
}
|
|
256
|
+
try {
|
|
257
|
+
const r = await session.append(pcm);
|
|
258
|
+
return Response.json({ id: session.id, samples: session.samples, duration: session.durationSeconds, speech: r.speech, segments: r.segments.map(segmentJson) });
|
|
259
|
+
} catch (e) {
|
|
260
|
+
return transcriptionFailure(e);
|
|
261
|
+
}
|
|
262
|
+
}
|
|
263
|
+
try {
|
|
264
|
+
const outcome = await session.finish();
|
|
265
|
+
sessionFormats.delete(route.id);
|
|
266
|
+
return transcriptionResponse(outcome, sessionFormats.get(route.id) ?? "json", session.params.task ?? "transcribe");
|
|
267
|
+
} catch (e) {
|
|
268
|
+
sessionFormats.delete(route.id);
|
|
269
|
+
return transcriptionFailure(e);
|
|
270
|
+
}
|
|
271
|
+
}
|
|
272
|
+
const parsed = await parseAudioRequest(request, route.task);
|
|
273
|
+
if (parsed instanceof Response) return parsed;
|
|
274
|
+
const taskName = route.task;
|
|
275
|
+
if (!parsed.stream) {
|
|
276
|
+
try {
|
|
277
|
+
const outcome = await service.transcribe(parsed.audio, { ...parsed.params, signal: request.signal });
|
|
278
|
+
return transcriptionResponse(outcome, parsed.format, taskName);
|
|
279
|
+
} catch (e) {
|
|
280
|
+
return transcriptionFailure(e);
|
|
281
|
+
}
|
|
282
|
+
}
|
|
283
|
+
// Streaming: one SSE event per finalized segment, then the final body.
|
|
284
|
+
const encoder = new TextEncoder();
|
|
285
|
+
const stream = new ReadableStream<Uint8Array>({
|
|
286
|
+
start: async (controller) => {
|
|
287
|
+
const send = (event: string, data: unknown) =>
|
|
288
|
+
controller.enqueue(encoder.encode(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`));
|
|
289
|
+
try {
|
|
290
|
+
const outcome = await service.transcribe(parsed.audio, {
|
|
291
|
+
...parsed.params,
|
|
292
|
+
signal: request.signal,
|
|
293
|
+
onSegment: (s: WhisperSegment) => {
|
|
294
|
+
if (s.text) send("transcript.text.delta", { delta: s.text, segment: { id: s.id, start: s.start, end: s.end } });
|
|
295
|
+
},
|
|
296
|
+
onProgress: (done, total) => send("transcript.progress", { done, total }),
|
|
297
|
+
});
|
|
298
|
+
const body = parsed.format === "verbose_json"
|
|
299
|
+
? toVerboseJson(outcome.result, outcome.durationSeconds, taskName)
|
|
300
|
+
: { text: outcome.result.text.trim() };
|
|
301
|
+
send("transcript.text.done", {
|
|
302
|
+
...body,
|
|
303
|
+
mlx_bun: { model: outcome.modelId, language: outcome.result.language, duration: outcome.durationSeconds, timings: outcome.timings, ...(outcome.vocabulary ? { vocabulary: outcome.vocabulary } : {}) },
|
|
304
|
+
});
|
|
305
|
+
} catch (e) {
|
|
306
|
+
send("error", { message: (e as Error).message });
|
|
307
|
+
} finally {
|
|
308
|
+
controller.close();
|
|
309
|
+
}
|
|
310
|
+
},
|
|
311
|
+
});
|
|
312
|
+
return new Response(stream, {
|
|
313
|
+
headers: { "content-type": "text/event-stream", "cache-control": "no-cache", connection: "keep-alive" },
|
|
314
|
+
});
|
|
315
|
+
}
|
|
316
|
+
|
|
317
|
+
function transcriptionFailure(e: unknown): Response {
|
|
318
|
+
if (e instanceof TranscriptionError) return errorJson(e.message, e.status);
|
|
319
|
+
const message = (e as Error).message ?? String(e);
|
|
320
|
+
if (/cancelled/.test(message)) return errorJson("request cancelled", 499);
|
|
321
|
+
if (/WAV|afconvert|transcode|audio|RIFF/i.test(message)) return errorJson(`undecodable audio: ${message}`, 400);
|
|
322
|
+
return errorJson(message, 500, "server_error");
|
|
323
|
+
}
|
package/src/serve/chat-stage.ts
CHANGED
|
@@ -12,7 +12,7 @@ import {
|
|
|
12
12
|
} from "./chat-request";
|
|
13
13
|
import type { InferenceRequest } from "./inference-request";
|
|
14
14
|
import type { ServingContext } from "./model-host";
|
|
15
|
-
import type { BuiltPrompt, ModelPromptBuilder } from "./model-binding";
|
|
15
|
+
import type { BuiltPrompt, ModelPromptBuilder, PromptNativeWork } from "./model-binding";
|
|
16
16
|
import { modelPromptBuilder } from "../backends/mlx/model-serving";
|
|
17
17
|
import { RequestError } from "./pipeline";
|
|
18
18
|
import { RequestOwnership } from "./request-plan";
|
|
@@ -23,7 +23,7 @@ export class ChatStage {
|
|
|
23
23
|
constructor(
|
|
24
24
|
private readonly ctx: ServingContext,
|
|
25
25
|
private readonly prep: RequestPrep,
|
|
26
|
-
private readonly promptCache: Pick<PromptCache, "peekPrefixLen"
|
|
26
|
+
private readonly promptCache: Pick<PromptCache, "peekPrefixLen"> & Partial<Pick<PromptCache, "objects">>,
|
|
27
27
|
/** Admission ceiling (fit() / the GLM memory plan), resolved once. */
|
|
28
28
|
private readonly contextLimit: number | null,
|
|
29
29
|
/** `serve --adapter <dir>` startup default; a request's explicit
|
|
@@ -40,9 +40,14 @@ export class ChatStage {
|
|
|
40
40
|
message.content.some((part) => part.type !== "text"));
|
|
41
41
|
const native = media || body.response_format != null || !!body.guided_grammar ||
|
|
42
42
|
!!body.guided_regex || !!body.guided_choice?.length || body.structured_outputs != null;
|
|
43
|
+
const nativeWork: PromptNativeWork = async (work) => {
|
|
44
|
+
signal?.throwIfAborted();
|
|
45
|
+
const entered = async () => { signal?.throwIfAborted(); return work(); };
|
|
46
|
+
return this.preparation ? this.preparation.run(entered, signal) : entered();
|
|
47
|
+
};
|
|
43
48
|
const build = async () => {
|
|
44
49
|
signal?.throwIfAborted();
|
|
45
|
-
const result = await this.runPrepared(request, requestId);
|
|
50
|
+
const result = await this.runPrepared(request, requestId, nativeWork);
|
|
46
51
|
if (signal?.aborted) {
|
|
47
52
|
result.plan.ownership.dispose();
|
|
48
53
|
signal.throwIfAborted();
|
|
@@ -52,14 +57,15 @@ export class ChatStage {
|
|
|
52
57
|
if (!native || !this.preparation) return build();
|
|
53
58
|
let reservation = await this.preparation.reserve?.(media ? "media" : "constraint", signal);
|
|
54
59
|
try {
|
|
55
|
-
const result = await
|
|
60
|
+
const result = await build();
|
|
56
61
|
if (reservation) result.plan.ownership.retain(reservation);
|
|
57
62
|
reservation = undefined;
|
|
58
63
|
return result;
|
|
59
64
|
} finally { reservation?.dispose(); }
|
|
60
65
|
}
|
|
61
66
|
|
|
62
|
-
private async runPrepared(request: ChatRequest, requestId
|
|
67
|
+
private async runPrepared(request: ChatRequest, requestId: string | undefined,
|
|
68
|
+
nativeWork: PromptNativeWork): Promise<InferenceRequest> {
|
|
63
69
|
const { ctx, prep } = this;
|
|
64
70
|
let body = request.params;
|
|
65
71
|
const id = requestId ?? `chatcmpl-${crypto.randomUUID()}`;
|
|
@@ -83,7 +89,7 @@ export class ChatStage {
|
|
|
83
89
|
!!body.guided_regex || !!body.guided_choice?.length ||
|
|
84
90
|
body.structured_outputs != null;
|
|
85
91
|
if (grammarReq) {
|
|
86
|
-
const g = await prep.compileGrammarForRequest(body);
|
|
92
|
+
const g = await nativeWork(() => prep.compileGrammarForRequest(body));
|
|
87
93
|
grammarCtrl = ownership.own(g.controller);
|
|
88
94
|
if (!g.controller && g.degradeHint) {
|
|
89
95
|
const degraded = applyGrammarDegrade(body, g.degradeHint);
|
|
@@ -94,7 +100,7 @@ export class ChatStage {
|
|
|
94
100
|
|
|
95
101
|
let built: BuiltPrompt;
|
|
96
102
|
try {
|
|
97
|
-
built = await this.buildPrompt(body, tools, ownership, prep);
|
|
103
|
+
built = await this.buildPrompt(body, tools, ownership, prep, nativeWork, this.promptCache.objects);
|
|
98
104
|
} catch (e) {
|
|
99
105
|
if (e instanceof RequestError) return reject(e.status, e.message);
|
|
100
106
|
return reject(400, `prompt build failed: ${(e as Error).message}`);
|
|
@@ -131,13 +137,11 @@ export class ChatStage {
|
|
|
131
137
|
// Token fast-forwarding (K3, MLX_BUN_FILL=strict): the determined-span
|
|
132
138
|
// table for this request's tools. Refused for the shapes only this stage
|
|
133
139
|
// can see — a compiled grammar (forced tokens are its job), media prompts
|
|
134
|
-
// (embeddings prefill / mRoPE)
|
|
135
|
-
//
|
|
136
|
-
// schemes (post-conversion multi-token append is unvalidated). The
|
|
140
|
+
// (embeddings prefill / mRoPE). Provider and cache-format support
|
|
141
|
+
// are declared by the engine's method binding. The
|
|
137
142
|
// body-level refusals live in prep.fillPlanFor.
|
|
138
143
|
let fillSession: import("../fill/fill-session").FillSession | null = null;
|
|
139
|
-
if (!grammarCtrl && !vision && !diffusionPixels
|
|
140
|
-
!options.kvBits && !options.kvConfig?.length && !options.turboQuant) {
|
|
144
|
+
if (!grammarCtrl && !vision && !diffusionPixels) {
|
|
141
145
|
fillSession = prep.fillPlanFor(body, tools, promptIds);
|
|
142
146
|
if (fillSession) options.fill = fillSession;
|
|
143
147
|
}
|
|
@@ -171,6 +175,7 @@ export class ChatStage {
|
|
|
171
175
|
topLogprobs,
|
|
172
176
|
adapterIds,
|
|
173
177
|
hasVision: !!vision,
|
|
178
|
+
hasPreparedPrefixIdentity: !!vision?.prefixIdentity,
|
|
174
179
|
userSeed: body.seed !== undefined,
|
|
175
180
|
hasGrammar: !!grammarCtrl,
|
|
176
181
|
hasDraft: !!ctx.draft,
|
|
@@ -280,9 +280,17 @@ export class CompletionExecutor {
|
|
|
280
280
|
control.signal?.throwIfAborted();
|
|
281
281
|
planned.transferOwnership();
|
|
282
282
|
|
|
283
|
+
// Bounded diagnostic of hidden markers versus visible events. No token
|
|
284
|
+
// readbacks are added: these IDs already crossed the sampling interface.
|
|
285
|
+
let traceTokens = control.trace ? 8 : 0;
|
|
283
286
|
const consumeToken: OnToken = (token, tokenLogprobs) => {
|
|
284
287
|
logprobs?.push(token, tokenLogprobs);
|
|
285
|
-
const
|
|
288
|
+
const route = traceTokens ? { token, index: 8 - traceTokens, events: 0 } : undefined;
|
|
289
|
+
const closeRoute = route ? control.trace!.begin("response.token_route", route) : undefined;
|
|
290
|
+
let pushed: ReturnType<CompletionSink["push"]>;
|
|
291
|
+
try { pushed = sink.push(token); if (route) route.events = pushed.events.length; }
|
|
292
|
+
finally { closeRoute?.(); }
|
|
293
|
+
if (traceTokens) traceTokens = pushed.events.length ? 0 : traceTokens - 1;
|
|
286
294
|
const observed = pushed.events.length
|
|
287
295
|
? control.onEvents?.(pushed.events)
|
|
288
296
|
: undefined;
|
|
@@ -30,10 +30,16 @@ export interface DiscoveryRoutes {
|
|
|
30
30
|
invalidateLibrary(): void;
|
|
31
31
|
}
|
|
32
32
|
|
|
33
|
+
export interface TranscriptionInfo {
|
|
34
|
+
id: string;
|
|
35
|
+
resident: boolean;
|
|
36
|
+
}
|
|
37
|
+
|
|
33
38
|
export function createDiscoveryRoutes(
|
|
34
39
|
ctx: ServingContext,
|
|
35
40
|
gateway: Pick<GenerationGateway, "batchMode">,
|
|
36
41
|
startedAt: number,
|
|
42
|
+
transcription: () => Promise<TranscriptionInfo | null> = async () => null,
|
|
37
43
|
): DiscoveryRoutes {
|
|
38
44
|
let libraryCache: { at: number; rows: unknown[] } | null = null;
|
|
39
45
|
|
|
@@ -106,6 +112,8 @@ export function createDiscoveryRoutes(
|
|
|
106
112
|
"POST /v1/messages",
|
|
107
113
|
"POST /v1/responses",
|
|
108
114
|
"POST /v1/embeddings",
|
|
115
|
+
"POST /v1/audio/transcriptions",
|
|
116
|
+
"POST /v1/audio/translations",
|
|
109
117
|
"GET /v1/models",
|
|
110
118
|
"GET/POST/DELETE /v1/adapters",
|
|
111
119
|
"GET /health",
|
|
@@ -132,6 +140,7 @@ export function createDiscoveryRoutes(
|
|
|
132
140
|
top_k: ctx.genDefaults.topK ?? null,
|
|
133
141
|
};
|
|
134
142
|
const capabilities = modelServingBinding(ctx).discovery;
|
|
143
|
+
const stt = await transcription();
|
|
135
144
|
const data: Array<Record<string, unknown>> = [{
|
|
136
145
|
id: ctx.modelId,
|
|
137
146
|
object: "model",
|
|
@@ -165,9 +174,16 @@ export function createDiscoveryRoutes(
|
|
|
165
174
|
audio: !!(ctx.audio || ctx.loadAudio),
|
|
166
175
|
adapters: capabilities.adapters,
|
|
167
176
|
training: capabilities.training,
|
|
177
|
+
transcription: stt !== null,
|
|
168
178
|
},
|
|
169
179
|
gen_defaults: genDefaults,
|
|
170
180
|
}];
|
|
181
|
+
if (stt)
|
|
182
|
+
data.push({
|
|
183
|
+
id: stt.id, object: "model", created, owned_by: "mlx-bun",
|
|
184
|
+
transcription: true, resident: stt.resident,
|
|
185
|
+
capabilities: { transcription: true, translation: true, chat_completions: false },
|
|
186
|
+
});
|
|
171
187
|
try {
|
|
172
188
|
const { Registry, visionCapable } = await import("../registry");
|
|
173
189
|
const { supportTier } = await import("../model/support");
|
|
@@ -175,7 +191,7 @@ export function createDiscoveryRoutes(
|
|
|
175
191
|
try {
|
|
176
192
|
if (registry.list().length === 0) await registry.scan();
|
|
177
193
|
for (const model of registry.listCanonical()) {
|
|
178
|
-
if (model.repoId === ctx.modelId) continue;
|
|
194
|
+
if (model.repoId === ctx.modelId || model.repoId === stt?.id) continue;
|
|
179
195
|
const tier = supportTier(model.modelType, model.repoId);
|
|
180
196
|
if (tier === null) continue;
|
|
181
197
|
data.push({
|
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
/** Token provenance for rendered multi-turn prompts. A generated BPE token
|
|
2
|
+
* sequence need not equal encode(decode(ids)). Preserve the sequence that
|
|
3
|
+
* produced the saved inference state when its rendered text is unchanged.
|
|
4
|
+
* This stores only text and IDs; KV ownership and persistence live elsewhere. */
|
|
5
|
+
export interface PromptTokenHistory {
|
|
6
|
+
remember(tokens: readonly number[]): void;
|
|
7
|
+
resolve(rendered: string, canonical: number[]): number[];
|
|
8
|
+
}
|
|
9
|
+
|
|
10
|
+
interface TokenCodec {
|
|
11
|
+
encode(text: string, addSpecialTokens?: boolean): number[];
|
|
12
|
+
decode(ids: number[], skipSpecialTokens?: boolean): string;
|
|
13
|
+
}
|
|
14
|
+
|
|
15
|
+
export class GeneratedTokenHistory implements PromptTokenHistory {
|
|
16
|
+
readonly #entries = new Map<string, number[]>();
|
|
17
|
+
#chars = 0;
|
|
18
|
+
constructor(readonly codec: TokenCodec, readonly maxChars = 2 ** 21, readonly maxEntries = 16) {}
|
|
19
|
+
|
|
20
|
+
remember(tokens: readonly number[]): void {
|
|
21
|
+
if (!tokens.length) return;
|
|
22
|
+
// Repeated requests can publish the same generated checkpoint. Its text
|
|
23
|
+
// and IDs are already owned here; refresh recency without decoding the
|
|
24
|
+
// entire conversation again on the request's completion path.
|
|
25
|
+
for (const [text, ids] of this.#entries) {
|
|
26
|
+
if (ids.length !== tokens.length || !ids.every((id, i) => id === tokens[i])) continue;
|
|
27
|
+
this.#entries.delete(text);
|
|
28
|
+
this.#entries.set(text, ids);
|
|
29
|
+
return;
|
|
30
|
+
}
|
|
31
|
+
const ids = [...tokens], text = this.codec.decode(ids, false);
|
|
32
|
+
// An unfinished UTF-8 token is not a text boundary. A later checkpoint
|
|
33
|
+
// can supply a complete prefix instead.
|
|
34
|
+
if (!text.length || text.endsWith("\uFFFD") || text.length > this.maxChars) return;
|
|
35
|
+
if (this.#entries.delete(text)) this.#chars -= text.length;
|
|
36
|
+
this.#entries.set(text, ids); this.#chars += text.length;
|
|
37
|
+
while (this.#entries.size > this.maxEntries || this.#chars > this.maxChars) {
|
|
38
|
+
const oldest = this.#entries.keys().next().value!;
|
|
39
|
+
this.#entries.delete(oldest); this.#chars -= oldest.length;
|
|
40
|
+
}
|
|
41
|
+
}
|
|
42
|
+
|
|
43
|
+
resolve(rendered: string, canonical: number[]): number[] {
|
|
44
|
+
let best: { text: string; ids: number[] } | undefined;
|
|
45
|
+
for (const [text, ids] of this.#entries) {
|
|
46
|
+
if (text.length >= rendered.length || (best && text.length <= best.text.length)) continue;
|
|
47
|
+
if (rendered.startsWith(text)) best = { text, ids };
|
|
48
|
+
}
|
|
49
|
+
if (!best) return canonical;
|
|
50
|
+
// Normal canonical output keeps the existing tokenizer fast path.
|
|
51
|
+
if (best.ids.every((id, i) => canonical[i] === id)) return canonical;
|
|
52
|
+
const ids = [...best.ids, ...this.codec.encode(rendered.slice(best.text.length), false)];
|
|
53
|
+
// Decoder cleanup/normalization (e.g. SentencePiece's leading-space
|
|
54
|
+
// rules) can make a standalone suffix unsafe to append.
|
|
55
|
+
return this.codec.decode(ids, false) === rendered ? ids : canonical;
|
|
56
|
+
}
|
|
57
|
+
}
|