Commit History
-
P4: MTP working — 75% accept, 15% net speedup (52.6 vs 45.6 tok/s)
anerjyThree fixes that took accept from 2-4% → 73-80%: 1. **Per-MTP-layer lm_head** (the big one). The 3 MTP layers each ship their OWN shared_head.output bf16 weights — verified by diffing model.layers.{45,46,47}.transformer.shared_head.output row-0 means in Hikari07jp's BF16 source (0.000029 vs 0.000008 vs 0.000034). Previously we routed the MTP forward through the backbone's 4-bit-quantized lm_head; the per-MTP-layer BF16 head gave the same correct logit signal vLLM gets. 2. **Per-MTP embed_tokens**. The MTP shard includes its own model.embed_tokens.weight (bf16); vLLM's Step3p5AMultiTokenPredictor keeps a separate embedding. Routed that to language_model.mtp.embed_tokens.weight and use it for the next_token_ids lookup in mtp_forward. 3. **+1 norm shift on ALL 7 Gemma norms** at offline extract time (vs only 4 in P3). vLLM's GemmaRMSNorm does `weight = stored + 1.0; rms_norm(x, weight, eps)` for every norm in Step3p5: enorm, hnorm, mtp_block.input…7818108 -
P3.7: Norm +1 shift on extract for Gemma-stored MTP weights
anerjyHikari07jp's BF16 MTP-draft shard stores 5 norm types in vLLM's GemmaRMSNorm zero-centered convention (mean ≈ 0): enorm, hnorm, mtp_block.input_layernorm, mtp_block.post_attention_layernorm. vLLM computes (1+weight)*normalize(x); MLX-side ZeroCenteredRMSNorm is plain mx.fast.rms_norm — no +1 baked in. Without shifting these at extract time, MTP activations are 10-20× too small and the head emits garbage. First bench post-shift hit 2.6-4.2% accept (vs 0% pre-shift) on greedy code/physics prompts at 200 tokens, confirming the shift direction is correct. Still well below break-even (~43% needed for net speedup vs MTP-off baseline on Step-3.7-Flash-4bit). shared_head_norm intentionally NOT shifted: its stored mean ≈ 1.70 already matches the order of magnitude of backbone model.norm.weight (~2.46), suggesting it was saved in plain RMSNorm convention. Shifting it to 2.70 made accept slightly worse (4.2% → 2.6%). Other tensors (q_norm, k_norm, eh_proj, all attention/MLP weights) load as-is…83aabfc -
P3: enable MTP shard loading (config + sanitize + shard rewrite script)
anerjy +1: Claude Opus 4.7Pre-stage of MTP speculative decoding wiring. Wraps up: 1. config.py: TextConfig gains `num_nextn_predict_layers: int = 0` so LanguageModel.__init__ can build MTPModule conditionally. 2. step3p7.py Model.sanitize: defensive route for raw `model.layers.X.*` and `model.embed_tokens.weight` into the text bucket. This path won't fire under mlx_vlm.utils.load() in practice — that loader skips sanitize when the first shard carries `format=mlx` metadata — but it keeps the contract clear if a user combines MTP weights without the pre-rewrite script. 3. language.py: already in P2 (MTPLayer + MTPModule classes). Now activated by config.num_nextn_predict_layers > 0. 4. scripts/rewrite_mtp_shard.py (new): offline rewrites a Hikari07jp BF16 shard keys → final MLX-side names so the shard drops into the model dir without needing Model.sanitize. Documents the mlx-vlm `format=mlx` skip-path that bit us during integration. Uses safetensors.…c3b003a -
Initial: Step 3.7 Flash MLX adapter (text + vision)
anerjy +1: Claude Opus 4.7First MLX-native runtime adapter for stepfun-ai/Step-3.7-Flash on Apple Silicon. mlx-vlm (multimodal) + mlx-lm (text-only fallback). Includes ImagePatcher port for multi-crop vision (728² global + N×504² sub-patches) and chat_template with reasoning_effort=none. 4 port bugs documented in README (MLX nn.Upsample NHWC, lazy serialize zeros, vision cache stale on multi-crop, vision weight provenance). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>98d164a