DFlash2 quantized-draft loader + M5 Max measurements

README.md

@@ -0,0 +1,121 @@
+---
+license: mit
+library_name: mlx
+tags:
+- mlx
+- speculative-decoding
+- dflash2
+- apple-silicon
+- quantization
+- benchmark
+base_model:
+- Qwen/Qwen3.8-27B
+---
+
+# DFlash2 on Apple Silicon — quantized-draft loader + M5 Max measurements
+
+Two things this repo provides, neither of which is model weights:
+
+1. **`load_draft_q4.py`** — a drop-in loader that lets the
+   [`cfontes/qwen3.8-27b-mlx-dflash2`](https://huggingface.co/cfontes/qwen3.8-27b-mlx-dflash2)
+   server and benchmarks use a **quantized** DFlash2 drafter. The upstream loader builds a
+   bf16 skeleton and calls `load_weights` directly, so a 4-bit draft fails on shape
+   mismatch.
+2. **`config.json`** — the draft config that
+   [`ProCreations/Qwen3.8-27B-DFlash2-MLXFast-Q4`](https://huggingface.co/ProCreations/Qwen3.8-27B-DFlash2-MLXFast-Q4)
+   ships without. That repo contains only `model.safetensors`, so no standard loader can
+   instantiate it.
+
+Plus the measurements that motivated both.
+
+## Headline: quantizing the *drafter* is close to free throughput
+
+Apple **M5 Max**, 137 GB, macOS 26.4.1, MLX 0.32.0 / mlx-lm 0.31.3. Target:
+`tozp/Qwen3.8-27B-OBLITERATED-V2-mlx-4bit`. Greedy, 300 tokens, K=4, interleaved A/B with a
+90 s cooldown before each measurement, two rounds.
+
+| drafter | size | tok/s | accept |
+|---|---:|---:|---:|
+| `z-lab/Qwen3.8-27B-DFlash2` (bf16) | 3.85 GB | 59.8 / 60.3 | 87.9% |
+| `ProCreations/...-MLXFast-Q4` (affine-4 g64) | **1.27 GB** | **76.1 / 76.0** | 86.9% |
+| | −2.6 GB | **+26%** | −1.0 pp |
+
+The drafter runs a forward pass every block, so shrinking it by 2.6 GB removes real
+bandwidth from the hot loop. The prediction quality it costs — one point of acceptance — is
+nowhere near enough to offset that.
+
+Stacked against no speculation at all, on the same target and machine:
+
+| configuration | tok/s | vs dense |
+|---|---:|---:|
+| dense (`mlx_lm.stream_generate`) | 33.3 | 1.00× |
+| + bf16 drafter | 61.4 | 1.84× |
+| **+ Q4 drafter** | **76.1** | **2.29×** |
+
+A 27B **dense** model at 76 tok/s in ~15.5 GB.
+
+## Target-model comparison (same protocol)
+
+| target | dense | DFlash2 (bf16 draft) | speedup | accept | greedy token-exact vs dense |
+|---|---:|---:|---:|---:|---|
+| `mlx-community/Qwen3.8-27B-mxfp4` | 35.4 | 54.0 | 1.52× | 82% | yes |
+| Qwen3.8-27B abliterated w/ [Heretic](https://github.com/p-e-w/heretic) | 33.3 | 62.2 | 1.87× | 89% | one near-tie flip |
+| `tozp/...-OBLITERATED-V2` | 33.3 | 61.4 | 1.84× | 88% | yes |
+
+Round-to-round spread under 2% on every cell.
+
+**The drafter transfers across quantization schemes and weight surgery.** It is trained
+against the stock model; both alternatives are abliterated by unrelated methods (Optuna
+search vs SVD+LEACE blend) *and* quantized differently (affine-4 g64 vs mxfp4 g32).
+Acceptance went **up** 6–7 points rather than down. Two independent methods landing in the
+same place suggests the cause is removing refusal directions as such — plausibly lower
+next-token entropy, which is testable and untested here.
+
+## Measurement traps on this hardware
+
+- **Thermal drift.** Running targets back-to-back without cooldowns, the same mxfp4
+  measurement read 34.9 t/s at the start and 26.8 t/s three minutes later — **−23%**, enough
+  to invert a model-vs-model conclusion. Interleaving with 90 s cooldowns took spread from
+  23% to under 2%.
+- **Cold start.** The first model load of a session measures far below steady state; one run
+  produced `dense 1.9 t/s / speedup 18.00x`. Warm up before timing.
+- **`ioreg` "Device Utilization %" is useless here** — it reads 100% at rest. Check GPU
+  memory in use and whether other engines hold models instead. A 37 GB workload on the GPU
+  went undetected this way and produced a full set of wrong numbers.
+
+## Usage
+
+```bash
+git clone https://huggingface.co/cfontes/qwen3.8-27b-mlx-dflash2 dflash2
+cd dflash2
+hf download ProCreations/Qwen3.8-27B-DFlash2-MLXFast-Q4 --local-dir models/draft-q4
+cp /path/to/this/repo/config.json          models/draft-q4/config.json
+cp /path/to/this/repo/load_draft_q4.py     bench/extra/load_draft_q4.py
+
+# point the server at the quantized draft
+python dflash2_mlx_server/server.py --main-dir models/<target> --draft-dir models/draft-q4
+```
+
+`server.py` calls `dflash_port.local_load.load_draft_from_dir`; swap that import for
+`load_draft_any` from `load_draft_q4.py`. The loader detects a `quantization` block and
+falls back to the original function when there isn't one, so bf16 drafts keep working.
+
+Verified end-to-end: server loads in 0.9 s and reports `accept_rate 0.917, tok_per_s 75.4`.
+
+## How the loader works
+
+`nn.quantize` is applied to the draft skeleton **before** `load_weights`, with a predicate
+that quantizes exactly the modules that have a `.scales` entry in the weight file — the same
+test `mlx_lm` uses. It also normalizes the candidate-selector codebook keys, which the two
+published drafts spell differently (`candidate_selector.predecessor_codebook` vs
+`...codebook.weight`); the upstream rename raises `KeyError` on the Q4 layout.
+
+## Credits
+
+- [z-lab](https://huggingface.co/z-lab/Qwen3.8-27B-DFlash2) — the DFlash2 drafter
+- [cfontes](https://huggingface.co/cfontes/qwen3.8-27b-mlx-dflash2) — the MLX port and server this patches (MIT)
+- [ProCreations](https://huggingface.co/ProCreations/Qwen3.8-27B-DFlash2-MLXFast-Q4) — the affine-4 quantization measured here
+- [tozp](https://huggingface.co/tozp/Qwen3.8-27B-OBLITERATED-V2-mlx-4bit) and [Heretic](https://github.com/p-e-w/heretic) — the abliterated targets
+
+No model weights are redistributed here. `config.json` is derived from z-lab's draft config
+with a quantization block added.

config.json

@@ -0,0 +1,65 @@
+{
+ "architectures": [
+  "DFlash2DraftModel"
+ ],
+ "attention_bias": false,
+ "attention_dropout": 0.0,
+ "bos_token_id": null,
+ "is_causal": false,
+ "dflash_config": {
+  "block_size": 8,
+  "conv_group_size": 16,
+  "conv_kernel_size": 2,
+  "mask_token_id": 248070,
+  "selector_rank": 256,
+  "selector_top_k": 16,
+  "target_layer_ids": [
+   5,
+   19,
+   33,
+   47,
+   61
+  ]
+ },
+ "dtype": "bfloat16",
+ "eos_token_id": 248044,
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 5120,
+ "initializer_range": 0.02,
+ "intermediate_size": 17408,
+ "layer_types": [
+  "sliding_attention",
+  "sliding_attention",
+  "sliding_attention",
+  "sliding_attention",
+  "sliding_attention"
+ ],
+ "max_position_embeddings": 262144,
+ "max_window_layers": 5,
+ "model_type": "qwen3",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 5,
+ "num_key_value_heads": 8,
+ "num_target_layers": 64,
+ "pad_token_id": 248044,
+ "rms_norm_eps": 1e-06,
+ "rope_parameters": {
+  "rope_theta": 10000000,
+  "rope_type": "default"
+ },
+ "sliding_window": 2048,
+ "tie_word_embeddings": false,
+ "transformers_version": "5.15.0",
+ "use_cache": true,
+ "use_sliding_window": true,
+ "vocab_size": 248320,
+ "quantization": {
+  "group_size": 64,
+  "bits": 4
+ },
+ "quantization_config": {
+  "group_size": 64,
+  "bits": 4
+ }
+}

load_draft_q4.py

@@ -0,0 +1,56 @@
+"""支援量化 draft 的載入器（原 local_load 只支援 bf16）。"""
+import json
+from pathlib import Path
+import mlx.core as mx
+import mlx.nn as nn
+from dflash_port.model_mlx import DFlash2DraftModel, DFlashDraftModel, DFlashConfig
+from dflash_port import local_load as _ll
+
+
+def load_draft_any(path):
+    path = Path(path)
+    cfg = json.loads((path / "config.json").read_text())
+    weights = {k: v for f in path.glob("*.safetensors") for k, v in mx.load(str(f)).items()}
+
+    # codebook 鍵名在兩種發佈裡不一致，統一成 .weight 結尾
+    for name in ("predecessor_codebook", "successor_codebook"):
+        bare = f"candidate_selector.{name}"
+        if bare in weights and f"{bare}.weight" not in weights:
+            weights[f"{bare}.weight"] = weights.pop(bare)
+
+    q = cfg.get("quantization") or {}
+    if not q:                                   # 非量化 → 走原本的路徑
+        return _ll.load_draft_from_dir(path)
+
+    # 重用原載入器的 config 組裝（複製其邏輯的關鍵欄位）
+    d = cfg.get("dflash_config", {})
+    rope = cfg.get("rope_parameters") or cfg.get("rope_scaling")
+    layer_types = tuple(cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"])
+    config = DFlashConfig(
+        hidden_size=cfg["hidden_size"], num_hidden_layers=cfg["num_hidden_layers"],
+        num_attention_heads=cfg["num_attention_heads"], num_key_value_heads=cfg["num_key_value_heads"],
+        head_dim=cfg["head_dim"], intermediate_size=cfg["intermediate_size"],
+        vocab_size=cfg["vocab_size"], rms_norm_eps=cfg["rms_norm_eps"],
+        rope_theta=cfg.get("rope_theta", (rope or {}).get("rope_theta", 10000.0)),
+        max_position_embeddings=cfg["max_position_embeddings"],
+        block_size=int(d.get("block_size", 16)), target_layer_ids=tuple(d["target_layer_ids"]),
+        num_target_layers=cfg["num_target_layers"], mask_token_id=d["mask_token_id"],
+        rope_scaling=rope, layer_types=layer_types, sliding_window=cfg.get("sliding_window"),
+        final_logit_softcapping=d.get("final_logit_softcapping", cfg.get("final_logit_softcapping")),
+        input_embedding_scale=float(d.get("input_embedding_scale", 1.0)),
+        output_multiplier=float(d.get("output_multiplier", 1.0)),
+        conv_kernel_size=int(d.get("conv_kernel_size", 0)), conv_group_size=int(d.get("conv_group_size", 0)),
+        selector_rank=int(d.get("selector_rank", 0)), selector_top_k=int(d.get("selector_top_k", 0)),
+        is_causal=cfg.get("is_causal"),
+    )
+    klass = DFlash2DraftModel if "DFlash2DraftModel" in (cfg.get("architectures") or []) else DFlashDraftModel
+    model = klass(config)
+
+    # 只量化「權重檔裡真的有 scales」的模組 —— 與 mlx-lm 的判斷方式一致
+    qset = {k[:-len(".scales")] for k in weights if k.endswith(".scales")}
+    nn.quantize(model, group_size=int(q["group_size"]), bits=int(q["bits"]),
+                class_predicate=lambda p, m: p in qset)
+    model.eval()
+    model.load_weights(list(weights.items()))
+    mx.eval(model.parameters())
+    return model

results.json

@@ -0,0 +1,94 @@
+{
+ "machine": "Apple M5 Max, 137 GB, macOS 26.4.1",
+ "stack": {
+  "mlx": "0.32.0",
+  "mlx_lm": "0.31.3"
+ },
+ "protocol": "interleaved A/B, 90s cooldown before each measurement, greedy, 300 tokens, K=4, GPU verified idle",
+ "draft_comparison": {
+  "target": "tozp/Qwen3.8-27B-OBLITERATED-V2-mlx-4bit",
+  "bf16_draft": {
+   "repo": "z-lab/Qwen3.8-27B-DFlash2",
+   "size_gb": 3.85,
+   "tok_per_s": [
+    59.8,
+    60.3
+   ],
+   "accept_pct": [
+    87.9,
+    87.9
+   ]
+  },
+  "q4_draft": {
+   "repo": "ProCreations/Qwen3.8-27B-DFlash2-MLXFast-Q4",
+   "size_gb": 1.27,
+   "tok_per_s": [
+    76.1,
+    76.0
+   ],
+   "accept_pct": [
+    86.9,
+    86.9
+   ]
+  },
+  "delta": {
+   "tok_per_s_pct": 26.4,
+   "accept_pp": -1.0,
+   "memory_gb": -2.58
+  }
+ },
+ "target_comparison_bf16_draft": {
+  "mxfp4": {
+   "dense": [
+    35.4,
+    35.0,
+    35.6
+   ],
+   "dflash": [
+    54.0,
+    54.1,
+    53.8
+   ],
+   "accept_pct": 82,
+   "token_exact": true
+  },
+  "heretic_ara": {
+   "dense": [
+    33.0,
+    33.3,
+    33.5
+   ],
+   "dflash": [
+    61.7,
+    62.2,
+    62.3
+   ],
+   "accept_pct": 89,
+   "token_exact": false
+  },
+  "obliterated_v2": {
+   "dense": [
+    33.0,
+    33.3,
+    33.6
+   ],
+   "dflash": [
+    61.2,
+    61.4,
+    61.6
+   ],
+   "accept_pct": 88,
+   "token_exact": true
+  }
+ },
+ "server_verification": {
+  "load_s": 0.9,
+  "accept_rate": 0.917,
+  "tok_per_s": 75.4
+ },
+ "thermal_drift_uncontrolled": {
+  "same_measurement_start": 34.9,
+  "after_3min": 26.8,
+  "pct": -23
+ }
+}