---
license: apache-2.0
base_model:
  - stepfun-ai/Step-3.7-Flash
  - stepfun-ai/Step-3.7-Flash-NVFP4
tags:
  - speculative-decoding
  - mtp
  - multi-token-prediction
  - vllm
  - nvfp4
  - step3
language:
  - en
  - zh
  - ja
library_name: vllm
pipeline_tag: text-generation
---

# Step-3.7-Flash MTP draft (for the NVFP4 checkpoint)

A tiny **Multi-Token-Prediction (MTP / nextn) draft** for **`stepfun-ai/Step-3.7-Flash-NVFP4`**, so you can run
**speculative decoding** on the NVFP4 checkpoint in vLLM.

> **Why this exists:** the official `Step-3.7-Flash-NVFP4` checkpoint **declares**
> `num_nextn_predict_layers: 3` in its config but **ships zero MTP weights** — the
> 3 nextn layers were dropped during quantization, and the per-layer config arrays
> were truncated to 45 (so even loading them would `IndexError`). The BF16 and FP8
> releases keep the MTP weights, but **the NVFP4 one — the SM120-friendly, smallest
> one — cannot do speculative decoding out of the box.** This repo is the missing
> piece: the 3 MTP layers extracted from the BF16 release, kept in BF16 (they're
> tiny), packaged as a vLLM-loadable draft.

- **~5.9 GB**, BF16. Base = NVFP4 (mixed precision is fine; the draft is small).
- **Lossless** in the speculative sense: vLLM's rejection sampling provably matches
  the target distribution; at `temperature=0` it follows the target's greedy tokens.
- Drop-in: point vLLM's `--speculative-config` at this directory.

## Usage (vLLM, stepfun37 image / vLLM ≥ the build with `Step3p5MTP`)

The draft is auto-routed to vLLM's native `Step3p5MTP` + `Step3p5MTPProposer`
because its config is `model_type: step3p7` with `num_nextn_predict_layers > 0`.

```bash
docker run -d --gpus all --ipc=host --shm-size=64g --network host \
  -v /path/to/Step-3.7-Flash-NVFP4:/model:ro \
  -v /path/to/Step-3.7-Flash-MTP-draft:/draft:ro \
  vllm/vllm-openai:stepfun37 \
  /model \
    --served-model-name step3p7 --port 8000 \
    --trust-remote-code --tensor-parallel-size 2 --enable-expert-parallel \
    --quantization modelopt --kv-cache-dtype fp8 \
    --max-model-len 262144 --gpu-memory-utilization 0.92 --async-scheduling \
    --speculative-config '{"method":"mtp","model":"/draft","num_speculative_tokens":1}'
```

JSON for `--speculative-config` must have **no spaces** (brace-expansion safety).
**`num_speculative_tokens: 1` (K=1) is the sweet spot** — see below.

## Benchmarks (2× RTX PRO 6000 Blackwell, SM120, TP=2)

Measured on the NVFP4 base + this draft, K=1, vs. NVFP4 with speculation off.
`per_req` = decode tok/s a single user feels (prefill excluded). Acceptance ≈ **0.80** in production traffic.

**Single-stream decode (short context):**

| workload | base | + MTP K=1 | speedup | accept |
|---|---|---|---|---|
| free-form | 106.8 | **125.5** | +17.5% | 0.77 |
| code | 106.7 | **133.7** | +25.3% | 0.88 |
| Japanese | 107.0 | **129.3** | +20.9% | 0.80 |
| tool-call | 106.9 | **135.4** | +26.6% | 0.90 |

**Decode speedup grows with context length** (longer KV → base is more
memory-bound → bigger speculative win):

| context | C=1 | C=2 | C=4 | C=8 |
|---|---|---|---|---|
| 1K | +20% | +8% | +32% | +34% |
| 8K | +22% | +24% | +25% | **+44%** |
| 32K | +22% | +26% | +20% | +17% |
| **128K** | **+28%** | **+33%** | **+38%** | — |

Net-positive across the whole concurrency range we tested (MoE stays memory-bound
to high batch). Best `K`: **K=1** (K=2/K=3 lose to draft cost — later positions
have lower acceptance and add forward cost). NaN-free on SM120 (Gate0 5/5).

## How it was built (reproducible)

The draft is **not retrained** — it's the original StepFun MTP layers, extracted verbatim:

1. From `stepfun-ai/Step-3.7-Flash` (BF16), take the 52 tensors of
   `model.layers.{45,46,47}.*` (the 3 nextn layers, dense-MLP, 17 tensors each)
   plus `model.embed_tokens.weight`. They all live in one shard
   (`model-00024.safetensors`).
2. Keep the **original BF16 weight names** — vLLM's `Step3p5MTP` loader does its own
   renaming (`.transformer.` strip, `shared_head.output→head`, `.mtp_block.` insert).
3. `config.json` = the **BF16 original** config (NOT the NVFP4 one): its per-layer
   arrays (`layer_types`, `partial_rotary_factors`, …) are length 48 and cover the
   MTP layer indices 45-47. **Strip `quantization_config`** so the draft loads as BF16.

Full scripts + benchmark harness: **[GitHub repo](#)** (`build_draft.py`,
`launch_mtp.sh`, `eval_mtp.py`, `bench_matrix.py`).

## License & attribution

Apache-2.0, inherited from the base model **`stepfun-ai/Step-3.7-Flash`**. These are
StepFun's weights, redistributed unchanged (only re-sharded/re-packaged as a draft).
All credit for the model and the MTP layers goes to StepFun.