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