diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py index ae1f7f5..35e1f41 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/nvidia/mtp.py @@ -110,10 +110,23 @@ class Glm5NextMultiTokenPredictorLayer(nn.Module): return hidden_states, hidden_states +def _mtp_hf_config(vllm_config: VllmConfig): + """The MTP head's own config. With a standalone draft checkpoint + (``--speculative-config '{"model": ...}'``) the draft config carries + num_nextn_predict_layers / n_routed_experts, while vllm_config.model_config + is still the *target* (which may have num_nextn_predict_layers == 0, e.g. an + checkpoint that ships without an MTP layer). Fall back to the + target config for in-checkpoint MTP layers.""" + spec = vllm_config.speculative_config + if spec is not None and spec.draft_model_config is not None: + return spec.draft_model_config.hf_config + return vllm_config.model_config.hf_config + + class Glm5NextMultiTokenPredictor(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - config = vllm_config.model_config.hf_config + config = _mtp_hf_config(vllm_config) self.mtp_start_layer_idx = config.num_hidden_layers self.num_mtp_layers = config.num_nextn_predict_layers self.layers = torch.nn.ModuleDict( @@ -210,7 +223,7 @@ class Glm5NextMultiTokenPredictor(nn.Module): class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - self.config = vllm_config.model_config.hf_config + self.config = _mtp_hf_config(vllm_config) self.quant_config = vllm_config.quant_config self.model = Glm5NextMultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")