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

[inureyes]

Summary

The MTP drafter is read once per drafted token, so its cost is weight traffic. qwen3.8-27b-mtp-bf16 is 810 MiB for one decoder layer plus an fc projection. Quantizing its eight 2-D projections to the scheme its own config declares (affine, group 64, 4-bit) takes that to 228 MiB, 0.281x.

Phase 3 of #1185. Full record in docs/benchmark_results/mtp-drafter-quantization-m5max-2026-08-17.md.

Result

M5 Max, qwen3.8-27b-4bit target, --draft-block-size 3 --temp 0, two reps per arm, 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. Against classic decode on the same cooled host (32.33, 31.87, 30.79 tok/s), MTP goes from 1.19x to about 1.5x.

Acceptance does not move. It falls 3.4% at 120 tokens and rises 1.4% at 300. Two lengths disagreeing in sign is noise; a real degradation would not reverse.

Output is unchanged, and that is structural

Classic decode, MTP with the dense drafter, and MTP with the quantized drafter produce byte-identical text at temperature 0.

That is not a tolerance argument. The drafter only proposes and the target verifies every proposal, so a worse draft is rejected and a better one accepted, and neither reaches the output. Quantization can cost acceptance and nothing else, which is why acceptance is the thing measured above.

Why a load-time conversion and not a checkpoint

Requiring a converted checkpoint means publishing one per drafter and keeping it version-matched. Instead Qwen35MtpDraftModel::from_path quantizes its own dense projections, so every existing bf16 drafter gets this without being re-downloaded and there is no artifact to ship.

Placed after sanitize_weights and before the bf16 to f16 pass. The ordering is load-bearing: that pass deliberately skips quantization auxiliaries, so running it first would leave scales as f16 where every shipped checkpoint keeps them bf16.

Two skips keep it from being a new failure mode. A tensor whose .scales sibling already exists is left alone, so a pre-converted checkpoint loads unchanged. 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 offline, for publishing a checkpoint or holding both on disk for an A/B. It reads the scheme from the target rather than choosing one, because the loader passes the drafter config's parameters to UnifiedLinear and a drafter quantized any other way would be read with the wrong ones.

The baseline this is measured against moved

#1185's numbers predate #1199, which buys back temperature-0 exactness on generation 15+ by forcing plain qmv on the verify path. That made the verify forward 13.6% more expensive, so the issue's own instruction to re-take Phase 0 applies. Re-taken on a cooled host with no overrides:

component	ms/round	share	before #1199
verify forward	40.02	64.0%	35.24 / 60.2%
draft_block	10.54	16.9%	10.67 / 18.2%
accept hook	10.33	16.5%	10.60 / 18.1%

The ceiling moved with it: a round with a free drafter is 41.6 ms emitting 2.33 tokens, so 56.0 tok/s or 1.77x classic, not the 1.93x in the issue.

Test plan

• -p mlxcel-core --lib is 1472 passed / 4 failed, matching baseline; the 22 qwen3_5_mtp drafter tests pass.
• cargo fmt --check and cargo clippy --all-targets --features metal,accelerate -- -D warnings clean.
• Byte-identity against classic decode confirmed with the switch on and off.
• The exactness gate passes with no override, so these are shipping-configuration numbers.

Not covered

Only the qwen3_5_mtp drafter. The Gemma 4 assistant drafter has the same shape of opportunity and is untouched here, deliberately: it needs its own acceptance A/B before it gets the same default, and this PR should not change two pairings on one measurement.

Phases 2a and 2b are worth re-ranking after this. They target the LM head, 21.4% of a drafter step that now costs 2.7 ms per round rather than 10.5, so the absolute prize has shrunk about fourfold and most of what remains is on the verify side.

Refs #1185