Step-3.7-p15-ream-shared8-head8

A pruned, merged, and mixed-precision-quantized MLX build of StepFun's Step-3.7-Flash (VLM, MoE), for Apple Silicon. 94 GB total.

What's different from the base model

1. REAP pruning, 15% — MoE expert count reduced 288 → 245 per layer, using REAP (router-weighted saliency: mean over routed tokens of gate_weight × ‖expert_activation‖) to rank experts and keep the highest-scoring ones.
2. REAM merge, not delete — instead of discarding the lowest-scoring 43 experts per layer, each is folded into its most-similar kept expert (router-row cosine similarity), with the kept expert becoming a saliency-weighted average of itself plus everything it absorbed. Output expert count is identical to a straight prune (245/288) — this changes the kept experts' values, not the count. 1,806 experts merged total (43/layer × 42 non-dense layers).
3. Mixed-precision quantization — 4-bit/group-size-64 affine base, with 8-bit/gs64 kept for the parts most sensitive to precision loss: the shared (always-active) expert on every MoE layer, the router gate itself, lm_head, embed_tokens, and the first 3 dense (non-MoE) layers. 179 modules total at 8-bit against a 4-bit base.
4. Native MTP (speculative decoding) enabled — see below. This is the only known working MTP implementation for Step-3.7 anywhere; upstream mlx-lm and StepFun's own reference code don't implement inference-time MTP for this architecture at all.

Architecture

Base	StepFun Step-3.7-Flash
Type	VLM (vision + text), MoE
Hidden size	4096
Layers	45 main + 3 MTP draft layers
Attention	hybrid sliding/full, per-layer layer_types
MoE experts	245 (post-REAP, from 288)
MoE top-k	8
Vocab	128,896
Context	262,144 tokens
Format	MLX safetensors, 21 shards

Benchmarks (measured on this exact build)

benchmark	score	n	time
GSM8K	76.0% (152/200)	200	1011.0 s
HumanEval	77.4% (127/164)	164 (full set)	2052.5 s
MMLU	31.5% (63/200)	200	456.4 s

GSM8K and HumanEval are strong for a 94 GB pruned+merged+quantized MoE. MMLU is notably low — close to the 25% random-guess floor for 4-choice multiple choice. At n=200 that's a real number (±~7% margin), not noise, but the cause isn't diagnosed yet: it could be genuine breadth-of-knowledge loss from the REAP/REAM step, or an eval-harness parsing mismatch with this model's output format. Worth independent verification if MMLU-style broad recall matters for your use case; the math and code results suggest the model's actual reasoning/generation capability held up well.

Earlier internal finding (background, not this exact file)

A related build from the same project ran a formal REAM-vs-straight-prune comparison and found perplexity gains didn't translate to accuracy gains: −0.194 NLL overall (including a −0.53 NLL / 41% swing on tool-use) produced a 24/24-vs-23/24 math result and an exact 14/15-vs-14/15 tie on tool-call accuracy. That test predates the benchmarks above and was run on a comparably-configured build, not this exact checkpoint — it's included here as methodology context (perplexity alone doesn't validate a merge; measure real task accuracy) rather than a verdict on this file, which the numbers above measure directly.

Native MTP (speculative decoding)

This build ships real, working MTP draft weights (layer index 45, extracted from the upstream checkpoint's model-00024.safetensors shard and quantized to match this build's base 4-bit/gs64 policy) and — if you're running it under oMLX — the model-side patch that makes them actually usable is upstreamed: jundot/omlx#2450.

Neither StepFun's own reference implementation nor stock mlx-lm implement MTP inference for Step-3.7 at all — the checkpoint carries the weights and config.json declares num_nextn_predict_layers, but nothing consumes them without that patch. The forward pass was ported directly from StepFun's own reference C++ (stepfun-ai/llama.cpp, branch step3p5-mtp), not inferred from tensor shapes.

Measured via oMLX's own throughput benchmark, mtp_enabled true vs false, same checkpoint/prompts, A/B across reloads:

test	baseline tgTPS	MTP tgTPS	speedup
pp 4096 / tg 128	52.7	64.4	1.22x
pp 8192 / tg 128	51.2	63.3	1.24x
pp 16384 / tg 128	48.8	61.1	1.25x
pp 32768 / tg 128	42.6	55.2	1.30x

Speedup grows with context, for +0.4GB peak memory. ~1.25x is the expected magnitude, not a ceiling — StepFun's own runtime is depth-1 (one draft token per verify cycle), which caps the theoretical best at 2x.

To use MTP: you need oMLX with the linked patch applied (merged upstream or manually installed — see the PR), then set mtp_enabled: true for this model in oMLX's settings and reload. Without that patch, this checkpoint still loads and runs completely normally as a standard (non-speculative) model — the extra weights are simply inert.

Quantization detail

base:      4-bit, group_size=64, affine
8-bit overrides (179 modules, same group_size/mode):
  - shared expert (share_expert.{gate,up,down}_proj) — every MoE layer, 126 modules
  - router gate (mlp.gate.gate) — every MoE layer, 42 modules
  - lm_head, embed_tokens — 2 modules
  - dense MLP, layers 0-2 (pre-MoE) — 9 modules

Running this model

oMLX (recommended — native MTP support, see above):

omlx serve True2456/Step-3.7-p15-ream-shared8-head8

mlx-vlm / mlx-lm (standard load, no MTP):

from mlx_vlm import load, generate
model, processor = load("True2456/Step-3.7-p15-ream-shared8-head8")

Context length defaults may need an explicit override in some servers (--context-length 262144 or your server's equivalent) — the config declares the model's real max, but some UIs cap the manual input lower.

License

Apache 2.0, inherited from the base model.