[Bug Fix] MTP Speculative Decoding with NVFP4: Weight Shape Mismatch

[jagguvarma15]

Fix: MTP Speculative Decoding with NVFP4 Weight Shape Mismatch

Fixes #35031

Summary

The eh_proj layer in multiple MTP (Multi-Token Prediction) model files was defined as a plain nn.Linear, which does not participate in vLLM's quantization framework. When NVFP4 quantized checkpoints are loaded, the weights have a different shape (packed uint8 with halved input dimension for FP4 packing), causing a weight shape mismatch error.

Root Cause

nn.Linear is not recognized by ModelOptQuantConfigBase.get_quant_method(), which only handles LinearBase subclasses. This means eh_proj never gets a quantization method assigned, so it expects unquantized weight shapes — but the checkpoint provides NVFP4-quantized weights with different dimensions.

Fix

Replaced nn.Linear with ReplicatedLinear (a LinearBase subclass) for the eh_proj layer across 6 MTP model files. ReplicatedLinear properly integrates with vLLM's quantization system, accepting a quant_config parameter so it correctly handles quantized weight formats including NVFP4.

The existing longcat_flash_mtp.py already used ReplicatedLinear correctly and served as the reference for this fix.

Files Changed

File	Change
vllm/model_executor/models/deepseek_mtp.py	nn.Linear → ReplicatedLinear for eh_proj
vllm/model_executor/models/step3p5_mtp.py	nn.Linear → ReplicatedLinear for eh_proj
vllm/model_executor/models/openpangu_mtp.py	nn.Linear → ReplicatedLinear for eh_proj
vllm/model_executor/models/glm_ocr_mtp.py	nn.Linear → ReplicatedLinear for eh_proj
vllm/model_executor/models/glm4_moe_mtp.py	nn.Linear → ReplicatedLinear for eh_proj
vllm/model_executor/models/glm4_moe_lite_mtp.py	nn.Linear → ReplicatedLinear for eh_proj

Changes per file

1. Added ReplicatedLinear import from vllm.model_executor.layers.linear
2. Replaced nn.Linear(hidden_size * 2, hidden_size, bias=False) with ReplicatedLinear(..., quant_config=quant_config, prefix=...)
3. Updated forward() to unpack the tuple return: hidden_states, _ = self.eh_proj(...)

---

Assisted by Claude Opus 4.6

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[dosubot]

Related Documentation

Checked 0 published document(s) in 1 knowledge base(s). No updates required.

^{How did I do? Any feedback?}  

[gemini-code-assist]

Code Review

This pull request correctly addresses the weight shape mismatch issue in MTP speculative decoding when using NVFP4 quantization. By replacing nn.Linear with ReplicatedLinear for the eh_proj layer across multiple MTP model files, the code now properly integrates with vLLM's quantization framework. The updates to the forward methods to handle the tuple return of ReplicatedLinear are also correctly implemented. The changes are consistent and follow established patterns in the repository.