Skip to content

v0.0.3

Choose a tag to compare

View all tags
@ARahim3 ARahim3 released this 30 Jun 22:16
v0.0.3
1d57b95

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

Assets 2
Loading