A high-throughput and memory-efficient inference and serving engine for LLMs
[SpecDecode] Reduce TP communication for large-vocab draft models speculative decoding (#39419)
Signed-off-by: EanWang211123 <wangyiheng@sangfor.com.cn>
yiheng committed
Jun 10, 2026 at 15:59 UTC
bd2d83ff31c58685cc4d2f6e171d7af36011fa02
7 files changed
+53
-39
vllm/model_executor/models/deepseek_eagle3.py
+2
-1
@@ -31,6 +31,7 @@ from vllm.model_executor.models.deepseek_v2 import (
31
)
32
from vllm.multimodal.inputs import NestedTensors
33
34
+from .interfaces import LocalArgmaxMixin
35
from .utils import (
36
AutoWeightsLoader,
37
get_draft_quant_config,
@@ -309,7 +310,7 @@ class DeepseekV2Eagle3Model(nn.Module):
310
return loaded_params
311
312
312
-class Eagle3DeepseekV2ForCausalLM(DeepseekV2ForCausalLM):
313
+class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):
314
"""Eagle3 speculative decoding model for DeepseekV2/V3."""
315
316
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
vllm/model_executor/models/interfaces.py
+35
@@ -1282,6 +1282,41 @@ def supports_any_eagle(
1282
return supports_eagle(model) or supports_eagle3(model)
1283
1284
1285
+class LocalArgmaxMixin:
1286
+ """Mixin for draft model heads in speculative decoding.
1287
+
1288
+ Provides a D2T-aware ``get_top_tokens`` that preserves the
1289
+ local-argmax communication reduction even when the draft vocabulary
1290
+ is smaller than the target vocabulary.
1291
+
1292
+ When ``draft_id_to_target_id`` is present (shape ``(draft_vocab_size,)``,
1293
+ containing per-token offset to target vocab id), the draft argmax index
1294
+ ``k`` is mapped to the target vocab id via::
1295
+
1296
+ target_id = k + draft_id_to_target_id[k]
1297
+
1298
+ This is mathematically equivalent to computing the full-vocab scatter
1299
+ logits and taking the global argmax, but requires only
1300
+ O(batch * 2 * tp_size) communication instead of O(batch * vocab_size).
1301
+
1302
+ Requires the subclass to expose:
1303
+ ``self.logits_processor``: LogitsProcessor
1304
+ ``self.lm_head``: ParallelLMHead
1305
+ ``self.draft_id_to_target_id`` (optional): nn.Parameter
1306
+ """
1307
+
1308
+ def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
1309
+ """Vocab-parallel argmax with optional D2T remapping."""
1310
+ top = self.logits_processor.get_top_tokens(
1311
+ self.lm_head,
1312
+ hidden_states,
1313
+ )
1314
+ d2t = getattr(self, "draft_id_to_target_id", None)
1315
+ if d2t is not None:
1316
+ top = top + d2t[top]
1317
+ return top
1318
+
1319
+
1320
class EagleModelMixin:
1321
aux_hidden_state_layers: tuple[int, ...] = ()
1322
vllm/model_executor/models/llama.py
+2
-1
@@ -62,6 +62,7 @@ from vllm.v1.attention.backend import AttentionType
62
from .adapters import as_embedding_model, as_seq_cls_model
63
from .interfaces import (
64
EagleModelMixin,
65
+ LocalArgmaxMixin,
66
SupportsEagle,
67
SupportsEagle3,
68
SupportsLoRA,
@@ -487,7 +488,7 @@ class LlamaModel(nn.Module, EagleModelMixin):
488
489
490
class LlamaForCausalLM(
490
- nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
491
+ LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
492
):
493
packed_modules_mapping = {
494
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
vllm/model_executor/models/llama4_eagle.py
-17
@@ -208,23 +208,6 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
208
) -> tuple[torch.Tensor, torch.Tensor]:
209
return self.model(input_ids, positions, hidden_states, inputs_embeds)
210
211
- def get_top_tokens(
212
- self,
213
- hidden_states: torch.Tensor,
214
- ) -> torch.Tensor:
215
- """Vocab-parallel argmax without all-gathering full logits.
216
-
217
- Falls back to full logits when draft_id_to_target_id remapping is
218
- active, since the shared lm_head covers the full target vocab but
219
- the draft model only predicts over a subset (draft_vocab_size).
220
- """
221
- if (
222
- hasattr(self, "draft_id_to_target_id")
223
- and self.draft_id_to_target_id is not None
224
- ):
225
- return self.compute_logits(hidden_states).argmax(dim=-1)
226
- return self.logits_processor.get_top_tokens(self.lm_head, hidden_states)
227
-
211
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
212
def transform(inputs):
213
name, loaded_weight = inputs
vllm/model_executor/models/qwen3.py
+8
-2
@@ -48,7 +48,13 @@ from vllm.sequence import IntermediateTensors
48
from vllm.transformers_utils.config import set_default_rope_theta
49
from vllm.v1.attention.backend import AttentionType
50
51
-from .interfaces import SupportsEagle, SupportsEagle3, SupportsLoRA, SupportsPP
51
+from .interfaces import (
52
+ LocalArgmaxMixin,
53
+ SupportsEagle,
54
+ SupportsEagle3,
55
+ SupportsLoRA,
56
+ SupportsPP,
57
+)
58
from .qwen2 import Qwen2MLP as Qwen3MLP
59
from .qwen2 import Qwen2Model
60
from .utils import AutoWeightsLoader, PPMissingLayer, extract_layer_index, maybe_prefix
@@ -259,7 +265,7 @@ class Qwen3Model(Qwen2Model):
265
266
267
class Qwen3ForCausalLM(
262
- nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
268
+ LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
269
):
270
packed_modules_mapping = {
271
"qkv_proj": [
vllm/model_executor/models/qwen3_5_mtp.py
+2
-1
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
22
VocabParallelEmbedding,
23
)
24
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
+from vllm.model_executor.models.interfaces import LocalArgmaxMixin
26
from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5RMSNorm
27
from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts
28
from vllm.sequence import IntermediateTensors
@@ -353,7 +354,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
354
"hidden_states": 0,
355
}
356
)
356
-class Qwen3_5MTP(nn.Module, SupportsMultiModal):
357
+class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):
358
packed_modules_mapping = {
359
"qkv_proj": [
360
"q_proj",
vllm/v1/spec_decode/llm_base_proposer.py
+4
-17
@@ -1464,23 +1464,10 @@ class SpecDecodeBaseProposer:
1464
f"{self.model.__class__.__name__} does not implement "
1465
"get_top_tokens()."
1466
)
1467
- # Warn if draft model has vocab remapping, which forces fallback
1468
- # to the full-logits path (negating the optimization).
1469
- if (
1470
- hasattr(self.model, "draft_id_to_target_id")
1471
- and self.model.draft_id_to_target_id is not None
1472
- ):
1473
- logger.warning(
1474
- "use_local_argmax_reduction is enabled but draft model "
1475
- "uses draft_id_to_target_id vocab remapping. The "
1476
- "optimization will be bypassed (falling back to full "
1477
- "logits gather + argmax)."
1478
- )
1479
- else:
1480
- logger.info(
1481
- "Using local argmax reduction for draft token generation "
1482
- "(communication: O(2*tp_size) vs O(vocab_size))."
1483
- )
1467
+ logger.info(
1468
+ "Using local argmax reduction for draft token generation "
1469
+ "(communication: O(2*tp_size) vs O(vocab_size))."
1470
+ )
1471
1472
@torch.inference_mode()
1473
def dummy_run(