mlx-bun 0.0.4

package/docs/reference/server-api.md

@@ -0,0 +1,224 @@
+# Server API
+
+`mlx-bun serve` (or `bun scripts/serve.ts`) exposes an OpenAI-compatible
+HTTP API on one model. The request's `model` field is ignored; the
+loaded model's id is echoed back. By default generation is serialized
+through a single queue (one GPU, batch = 1); `--batch N` switches the
+server into bf16 continuous batching.
+
+For the full set of start flags (`--port`, `--memory-budget`,
+`--prompt-cache`, `--batch`, `--kv-quant`, `--thinking`, sampling
+defaults, and the perf levers) and a compatibility matrix of which
+combinations compose, see [server-config.md](./server-config.md). This
+doc covers the request/response wire format. LoRA adapters are mounted at
+runtime via `POST /v1/adapters` (below), not a start flag.
+
+## POST /v1/chat/completions
+
+Request body (OpenAI chat schema; unknown fields ignored):
+
+```jsonc
+{
+  "messages": [ /* role: system | user | assistant | tool */ ],
+  "stream": false,
+  "max_tokens": 1024,            // or max_completion_tokens (wins)
+  "temperature": 0.7,            // 0 = greedy
+  "top_p": 0, "top_k": 0,        // 0 = off
+  "seed": 1234,                  // omit for time-derived
+  "repetition_penalty": 1.1,     // optional
+  "stop": "\n\n",                // or ["###", "\n\n"] (spec: up to 4)
+  "tools": [ /* OpenAI function tools */ ],
+  "tool_choice": "auto",         // "none" disables tools
+  "chat_template_kwargs": {      // forwarded to the chat template
+    "enable_thinking": false     // MiniCPM5: <think> channel on/off
+  },
+  "adapter": "id"                // LoRA: "id", stacked "a+b", or "none"
+}
+```
+
+Sampling defaults follow the model author's `generation_config.json`
+when a field is omitted (optiq serve's gen_config behavior); explicit
+request values always win. MiniCPM5 defaults to the no-think direct
+answer mode unless `chat_template_kwargs.enable_thinking` is `true`.
+
+`stop` sequences are matched on **decoded text**, not token ids, so a
+sequence that spans token boundaries still fires. Generation halts at
+the first match; the stop sequence itself is excluded from the content
+and `finish_reason` is `"stop"`.
+
+Message `content` is a string or an array of parts:
+`{ "type": "text", "text": ... }` and
+`{ "type": "image_url", "image_url": { "url": "data:image/png;base64,..." } }`
+(http/https URLs also accepted; PNG, JPEG, HEIC, AVIF, WebP, TIFF, GIF,
+BMP via native OS codecs; requires a model with the vision sidecar).
+
+Non-streaming response:
+
+```jsonc
+{
+  "id": "chatcmpl-…", "object": "chat.completion", "created": 1760000000,
+  "model": "<loaded model id>",
+  "choices": [{
+    "index": 0,
+    "message": {
+      "role": "assistant",
+      "content": "…",                      // "" when only tool calls
+      "tool_calls": [{                     // present when the model called tools
+        "id": "call_…", "type": "function",
+        "function": { "name": "…", "arguments": "{…json…}" }
+      }]
+    },
+    "finish_reason": "stop" | "length" | "tool_calls"
+  }],
+  "usage": {
+    "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0,
+    "prompt_tokens_details": { "cached_tokens": 0 }   // prompt-cache reuse
+  }
+}
+```
+
+Streaming (`"stream": true`) is SSE: `data: <chunk>\n\n` per event,
+terminated by `data: "[DONE]"`. Chunks are `chat.completion.chunk`
+objects whose `choices[0].delta` carries `{role}`, then `{content}`
+increments (multi-byte sequences are held back until decodable; text
+that could begin a `stop` sequence is held back until disambiguated,
+so no part of a stop sequence is ever streamed), then
+for tool calls a final `{tool_calls: [{index, id, type, function}]}`
+delta; the last chunk carries `finish_reason` and `usage`.
+
+Tool round-trip: send the assistant message with its `tool_calls` back,
+followed by `{ "role": "tool", "tool_call_id": …, "content": … }`
+messages; multi-turn prompt prefixes reuse the KV prompt cache
+automatically.
+
+Tool-call parsing is per model family. Gemma 4 uses its native
+`<|tool_call>`…`<tool_call|>` sentinel tokens. MiniCPM5 emits XML in
+decoded text (`<function name="…"><param name="…">…`, CDATA-wrapped
+values supported); content before the tool markup still streams live,
+only the markup is withheld and converted to `tool_calls`. Argument
+values are decoded against the tool's JSON schema (string-typed params
+stay strings); markup that fails to parse falls back to plain content.
+
+### Errors
+
+All errors are `{ "error": { "message": …, ... } }`.
+
+- `400` — malformed JSON, empty `messages`, unknown adapter id, vision
+  request on a model without a sidecar, prompt build failures.
+- `400` with `"type": "memory_admission"`, `"code":
+  "context_over_budget"` — `prompt + max_tokens` exceeds the memory
+  budget's max safe context (only when serving with `--memory-budget`;
+  the GPU OOM this prevents would kill the process, so it is refused
+  up front). Lower `max_tokens` or shorten the prompt; the ceiling is
+  visible at `/stats`.
+
+## POST /v1/messages (Anthropic Messages API)
+
+Anthropic-protocol surface over the same engine — on by default, like
+`optiq serve`. Point any Anthropic-SDK tool at the server
+(`ANTHROPIC_BASE_URL=http://localhost:8090`, any `x-api-key`) — Claude
+Code works as a client this way.
+
+- `system` (string or text blocks), `messages` with string or
+  content-block arrays; `tool_use` / `tool_result` blocks map to the
+  native gemma tool-calling path (better than the optiq shim, which
+  inlines them as text); `image` blocks (base64 or url source) hit the
+  vision path on sidecar models.
+- `tools` (`{name, description, input_schema}`) map to function tools;
+  server-tool types (web_search, …) are dropped silently.
+- `max_tokens`, `temperature`, `top_p`, `top_k`, `stop_sequences`,
+  `stream` as in the Anthropic spec.
+- Response: `{id: "msg_…", type: "message", content: [{type: "text"} |
+  {type: "tool_use"}…], stop_reason, usage: {input_tokens,
+  output_tokens, cache_read_input_tokens}}` —
+  `cache_read_input_tokens` comes from the prompt cache.
+- Streaming follows the Anthropic event grammar exactly:
+  `message_start → content_block_start/delta/stop (text_delta,
+  input_json_delta) → message_delta (stop_reason + usage) →
+  message_stop`. Errors are `event: error` frames.
+- Errors: `{type: "error", error: {type: "invalid_request_error" |
+  "api_error", message}}`.
+
+## POST /v1/responses (OpenAI Responses API)
+
+Responses-protocol surface (Codex, Cursor, Continue, Cline, and the
+OpenAI SDK speak this now). Oracle: optiq responses shim.
+
+- `input` (string or item array: `message`, `function_call`,
+  `function_call_output`), `instructions` (merged with any
+  system/developer items into one leading system message),
+  `max_output_tokens`, `temperature`, `top_p`, `top_k`, flat
+  `tools`/`tool_choice` (built-in tool types dropped), `stream`.
+- **`previous_response_id` resumption**: pass a prior response id
+  instead of resending the conversation; the server splices the stored
+  input + output back in (instructions carry forward when omitted).
+  Store is per-process, 1 h TTL, 32 MiB byte-capped LRU — observable
+  at `GET /stats` (`response_store`). Unknown/expired id → 404.
+- Response: `{id: "resp_…", object: "response", status: "completed" |
+  "incomplete", output: [{type: "message"|"function_call"…}], usage}`.
+- Streaming event chain: `response.created → response.in_progress →
+  response.output_item.added → response.content_part.added →
+  response.output_text.delta… → response.output_text.done →
+  response.content_part.done → response.output_item.done →
+  response.completed` (+ `response.function_call_arguments.delta/.done`
+  for tool calls).
+
+## GET /v1/models
+
+`{ "object": "list", "data": [{ "id": "<model id>", "object": "model", … }] }`
+
+## GET /stats
+
+```jsonc
+{
+  "server": { "owner": "serve" | "pi-session" | "embedded", "model": "...", "started_at": 0 },
+  "prompt_cache": { "entries": 0, "bytes": 0, "max_bytes": 0, "hits": 0, "misses": 0 },
+  "response_store": { "entries": 0, "bytes": 0, "max_bytes": 33554432, "ttl_ms": 3600000 },
+  "kv_quant": { "mode": "mixed (kv_config.json)" | "uniform-kv8" | "bf16",
+                 "layers": { "kv4": 8, "bf16": 40 },
+                 "attention": { "global": 10, "sliding_window": 38 } },
+  "admission": {
+    "max_safe_context": 0,            // tokens; requests above this 400
+    "memory_budget_bytes": null,      // explicit budget, or null (machine default)
+    "usable_bytes": 0,
+    "weights_bytes": 0
+  },
+  "batch": {
+    "configured": 1,                  // the --batch N value
+    "batched": false,                 // batching enabled (N>1)
+    "active_rows": 0                  // rows currently decoding in the batch
+  }
+}
+```
+
+## Adapters (LoRA hot-swap)
+
+- `GET /v1/adapters` — `{ adapters: [{ id, path, rank, scale, size_bytes, mounted_layers }] }`
+- `POST /v1/adapters` — `{ "id": "...", "path": "/dir" }`; mounts
+  through the generation queue (never races a forward pass). 400 on
+  shape/compat mismatch — validation is all-or-nothing.
+- `DELETE /v1/adapters/<id>` — unmount; 404 if not mounted.
+
+Select per request with the `adapter` body field. Prompt-cache entries
+are namespaced per adapter spec, so switching adapters never reuses
+another adapter's KV.
+
+## Client setup: pi
+
+`~/.pi/agent/models.json`:
+
+```jsonc
+{
+  "providers": {
+    "mlx-bun": {
+      "baseUrl": "http://localhost:8090/v1",
+      "api": "openai-completions",
+      "apiKey": "sk-anything-nonempty",
+      "models": [{ "id": "<model id from /v1/models>" }]
+    }
+  }
+}
+```
+
+Any OpenAI SDK works the same way: `baseURL: "http://localhost:8090/v1"`,
+any non-empty `apiKey`.

package/docs/reference/server-config.md

@@ -0,0 +1,243 @@
+# Server configuration
+
+Every start flag for `mlx-bun serve` (and `mlx-bun pi`, which shares the
+same flag set), what each does, its default, and — the part that bites —
+**which combinations compose and which don't**. For the request/response
+wire format (chat/messages/responses bodies, SSE grammar, tool calling,
+adapters admin) see [server-api.md](./server-api.md); this doc is about
+*starting* the server.
+
+The flags are parsed in `serverRuntimeFlags()`
+([src/cli.ts](../../src/cli.ts)); the runtime knobs live in
+`ServerOptions` ([src/server.ts](../../src/server.ts)). The two execution
+lanes (serial vs. `--batch N`) and the lane picker are in
+[src/serve/generation-gateway.ts](../../src/serve/generation-gateway.ts);
+the design rationale is in
+[docs/design/parallel-slots.md](../design/parallel-slots.md).
+
+## Start flags
+
+| Flag | Arg | Default | What it does |
+| --- | --- | --- | --- |
+| `--host` | addr | all interfaces | Interface to bind. `127.0.0.1` = loopback only. |
+| `--port` | n | `8090` | Listen port. A pre-flight probe refuses to start if the port is already serving. |
+| `--memory-budget` | GB | machine RAM × 0.75 | Admission ceiling. Requests whose `prompt + max_tokens` exceed the budget's max safe context are **rejected with 400** (`type: memory_admission`) instead of risking an uncatchable GPU OOM. Also caps the mlx allocator (`mlx_set_memory_limit`) as defense in depth. **Decimal GB (×10⁹).** |
+| `--prompt-cache` | GB | `2` (≈2e9 bytes) | Byte cap for the prompt (prefix-reuse KV) cache. **Binary GiB (×2³⁰)** on input. `0`/unset ⇒ default; there is no flag to disable the cache outright. |
+| `--batch` | n | `1` (serial) | Max concurrent requests batched through the mlx-lm-parity engine. `>1` switches the **whole server** into bf16 continuous batching — a *mode*, not a load fallback. See [Execution modes](#execution-modes-serial-vs---batch-n). Alias: `--decode-concurrency` (mlx_lm.server drop-in). |
+| `--kv-quant` | `config`\|`off`\|`4`\|`8` | `config` serial / `off`(bf16) under `--batch N` | KV-cache quantization. `config` = per-layer `kv_config.json` (optiq parity); `off` = bf16; `4`/`8` = uniform bits (group 64, start 0). Under `--batch N`, an explicit value routes those requests to the **serial** lane (batched is bf16-only). |
+| `--thinking` | `true`\|`false` | model's own (false for CPM) | Server-wide default for the chat template's `enable_thinking` (MiniCPM5/CPM hybrid reasoning). A request's `chat_template_kwargs.enable_thinking` overrides it. |
+| `--temperature` | n ∈ [0,5] | `generation_config.json` | Server-wide sampling default. Per-request `temperature` still wins; the browser chat (sends none) inherits this. |
+| `--top-p` | n ∈ [0,1] | `generation_config.json` | Server-wide top-p default (per-request `top_p` wins). |
+| `--top-k` | n ∈ [0,1e6] | `generation_config.json` | Server-wide top-k default (per-request `top_k` wins). |
+| `--compiled-decode` | on\|off | on | Replay the per-step decode graph in C++ (`MLX_BUN_COMPILED_DECODE`). Bit-exact A/B lever. **Serial lane only** (see note below). |
+| `--perf-kernel` | on\|off | **on** | Fused quantized-KV decode-SDPA Metal kernel (`MLX_BUN_PERF_KERNEL`), the perf side of the compat A/B. Engages on quantized caches at decode. **Serial lane only.** |
+| `--fused-decode` | on\|off | off | Experimental: tile the quantized decode SDPA (`MLX_BUN_FUSED_DECODE`). **Serial lane only.** |
+| `--fused-sdpa` | on\|off | on | Fused SDPA path for quantized prefill/continuation (inverted env `MLX_BUN_NO_FUSED_SDPA`). **Serial lane only.** |
+| `--force-wire` | (bool) | off | Wire weights into memory for the whole generation (`MLX_BUN_FORCE_WIRE`). Near-ceiling models (e.g. 26B) need it. **Serial lane only.** |
+
+The five performance levers are A/B knobs whose defaults are the measured
+winners; flip them to compare. They are set as `MLX_BUN_*` env vars
+before the model loads, so they apply to `mlx-bun pi` too. They affect
+the **serial** decode path (`generate()`); the batched scheduler calls
+`model.forwardHidden` directly and so is unaffected by all of them — see
+[Levers that don't reach the batched lane](#levers-that-dont-reach-the-batched-lane).
+
+> **Note — `--perf-kernel` default.** The code defaults it **on**
+> (`perfKernelEnabled()` returns true unless `MLX_BUN_PERF_KERNEL=0`, and
+> the ready card / `/stats` report it on). `STATUS.md` still lists a
+> "default flip" as pending; treat that as stale — the engaged default
+> today is on. (If the intent is genuinely off-until-clean-machine-pass,
+> the *code* is what needs changing, not this doc.)
+
+## Per-request overrides
+
+Most quality knobs can be set per request in the chat body and override
+the server-wide default. Precedence, highest first:
+
+1. explicit request field (`temperature`, `top_p`, `top_k`, `seed`,
+   `repetition_penalty`, `max_tokens`/`max_completion_tokens`, `stop`,
+   `chat_template_kwargs.enable_thinking`, `adapter`, `tools`)
+2. the matching `--temperature`/`--top-p`/`--top-k`/`--thinking` server default
+3. the model's `generation_config.json`
+4. the built-in fallback (temp 0.7, top-p/top-k 0 = off)
+
+These ride all three protocol surfaces (`/v1/chat/completions`,
+`/v1/messages`, `/v1/responses`) because each funnels through one
+`handleChat`. Full field list: [server-api.md](./server-api.md).
+
+## Environment levers
+
+The CLI flags above set these; you can also export them directly (e.g.
+for `bun scripts/serve.ts` or paired A/B harnesses). One has no CLI flag.
+
+| Env var | CLI flag | Default | Effect |
+| --- | --- | --- | --- |
+| `MLX_BUN_COMPILED_DECODE` | `--compiled-decode` | on (`!=="0"`) | Compiled decode graph replay. |
+| `MLX_BUN_PERF_KERNEL` | `--perf-kernel` | on (`!=="0"`) | Fused quantized-KV decode kernel (not bit-exact; perf A/B). |
+| `MLX_BUN_FUSED_DECODE` | `--fused-decode` | off (`==="1"`) | Tile quantized decode SDPA. |
+| `MLX_BUN_NO_FUSED_SDPA` | `--fused-sdpa` (inverted) | fused on | `=1` forces the stock unfused SDPA everywhere. |
+| `MLX_BUN_FUSED_GELU` | *(none)* | on (`!=="0"`) | Fused GeGLU MLP kernel. Bit-exact, so it stays on both lanes; env-only opt-out. |
+| `MLX_BUN_FORCE_WIRE` | `--force-wire` | off (`==="1"`) | Wire weights for the generation. |
+
+## Execution modes: serial vs. `--batch N`
+
+`--batch 1` (default) is the serialized single-queue path: one GPU, one
+generation at a time, prompt-cache prefix reuse, mixed-precision KV by
+default. Untouched and unconditional.
+
+`--batch N` (N>1) opts the **whole server** into a continuous-batching
+engine that is **bit-parity with `mlx_lm.server` at `B=N`**. It is a
+*mode switch*, not a load-dependent fallback (results must not depend on
+concurrency).
+
+Parity is the guarantee, so the batched lane runs in **compat mode**: it
+exposes **none** of the serial lane's perf flags and runs the plain
+bit-exact forward — the same path mlx-lm runs, never the optional
+parity-breaking kernels. This is the mechanism of the guarantee, not a
+missing feature (see
+[compat mode](#--batch-n-is-compat-mode--perf-flags-dont-apply-by-design)).
+Because mlx-lm's batched path is bf16 (its quantized batching is NYI),
+bf16 continuous batching **is** the drop-in:
+
+- **KV quant unset ⇒ bf16** so the batch path engages out of the box
+  ("Option B"). The serial default stays mixed-precision (optiq parity).
+- **Explicit `--kv-quant config|4|8` ⇒** those requests route to the
+  serial lane (batched is bf16-only; a startup warning is printed). With
+  an explicit `--kv-quant`, *every* request carries a quant scheme, so
+  **nothing batches** — `--batch N --kv-quant config` is effectively
+  serial-with-quant. Omit `--kv-quant` to actually batch.
+
+### The lane picker (`GenerationGateway.willBatch`)
+
+Under `--batch N`, each request is routed per-request. It joins the
+batch only if **all** of these hold; otherwise it drains the batch and
+runs solo (mlx-lm's `_is_batchable` behavior). The two lanes are
+mutually exclusive on the GPU (one `AsyncMutex`), so a serial-fallback
+request never runs alongside a batched step — but batched requests run
+concurrently with each other.
+
+| Request property | Batches? |
+| --- | --- |
+| vision (image parts) | ❌ serial — needs offset-0 single-seq prefill + bidirectional image mask |
+| LoRA `adapter` (resolves to ≥1) | ❌ serial — `loraState.active` is one per-generation field; per-row adapters unsupported |
+| `repetition_penalty` | ❌ serial — per-row logits processors are a later refinement |
+| explicit `seed` | ❌ serial — reproducibility ⇒ solo (matches mlx-lm) |
+| KV quant active (explicit `--kv-quant`) | ❌ serial — batched is bf16-only in v1 |
+| `temperature` / `top_p` / `top_k` | ✅ batches (each row samples with its own seed) |
+| `stop` sequences | ✅ batches (per-row `StopMatcher` in the onToken closure) |
+| `tools` / `tool_choice` | ✅ batches (per-row tool router; decode-layer parse) |
+| `--thinking` / `enable_thinking` | ✅ batches (template-render concern, lane-independent) |
+| multi-turn / long prompt | ✅ batches, but **no prompt-cache reuse** (`cached_tokens=0`) |
+
+Both full-attention (CPM) and sliding-window (Gemma) models batch; the
+scheduler assembles each layer's cache by attention type.
+
+## Compatibility matrix
+
+How each option behaves in each mode. "serial" = `--batch 1`; "`--batch N`"
+= the batched mode (a given request may still take the serial *lane*
+inside it per the table above).
+
+| Option | serial (`--batch 1`) | `--batch N` (N>1) |
+| --- | --- | --- |
+| `--kv-quant config`/`4`/`8` | ✅ applied to all requests | ⚠️ applied, but forces **all** requests to the serial lane (no batching) |
+| `--kv-quant off` | ✅ bf16 | ✅ bf16 (same as the implicit batch default) |
+| *(kv-quant unset)* | mixed-precision `config` | **bf16** (Option B) — incl. serial-lane fallback requests |
+| `--memory-budget` | ✅ per-request admission | ✅ per-request admission — but **not aggregate** across rows (see limitations) |
+| `--prompt-cache` | ✅ prefix reuse | ⚠️ bypassed for batched requests (`cached_tokens=0`); serial-lane requests still reuse |
+| `--temperature`/`--top-p`/`--top-k` | ✅ | ✅ (per-row) |
+| `--thinking` | ✅ | ✅ |
+| vision request | ✅ | ✅ via serial lane (in bf16 under Option B) |
+| LoRA `adapter` | ✅ | ✅ via serial lane |
+| `repetition_penalty` | ✅ | ✅ via serial lane |
+| `seed` | ✅ | ✅ via serial lane |
+| `tools` / `stop` | ✅ | ✅ (batches) |
+| `--compiled-decode`/`--perf-kernel`/`--fused-*`/`--force-wire` | ✅ (serial perf tree) | **n/a — compat mode, no perf flags by design** |
+
+### `--batch N` is compat mode — perf flags don't apply by design
+
+The bit-parity guarantee (mlx-lm B=N) is the *whole point* of `--batch N`,
+and it requires running the plain forward path. So the batched lane
+deliberately exposes **no** perf knobs: the scheduler
+([batch-scheduler.ts](../../src/serve/batch-scheduler.ts)) drives the
+model through `forwardHidden`/`logitsFromHidden` directly (not
+`generate()`), running the same bit-exact kernels mlx-lm runs — never the
+optional, parity-breaking ones. Flagging the batched lane would defeat
+the guarantee, so it's intentionally not wired.
+
+- **`--perf-kernel` / `--fused-decode` / `--compiled-decode`** — never
+  engage in the batched lane. They diverge from the -O0 reference (or, for
+  the quantized-KV kernels, are moot since batched is bf16). They're the
+  L3 perf row, a separate KL-gated path — see the validation matrix in
+  [parallel-slots.md](../design/parallel-slots.md).
+- **`--force-wire`** — doesn't wire (the scheduler bypasses `generate()`'s
+  wired scope). A model that needs wiring for speed (e.g. 26B: 8.6 → 32.3
+  tok/s wired) runs at mlx-lm-equivalent unwired speed under `--batch N`;
+  in practice such a model has no headroom for B>1 KV anyway.
+- **Always-on bit-exact kernels still run.** Fused GeGLU is bit-exact with
+  the spelled-out MLP, so it stays on in both lanes without breaking
+  parity — "compat mode" means *no parity-breaking optionality*, not
+  -O0.
+
+The serial lane is where the perf/optimization tree lives: `--perf-kernel`
+/ `--fused-*` / `--compiled-decode` engage there, and the
+mixed-precision-KV (optiq) default applies there.
+
+## Known limitations under `--batch N`
+
+These are deliberate v1 scope, not bugs — but they change behavior, so
+know them:
+
+1. **Prompt cache bypassed.** Batched requests solo-prefill every row;
+   `cached_tokens=0`. Wiring `PromptCache` into the scheduler is a
+   follow-up.
+2. **Admission is per-request, not aggregate.** `--memory-budget` checks
+   each request against single-sequence max-safe-context, but N
+   concurrent rows can collectively exceed the budget (the `B×S_max`
+   KV-budget admission is TODO). With a tight budget and several
+   large-context requests, the allocator cap is the only backstop — and
+   a true GPU OOM here is uncatchable. Size the budget with headroom for
+   N rows, or keep N small.
+3. **Short-context only.** Verified pre-ring-wrap (rows < the 1024
+   sliding window). Long-context (context > window) batched decode is a
+   separate validation.
+4. **bf16 only — by contract.** mlx-lm's batched path *is* bf16, so
+   bf16-only batching is exactly what mlx-lm-parity means (not a
+   shortcoming). Going further — batched + mixed-precision KV quant — is
+   novel territory with no mlx-lm/optiq oracle, so it's a deferred,
+   KL-gated extension.
+5. **`extend` join not yet used.** A joining request re-merges the whole
+   batch (O(B·S)); mlx-lm's keep-the-running-batch `extend` is a later
+   optimization (numerically equivalent).
+
+## Observability — `GET /stats`
+
+The live config and batch state:
+
+```jsonc
+{
+  "server":  { "owner": "serve", "model": "...", "started_at": 0 },
+  "prompt_cache":  { "entries": 0, "bytes": 0, "max_bytes": 2000000000, "hits": 0, "misses": 0 },
+  "response_store": { "entries": 0, "bytes": 0, "max_bytes": 33554432, "ttl_ms": 3600000 },
+  "kv_quant": {
+    "mode": "mixed (kv_config.json)" | "uniform-kv8" | "bf16",
+    "layers": { "kv4": 8, "bf16": 40 },
+    "attention": { "global": 10, "sliding_window": 38 }
+  },
+  "admission": {
+    "max_safe_context": 0,          // requests above this 400
+    "memory_budget_bytes": null,    // explicit budget, or null = machine default
+    "usable_bytes": 0, "weights_bytes": 0
+  },
+  "batch": {
+    "configured": 1,                // the --batch N value
+    "batched": false,               // batching enabled (N>1) for this server
+    "active_rows": 0                // rows currently decoding in the batch
+  }
+}
+```
+
+`batch.batched` reflects only whether `--batch N` (N>1) is configured;
+with an explicit `--kv-quant` it can read `true` while `active_rows`
+never exceeds 1 (every request routes serial). `active_rows` is the
+honest signal of whether anything is actually batching.

package/docs/segmented-backward-handoff.md

@@ -0,0 +1,148 @@
+# Handoff — segmented backward training (2026-06-16)
+
+Current-state report for the segmented-backward LoRA training work. Deep dossier
+(design, proofs, every measurement): [`docs/design/segmented-backward-training.md`](design/segmented-backward-training.md)
+§9 (MiniCPM5) and §10 (e4b). Branch: `segmented-backward-training`,
+PR [#9](https://github.com/joshuarossi/mlx-bun/pull/9).
+
+## TL;DR
+
+Long-context LoRA SFT that **streams the backward segment-by-segment** so only one
+segment's activations are live at a time — training fits at context lengths where
+the full backward spikes. Two phases, both working:
+
+- **MiniCPM5 (Phase A): done, production-ready.** Bit-exact mechanism, no memory
+  leak, trains end-to-end. Real 300-iter chunk fine-tune → **chunk-eval 95.10/100**
+  (beats the 91.70 non-segmented baseline; pre-tune 11.89).
+- **e4b / gemma-3n (Phase B): built + validated.** Forward bit-exact; grads
+  bf16-class; trains **all 42 layers at 8K context** where `mlx_lm.lora` OOMs.
+
+**Parity:** our **non-segmented** training is **bit-exact with mlx-lm** at 2k
+(forward, loss, AND gradient — verified by running mlx-lm's own `default_loss` /
+`nn.value_and_grad`). The **segmented** path is **bf16-class (~2%)** off, a
+grouping-controllable bf16 non-associativity (not a logic bug).
+
+## What was built
+
+| file | what |
+|---|---|
+| `src/train/segmented.ts` | `SegmentedBackward` (MiniCPM5) + `SegmentedBackwardGemma4` (e4b) + `planSegmentsBySize`; the per-segment vjp driver, `detachLeaf` |
+| `src/mlx/autograd.ts` | `Vjp` — `mlx_vjp` binding (the backward primitive; see "leak" below) |
+| `src/mlx/ffi.ts` | `mlx_vjp` FFI decl |
+| `src/mlx/array.ts` | `MlxArray.fromBytesCopy` — graph-free leaf copy (row-major) |
+| `src/model/minicpm5.ts` | `runLayerRange` (forwardLayers delegates) |
+| `src/model/gemma4.ts` | `runLayerRange` + `makeTrainingMasks` + `embedForSegmented` + `reusedDonors` (additive; forwardLayers untouched) |
+| `src/train/trainer.ts` | `TrainConfig.segmentSize`; gated segmented path (MiniCPM5 + Gemma4) |
+| `src/train/loss.ts` | exported `responseOnlyCe`, `maskedCe` |
+
+**Enable:** set `segmentSize` (layers per segment) > 0 in the train config, or
+`SEG=n` in `scripts/chunk-finetune.ts`. 0 = the normal full backward.
+
+How it works (B=1 SFT): forward saves a DETACHED boundary (graph-free leaf) at each
+segment edge → a loss-head `mlx_vjp` gives `dh` for the last boundary → reverse
+per-segment `mlx_vjp`s thread `dh` and the segment's LoRA params, accumulating
+grads. Only one segment's activations are live. e4b adds two boundary streams: the
+per-layer-input tensor (a pure constant) and the **KV-shared donor K/V** (a second
+boundary stream WITH cotangent accumulation — the intricate part; docs §10).
+
+## What was validated (honest numbers)
+
+**Correctness**
+- MiniCPM5 segmented grads vs the full value_and_grad: **bit-exact under flash**
+  (relNorm 0.0000%). Under `ops.sdpa`, ~6% — a *pre-existing* `ops.sdpa`
+  eager-vs-autograd forward divergence, not segmentation.
+- e4b: forward **bit-exact** vs full; single-consumer donor reuse **bit-exact**;
+  natural-cut donor-KV grads **~1% (bf16 non-associativity)** — single-consumer
+  exact, flat in consumer count, tracks summand count (donor 22's 15 sharers
+  ≈0.46% vs donor 23's 3 ≈0%), grouping-dependent for both bf16 AND fp32. fp32
+  accumulation does NOT help (the term is mlx's within-vjp bf16 sum). **Controllable
+  by the segment cut** (coarser sharer grouping → tighter, down to bit-exact at
+  1 consumer), not the dtype.
+- **Parity vs mlx-lm (ran their code):** forward/loss — denominator (`M` vs
+  `ntoks`) bit-exact across 6 mask spans incl. edges; response-only CE == our
+  full-masked CE (a pure memory optimization); loss bit-exact on 4/6 spans, 0.5
+  float32-ulp on 2/6 (a logsumexp/reduction f32 effect, ~5e-5%, harmless).
+  Gradient — our FULL (non-segmented) `dB` over all 42 q_proj **bit-exact** with
+  mlx-lm; segmented `dB` ~2.3% (the bf16 donor-KV effect).
+
+**Memory leak (found + fixed).** The first cut used a surrogate-loss
+`value_and_grad`; it leaked ~32 MB/segment/step (mlx-level, not GC/cache/sync
+reclaimable). **Fixed by using `mlx_vjp`** (the natural primitive — no surrogate).
+End-to-end now flat (active constant across a 300-iter run). Repro:
+`scripts/vag-leak-test.ts` (MODE=layers leaks via value_and_grad; MODE=vjp flat).
+
+**Memory vs the real reference** (`mlx_lm.lora --grad-checkpoint`, the way a user
+trains — it loads the OptiQ e4b directly). mlx-lm checkpoints **per-layer**; ours
+groups `segmentSize` layers. Matched grouping (both per-layer = segSize 1, all 42
+layers, batch 1):
+
+| L | mlx-lm (per-layer) | ours (segSize=1) |
+|---|---|---|
+| 2048 | 12.84 GB | 8.76 GB |
+| 4096 | 20.87 GB | 10.93 GB |
+| 8192 | **OOM** (25.7 GB @16 layers) | 15.29 GB |
+
+Ours lower at every length AND trains all 42 layers at 8K where mlx-lm OOMs. **Two
+confounds, not isolated:** (1) mlx-lm materializes the full `[1,L,262k]` logits
+(~4 GB @4K) while we use response-only CE; (2) our full backward streaming vs
+mlx-lm's per-layer checkpoint. So the gap is part CE-optimization, part mechanism —
+not purely "segmentation beats checkpointing." (Earlier "reference crashes at 4K /
+~70 GB" was WRONG — it used mlx-bun's *own* ineffective checkpoint as the baseline.)
+
+## How to run an adapter
+
+The fine-tune saves a hot-swap adapter dir (`adapters.safetensors` +
+`optiq_lora_config.json` + `adapter_config.json`).
+
+**One-off generate** (verified working) — `scripts/run-adapter.ts`:
+```sh
+ADAPTER=~/.cache/mlx-bun-finetunes/minicpm5-chunk-segmented \
+  PROMPT="Split into chunks and return JSON: ..." \
+  MAXTOK=256 bun scripts/run-adapter.ts
+# omit ADAPTER for the base model; MODEL=<dir> to change base; SYSTEM=... TEMP=...
+```
+
+**Score it on the chunk task:**
+```sh
+ADAPTER=~/.cache/mlx-bun-finetunes/minicpm5-chunk-segmented bun scripts/chunk-eval.ts
+```
+
+**Serve it** (you run the server; mount via the API):
+```sh
+# 1) start: mlx-bun serve <model-dir>
+# 2) mount:  curl -X POST localhost:PORT/v1/adapters -d '{"id":"chunk","path":"<adapter-dir>"}'
+# 3) list:   curl localhost:PORT/v1/adapters     # unmount: DELETE /v1/adapters/chunk
+```
+The adapter activates per request through `loraState`; in code the pattern is
+`new AdapterManager(model).mount(id, dir)` then `model.loraState.active = [id]`.
+
+**Train a new adapter** (segmented):
+```sh
+# MiniCPM5, seq 4096, segments of 4 layers:
+MLX_BUN_PERF_KERNEL=0 MLX_BUN_FUSED_GELU=0 \
+  SEQ=4096 ITERS=300 SEG=4 DATA=<data-dir> ADAPTER=<out-dir> \
+  bun scripts/chunk-finetune.ts
+# e4b: add MODEL=<e4b-dir> and use SEG=2 (or 1) for ≥8K to keep peak down.
+```
+Training REQUIRES `MLX_BUN_PERF_KERNEL=0 MLX_BUN_FUSED_GELU=0` (those fused kernels
+have no vjp). Use `MLX_BUN_TRAIN_ATTN=flash` for exact grads (slower).
+
+## What's next / open
+
+1. **e4b @8K toward ~10 GB** (currently 15.3 GB segSize=1): the §5 full-attention
+   ISOLATION `planSegments` — put each O(L²) full-attn layer (5,11,17,23,29,35,41)
+   in its own short segment. Mechanism + tuning knob are proven; this is peak-shaving.
+2. **Isolate the memory confound** (optional): run ours at segSize=1 with full-logits
+   CE to separate the response-only-CE saving from the streaming mechanism.
+3. **mlx-bun's own gradient checkpoint is ineffective** (23 GB @2048 vs mlx-lm's
+   12.8) — orthogonal bug; segmented sidesteps it.
+4. Scale knobs: `segmentSize` is the peak↔grad-fidelity knob (finer = lower peak,
+   more donor-KV bf16 deviation).
+
+## Diagnostic scripts (kept)
+
+`scripts/segmented-grad-test.ts` (MiniCPM5 grads + peak + leak loop),
+`segmented-grad-test-e4b.ts` (e4b grads + peak; RANGES/SEG_ONLY/SEG_DUMP),
+`parity-vs-mlxlm.ts` + `grad-parity-vs-mlxlm.ts` (mlx-lm parity; Python sides in
+`/tmp/mlxlm-*.py`), `vag-leak-test.ts`, `seg-debug.ts`, `gemma4-runlayerrange-test.ts`,
+`run-adapter.ts`, `chunk-{finetune,eval}.ts`.