MTP fails (vLLM v0.22), improper MTP model by author
The base model loaded fine. Everything blows up only when loading the draft (MTP) model:
gemma4_mtp.py", line 440, in load_weights
weight_loader(param, loaded_weight)
...
load_row_parallel_weight
loaded_weight = loaded_weight.narrow(...)
RuntimeError: start (0) + length (8192) exceeds dimension size (4096).
This is a shape mismatch in the speculative draft checkpoint, not a config/flag problem. TP size is 1, so this isn't a sharding artifact.
The 8192 vs 4096 (exactly 2:1) ratio is the tell. MTP layers have a fusion/embed-projection (eh_proj / fc) that takes the concatenation of [token_embedding ; previous_hidden_state], i.e. 2 × hidden_size = 8192, and projects it back to hidden_size = 4096. vLLM's Gemma4MTPModel allocates that layer expecting an 8192-wide input dimension, then tries to narrow() 8192 rows out of the tensor in the checkpoint — but the checkpoint's tensor only has 4096.
So melcheikh/gemma-4-31B-it-qat-assistant-NVFP4-Blackwell was exported with a projection sized for a single hidden_size input (4096) rather than the concatenated 2*hidden_size. The draft checkpoint is incompatible with vLLM 0.22.0's Gemma4 MTP implementation. Either it was built against a different MTP layout, or the NVFP4 quantization/export of that fusion layer dropped/halved a dimension.
Hi @AQLabs ,
Thank you very much for pointing this out and for the precise diagnosis! You are 100% correct.
The issue is indeed a shape mismatch caused by the MTP projection layers (pre_projection and post_projection) being incorrectly quantized to FP4. Because they are standard Linear layers, the default ModelOpt configuration matched and quantized/packed them, which halved their dimensions (reducing the input size from 8192 to 4096 in the 12B case, or 10752 to 5376 in our 31B case). Since vLLM expects these handshake layers to remain in full precision (BF16), it failed with the narrow out-of-bounds error.
We have updated our quantization scripts to explicitly exclude pre_projection and post_projection from quantization. We are also increasing the calibration dataset size (to 512 samples) to ensure optimal perplexity/quality.
We are currently recompiling the corrected weights and will push the updated assistant (and base) checkpoints shortly. We'll update this thread as soon as the fixed models are uploaded. Thanks again for the help!
Update: The corrected assistant weights are now uploaded. We have explicitly excluded the MTP projection layers (pre_projection and post_projection) from quantization to resolve the shape mismatch, and calibrated the rest of the model with 512 samples.
vLLM should now load and run speculative decoding without issues. Thank you again for the precise diagnosis!
vllm serve melcheikh/gemma-4-31B-it-qat-NVFP4-Blackwell --served-model-name Qwen3.5 --trust-remote-code --gpu-memory-utilization 0.9 --tensor-parallel-size 2 --enable-prefix-caching --block-size 128 --disable-log-stats --max-num-seqs 256 --performance-mode throughput --optimization-level 2 --distributed-executor-backend mp --dtype auto --kv-cache-dtype fp8 --max-model-len 262144 --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4 --speculative-config '{"method":"mtp", "model": "melcheikh/gemma-4-31B-it-qat-assistant-NVFP4-Blackwell", "num_speculative_tokens": 8}'
RuntimeError: start (4096) + length (4096) exceeds dimension size (4096).
AssertionError: param_data.shape == loaded_weight.shape
vllm 0.23.0, --tensor-parallel-size 2, 2 x RTX PRO 6000， The service runs normally when the melcheikh/gemma-4-31B-it-qat-assistant-NVFP4-Blackwell model is not used.
Hi @wangaocheng ,
Thank you for reporting this.
This is an issue inherent to vLLM's current implementation for speculative/MTP models when running under tensor parallelism (tensor-parallel-size 2). We have already forwarded this issue for a fix upstream.
Once the fix is integrated into vLLM, you should be able to run speculative decoding with tensor parallelism successfully. We'll update this thread as soon as the upstream resolution is available.
Hi @wangaocheng ,
Update: The fix has been officially merged upstream in vLLM main (commit 7a90eb98a).
Once you upgrade your vLLM installation to the latest state on the main branch (or in the upcoming release), speculative decoding under tensor parallelism will work out of the box.
Please note that the speculative decoding CLI flags have changed in the latest main, so you should run the server using --spec-model and --spec-tokens like this:
python -m vllm.entrypoints.openai.api_server \
--model google/gemma-4-31b-it \
--spec-model melcheikh/gemma-4-31B-it-qat-assistant-NVFP4-Blackwell \
--spec-tokens 8 \
--tensor-parallel-size 2