From 4e52eb8c9c9a4a37683a5aed8dc89bb7f29b8a42 Mon Sep 17 00:00:00 2001 From: malaiwah Date: Sun, 26 Jul 2026 02:08:32 +0000 Subject: [PATCH] [SpecDecode] deepseek_mtp: normalize rank-sliced weight names in load_weights DeepSeekMTP.load_weights does not apply normalize_rank_sliced_weight_name(), which DeepseekV2ForCausalLM.load_weights already does. Quant configs that serialize per-TP-rank payloads (EXL3 rank-sliced trellis) therefore fail to load an MTP layer stored in that schema: KeyError: 'model.layers.78.mtp_block.mlp.experts.routed_experts.w2_rank0.mcg' The expert mapping rewrites the projection but the '.rank{r}.{suffix}' tail is never normalized. This mirrors the deepseek_v2 hunk and is a no-op (getattr returns None) for quant configs without rank slicing. Validated on GLM-5.2 753B, TP4, 4x RTX PRO 6000 Blackwell: with this fix a 3.0bpw EXL3 trellis MTP-layer overlay loads and serves at BF16-parity draft acceptance (GSM8K MAL 3.517 vs 3.528, acceptance 83.9% vs 84.3%) while the draft's footprint drops from 19.3 GB to 3.7 GB, freeing ~3.65 GiB/GPU of KV. Signed-off-by: Michel Belleau --- vllm/model_executor/models/deepseek_mtp.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index e37303a1d36f..b78459daf6ba 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -741,7 +741,25 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded_params: set[str] = set() _pending_wk_fp8: dict = {} # FP8 indexer wk dequant buffer _pending_fp8_linear: dict = {} + # Normalize rank-sliced checkpoint names the same way + # DeepseekV2ForCausalLM.load_weights does. Quant configs that serialize + # per-TP-rank payloads (EXL3 rank-sliced trellis) expose + # normalize_rank_sliced_weight_name(); without applying it here the MTP + # layer's per-expert `.rank{r}.{suffix}` tensors never map onto the fused + # RoutedExperts parameters and loading dies with e.g. + # KeyError: '...mtp_block.mlp.experts.routed_experts.w2_rank0.mcg'. + # No-op (getattr -> None) for quant configs without rank slicing. + quant_config = getattr(self, "quant_config", None) + rank_sliced_name = ( + getattr(quant_config, "normalize_rank_sliced_weight_name", None) + if quant_config is not None + else None + ) for name, loaded_weight in weights: + if rank_sliced_name is not None: + name = rank_sliced_name(name) + if name is None: + continue if "rotary_emb.inv_freq" in name: continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)