[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 <movzx@yahoo.com>

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"],

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"],