vllm_patch/README.md

vLLM runtime patch (ships with this MXFP4 checkpoint)

Why

This checkpoint quantizes the routed MoE experts to MXFP4; everything else (attention, dense MLP, shared expert, router, the MTP/next-token-prediction layers) stays BF16. Step-3.7-Flash includes 3 MTP draft layers (num_nextn_predict_layers: 3, layers 45/46/47) used for speculative decoding.

When you enable MTP:

--speculative_config '{"method": "mtp", "num_speculative_tokens": 3}'

the stock vllm/vllm-openai:stepfun37 image fails at draft-model load with:

RuntimeError: Some parameters like
  model.layers.46.mtp_block.self_attn.attn.k_zero_point
  ( / v_zero_point / v_scale ) are not in the checkpoint and will falsely
  use random initialization

This is not a quantization defect. The MTP block is BF16 in every release — the official FP8 and NVFP4 Step-3.7 checkpoints also ship MTP attention as plain .weight, with no scales or zero-points. The problem is purely loader-side:

• The compressed-tensors mixed-precision path makes vLLM's inner Attention layer allocate k/v_scale + *_zero_point buffers that are never read (the KV cache is not quantized here).
• The strict MTP weight loader in step3p5_mtp.py then demands those buffers be present in the checkpoint. It already excuses scalar scales (.k_scale/.v_scale/... with numel == 1), but Step-3.7's are non-scalar (per-head) and there is no excuse for the zero-points at all.

The base model (no MTP) loads and runs correctly via a lenient loader — proving these params are inert (it scores ~95% GSM8K with the same Attention layer).

The official FP8/NVFP4 releases avoid the allocation a different way — through their own config dialect (modules_to_not_convert for FP8, modelopt exclude_modules for NVFP4). The compressed-tensors path needs the loader fix below instead.

The fix

step3p5_mtp.py here is the patched file. It replaces the optional_params set in load_weights() with an unconditional suffix match (no numel == 1 gate) over the full KV-attention quant-param family:

optional_params = {
    name
    for name in params_dict
    if name.endswith(
        (".k_scale", ".v_scale", ".q_scale", ".prob_scale",
         ".k_zero_point", ".v_zero_point", ".q_zero_point", ".prob_zero_point")
    )
}

So the strict loader treats these inert KV-cache buffers as optional and falls back to defaults, exactly as it already intended for the scalar scales.

Send the same change upstream to the fork.

How to apply (no image rebuild)

Bind-mount the file over the stock one. Path is for the image's Python 3.12:

docker run -d --name step37 --gpus all --privileged --ipc=host -p 8000:8000 \
  -e VLLM_MXFP4_USE_MARLIN=1 \
  -v $(pwd):/model \
  -v $(pwd)/vllm_patch/step3p5_mtp.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/step3p5_mtp.py \
  vllm/vllm-openai:stepfun37 /model \
  --served-model-name step3p7-flash \
  --tensor-parallel-size 1 --disable-cascade-attn \
  --reasoning-parser step3p5 --enable-auto-tool-choice --tool-call-parser step3p5 \
  --speculative_config '{"method": "mtp", "num_speculative_tokens": 3}' \
  --gpu-memory-utilization 0.97 --enforce-eager \
  --max-model-len 8192 --max-num-batched-tokens 2048 \
  --linear-backend marlin --trust-remote-code

(For multi-GPU, use the official --tensor-parallel-size 8 --enable-expert-parallel instead of --tensor-parallel-size 1.)

Verify after startup: the log shows Detected MTP model. Sharing target model embedding weights with the draft model and Application startup complete, and SpecDecoding metrics report a healthy mean acceptance length (~3.0 here).

If you serve without MTP, this patch is not needed.

Patched against image vllm/vllm-openai:stepfun37 (vllm 0.1.dev16944+ge9c8946e7).