From 27a02df6be184d3f02210104265c5783736c314a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 17 Aug 2026 21:33:18 +0900 Subject: [PATCH] perf(speculative): quantize the MTP drafter's projections at load The drafter is read once per drafted token, so its cost is weight traffic: 810 MiB of bf16 for one decoder layer plus an fc projection. Quantizing its eight 2-D projections to the scheme its config declares (affine, group 64, 4-bit) takes that to 228 MiB. Measured on M5 Max against qwen3.8-27b-4bit, two reps per arm, alternated: - draft_block 10.40/10.71 -> 2.70/2.69 ms per round, 3.9x cheaper - accept hook 10.35/10.60 -> 2.74/2.71 ms per round - verify forward unchanged, as it must be - 37.90/35.57 -> 49.71/50.33 tok/s at n=120, and 35.24/34.76 -> 47.68/45.68 at n=300, so about 1.19x to 1.5x against classic decode Acceptance does not move: 0.6831 to 0.6601 at 120 tokens and 0.6500 to 0.6589 at 300. Two lengths disagreeing in sign is noise, not degradation. Output is byte-identical to classic decode with the drafter dense or quantized. That is structural: the drafter proposes and the target verifies every proposal, so drafter numerics cannot reach the output and the only exposure is acceptance. Done at load rather than as a second checkpoint, so every existing bf16 drafter gets it without being re-downloaded and there is no artifact to publish and version-match. A tensor whose .scales sibling already exists is left alone, so a pre-converted checkpoint still loads unchanged, and one whose contraction axis is not a multiple of the group size stays dense rather than failing the load. Verified to reproduce offline conversion exactly: same acceptance, same per-round drafter cost. MLXCEL_MTP_QUANTIZE_DRAFTER=0 keeps the checkpoint's precision, for an acceptance A/B on a pairing this has not been measured on. scripts/tools/quantize_mtp_drafter.py converts offline, for publishing a checkpoint or holding both on disk at once. Refs #1185 --- ...p-drafter-quantization-m5max-2026-08-17.md | 107 ++++++++++++++++++ scripts/tools/quantize_mtp_drafter.py | 52 +++++++++ .../src/drafter/qwen3_5_mtp/model.rs | 99 ++++++++++++++++ 3 files changed, 258 insertions(+) create mode 100644 docs/benchmark_results/mtp-drafter-quantization-m5max-2026-08-17.md create mode 100755 scripts/tools/quantize_mtp_drafter.py diff --git a/docs/benchmark_results/mtp-drafter-quantization-m5max-2026-08-17.md b/docs/benchmark_results/mtp-drafter-quantization-m5max-2026-08-17.md new file mode 100644 index 000000000..ea95b664c --- /dev/null +++ b/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. diff --git a/scripts/tools/quantize_mtp_drafter.py b/scripts/tools/quantize_mtp_drafter.py new file mode 100755 index 000000000..39206ed50 --- /dev/null +++ b/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)") diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs index 6e641ccb4..6cfe04203 100644 --- a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs +++ b/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 = 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 = 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(