Releases: ARahim3/mlx-dspark
Release list
v0.2.0
v0.2.0 — continuous batching, auto-calibration, prompt-lookup, penalties & logprobs
mlx-dspark now scales past single-user: serve concurrent requests in one batched pass, speculate on any model with a new drafter-free prompt-lookup mode (for the models without a matching DSpark/DFlash drafter), and let it tune the draft length to your Mac — plus OpenAI penalties, logprobs, and a fix for serving Gemma-4.
Highlights
- Continuous batching (
serve --max-batch N) — run up to N concurrently-queued requests through one batched target forward so they share a single weight-read per step (ideal for a local agent swarm). Both the target verify and the DSpark drafter are batched; a lone request — or one using penalties / logprobs /temperature > 0dspark — takes the serial path, so B=1 latency never regresses. Measured Qwen3-4B-8bit on an M4 Pro, 4 concurrent: baseline B=4 2.48× aggregate, batched dspark B=4 2.51× vs serialized baseline (130 tok/s; 1.73× over serialized dspark). Dense mlx-lm targets (Qwen3 / Llama / Mistral-class); Gemma-4 (mlx-vlm) falls back to serialized. --mode auto— picks the best available speculation for any target (a known DSpark drafter → else DFlash → else drafter-free n-gram lookup), so any repo serves with some speedup and no extra flags.--mode lookup+ hybrid drafting — prompt-lookup (n-gram) speculation with no drafter at all, for any model; great when output copies its own context (RAG quotes, code edits, repeat/refine turns). Hybrid drafting (dspark, on by default) verifies a free n-gram continuation instead of running the drafter on rounds where the suffix already occurred, reaching ~2.4× on copy-heavy content while general chat stays ≈neutral (--no-lookup-draftsto disable).--max-draft auto— measures this machine + model's verify/drafter cost curves once (a few seconds, cached on disk) and picks the cap per round from the curves + a live acceptance estimate, so the cap tracks the hardware (M1→M5) instead of a hard-coded default. Lossless — the cap only sets how many drafts get verified.presence_penalty/frequency_penalty— OpenAI penalties applied to the target logits so speculative/greedy output equals sequential decoding of the penalized target (works attemperature > 0too).logprobs/top_logprobs— chosen + top-k target log-probabilities per token, for both chat and completions, computed only when requested.mlx-dspark benchmark— a warm, device-stamped, reproducible sweep (--json) for the community M1→M5 matrix.
Fixes & performance
import mlx_dsparkcrashed on transformers ≥5.13 — which fresh installs now resolve to — withAttributeError: 'str' object has no attribute '__module__'(mlx_lm's string-keyed tokenizer registration hitting transformers 5.13's stricterregister, at import time). Fixed with a scoped, version-agnostic import-time compat shim (notransformerspin). Thanks to @zboyles for the report (#1).- Serving Gemma-4 (mlx-vlm targets) was broken since 0.1.0 — every request failed with
There is no Stream(gpu, 1) in current thread(mlx-vlm's load switches the loading thread's default stream). The engine now loads and generates on one thread; Gemma multi-turn also gets prefix caching now (reused while under the sliding window). - Decode-path perf pass — incremental streaming detokenization (was O(n²) on long / thinking outputs), one device sync per spec round, and a pipelined baseline: Qwen3-4B baseline now matches
mlx_lm.generate(~52 tok/s), dspark ~1.46×. - Server polish — sampling defaults from the model's
generation_config.json,--default-max-tokens2048 /--max-tokens-cap32768 (was 512/8192, which truncated thinking traces), a client disconnect keeps the prefix cache, LRU prefix-cache slots, and chunked prefill + wired-memory limit for small Macs.
Upgrade
pip install -U mlx-dspark # or: uv pip install -U mlx-dsparkNo breaking changes — every new behavior is opt-in or a transparent default; existing --model / --mode dspark|dflash|baseline usage is unchanged.
Speedup summary (all lossless)
| scenario | speedup vs baseline |
|---|---|
| single request, general content (dspark) | ~1.4× (Qwen3-4B) … ~1.75× (Gemma-4 12B) |
| single request, copy-heavy content (hybrid / lookup) | up to ~2.4× |
concurrent, 4 requests (--max-batch 4) |
~2.5× aggregate |
| multi-turn follow-up (prefix caching, long shared context) | ~13× faster follow-up turns |
Full diff: v0.1.0...v0.2.0
v0.1.0
v0.1.0 — serve it: OpenAI-compatible API, tool calls, prefix caching
mlx-dspark grew up from a library + demo CLI into a usable local tool. You can now serve a DSpark / DFlash model to LM Studio, the openai SDK, or any OpenAI-compatible client — with the speculative speedup transparent behind the API. Everything stays lossless (greedy is byte-identical to plain decoding up to fp ties; --temperature > 0 is an exact sample from the target at T). No heavy dependency added — the server is stdlib-only.
Highlights
- OpenAI-compatible API server —
mlx-dspark serve --model <repo>→http://127.0.0.1:8080/v1. Point any OpenAI client (LM Studio, theopenaiSDK, curl, LangChain) at it:POST /v1/chat/completions(streaming SSE and non-stream, multi-turn),POST /v1/completions,GET /v1/models,GET /health,GET /metrics.- Serves
dspark/dflash/baselineon one target.temperature,top_p,top_k,max_tokens,stop,seed, optional--api-key, CORS. Each response carries anx_mlx_dsparkblock (accept length + tok/s) so the spec-decode gain is visible.
- Prefix caching (in-memory + optional SSD spill) — reuse the shared conversation prefix's KV across turns instead of re-prefilling it. ~13× faster follow-up turns on a ~750-token shared context (measured: 87 ms vs 1132 ms). Lossless to the same fp-tie standard; invalidated on any error so it can't desync. On by default for dspark/baseline on dense targets (Qwen3); falls back for Gemma-4's rotating caches and DFlash.
- Tool calling — OpenAI
tools/tool_calls, parsed from both native formats (Qwen3 Hermes-JSON and Gemma-4's<|tool_call>call:…), streamed asdelta.tool_calls, with a full request → tool-call → result → answer round-trip. - Lossless top-p / top-k sampling — nucleus / top-k truncation applied to both the draft and the target, so temperature sampling stays an exact sample from the (truncated) target. Validated model-free.
- Model-centric interface — name the target with
--model <hf-repo | local-path>(like mlx-lm); the matched drafter auto-resolves (quantization-agnostic), or pass--drafter. This makes Qwen3-8B a first-class target and replaces the old 2-value--family.mlx-dspark modelslists targets with a known drafter. - Thinking toggle — per-request
enable_thinking/chat_template_kwargsand a server--no-thinkingdefault (silences Qwen3<think>blocks for a served endpoint). - Subcommand CLI + tests —
serve/generate/models/doctor, amlx-dsparkconsole-script entry point, and a 45-test model-free suite (server protocol, streaming, stop sequences, tool-call parsing, top-p losslessness, prefix-cache manager, drafter resolver).
Quickstart
pip install -U mlx-dspark
mlx-dspark serve --model mlx-community/Qwen3-8B-8bit # → http://127.0.0.1:8080/v1from openai import OpenAI
c = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="not-needed")
print(c.chat.completions.create(model="Qwen3-8B-8bit",
messages=[{"role": "user", "content": "Explain rainbows briefly."}]).choices[0].message.content)API
- New:
mlx_dspark.server(Engine,run_server),encode_messages(multi-turn),resolve+REGISTRY(target→drafter),mlx_dspark.tools,mlx_dspark.sampling,mlx_dspark.prefix_cache. speculative_generate/dflash_generate/greedy_generategainedprompt_ids=,stop=,top_p=/top_k=,cache=/ctx_caches=/reuse_len=(prefix reuse), and afinish_reasononGenResult.- No breakage:
--family/--target/load_pair("qwen3")are kept as deprecated aliases, and the old flatpython -m mlx_dspark --prompt …form still works.
Notes
- Prefix caching is exact for dense trimmable KV caches (Qwen3). Gemma-4's sliding-window caches and the DFlash drafter cache can't be safely rolled back to an arbitrary prefix, so they fall back to a fresh prefill (correct, just no reuse).
- MoE / linear-attention targets (
Qwen3.5-*,gpt-oss-*) still need the gated-delta KV rollback that isn't wired yet — PRs welcome. - The benchmark numbers from v0.0.3 (DSpark vs DFlash head-to-head) are unchanged and still reproducible.
Full diff: v0.0.3...v0.1.0
v0.0.3
v0.0.3 — z-lab DFlash support + DSpark-vs-DFlash head-to-head
mlx-dspark now runs z-lab's original DFlash drafters (block diffusion) alongside DeepSeek's DSpark — under the same lossless verify loop — so you can benchmark the two on one target/Mac. Output stays lossless on both paths (greedy is byte-identical to plain decoding up to fp ties; --temperature > 0 is an exact sample from the target at temperature T).
Highlights
- Run z-lab DFlash natively (
--mode dflash): block-diffusion drafter that denoises a whole 16-token block in one parallel pass and reuses the target's embed + lm-head (so the checkpoints are tiny — Gemma-4 is 1.45 GB). Loads z-lab's published checkpoints as-is. Presets:gemma4(z-lab/gemma4-12B-it-DFlash),qwen3(z-lab/Qwen3-4B-DFlash-b16). Any other z-lab adapter loads the same way:load_dflash(repo)+load_target(matched_target)+drafter.bind(target.model)(e.g.Qwen3-8B-DFlash-b16).--max-draft 0runs the full 16-block (DFlash's native operating point — strongest on code/math). - DSpark vs DFlash, measured on-device (the part nobody else has — same target, same Mac, one lossless loop): on Gemma-4 12B, DFlash's block-16 wins code/math (accept ~6.0, ~2.1×); DSpark's Markov head wins open chat (1.65×). They're complementary, exactly as the paper frames it.
- Temperature speculative sampling for DFlash too —
--temperature > 0is lossless wrt the target at T, via the same acceptance rule as DSpark.
API
- New:
load_dflash,load_dflash_pair,dflash_generate,DFLASH_PRESETS,DFlashDraftModel,DFlashConfig(all exported frommlx_dspark). - New CLI mode:
python -m mlx_dspark --mode dflash [--family gemma4|qwen3] [--max-draft 0] [--temperature T]. - The DSpark path (
load_pair,speculative_generate,--mode dspark) is unchanged — no breakage.
Head-to-head (gemma-4-12B-it-8bit, M4 Pro, warm, greedy/lossless, 4 prompts/domain — accept / tok·s)
| method | chat | code | math |
|---|---|---|---|
| DSpark (cap 2) | 2.45 / 28.5 | 2.78 / 32.8 | 2.86 / 32.4 |
| DFlash (cap 2) | 2.15 / 24.2 | 2.76 / 31.3 | 2.71 / 29.6 |
| DFlash (full 16) | 2.68 / 16.9 | 5.95 / 36.6 | 6.20 / 36.3 |
Greedy baseline ≈ 17.3 tok/s. Qwen3-4B DFlash also runs (full-block accept ~3.3 on code) with a smaller speedup — its cheap verify leaves little to amortize.
Attribution
DFlash is by z-lab (Chen et al., DFlash: Block Diffusion for Flash Speculative Decoding, arXiv:2602.06036, MIT). The drafter model classes in src/mlx_dspark/dflash_model.py are vendored from z-lab/dflash under the MIT License (full notice in the file header + NOTICE); the verification loop is mlx-dspark's own.
Notes
- MoE / linear-attention DFlash targets (
Qwen3.5-*,gpt-oss-*) use a gated-delta KV rollback not wired in this release — they may need extra cache handling. PRs welcome. - Drafter quantization does not affect acceptance (4-bit stays the default for the DFlash backbone too).
Full diff: v0.0.2...v0.0.3
v0.0.2
v0.0.2 — speed + lossless sampling + an honest performance model
Faster drafting, a new sampling mode, and measured guidance on what to expect on Apple Silicon. Output stays lossless: greedy is byte-identical to plain greedy decoding, and --temperature > 0 is an exact sample from the target at temperature T.
Highlights
- Lossless temperature speculative sampling (
--temperature,--seed) — the paper's actual method (§2.1): sample the draft, accept with probmin(1, p/q), resample the residual on rejection. Validated that the output distribution exactly equals the target's. (Default stays greedy /temperature 0.) - ~9–10% faster at the default cap from a drafter-slice fix: the block's
lm_head+ Markov head now run only over the verified positions instead of the full 7 every round. Output-neutral. --max-draft 2is the new default — the measured optimum on Apple Silicon (the per-token verify cost makes longer blocks slower for typical acceptance). Was 4.- 4-bit target option (
--target mlx-community/gemma-4-12B-it-4bit) — highest absolute throughput and fits ≤24 GB Macs (at some quality cost). 8-bit remains the default sweet spot.
Changed
- Default
max_draft_tokens4 → 2. - Confidence-based truncation now uses the cumulative prefix-survival product
∏ c_i(paper Eq 7-8) instead of a per-position threshold. Kept as an option; not the default. greedy_generategainstemperature/seed(so--mode baselinecan sample for fair A/Bs).
Performance (M4 Pro, warm, vs the official mlx_lm / mlx_vlm tools)
| model | baseline | mlx-dspark | speedup |
|---|---|---|---|
| Gemma-4 12B | 18.4 tok/s (mlx_vlm) | ~30 tok/s | ~1.6× (up to ~2× on code/math) |
| Qwen3-4B | 52.9 tok/s (mlx_lm) | ~73 tok/s | ~1.4× |
This matches DSpark's own per-user serving claim (~1.57–1.85×). The README documents the
Apple-Silicon cost model / ceiling and an on-device DSpark-vs-DFlash comparison.
Notes
- Greedy decoding path is unchanged and remains greedy-correct.
- Drafter quantization does not affect acceptance (4/8-bit/bf16 identical); 4-bit stays default.
- No API breakage; new arguments are optional with greedy-preserving defaults.
Full diff: v0.0.1...v0.0.2