# 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:

```python
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:

```bash
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).