mlx-bun 0.3.0 → 0.5.0

package/docs/design/generic-model-support.md

@@ -23,8 +23,12 @@ Two tiers, one contract:
 
 - **Tier-0 GENERIC** — one config-driven "universal" module per architecture
   *family*. Runs any covered checkpoint at **L1 parity** (bit-exact vs mlx-lm
-  on this machine's GPU). Monolith path only: no compiled decode, no
-  kv-quant. Slow, never broken. Built: `src/model/universal/` (dense
+  on this machine's GPU). Monolith path only, without compiled decode.
+  Standard attention now uses the shared plain/affine KV port; the softcap
+  branch retains plain KV. Qualified full-attention targets compose delayed
+  affine conversion with standalone and lookup drafting through the same
+  cache transaction interface; rotating layouts use that interface for
+  per-layer conversion and assistant donors. Built: `src/model/universal/` (dense
   llama-family). Design only: UniversalMoE.
 - **Tier-1 TARGETED** — a model *graduates*: dedicated or generated forward,
   L2 (mlx-optiq) composition where an optiq artifact exists, Lab experiments
@@ -200,7 +204,9 @@ Key choices:
   or missing tensors are a LOAD ERROR naming them, never a silently-wrong
   model.
 - **KV cache** — `KVCache`/`RotatingKVCache` unchanged; `layer_types` /
-  `sliding_window` already normalized. No kv-quant at Tier-0.
+  `sliding_window` already normalized. Standard attention consumes the shared
+  cache-attention interface; qualified Llama uniform KV4/KV8 serving and
+  generated-cache coverage are described in batching.md.
 - **RuntimeModel surface** — the minimal contract `Qwen3Model` implements
   (config, weightsBytes, loraState, makeCache, forwardHidden,
   logitsFromHidden, forward, generate, loraTargets) so serving/eval/LoRA
@@ -336,6 +342,72 @@ design-relevant facts; phase logs live in PLAN.md at the named anchor.
 - Training reference model for segmented backward and ORPO
   (docs/design/orpo-training.md, orpo-training.md).
 
+#### 6.2.1 MiniCPM5-2B investigation
+
+Metadata and template inspection completed 2026-09-10. No weights were
+downloaded and no GPU inference was run. This checkpoint is a candidate,
+not an artifact with established parity or local performance results.
+
+Start with the official [OpenBMB MLX artifact](https://huggingface.co/openbmb/MiniCPM5-2B-MLX),
+revision `3d00c3da500debfe345e60e25a87ed2669f5b1ee`. Its config uses standard
+Llama blocks, 42 layers, hidden size 2048, intermediate size 6144, 16 query
+and 2 KV heads of dimension 128, untied embeddings, vocab 130560 and
+RoPE theta 5000000. Quantization is affine 4-bit with group size 64.
+Preserve both EOS IDs, `1` and `130073`, and the separate
+`chat_template.jinja`. The release is text-only and Apache-2.0.
+
+Current integration findings:
+
+- `isMiniCPM5Config` in `src/model/support.ts` matches the older 24-layer,
+  hidden-1536 artifact. The 2B config fails that predicate and qualifies for
+  generic Llama through `genericArgsFor`; `familyProfile` therefore selects
+  the universal graph. Do not widen the predicate before auditing dedicated
+  execution, training and cache assumptions against this shape.
+- A CPU-only smoke using the pinned config and template passed
+  `loadModelConfig`, `genericArgsFor`, `ChatTemplate.load` and rendering a
+  user → assistant tool call → tool result → user conversation. Explicit
+  thinking true/false both render correctly; the XML function call survives.
+  This checks rendering, not generated tool-call parsing or tokenizer parity.
+- `resolveEnableThinking` applies the MiniCPM no-thinking default through
+  the same exact-shape predicate. It does not apply to 2B. With no explicit
+  option, this artifact's prompt ends at the assistant header; true appends
+  an open think tag, false appends a closed empty think block. Set this
+  explicitly in comparisons. Template support for disabling thinking does
+  not establish no-thinking quality.
+
+Evaluation sequence:
+
+1. Pin an 8-bit or BF16 control alongside the official 4-bit candidate.
+   Establish tokenizer/prompt identity and same-machine logits against the
+   pinned mlx-lm oracle, then test multi-turn tool calls, reasoning splits,
+   EOS, cache reuse and batched continuation. Keep quantization quality
+   comparisons separate from same-artifact runtime parity.
+2. Compare with MiniCPM5-1B and Gemma e4b on held-out Dreaming segmentation,
+   extraction and source-grounded synthesis, plus bounded repository/tool
+   tasks. Score boundary/label accuracy and completed tasks, not just valid
+   JSON or XML. Use explicit thinking modes and equal token/time budgets.
+3. Record first-answer latency, complete task time, output/reasoning tokens,
+   repetition and truncation rates, memory and context scaling. Local
+   measurements belong only in `docs/reference/benchmarks.md`, with a named
+   quiet machine and the required serving benchmark procedure.
+4. Investigate the [released DSpark draft](https://huggingface.co/openbmb/MiniCPM5-2B-DSpark)
+   after ordinary inference passes. Its config declares `Qwen3DSparkModel`,
+   five draft layers, block size seven, and target taps `[1,10,20,30,39]`.
+   Existing Gemma DSpark support does not establish compatibility. Audit
+   the draft's architecture, target taps, weights and verifier before a
+   paired task-time comparison.
+
+Evidence to use when designing the quality checks:
+[OpenBMB's model card](https://huggingface.co/openbmb/MiniCPM5-2B) distinguishes
+internally reproduced scores from Artificial Analysis results. The
+[independent September 7 evaluation](https://artificialanalysis.ai/articles/openbmb-releases-minicpm5-2b)
+supports strong capability for its size but reports substantial reasoning
+token use and weak absolute terminal-task success. These evaluations do not
+validate MLX 4-bit quality. [Upstream issue 374](https://github.com/OpenBMB/MiniCPM/issues/374)
+reports repetition sensitivity in GGUF coding runs. That is an unreproduced
+cross-runtime hypothesis for our sweep, not evidence of the same MLX failure
+or a reason to change repetition defaults.
+
 ### 6.3 Qwen3 / Qwen3-MoE / Qwen3.5-family (3.6, 3.8) — `qwen3.ts`, `qwen3-moe.ts`, `qwen3_5.ts`
 
 - **Qwen3** (`model_type: qwen3`): dense GQA with per-head q/k RMSNorm
@@ -581,6 +653,122 @@ non-goals.
 
 ---
 
+### 6.7 Whisper (speech-to-text) — `src/model/whisper.ts`, `src/audio/whisper-*.ts`, `src/serve/transcription-*.ts`
+
+Audio in → text out, the first encoder-decoder family. Target artifact:
+`mlx-community/whisper-large-v3-turbo` (fp16, 32 encoder / 4 decoder layers,
+128 mel bins, 51 866 vocab). Plan/evidence anchor: PLAN.md "Phase: Whisper
+speech-to-text". Served surface: [server-api.md](../reference/server-api.md#post-v1audiotranscriptions--post-v1audiotranslations-speech-to-text),
+[cli.md `transcribe`](../reference/cli.md#transcribe--speech-to-text).
+
+- **Oracle**: mlx-whisper 0.4.3 (`ml-explore/mlx-examples/whisper`) in its
+  own venv (`~/Code/mlx-whisper-oracle/.venv`, mlx pinned to
+  `MLX_CORE_VERSION`; it drags torch + numba, which must not enter the
+  pinned mlx-lm venv). The profile declares `{ l1, mlx-whisper, bit-exact }`
+  — no mlx-lm arm exists. Generator `scripts/oracle/gen-whisper-golden.py`
+  (`bun scripts/regen.ts whisper`) → `goldens/whisper.json` + untracked
+  blobs; gate `tests/parity/whisper.test.ts`: mel (f32), decode-window
+  encoder output (f16 raw) and every per-step PRE-filter logit vector (f32)
+  bit-exact; text / segment tokens / timestamps identical. Clips: the
+  tracked `speech-fox.wav`, whisper.cpp's public-domain `jfk.wav` (sotto's
+  vendored checkout, optional), and a synthesized 33 s concatenation that
+  crosses a 30 s window.
+- **Front end** (`whisper-mel.ts`): op-for-op `log_mel_spectrogram` on the
+  GPU — host reflect pad, `as_strided` frames, periodic Hann, `rfft` 400,
+  `|·|²`, slaney filter bank (librosa's float64 → float32 rounding order
+  reproduced, bit-exact vs the shipped `mel_filters.npz`), `log10`, 8-decade
+  clamp, `(x+4)/4`. 1.5 ms warm for a 41 s padded clip (M1 Max).
+- **Model** (`whisper.ts`): faithful path = the oracle's graph, including
+  its quirks that change bits: `nn.gelu` is `mx.compile`d in mlx, and the
+  FUSED kernel rounds differently from four separate ops (2 % of conv1
+  outputs differ by an ulp) — so the same graph is compiled here
+  (`CompiledFunction`, shapeless); attention scales q and k by
+  `(D/H)^-0.25` each with fp16 `q@kᵀ`, additive `-inf` fp16 causal mask,
+  precise softmax; self-attention K/V are cached un-split `[B, T, D]` and
+  concatenated per step; cross-attention K/V are projected once per window
+  and kept `[1, 1500, D]` — beam rows attend to them by broadcasting, never
+  by copying.
+- **Decoding** (`whisper-decode.ts`): DecodingTask port — sot sequence,
+  prompt/prefix budgets (`n_ctx/2 − 1 = 223` prompt tokens), suppress-blank,
+  non-speech suppression, timestamp rules, language detection, greedy /
+  sampled, MaximumLikelihoodRanker. **Two deliberate divergences from
+  mlx-whisper**, both documented in code: (1) the monotonic-timestamp rule
+  follows openai-whisper (mlx-whisper's port collects sequence indices and
+  slices an empty range, so its rule never fires) — pre-filter logits stay
+  bit-exact, `avg_logprob` moves ~1e-4; (2) `compression_ratio` uses Bun's
+  zlib, which emits one byte more than CPython's 1.2.12 for the same
+  level-6 stream — the 2.4 fallback threshold decision is asserted equal,
+  the value within 5 %. Beam search is ours (mlx-whisper raises
+  NotImplemented): openai-whisper's BeamSearchDecoder over cache rows, top-k
+  on device (`argpartition` on negated log-probs; the `[B, k]` slice is
+  materialized before host readback — reading a strided view linearly was a
+  real bug), finished sequences collected per audio, patience honoured.
+  Checked against whisper.cpp v1.9.3 beam-5 transcripts (identical on the
+  JFK clip with and without vocabulary hints).
+- **Transcribe loop** (`whisper-transcribe.ts`): transcribe.py's seek loop —
+  30 s windows over the zero-padded mel, temperature fallback ladder gated by
+  compression ratio / avg log-prob / no-speech probability, timestamp-driven
+  seeking, previous-text conditioning, `clip_timestamps`. Async only to
+  yield between windows (streaming, cancellation). Word timestamps (DTW on
+  the checkpoint's `alignment_heads`) are not ported yet.
+- **Serving** (`transcription-service.ts`, `audio-routes.ts`,
+  `transcription-server.ts`): residency-managed service (load on first
+  request, dispose after `--whisper-idle-unload` seconds, FIFO queue, runs
+  under the gateway's exclusive lock beside chat), OpenAI multipart/JSON
+  routes with `json/verbose_json/text/srt/vtt` and SSE streaming, sotto's
+  vocabulary-hint fitting with included/omitted reporting, and a
+  transcription-only server for `serve <whisper checkpoint>`. Page-in after
+  an explicit unload: 37–45 ms on the M1 Max (weights are mmap'd; the OS
+  file cache still holds them).
+- **Tokenizer**: the mlx-community repos ship no tokenizer; ids are identical
+  to HF's `openai/whisper-large-v3-turbo` tokenizer.json (checked against
+  tiktoken), which the loader resolves from the HF cache by `n_vocab`.
+- **Fast path** (`whisper-fast.ts`, default; `faithful`/`--faithful` selects
+  the oracle graph): fused SDPA everywhere (whisper.cpp flash-attn); one
+  fused q/k/v gemm per encoder layer; cross-K/V for all four decoder layers
+  from ONE matmul pair per window, head-split once (whisper.cpp `kv_cross`,
+  vLLM encoder-prefill cross KV); the single-token decoder step — embed, 4
+  layers, ln, logits, every logit filter, argmax (greedy) or top-k (beam) —
+  is one shapeless `mx.compile`d closure with per-step state as int32 arrays
+  (the LLM compiled-decode pattern); greedy runs device-resident with the
+  token array flowing into the next step and a one-step-lag host read
+  (mlx-whisper's async_eval); beam reorders caches with `take`. Gate:
+  `tests/parity/whisper-fast.test.ts` — at most two differing tokens per
+  clip vs the oracle/faithful path (one near-tie flips "dog"→"dog." plus
+  its timestamp on the 33 s no-conditioning clip). Word timestamps
+  (`whisper-timing.ts`, timing.py port: alignment-head cross-attention →
+  median filter → DTW in JS → punctuation merge and duration heuristics)
+  ride the faithful decoder once per window.
+- **Silero VAD** (`silero-vad.ts`): the v6.2 TorchScript graph ported
+  exactly (64-sample inter-chunk context, right reflect pad, STFT-as-conv,
+  4 conv blocks, LSTMCell, sigmoid) — whisper.cpp's port drops the context
+  and pads both sides — plus `get_speech_timestamps`; STFT/convs/input
+  gates batched over all chunks in one MLX graph, the recurrence on the
+  host; streaming state carried across chunks (streaming == batch within
+  1e-6). Gate: `tests/parity/silero-vad.test.ts` vs the pip reference
+  (probabilities within 5e-3 because the ggml weights are fp16; segments
+  identical). Used as the service's gate (`vad`) and inside sessions.
+- **Streaming run** (`WhisperStreamingRun`): the seek loop as a resumable
+  object; `feed()` transcribes each window with 30 s of audio past the seek
+  point, `finish()` drains. Batch = one feed + finish, so parity covers it.
+  Sessions (`/v1/audio/sessions`) wrap it with the VAD stream and the
+  exclusive lock.
+- **Voice input**: `mlx-bun dictate` (AVAudioEngine sidecar
+  `src/native/mic_capture.swift`, hold a keycode or Enter-toggle, print /
+  clipboard / System Events keystrokes) and the web composer's hold-to-talk
+  mic (`src/web/src/voice.ts`: 16 kHz AudioContext → session chunks → caret
+  insert; shown when the `ready` frame carries `transcription`). Both are
+  the streaming session over the same server code.
+- **Decode input**: AudioToolbox `ExtAudioFile` over `bun:ffi`
+  (`audiotoolbox.ts`) for every non-WAV container — CoreAudio's decoder and
+  resampler in-process; WAV keeps the exact PCM parser.
+- **What the measurements say (M1 Max 32 GB, 2026-09-15, benchmarks.md):**
+  the encoder is ~85 % of a short clip's time and is bound by its fp16
+  gemms at ~7–8 TFLOPS (padding rows to 1536 for tile alignment, fused
+  vs unfused GELU, batching two windows: all 0–5 %); decode is ~1.5 ms per
+  token. Whisper.cpp's `audio_ctx` truncation is the only lever left on
+  the encoder and it is a quality trade (`audio_ctx`, Lab, off).
+
 ## 7. Top design decisions
 
 1. **Explicit descriptor table keyed by model_type (+ vendored
@@ -627,3 +815,7 @@ non-goals.
 - 2026-08-23 — This doc consolidated: audio-input-plan,
   diffusion-gemma-port, minicpm5-decode-megakernel, and the Colibri archive
   summarized here; ports become sections.
+
+
+Non-Qwen grouped TurboQuant drafting, including assistant donor attention, uses
+the shared cache interfaces. See [composition coverage](batching.md#turboquant-donor-attention-and-grouped-draft-composition).