diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index dc153ac9e0b..492081fd66c 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/vllm/model_executor/models/deepseek_eagle3.py @@ -31,6 +31,7 @@ from vllm.model_executor.models.deepseek_v2 import ( ) from vllm.multimodal.inputs import NestedTensors +from .interfaces import LocalArgmaxMixin from .utils import ( AutoWeightsLoader, get_draft_quant_config, @@ -309,7 +310,7 @@ class DeepseekV2Eagle3Model(nn.Module): return loaded_params -class Eagle3DeepseekV2ForCausalLM(DeepseekV2ForCausalLM): +class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM): """Eagle3 speculative decoding model for DeepseekV2/V3.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 46be838f8ac..68dbcf90f87 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1282,6 +1282,41 @@ def supports_any_eagle( return supports_eagle(model) or supports_eagle3(model) +class LocalArgmaxMixin: + """Mixin for draft model heads in speculative decoding. + + Provides a D2T-aware ``get_top_tokens`` that preserves the + local-argmax communication reduction even when the draft vocabulary + is smaller than the target vocabulary. + + When ``draft_id_to_target_id`` is present (shape ``(draft_vocab_size,)``, + containing per-token offset to target vocab id), the draft argmax index + ``k`` is mapped to the target vocab id via:: + + target_id = k + draft_id_to_target_id[k] + + This is mathematically equivalent to computing the full-vocab scatter + logits and taking the global argmax, but requires only + O(batch * 2 * tp_size) communication instead of O(batch * vocab_size). + + Requires the subclass to expose: + ``self.logits_processor``: LogitsProcessor + ``self.lm_head``: ParallelLMHead + ``self.draft_id_to_target_id`` (optional): nn.Parameter + """ + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Vocab-parallel argmax with optional D2T remapping.""" + top = self.logits_processor.get_top_tokens( + self.lm_head, + hidden_states, + ) + d2t = getattr(self, "draft_id_to_target_id", None) + if d2t is not None: + top = top + d2t[top] + return top + + class EagleModelMixin: aux_hidden_state_layers: tuple[int, ...] = () diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index cf59ccdf750..39044f5e8b4 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -62,6 +62,7 @@ from vllm.v1.attention.backend import AttentionType from .adapters import as_embedding_model, as_seq_cls_model from .interfaces import ( EagleModelMixin, + LocalArgmaxMixin, SupportsEagle, SupportsEagle3, SupportsLoRA, @@ -487,7 +488,7 @@ class LlamaModel(nn.Module, EagleModelMixin): class LlamaForCausalLM( - nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 + LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 ): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], diff --git a/vllm/model_executor/models/llama4_eagle.py b/vllm/model_executor/models/llama4_eagle.py index 962377fd178..068a15b6254 100644 --- a/vllm/model_executor/models/llama4_eagle.py +++ b/vllm/model_executor/models/llama4_eagle.py @@ -208,23 +208,6 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM): ) -> tuple[torch.Tensor, torch.Tensor]: return self.model(input_ids, positions, hidden_states, inputs_embeds) - def get_top_tokens( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - """Vocab-parallel argmax without all-gathering full logits. - - Falls back to full logits when draft_id_to_target_id remapping is - active, since the shared lm_head covers the full target vocab but - the draft model only predicts over a subset (draft_vocab_size). - """ - if ( - hasattr(self, "draft_id_to_target_id") - and self.draft_id_to_target_id is not None - ): - return self.compute_logits(hidden_states).argmax(dim=-1) - return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: def transform(inputs): name, loaded_weight = inputs diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 6dec60232b1..b070eac3255 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -48,7 +48,13 @@ from vllm.sequence import IntermediateTensors from vllm.transformers_utils.config import set_default_rope_theta from vllm.v1.attention.backend import AttentionType -from .interfaces import SupportsEagle, SupportsEagle3, SupportsLoRA, SupportsPP +from .interfaces import ( + LocalArgmaxMixin, + SupportsEagle, + SupportsEagle3, + SupportsLoRA, + SupportsPP, +) from .qwen2 import Qwen2MLP as Qwen3MLP from .qwen2 import Qwen2Model from .utils import AutoWeightsLoader, PPMissingLayer, extract_layer_index, maybe_prefix @@ -259,7 +265,7 @@ class Qwen3Model(Qwen2Model): class Qwen3ForCausalLM( - nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 + LocalArgmaxMixin, nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3 ): packed_modules_mapping = { "qkv_proj": [ diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 7dd478d4243..021462f3ee5 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import LocalArgmaxMixin from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5RMSNorm from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts from vllm.sequence import IntermediateTensors @@ -353,7 +354,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module): "hidden_states": 0, } ) -class Qwen3_5MTP(nn.Module, SupportsMultiModal): +class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): packed_modules_mapping = { "qkv_proj": [ "q_proj", diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index cf9b70a7ca6..88e3030d2e0 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1464,23 +1464,10 @@ class SpecDecodeBaseProposer: f"{self.model.__class__.__name__} does not implement " "get_top_tokens()." ) - # Warn if draft model has vocab remapping, which forces fallback - # to the full-logits path (negating the optimization). - if ( - hasattr(self.model, "draft_id_to_target_id") - and self.model.draft_id_to_target_id is not None - ): - logger.warning( - "use_local_argmax_reduction is enabled but draft model " - "uses draft_id_to_target_id vocab remapping. The " - "optimization will be bypassed (falling back to full " - "logits gather + argmax)." - ) - else: - logger.info( - "Using local argmax reduction for draft token generation " - "(communication: O(2*tp_size) vs O(vocab_size))." - ) + logger.info( + "Using local argmax reduction for draft token generation " + "(communication: O(2*tp_size) vs O(vocab_size))." + ) @torch.inference_mode() def dummy_run(