80 lines
3.2 KiB
Diff
80 lines
3.2 KiB
Diff
From cf89be37d15d2b60d08f6be1b032f8a1dd71bce9 Mon Sep 17 00:00:00 2001
|
|
From: Hinata <hinata@idihore.id>
|
|
Date: Sat, 10 Oct 2026 13:35:32 +0700
|
|
Subject: [PATCH 099/103] qwen4_exp: mix LocalArgmaxMixin into the MTP draft
|
|
head
|
|
|
|
Qwen4ExpMTP picks each MTP draft token by materialising full-vocab logits,
|
|
even though the tree already provides the O(2 * tp_size) path for exactly
|
|
this case: LocalArgmaxMixin builds the argmax from per-rank (value, index)
|
|
pairs instead of all-gathering a vocab-sized tensor.
|
|
|
|
The mixin needs only self.logits_processor and self.lm_head, both already
|
|
present on this class, and it is the established idiom for draft heads
|
|
(Qwen3_5MTP, the llama/qwen3 Eagle drafters, deepseek_eagle3). Wiring it in
|
|
is what makes use_local_argmax_reduction usable with this model: without
|
|
get_top_tokens the speculative-config validator refuses to start.
|
|
|
|
Applied to both platform implementations (nvidia, amd) to keep them in step.
|
|
No functional change unless the flag is enabled.
|
|
---
|
|
vllm/models/qwen4_exp/amd/mtp.py | 4 ++--
|
|
vllm/models/qwen4_exp/nvidia/mtp.py | 12 ++++++++++--
|
|
2 files changed, 12 insertions(+), 4 deletions(-)
|
|
|
|
diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py
|
|
index 40ee6a6..f9637ae 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,
|
|
@@ -383,7 +383,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 76c0d7c..01fb69c 100644
|
|
--- a/vllm/models/qwen4_exp/nvidia/mtp.py
|
|
+++ b/vllm/models/qwen4_exp/nvidia/mtp.py
|
|
@@ -35,7 +35,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
should_use_host_embed_table,
|
|
)
|
|
from vllm.model_executor.model_loader.utils import configure_quant_config
|
|
-from vllm.model_executor.models.interfaces import SupportsPP, SupportsReplaySSM
|
|
+from vllm.model_executor.models.interfaces import (
|
|
+ LocalArgmaxMixin,
|
|
+ SupportsPP,
|
|
+ SupportsReplaySSM,
|
|
+)
|
|
from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
|
|
from vllm.model_executor.models.utils import (
|
|
AutoWeightsLoader,
|
|
@@ -409,7 +413,11 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):
|
|
}
|
|
)
|
|
class Qwen4ExpMTP(
|
|
- nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts, SupportsReplaySSM
|
|
+ LocalArgmaxMixin,
|
|
+ nn.Module,
|
|
+ SupportsPP,
|
|
+ Qwen4ExpMixtureOfExperts,
|
|
+ SupportsReplaySSM,
|
|
):
|
|
packed_modules_mapping = {
|
|
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
|
--
|
|
2.43.0
|
|
|