Skip to content

Releases: ARahim3/mlx-dspark

v0.2.0

Choose a tag to compare

@ARahim3 ARahim3 released this 04 Jul 21:24

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 > 0 dspark — 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-drafts to 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 at temperature > 0 too).
  • 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_dspark crashed on transformers ≥5.13 — which fresh installs now resolve to — with AttributeError: 'str' object has no attribute '__module__' (mlx_lm's string-keyed tokenizer registration hitting transformers 5.13's stricter register, at import time). Fixed with a scoped, version-agnostic import-time compat shim (no transformers pin). 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-tokens 2048 / --max-tokens-cap 32768 (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-dspark

No 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

Choose a tag to compare

@ARahim3 ARahim3 released this 02 Jul 18:59

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 servermlx-dspark serve --model <repo>http://127.0.0.1:8080/v1. Point any OpenAI client (LM Studio, the openai SDK, 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 / baseline on one target. temperature, top_p, top_k, max_tokens, stop, seed, optional --api-key, CORS. Each response carries an x_mlx_dspark block (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 as delta.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 models lists targets with a known drafter.
  • Thinking toggle — per-request enable_thinking / chat_template_kwargs and a server --no-thinking default (silences Qwen3 <think> blocks for a served endpoint).
  • Subcommand CLI + testsserve / generate / models / doctor, a mlx-dspark console-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/v1
from 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_generate gained prompt_ids=, stop=, top_p= / top_k=, cache= / ctx_caches= / reuse_len= (prefix reuse), and a finish_reason on GenResult.
  • No breakage: --family / --target / load_pair("qwen3") are kept as deprecated aliases, and the old flat python -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

Choose a tag to compare

@ARahim3 ARahim3 released this 30 Jun 22:16

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 0 runs 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 > 0 is 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 from mlx_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

Choose a tag to compare

@ARahim3 ARahim3 released this 30 Jun 06:21

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 prob min(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 2 is 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_tokens 4 → 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_generate gains temperature / seed (so --mode baseline can 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

v0.0.1

Choose a tag to compare

@ARahim3 ARahim3 released this 29 Jun 21:00

First release: Gemma-4 12B + Qwen3-4B DSpark on MLX.