From 5e6656e7a044e3144035b8e415f710787246ce58 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Sun, 27 Sep 2026 10:31:47 +0700 Subject: [PATCH] [Bugfix][Spec Decode] Add get_top_tokens() to Qwen4ExpMTP draft model use_local_argmax_reduction hard-fails at worker startup when the draft model lacks get_top_tokens(): the V2 speculator probes for the method in _validate_local_argmax_reduction and raises, with no fallback to the full-vocab logits path. Qwen4ExpMTP (nvidia and amd variants) has compute_logits but no get_top_tokens, so any Qwen4Exp MTP drafter with the flag enabled crashed both TP ranks during load_model. Qwen4ExpMTP owns a flat self.lm_head and its compute_logits is exactly self.logits_processor(self.lm_head, hidden_states), the same shape as Qwen3_5MTP, so LocalArgmaxMixin applies unchanged: vocab-parallel argmax with a (value, index) gather, O(2*tp_size) vs O(vocab_size) communication per draft token. No D2T remapping needed since the draft vocabulary equals the target vocabulary. Signed-off-by: MOVZX --- vllm/models/qwen4_exp/amd/mtp.py | 4 ++-- vllm/models/qwen4_exp/nvidia/mtp.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index a2b420603f..d103f17d0a 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -33,7 +33,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.utils import configure_quant_config -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import LocalArgmaxMixin, SupportsPP from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -379,7 +379,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): "hidden_states": 0, } ) -class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): +class Qwen4ExpMTP(LocalArgmaxMixin, nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index af3bcd4659..1fbfb0b25c 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -32,7 +32,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.utils import configure_quant_config -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import LocalArgmaxMixin, SupportsPP from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -359,7 +359,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): return loader.load_weights(weights, mapper=mapper) -class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): +class Qwen4ExpMTP(LocalArgmaxMixin, nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"],