Embedded MTP keys make mlx-lm shift every TRUNK norm by +1.0, collapsing acceptance on MLX builds (separate cause from #301's sidecar bug)

[elkaix]

Summary

Draft acceptance collapses to 0–6% on any Qwen3.5/3.8-family checkpoint that is both
MLX-quantized and embeds its MTP head in the trunk shards. The cause is not in the MTP
sidecar — it is that mlx-lm's qwen3_5 sanitizer adds +1.0 to every trunk RMSNorm
gain, because it infers "this checkpoint needs the delta→absolute restore" from the mere
presence of mtp.* keys rather than from the values.

The trunk is corrupted before the drafter ever runs, so no sidecar can match it.

This is distinct from the sidecar-extraction bug also filed in #301. That one was real and
is fixed (Refs #301), but fixing it does not restore acceptance, and #301 should not
be closed on the strength of it.

Measured

Loading a forged artifact whose shards embed the head, then comparing one trunk norm on disk
against the same tensor after mlx_lm.utils.load:

trunk norm  language_model.model.layers.0.input_layernorm.weight
  on disk : 0.9666
  runtime : 1.9666
  delta   : +1.0000

Verification of that artifact:

AR 15.04 tok/s  |  D1 13.46 (0.895x) acc 6.7%
D2 11.47 (0.763x) acc 0.9/0.0%  |  D3 11.18 (0.743x) acc 2.9/0.0/0.0%
verdict: mtp_acceptance_collapsed

Contract calibration tested 64 candidates and returned best_agreement: 0.0,
agreement_by_depth: [0.0, 0.0, 0.0], calibration_status: no_agreement_signal.

A control artifact built from a trunk of the same base model, same
{"bits": 8, "group_size": 64, "mode": "affine"} quantization, but with no mtp.* keys
in its shards (sidecar supplied separately) reaches best_agreement: 1.0 and:

AR 17.23 tok/s  |  D1 39.14 (2.271x) acc 97.8%
D2 47.15 (2.736x) acc 95.2/89.6%  |  D3 54.12 (3.140x) acc 94.8/90.7/79.3%
verdict: mtp_depth_wins

The two trunks' norms are numerically identical (layers.0.input_layernorm 0.9666 vs
0.9666, post_attention_layernorm 0.7827 vs 0.7827, model.norm 1.9441 vs 1.9441). The
only structural difference is the presence of the 29 embedded mtp.* keys.

Mechanism

mlx_lm/models/qwen3_5.py, Model.sanitize:

def sanitize(self, weights):
    has_mtp_weights = any("mtp." in k for k in weights)
    has_unsanitized_conv1d = any(
        "conv1d.weight" in k and v.shape[-1] != 1 for k, v in weights.items()
    )
    should_shift_norm_weights = has_mtp_weights or has_unsanitized_conv1d
    weights = {k: v for k, v in weights.items() if "mtp." not in k}
    ...
        if should_shift_norm_weights and any(k.endswith(sfx) for sfx in norm_keys):
            if v.ndim == 1:
                weights[k] = v + 1.0

has_unsanitized_conv1d is a genuine value-based signal (shape[-1] != 1).
has_mtp_weights is a bare presence check. A raw HF export happens to satisfy both, so the
heuristic looks correct there — but an MLX build that already stores gains absolute and
still ships its head trips the first condition and gets shifted a second time.

Note the MTP tensors are discarded on the next line, so the keys that trigger the shift are
not even the ones being shifted.

Pruning model.safetensors.index.json does not help: mlx_lm/utils.py loads via
glob.glob(str(model_path / "model*.safetensors")) and ignores the index entirely.

Why the sidecar is not implicated

MTPLX loads mtp.safetensors exclusively when present, treating embedded weights as a
fallback. Substituting a known-good sidecar (sha256-identical to the one driving the 3.140x
control above) into the failing artifact and re-verifying against the same trunk changed
nothing — D3 acceptance 2.0/0.0/0.0, 11.20 tok/s, 0.735x. The sidecar is inert here.

I also confirmed the extracted sidecar is byte-faithful after the #301 extraction fix:
input_layernorm 1.0361, post_attention_layernorm 1.2063, k_norm 1.7795,
q_norm 1.7906, mtp.norm 2.2520 — source and sidecar identical to within 1e-6.

Reproduce

1. hf download chimingw/Qwen3.8-27B-Uncensored-OrcaRouter-MLX-8bit (31.2 GB, 29 embedded
   language_model.mtp.* tensors).
2. mtplx forge build --repo <path> --recipe '{"body_bits":8,"body_group_size":64,"body_mode":"affine","mtp_policy":"keep_bf16"}' ...
3. Compare any trunk norm on disk against its value after mlx_lm.utils.load on the forged
   artifact — it is exactly +1.0 higher.

Suggested direction

has_mtp_weights should be replaced with a value-based test, the same way
has_unsanitized_conv1d already is. The gains separate cleanly: on this family
delta-encoded q/k norms sit near 0.78 and absolute ones near 1.78, and MTPLX already ships a
two-signal predicate (qk_max >= 1.25 or low_min >= 0.5 ⇒ already absolute) that classifies
both conventions correctly with wide margins.

Since the fix is upstream, MTPLX may want an interim guard: detect at forge time that a
source embeds mtp.* keys and stores gains absolute, and either refuse to brand it or
strip the embedded keys from the physical shards rather than only from the index.

Environment

• mtplx 2.9.0
• mlx_lm 0.31.3, mlx 0.32.1
• macOS 26.5.2, Apple M5 Max, 128 GB
• Source: chimingw/Qwen3.8-27B-Uncensored-OrcaRouter-MLX-8bit
• Control: orcarouter/Qwen3.8-27B-Uncensored-MLX (8-bit/) + an externally supplied sidecar

[youssofal]

Reproduced, and it is worse than filed: mtplx serve loads the trunk through mlx-lm too, so a forged artifact with embedded absolute mtp keys gets its trunk double shifted inside mtplx itself, which is exactly your acceptance collapse. The forge side convention bug is fixed in the next release (#301) but that does not close this. The real fix, omitting mtp keys from written shards when the standalone sidecar exists, touches the artifact integrity path, so it is being done deliberately with a load time guard first. Staying open until that lands.