perf(speculative): quantize the MTP drafter's projections at load

docs/benchmark_results/mtp-drafter-quantization-m5max-2026-08-17.md

@@ -0,0 +1,107 @@
+# Quantizing the MTP drafter, Qwen 3.8 27B on M5 Max
+
+Issue #1185, Phase 3. The drafter ships bf16 and is read once per drafted
+token, so its cost is weight traffic. This measures what 4-bit affine
+quantization does to that cost, to acceptance, and to output.
+
+## Setup
+
+| Field | Value |
+|---|---|
+| Host | Apple M5 Max, 40 GPU cores, 128 GB, macOS 26.6.1, Xcode 26.5 |
+| Build | `cargo build --release --features metal,accelerate`, `main` at `1e7b1d13` |
+| Target | `models/qwen3.8-27b-4bit` (affine, group 64, 4-bit) |
+| Drafter | `models/qwen3.8-27b-mtp-bf16`, 810 MiB, 15 tensors |
+| Command | `mlxcel generate --draft-kind mtp --draft-block-size 3 --temp 0` |
+| Exactness gate | Passed with no override. The #1199 retry disables `qmv_wide` and the probe reports the verify block byte-identical to the single-token chain, so these are shipping-configuration numbers. |
+
+Load averages stayed between 1.5 and 2.1 across the run, and the arms were
+alternated, because this host loses throughput on repeated 27B work faster
+than it loses it to anything else.
+
+## What gets quantized
+
+The drafter is one decoder layer plus an `fc` projection and seven norms. It
+borrows `embed_tokens` and the LM head from the target, so those are not its
+to convert. Eight 2-D projections are, all with the contraction axis last and
+all divisible by the group size:
+
+| tensor | shape |
+|---|---|
+| `fc.weight` | 5120 x 10240 |
+| `layers.0.mlp.{gate,up}_proj.weight` | 17408 x 5120 |
+| `layers.0.mlp.down_proj.weight` | 5120 x 17408 |
+| `layers.0.self_attn.q_proj.weight` | 12288 x 5120 |
+| `layers.0.self_attn.{k,v}_proj.weight` | 1024 x 5120 |
+| `layers.0.self_attn.o_proj.weight` | 5120 x 6144 |
+
+810.0 MiB becomes 227.9 MiB, 0.281x, including scales and biases. The seven
+1-D norms stay bf16.
+
+## Result
+
+Two reps per arm at each length, alternated.
+
+| | bf16 drafter | 4-bit drafter |
+|---|---:|---:|
+| `draft_block` | 10.40, 10.71 ms/round | **2.70, 2.69** |
+| accept hook | 10.35, 10.60 ms/round | **2.74, 2.71** |
+| verify forward | 40.08, 43.50 ms/round | 40.58, 40.05 |
+| acceptance, n=120 | 0.6831 | 0.6601 |
+| acceptance, n=300 | 0.6500 | 0.6589 |
+| tok/s, n=120 | 37.90, 35.57 | **49.71, 50.33** |
+| tok/s, n=300 | 35.24, 34.76 | **47.68, 45.68** |
+
+**The drafter step is 3.9x cheaper and the verify forward is untouched**, which
+is the shape the cost model predicted: the drafter is memory-bound and the
+target is not affected by anything the drafter does.
+
+**Acceptance does not move.** It falls 3.4% at 120 tokens and rises 1.4% at
+300. Two lengths disagreeing in sign is what noise looks like; a real
+degradation would not reverse.
+
+Against classic decode measured on the same cooled host (32.33, 31.87, 30.79
+tok/s, mean 31.66), MTP goes from **1.19x to about 1.5x**.
+
+## Output is unchanged, and that is not a tolerance claim
+
+Classic decode, MTP with the bf16 drafter, and MTP with the 4-bit drafter
+produce **byte-identical** text at temperature 0. That is structural rather
+than fortunate: the drafter only proposes and the target verifies every
+proposal, so a worse draft is rejected and a better one is accepted, and
+neither reaches the output. Quantization can cost acceptance and nothing else.
+
+## Shipped as a load-time conversion, not a second checkpoint
+
+Requiring a converted checkpoint would mean publishing and version-matching
+one per drafter. Instead `Qwen35MtpDraftModel::from_path` quantizes its own
+dense projections after sanitizing weights and before the bf16 to f16 pass,
+so every existing bf16 drafter gets this without being re-downloaded.
+
+A tensor whose `.scales` sibling already exists is left alone, so a
+pre-converted checkpoint loads unchanged, and one whose contraction axis is
+not a multiple of the group size stays dense rather than failing the load.
+
+Verified equivalent to converting offline:
+
+| arm | acceptance | `draft_block` |
+|---|---:|---:|
+| bf16 checkpoint, load-time conversion | 0.6601 | 2.74 ms/round |
+| pre-converted 4-bit checkpoint | 0.6601 | 2.70 ms/round |
+| bf16 checkpoint, `MLXCEL_MTP_QUANTIZE_DRAFTER=0` | 0.6831 | 10.43 ms/round |
+
+`scripts/tools/quantize_mtp_drafter.py` still converts a checkpoint offline,
+for publishing one or for A/B work that wants the two on disk side by side.
+
+## What this leaves for the rest of Phase 2
+
+The drafter side was 33.4% of the round before this and is about 11.8% after.
+Phases 2a and 2b target the LM head, which is 21.4% of a drafter step that now
+costs 2.7 ms per round rather than 10.5, so the absolute prize has shrunk by
+roughly a factor of four.
+
+The ceiling has also moved. With the #1199 `qmv_wide` retry the verify forward
+costs 40.02 ms per round, so a round with a free drafter is 41.6 ms and emits
+2.33 tokens: 56.0 tok/s, or 1.77x classic. The 1.93x in #1185 was computed
+against the cheaper, non-exact verify. At about 1.5x today, most of what
+remains is on the verify side, which is where tree drafting would act.

scripts/tools/quantize_mtp_drafter.py

@@ -0,0 +1,52 @@
+"""Quantize a qwen3_5_mtp drafter checkpoint to match its target's scheme.
+
+Out-of-band tooling (issue #1185 Phase 3). The drafter borrows embed_tokens and
+the LM head from the target, so only its own 2D projections are converted; the
+1D norms stay bf16, matching what every shipped mlx-community checkpoint does.
+
+The scheme is taken from the target rather than chosen here, because the
+mlxcel loader passes the *target's* group_size/bits to UnifiedLinear when it
+builds the drafter. A drafter quantized any other way would be read with the
+wrong parameters.
+"""
+import json, shutil, sys
+from pathlib import Path
+import mlx.core as mx
+
+src, tgt, dst = (Path(p) for p in sys.argv[1:4])
+scheme = json.loads((tgt / "config.json").read_text())["quantization"]
+group_size, bits, mode = scheme["group_size"], scheme["bits"], scheme.get("mode", "affine")
+print(f"target scheme: group_size={group_size} bits={bits} mode={mode}")
+
+w = mx.load(str(src / "model.safetensors"))
+out, converted, kept = {}, 0, 0
+before = sum(v.nbytes for v in w.values())
+for key in sorted(w):
+    v = w[key]
+    if v.ndim == 2 and key.endswith(".weight"):
+        prefix = key[: -len(".weight")]
+        if v.shape[-1] % group_size:
+            raise SystemExit(f"{key}: last dim {v.shape[-1]} not divisible by {group_size}")
+        packed, scales, biases = mx.quantize(v, group_size=group_size, bits=bits)
+        out[key] = packed
+        out[f"{prefix}.scales"] = scales
+        out[f"{prefix}.biases"] = biases
+        converted += 1
+        print(f"  quantized {key:48s} {tuple(v.shape)} -> {tuple(packed.shape)}")
+    else:
+        out[key] = v
+        kept += 1
+
+dst.mkdir(parents=True, exist_ok=True)
+mx.eval(list(out.values()))
+mx.save_safetensors(str(dst / "model.safetensors"), out)
+cfg = json.loads((src / "config.json").read_text())
+cfg["quantization"] = {"group_size": group_size, "bits": bits, "mode": mode}
+(dst / "config.json").write_text(json.dumps(cfg, indent=2) + "\n")
+for extra in ("tokenizer.json", "tokenizer_config.json", "vocab.json", "README.md"):
+    if (src / extra).exists():
+        shutil.copy2(src / extra, dst / extra)
+
+after = sum(v.nbytes for v in out.values())
+print(f"\nconverted {converted} projections, kept {kept} tensors as-is")
+print(f"{before/2**20:.1f} MiB -> {after/2**20:.1f} MiB ({after/before:.3f}x)")

src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs

@@ -165,6 +165,24 @@ impl std::fmt::Debug for Qwen35MtpDraftModel {
     }
 }
 
+/// `MLXCEL_MTP_QUANTIZE_DRAFTER=0` keeps a dense drafter dense.
+///
+/// On by default because the measurement says the cost is throughput-only:
+/// the target verifies every proposal, so the switch trades acceptance, not
+/// correctness, and acceptance did not move measurably (#1185 Phase 3). Read
+/// once per process, matching the other `MLXCEL_*` switches.
+fn quantize_drafter_enabled() -> bool {
+    static FLAG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *FLAG.get_or_init(|| {
+        std::env::var("MLXCEL_MTP_QUANTIZE_DRAFTER")
+            .map(|v| {
+                let v = v.trim().to_ascii_lowercase();
+                !matches!(v.as_str(), "0" | "false" | "no" | "off")
+            })
+            .unwrap_or(true)
+    })
+}
+
 impl Qwen35MtpDraftModel {
     /// Construct from a checkpoint directory containing `config.json` and
     /// safetensors shards. Used by [`crate::drafter::load_drafter`]'s `Mtp`
@@ -185,6 +203,13 @@ impl Qwen35MtpDraftModel {
         let mut weights = crate::weights::load_weights_from_dir(path)
             .map_err(|reason| DrafterError::WeightLoad { reason })?;
         Self::sanitize_weights(&mut weights);
+        // Quantize the drafter's own projections before the dtype pass, so a
+        // dense checkpoint costs what a 4-bit one costs (issue #1185 Phase 3).
+        // Must run first: the bf16 -> f16 conversion below deliberately skips
+        // quantization auxiliaries, and the packed payload is uint32 either
+        // way, so ordering it after would only leave scales as f16 where every
+        // shipped checkpoint keeps them bf16.
+        Self::quantize_dense_projections(&mut weights, &config);
         // Apple Silicon precision: bf16 → f16 on non-quantized tensors at
         // the weight-loading boundary, matching the target model loaders and
         // the DFlash drafter loader. The published drafter is bf16; the
@@ -194,6 +219,80 @@ impl Qwen35MtpDraftModel {
         Self::from_weights(&weights, config)
     }
 
+    /// Quantize the drafter's dense 2-D projections in place, to the scheme
+    /// its own config declares (group 64, 4-bit affine by default).
+    ///
+    /// The drafter is memory-bound and ships bf16: 810 MiB of weights read
+    /// once per drafted token. Quantizing to 4-bit takes that to 228 MiB, and
+    /// measured on an M5 Max the drafter step went from 10.5 to 2.7 ms per
+    /// round while the target verify forward was unchanged (#1185).
+    ///
+    /// Safe by construction rather than by tolerance: the drafter only
+    /// proposes and the target verifies every proposal, so drafter numerics
+    /// cannot reach the output. The exposure is acceptance rate, and that was
+    /// measured across two generation lengths: 0.683 to 0.660 at 120 tokens
+    /// and 0.650 to 0.659 at 300, which is noise in both directions. Output
+    /// stayed byte-identical to classic decode.
+    ///
+    /// Skips a tensor whose `.scales` sibling already exists (a pre-converted
+    /// checkpoint) and one whose contraction axis is not a multiple of the
+    /// group size, leaving those dense rather than failing the load.
+    ///
+    /// `MLXCEL_MTP_QUANTIZE_DRAFTER=0` keeps the checkpoint's own precision,
+    /// for an acceptance A/B on a pairing this has not been measured on.
+    fn quantize_dense_projections(weights: &mut WeightMap, config: &Qwen35MtpConfig) {
+        if !quantize_drafter_enabled() {
+            return;
+        }
+        let text_cfg = config.text_config();
+        let (group_size, bits) = (text_cfg.group_size(), text_cfg.bits());
+
+        let candidates: Vec<String> = weights
+            .iter()
+            .filter(|(key, value)| {
+                key.ends_with(".weight")
+                    && ffi::array_shape(value).len() == 2
+                    && !weights.contains_key(&format!("{}.scales", key.trim_end_matches(".weight")))
+            })
+            .map(|(key, _)| key.clone())
+            .collect();
+
+        let mut converted = 0usize;
+        for key in candidates {
+            let prefix = key.trim_end_matches(".weight").to_string();
+            let shape = weights.get(&key).map(|w| ffi::array_shape(w));
+            let Some(shape) = shape else { continue };
+            if shape[shape.len() - 1] % group_size != 0 {
+                continue;
+            }
+            let quantized = {
+                let Some(w) = weights.get(&key) else { continue };
+                ffi::quantize_weights_with_mode(w, group_size, bits, "affine")
+            };
+            let packed = ffi::quantized_weights_w(&quantized);
+            let scales = ffi::quantized_weights_scales(&quantized);
+            let has_biases = ffi::quantized_weights_has_biases(&quantized);
+            weights.insert(key.clone(), packed);
+            weights.insert(format!("{prefix}.scales"), scales);
+            if has_biases {
+                weights.insert(
+                    format!("{prefix}.biases"),
+                    ffi::quantized_weights_biases(&quantized),
+                );
+            }
+            converted += 1;
+        }
+
+        if converted > 0 {
+            tracing::debug!(
+                converted,
+                group_size,
+                bits,
+                "quantized the MTP drafter's dense projections at load"
+            );
+        }
+    }
+
     /// Construct from an in-memory weight map (already sanitized). Used by
     /// `from_path` and unit tests with synthetic fixtures.
     pub fn from_weights(