Add MTP support for Step 3.5 Flash

[janhilgard]

Summary

• Add MTP (Multi-Token Prediction) speculative decoding support to the existing Step 3.5 Flash model
• 3 MTP prediction layers with per-layer shared_head (dense MLP, not MoE)
• Backward compatible — MTP only instantiated when num_nextn_predict_layers > 0

Changes

Component	Description
Step3p5SharedHead	Per-layer prediction head: norm + output projection
Step3p5MTPLayer	hnorm + enorm → eh_proj → sliding attention + dense MLP → shared_head
Step3p5MTP	Container for multiple prediction layers
Model.__call__	Added return_hidden parameter for MTP integration
Model.mtp_forward	Single-step MTP prediction (hidden + next token → logits)
Model.make_mtp_cache	KV cache factory for MTP layers
sanitize	Handle MTP weight loading (instead of filtering)
quant_predicate	Exclude MTP norm layers from quantization

Also fixes two pre-existing ruff E741 lint warnings (l → layer in generator expressions).

MTP Architecture

hidden_states ──→ hnorm ──→ ┐
                             ├──→ concat [e, h] ──→ eh_proj ──→ attn + MLP ──→ shared_head ──→ logits
next_token_ids ──→ embed ──→ enorm ──→ ┘

Each MTP layer uses sliding attention with dense SwiGLU MLP (not MoE), matching the original StepFun architecture.

Test plan

• Basic inference without MTP (backward compatibility)
• return_hidden=True returns prenorm hidden states
• mtp_forward() produces correct-shaped logits
• make_mtp_cache() returns KVCache list matching MTP layer count
• ruff format and ruff check pass cleanly
• End-to-end test with mlx-community/Step-3.5-Flash-4bit (111 GB, requires Apple Silicon)

🤖 Generated with Claude Code

[janhilgard]

Rebased onto current main and refreshed. Status update, since this has been sitting a while and the ground has moved under it.

The base model landed upstream in the meantime — mlx_lm/models/step3p5.py is in 0.31.3. MTP still is not: there is no Step3p5MTPLayer or Step3p5SharedHead, and sanitize() still drops the head outright:

for k, v in weights.items():
    if ".mtp" in k:
        continue

So the checkpoint ships MTP weights that are discarded at load, and speculative decoding has nothing to draft with. This PR is not stale — it just never got picked up.

One real conflict the rebase surfaced, and it was semantic rather than textual. #949 gave sliding attention layers a RotatingKVCache. The MTP layer runs sliding attention (it borrows the first sliding_attention layer's config so RoPE matches the checkpoint), but make_mtp_cache() was handing it a plain KVCache. The rebase applied cleanly and hid that completely — the MTP cache would have grown without bound across a generation while the main model's equivalent layers stayed capped at sliding_window.

Fixed by following the same policy, keyed off self_attn.is_sliding rather than assumed, since which layer the head borrows depends on layer_types:

return [
    RotatingKVCache(max_size=self.args.sliding_window)
    if layer.self_attn.is_sliding else KVCache()
    for layer in self.mtp.layers
]

test_step3p5_mtp_cache_follows_the_main_model_policy covers it — class, window size, and that the cache actually stays bounded over 4× sliding_window steps, plus the full-attention config keeping a plain KVCache. A plain KVCache fails it, which I checked by reverting the change.

tests/test_models.py: 78 passed, 1 skipped, 52 subtests.

One caveat I would rather state than have found in review: I no longer run Step 3.5 Flash, so this round is structural verification only — I cannot re-measure MTP acceptance against real weights the way the original numbers were taken. The cache fix is checked against behaviour, not just types, but if you want live acceptance figures before merging, someone with the model would need to take them.

[janhilgard]

Closing this — not because it went stale, but because I can no longer stand behind it.

State as of today, for anyone who finds this later:

• The base model landed upstream in the meantime; mlx_lm/models/step3p5.py is in 0.31.3.
• MTP is still not there. There is no Step3p5MTPLayer or Step3p5SharedHead, and sanitize() still drops the head outright (if ".mtp" in k: continue), so checkpoints ship MTP weights that are discarded at load and speculative decoding has nothing to draft with.
• The branch is rebased onto current main and green: tests/test_models.py 78 passed, 1 skipped, 52 subtests.
• The rebase surfaced a semantic conflict worth recording. step3p5: use rotating cache for sliding attention layers #949 gave sliding-attention layers a RotatingKVCache, and the MTP layer runs sliding attention — it borrows the first sliding_attention layer's config so RoPE matches the checkpoint — while make_mtp_cache() was still handing it a plain KVCache. Textually the rebase was clean, so nothing flagged it; the MTP cache would have grown without bound across a generation while the main model's equivalent layers stayed capped at sliding_window. Fixed by keying off self_attn.is_sliding, with a regression test that checks the cache actually stays bounded rather than just having the right class.

What I cannot do is the part that matters for merging it: I no longer run Step 3.5 Flash, so I cannot re-measure MTP acceptance against real weights. The original numbers came from a live deployment that no longer exists, and shipping a speculative-decoding path whose acceptance rate nobody can currently verify is not something I want to advocate for.

The branch stays up. If someone with the model wants to pick it up, the code and tests are here and the cache fix above is the non-obvious part.