The MTP Sidecar Mismatch — Plan to Fix It

Author: Hakim Ghelab, VegaLaboratories LTD

Status: Implemented, built into a real full output model, and benchmarked end to end (2026-09-08). --correct-mtp/--mtp-bits flags added to 05_full_model_quantize.py. All 7 MTP tensors corrected, every one passed the safety gate on the first attempt, real Hessian-weighted error vs. naive ranges from ~185x better (k_proj) to ~9,675x better (o_proj). Two real bugs found and fixed by running it, not by review — see CHANGELOG.md’s matching entry. Real, isolated mtplx tune comparison (same trunk, same bit-allocation plan — the only variable that changed is naive vs. YAQA-corrected MTP sidecar) confirms the correction is a real, measurable win, not just a theoretical one:

Real numbers, mtplx tune, isolated comparison (14_compare_mtplx_tune.py, same source/plan, only the MTP sidecar’s correction method differs): rejected drafts at D3 more than halved (26 → 10); D1/D2 also faster (+34.9%/+9.0%) with acceptance essentially unchanged (-0.77pp/-0.53pp). One real, honest caveat, not swept under the rug: the same comparison run also flagged an output-validation failure (balanced_delimiters, an unclosed parenthesis) on the corrected model at D1 and D3 — plausibly sampling noise from a small validation set, not yet confirmed either way. Separately, a different comparison against this project’s own best existing product (HybridPareto5bpw-V3.3_Stratified, a different bit-allocation plan entirely, also naive at both trunk and MTP) showed a D3 regression — that comparison changes two variables at once (plan and correction method), so it is real but does not isolate the MTP fix; that thread is open separately, not part of this result.

The real gap, confirmed with evidence

The YAQA build ships a trunk where 362/363 tensors received real, two-sided Hessian-weighted correction — and an MTP (speculative-decode draft head) sidecar quantized with plain mx.quantize, zero correction, identical to what the fully-naive reference baseline does. Confirmed directly:

• YAQA resume manifest: 0 MTP-related entries out of 363 real corrected tensors — MTP was never inside YAQA’s own correction loop.
• optiq.runtime.mtp_convert.preserve_mtp() → _quantize_mtp(): calls mx.quantize(arr, group_size, bits) directly on the MTP projection tensors, no Hessian, no correction of any kind.
• The reference model’s own real build manifest (V5_BUILD_MANIFEST.json inside Qwen3.8-27B-heretic-ara-HybridPareto5bpw-V3.3_Stratified/): "language_trunk_source": "original_source + exact plan predicate", "elapsed_build_s": 35.77 — the whole trunk build took under 36 seconds, consistent only with plain per-tensor quantization straight from the plan, never with real Hessian correction. Confirmed further: zero Hessian/GPTQ/correction code anywhere in 05_build_final_model_V5.py, the actual script that built this reference model.

What this means: the reference model is naive trunk + naive sidecar — internally consistent. The YAQA model is corrected trunk + naive sidecar — internally inconsistent. The two builds were never a clean apples-to-apples comparison. A trunk that moved substantially closer to true model behavior, sitting next to a draft head that didn’t move at all, is a plausible real contributor to the D1/D3 speculative-decode degradation found earlier today (deeper depths depend most on trunk/draft agreement).

The fix, not a workaround

Not “control for this in the comparison” — finish applying the technology. Build a genuinely YAQA-corrected MTP sidecar, at the same bit width as the existing naive one (Q4, confirmed from the real build manifest), so the comparison becomes naive-Q4-sidecar vs. YAQA-corrected-Q4-sidecar on the exact same tensors — the same “same plan, different rounding method” design already used for the trunk comparison.

Why this is actually feasible, not a from-scratch rebuild

Checked optiq.runtime.mtp.mtp_patch.inject_mtp_support() directly. The real MTP module it constructs:

class _MTPModule(nn.Module):
    def __init__(self, args, n_layers):
        self.pre_fc_norm_hidden = nn.RMSNorm(...)
        self.pre_fc_norm_embedding = nn.RMSNorm(...)
        self.fc = nn.Linear(args.hidden_size * 2, args.hidden_size, bias=False)
        self.layers = [DecoderLayer(args, layer_idx=fa_idx) for _ in range(n_layers)]
        self.norm = nn.RMSNorm(...)

DecoderLayer here is mlx_lm.models.qwen3_5.DecoderLayer — the exact same class the trunk’s own decoder layers are built from. The MTP layer is not a bespoke architecture; it is architecturally identical to a trunk layer (same self_attn/mlp sub-structure), just one additional instance with its own weights, wrapped by a small fc-combine + norm stage that stays BF16 either way. inject_mtp_support() also already supports loading the MTP weights before quantizing them (contract.mtp_prequantized=False path) — i.e. a real, full-precision, callable MTP module is already something this codebase knows how to construct.

This means YAQA’s existing Hessian-collection and two-sided correction machinery — already proven correct on 362 real trunk DecoderLayer instances — should attach to this one additional DecoderLayer instance directly, not require new correction math.

What actually needs building

1. Construct the real MTP module in full precision (reusing inject_mtp_support()’s construction logic, or a close adaptation of it) and attach it to the loaded model.
2. Run a real forward pass over real calibration data, capturing the hidden states at the point they’d feed into the MTP module, to collect real Hessians for its self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj tensors — the one genuinely new piece, since trunk Hessian collection captures at a different point in the graph.
3. Apply the same adaptive_safe_correction() (yaqa_core.py) already proven on the trunk, per-tensor, to those 7 tensors.
4. Write the corrected sidecar in the same optiq/mtp.safetensors format the naive path already produces, keeping mtp.fc.weight and all norms at BF16 unchanged (matching the existing naive recipe exactly).

Design requirement — integrated, not a new standalone script

This does not ship as a separate one-off script. It’s a new flag on 05_full_model_quantize.py (--correct-mtp, --mtp-bits {4,5,6,8}), the same entry point already used for the trunk, so there is one build command, not a soup of scripts. Written to detect MTP layer count/architecture from the model’s own config rather than hardcoding this project’s specific model, so the same flag works for any future MTP-capable model this pipeline is pointed at (Qwen3.6, or a different model family entirely) — matching this project’s own established convention of writing tools that generalize past the one model they were first built against (see 08_gptq_apply_plan.py’s own native_mtp_defaults() for the precedent).

First real experiment

Build the hybrid model — YAQA-corrected trunk (already exists, untouched) + newly YAQA-corrected MTP sidecar at --mtp-bits 4 (bit-for-bit matching the existing naive sidecar) — into a new output directory, leaving the current YAQA build (naive sidecar) intact for comparison. Re-run the same mtplx tune + 5-task quality benchmark suite already run today. If this closes some or all of the D1/D3 gap, that’s real, direct evidence the sidecar mismatch was a real contributor, not just a plausible story.

© 2026 Hakim Ghelab, VegaLaboratories LTD. All rights reserved.