From abaf62a5e1f537cd4dc60becbd4032e2ba275d66 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Tue, 30 Sep 2025 00:57:34 +0000 Subject: [PATCH 01/13] micro opt and refactor Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/model_executor/models/llama_eagle.py | 3 +- vllm/model_executor/models/llama_eagle3.py | 23 +---------- vllm/v1/spec_decode/eagle.py | 10 ++--- vllm/v1/spec_decode/eagle_mixin.py | 45 ++++++++++++++++++++++ 4 files changed, 54 insertions(+), 27 deletions(-) create mode 100644 vllm/v1/spec_decode/eagle_mixin.py diff --git a/vllm/model_executor/models/llama_eagle.py b/vllm/model_executor/models/llama_eagle.py index d7d6b1745fc8..2213173dff63 100644 --- a/vllm/model_executor/models/llama_eagle.py +++ b/vllm/model_executor/models/llama_eagle.py @@ -18,6 +18,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.llama import (LlamaDecoderLayer, LlamaForCausalLM) +from vllm.v1.spec_decode.eagle_mixin import EagleChainMixin from .utils import AutoWeightsLoader, maybe_prefix @@ -133,7 +134,7 @@ def load_weights(self, weights: Iterable[tuple[str, return loaded_params -class EagleLlamaForCausalLM(LlamaForCausalLM): +class EagleLlamaForCausalLM(EagleChainMixin, LlamaForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 3fb6f2f8d5ec..96b43e8a436c 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -20,6 +20,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.llama import (LlamaDecoderLayer, LlamaForCausalLM) +from vllm.v1.spec_decode.eagle_mixin import Eagle3ChainMixin from .utils import AutoWeightsLoader, maybe_prefix @@ -207,7 +208,7 @@ def load_weights(self, weights: Iterable[tuple[str, return loaded_params -class Eagle3LlamaForCausalLM(LlamaForCausalLM): +class Eagle3LlamaForCausalLM(Eagle3ChainMixin, LlamaForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) @@ -254,26 +255,6 @@ def forward( ) -> tuple[torch.Tensor, torch.Tensor]: return self.model(input_ids, positions, hidden_states, inputs_embeds) - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> Optional[torch.Tensor]: - logits = self.logits_processor(self.lm_head, hidden_states) - if self.draft_id_to_target_id is None: - assert logits.shape[1] == self.config.vocab_size, \ - "Expected logits to have shape " \ - f"(*, {self.config.vocab_size}), but got {logits.shape}" - return logits - - base = torch.arange(self.config.draft_vocab_size, device=logits.device) - targets = base + self.draft_id_to_target_id - logits_new = logits.new_full(( - logits.shape[0], - self.config.vocab_size, - ), float('-inf')) - logits_new[:, targets] = logits - return logits_new - def combine_hidden_states( self, hidden_states: torch.Tensor, diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index a2f7dbe5703f..1b0ca31a2eec 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -255,11 +255,10 @@ def propose( else: last_hidden_states, hidden_states = ret_hidden_states sample_hidden_states = last_hidden_states[last_token_indices] - logits = self.model.compute_logits(sample_hidden_states) # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1: - draft_token_ids = logits.argmax(dim=-1) + draft_token_ids = self.model.propose_chain(sample_hidden_states) return draft_token_ids.view(-1, 1) if self.uses_mrope: @@ -273,6 +272,7 @@ def propose( if isinstance(attn_metadata, TreeAttentionMetadata): # Draft using tree attention. + logits = self.model.compute_logits(sample_hidden_states, None) draft_token_ids_list = self.propose_tree( batch_size=batch_size, logits=logits, @@ -283,7 +283,7 @@ def propose( # [batch_size, num_tree_tokens] return torch.cat(draft_token_ids_list, dim=1) - draft_token_ids = logits.argmax(dim=-1) + draft_token_ids = self.model.propose_chain(sample_hidden_states) if self.allowed_attn_types is not None and \ not isinstance(attn_metadata, self.allowed_attn_types): @@ -406,8 +406,8 @@ def propose( else: last_hidden_states, hidden_states = ret_hidden_states hidden_states = hidden_states[:batch_size] - logits = self.model.compute_logits(last_hidden_states[:batch_size]) - draft_token_ids = logits.argmax(dim=-1) + draft_token_ids = self.model.propose_chain( + last_hidden_states[:batch_size]) draft_token_ids_list.append(draft_token_ids) # [batch_size, num_speculative_tokens] diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py new file mode 100644 index 000000000000..bcc5a1d111c0 --- /dev/null +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -0,0 +1,45 @@ +import torch +from typing import Optional + +class EagleMixin: + def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.compute_logits(hidden_states) + draft_token_ids = logits.argmax(dim=-1) + return draft_token_ids + +class Eagle3Mixin: + def compute_draft_logits( + self, + hidden_states: torch.Tensor, + ) -> Optional[torch.Tensor]: + # Computes draft logits from hidden states. + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> Optional[torch.Tensor]: + # Maps draft logits to target logits. + logits = self.compute_draft_logits(hidden_states) + if self.draft_id_to_target_id is None: + assert logits.shape[1] == self.config.vocab_size, \ + "Expected logits to have shape " \ + f"(*, {self.config.vocab_size}), but got {logits.shape}" + return logits + + base = torch.arange(self.config.draft_vocab_size, device=logits.device) + targets = base + self.draft_id_to_target_id + logits_new = logits.new_full(( + logits.shape[0], + self.config.vocab_size, + ), float('-inf')) + logits_new[:, targets] = logits + return logits_new + + def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.compute_draft_logits(hidden_states) + draft_token_ids = logits.argmax(dim=-1) + offset = self.draft_id_to_target_id[draft_token_ids] + draft_token_ids += offset + return draft_token_ids \ No newline at end of file From 65abf203cd63587e0d0132ccf0813ec6c8bbae58 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Tue, 30 Sep 2025 01:08:02 +0000 Subject: [PATCH 02/13] fix Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/model_executor/models/llama_eagle.py | 4 ++-- vllm/model_executor/models/llama_eagle3.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/llama_eagle.py b/vllm/model_executor/models/llama_eagle.py index 2213173dff63..fbfc291ad58e 100644 --- a/vllm/model_executor/models/llama_eagle.py +++ b/vllm/model_executor/models/llama_eagle.py @@ -18,7 +18,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.llama import (LlamaDecoderLayer, LlamaForCausalLM) -from vllm.v1.spec_decode.eagle_mixin import EagleChainMixin +from vllm.v1.spec_decode.eagle_mixin import EagleMixin from .utils import AutoWeightsLoader, maybe_prefix @@ -134,7 +134,7 @@ def load_weights(self, weights: Iterable[tuple[str, return loaded_params -class EagleLlamaForCausalLM(EagleChainMixin, LlamaForCausalLM): +class EagleLlamaForCausalLM(EagleMixin, LlamaForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 96b43e8a436c..2e4f81de7e7b 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -20,7 +20,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.llama import (LlamaDecoderLayer, LlamaForCausalLM) -from vllm.v1.spec_decode.eagle_mixin import Eagle3ChainMixin +from vllm.v1.spec_decode.eagle_mixin import Eagle3Mixin from .utils import AutoWeightsLoader, maybe_prefix @@ -208,7 +208,7 @@ def load_weights(self, weights: Iterable[tuple[str, return loaded_params -class Eagle3LlamaForCausalLM(Eagle3ChainMixin, LlamaForCausalLM): +class Eagle3LlamaForCausalLM(Eagle3Mixin, LlamaForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) From e827152ef2d0a5f8d0bc26f992b726a1fdd0ff5d Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 3 Oct 2025 03:41:01 +0000 Subject: [PATCH 03/13] lint Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle_mixin.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py index bcc5a1d111c0..08003338e87d 100644 --- a/vllm/v1/spec_decode/eagle_mixin.py +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -1,17 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any, Optional + import torch -from typing import Optional + class EagleMixin: + compute_logits: Any + def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: logits = self.compute_logits(hidden_states) draft_token_ids = logits.argmax(dim=-1) return draft_token_ids - + + class Eagle3Mixin: + lm_head: Any + logits_processor: Any + draft_id_to_target_id: Optional[Any] + config: Any + def compute_draft_logits( self, hidden_states: torch.Tensor, - ) -> Optional[torch.Tensor]: + ) -> torch.Tensor: # Computes draft logits from hidden states. logits = self.logits_processor(self.lm_head, hidden_states) return logits @@ -36,10 +48,10 @@ def compute_logits( ), float('-inf')) logits_new[:, targets] = logits return logits_new - + def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: logits = self.compute_draft_logits(hidden_states) draft_token_ids = logits.argmax(dim=-1) offset = self.draft_id_to_target_id[draft_token_ids] draft_token_ids += offset - return draft_token_ids \ No newline at end of file + return draft_token_ids From e45518b51c15df26c2e21e22b7e549837ea6c1ab Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 3 Oct 2025 03:44:17 +0000 Subject: [PATCH 04/13] gemini Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 5d1c4431f34c..feffed973d78 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -291,7 +291,7 @@ def propose( if isinstance(attn_metadata, TreeAttentionMetadata): # Draft using tree attention. - logits = self.model.compute_logits(sample_hidden_states, None) + logits = self.model.compute_logits(sample_hidden_states) draft_token_ids_list = self.propose_tree( batch_size=batch_size, logits=logits, From b60c9d4588aa5edf912f573e0e4dc96ccfb4a0e7 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Wed, 31 Dec 2025 15:59:15 +0000 Subject: [PATCH 05/13] address comments Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/model_executor/models/llama_eagle.py | 1 - vllm/model_executor/models/llama_eagle3.py | 1 - vllm/v1/spec_decode/eagle.py | 7 +++-- vllm/v1/spec_decode/eagle_mixin.py | 30 ++++++++++++++-------- 4 files changed, 22 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/models/llama_eagle.py b/vllm/model_executor/models/llama_eagle.py index ad230ffc11eb..4138d7307a68 100644 --- a/vllm/model_executor/models/llama_eagle.py +++ b/vllm/model_executor/models/llama_eagle.py @@ -163,7 +163,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class EagleLlamaForCausalLM(EagleMixin, LlamaForCausalLM): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) self.config = vllm_config.speculative_config.draft_model_config.hf_config diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index e64ea932d782..e3fe57b82adb 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -261,7 +261,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Eagle3LlamaForCausalLM(Eagle3Mixin, LlamaForCausalLM): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) self.config = vllm_config.speculative_config.draft_model_config.hf_config diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 6ddcc84b4033..c6711dd7d93d 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -346,7 +346,7 @@ def propose( # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1: - draft_token_ids = self.model.propose_chain(sample_hidden_states) + draft_token_ids = self.model.sample_chain(sample_hidden_states) return draft_token_ids.view(-1, 1) if self.uses_mrope: @@ -376,7 +376,7 @@ def propose( # [batch_size, num_tree_tokens] return torch.cat(draft_token_ids_list, dim=1) - draft_token_ids = self.model.propose_chain(sample_hidden_states) + draft_token_ids = self.model.sample_chain(sample_hidden_states) if self.allowed_attn_types is not None and not isinstance( attn_metadata, self.allowed_attn_types @@ -534,8 +534,7 @@ def propose( else: last_hidden_states, hidden_states = ret_hidden_states hidden_states = hidden_states[:batch_size] - draft_token_ids = self.model.propose_chain( - last_hidden_states[:batch_size]) + draft_token_ids = self.model.sample_chain(last_hidden_states[:batch_size]) draft_token_ids_list.append(draft_token_ids) # [batch_size, num_speculative_tokens] diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py index 08003338e87d..2bc937046342 100644 --- a/vllm/v1/spec_decode/eagle_mixin.py +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any, Optional +from typing import Any import torch @@ -8,7 +8,7 @@ class EagleMixin: compute_logits: Any - def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: + def sample_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: logits = self.compute_logits(hidden_states) draft_token_ids = logits.argmax(dim=-1) return draft_token_ids @@ -17,7 +17,7 @@ def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: class Eagle3Mixin: lm_head: Any logits_processor: Any - draft_id_to_target_id: Optional[Any] + draft_id_to_target_id: Any | None config: Any def compute_draft_logits( @@ -31,27 +31,35 @@ def compute_draft_logits( def compute_logits( self, hidden_states: torch.Tensor, - ) -> Optional[torch.Tensor]: + ) -> torch.Tensor | None: # Maps draft logits to target logits. logits = self.compute_draft_logits(hidden_states) if self.draft_id_to_target_id is None: - assert logits.shape[1] == self.config.vocab_size, \ - "Expected logits to have shape " \ + assert logits.shape[1] == self.config.vocab_size, ( + "Expected logits to have shape " f"(*, {self.config.vocab_size}), but got {logits.shape}" + ) return logits base = torch.arange(self.config.draft_vocab_size, device=logits.device) targets = base + self.draft_id_to_target_id - logits_new = logits.new_full(( - logits.shape[0], - self.config.vocab_size, - ), float('-inf')) + logits_new = logits.new_full( + ( + logits.shape[0], + self.config.vocab_size, + ), + float("-inf"), + ) logits_new[:, targets] = logits return logits_new - def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: + def sample_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: logits = self.compute_draft_logits(hidden_states) draft_token_ids = logits.argmax(dim=-1) + + if self.draft_id_to_target_id is None: + return draft_token_ids + offset = self.draft_id_to_target_id[draft_token_ids] draft_token_ids += offset return draft_token_ids From 58a84e41574992e2ad5ddfae6e14bcf0ebe8cb6e Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Wed, 31 Dec 2025 17:46:04 +0000 Subject: [PATCH 06/13] ruff Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle_mixin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py index 2bc937046342..ee759f7f9375 100644 --- a/vllm/v1/spec_decode/eagle_mixin.py +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -56,10 +56,10 @@ def compute_logits( def sample_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: logits = self.compute_draft_logits(hidden_states) draft_token_ids = logits.argmax(dim=-1) - + if self.draft_id_to_target_id is None: return draft_token_ids - + offset = self.draft_id_to_target_id[draft_token_ids] draft_token_ids += offset return draft_token_ids From 85449d3eee978cad93f535c3e8b52e928710bc5a Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Mon, 9 Feb 2026 01:00:57 +0000 Subject: [PATCH 07/13] fix Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 1841229346ee..e6a7be22f471 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -494,8 +494,7 @@ def propose( last_hidden_states, hidden_states = ret_hidden_states sample_hidden_states = last_hidden_states[token_indices_to_sample] - logits = self.model.compute_logits(sample_hidden_states) - + # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1 or self.parallel_drafting: draft_token_ids = self.model.sample_chain(sample_hidden_states) From efbbb8611087e86e622ec0522d2461f89585f573 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Mon, 9 Feb 2026 01:02:15 +0000 Subject: [PATCH 08/13] fix Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index e6a7be22f471..69018f474fb2 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -494,7 +494,7 @@ def propose( last_hidden_states, hidden_states = ret_hidden_states sample_hidden_states = last_hidden_states[token_indices_to_sample] - + # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1 or self.parallel_drafting: draft_token_ids = self.model.sample_chain(sample_hidden_states) From 0091604f987a44bd58a9b41fab055881712259b6 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 27 Feb 2026 20:29:53 +0000 Subject: [PATCH 09/13] resolve conflicy Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle_mixin.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py index ee759f7f9375..f6b097d47f87 100644 --- a/vllm/v1/spec_decode/eagle_mixin.py +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -9,9 +9,10 @@ class EagleMixin: compute_logits: Any def sample_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: - logits = self.compute_logits(hidden_states) - draft_token_ids = logits.argmax(dim=-1) - return draft_token_ids + """Greedy-sample draft tokens from hidden states.""" + if self.use_local_argmax_reduction: + return self.get_top_tokens(hidden_states) + return self.compute_logits(hidden_states).argmax(dim=-1) class Eagle3Mixin: @@ -54,12 +55,19 @@ def compute_logits( return logits_new def sample_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: - logits = self.compute_draft_logits(hidden_states) - draft_token_ids = logits.argmax(dim=-1) + if ( + self.use_local_argmax_reduction + and not (hasattr(self, "draft_id_to_target_id") + and self.draft_id_to_target_id is not None) + ): + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + else: + logits = self.compute_draft_logits(hidden_states) + draft_token_ids = logits.argmax(dim=-1) - if self.draft_id_to_target_id is None: - return draft_token_ids + if self.draft_id_to_target_id is None: + return draft_token_ids - offset = self.draft_id_to_target_id[draft_token_ids] - draft_token_ids += offset - return draft_token_ids + offset = self.draft_id_to_target_id[draft_token_ids] + draft_token_ids += offset + return draft_token_ids From 888d041fd4b662adf47a84e5dec9b38800af42f5 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 27 Feb 2026 20:39:19 +0000 Subject: [PATCH 10/13] resolve conflicy Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index ab56e8535d7e..0368e05829a0 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -497,11 +497,7 @@ def propose( # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1 or self.parallel_drafting: -<<<<<<< HEAD draft_token_ids = self.model.sample_chain(sample_hidden_states) -======= - draft_token_ids = self._greedy_sample(sample_hidden_states) ->>>>>>> 2decec9856033347f3129f1d1b2ec015e1ad88ea return draft_token_ids.view(-1, self.num_speculative_tokens) if self.uses_mrope: @@ -519,11 +515,7 @@ def propose( hidden_states = hidden_states[token_indices_to_sample] if isinstance(attn_metadata, TreeAttentionMetadata): -<<<<<<< HEAD - # Draft using tree attention. -======= # Draft using tree attention - requires full logits for top-k ->>>>>>> 2decec9856033347f3129f1d1b2ec015e1ad88ea logits = self.model.compute_logits(sample_hidden_states) draft_token_ids_list = self.propose_tree( batch_size=batch_size, @@ -536,11 +528,7 @@ def propose( # [batch_size, num_tree_tokens] return torch.cat(draft_token_ids_list, dim=1) -<<<<<<< HEAD - draft_token_ids = self.model.sample_chain(sample_hidden_states) -======= draft_token_ids = self._greedy_sample(sample_hidden_states) ->>>>>>> 2decec9856033347f3129f1d1b2ec015e1ad88ea if self.allowed_attn_types is not None and not isinstance( attn_metadata, self.allowed_attn_types @@ -705,11 +693,7 @@ def propose( last_hidden_states, hidden_states = ret_hidden_states hidden_states = hidden_states[:batch_size] -<<<<<<< HEAD - draft_token_ids = self.model.sample_chain(last_hidden_states[:batch_size]) -======= draft_token_ids = self._greedy_sample(last_hidden_states[:batch_size]) ->>>>>>> 2decec9856033347f3129f1d1b2ec015e1ad88ea draft_token_ids_list.append(draft_token_ids) # [batch_size, num_speculative_tokens] From ad8586ff2ca6b2f88be5f83af35358f0766755c2 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 27 Feb 2026 20:46:48 +0000 Subject: [PATCH 11/13] resolve conflicy Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle.py | 6 +++--- vllm/v1/spec_decode/eagle_mixin.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 0368e05829a0..b98de2fa7de2 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -497,7 +497,7 @@ def propose( # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1 or self.parallel_drafting: - draft_token_ids = self.model.sample_chain(sample_hidden_states) + draft_token_ids = self.model.sample_chain(sample_hidden_states, use_local_argmax_reduction=self.use_local_argmax_reduction) return draft_token_ids.view(-1, self.num_speculative_tokens) if self.uses_mrope: @@ -528,7 +528,7 @@ def propose( # [batch_size, num_tree_tokens] return torch.cat(draft_token_ids_list, dim=1) - draft_token_ids = self._greedy_sample(sample_hidden_states) + draft_token_ids = self.model.sample_chain(sample_hidden_states, use_local_argmax_reduction=self.use_local_argmax_reduction) if self.allowed_attn_types is not None and not isinstance( attn_metadata, self.allowed_attn_types @@ -693,7 +693,7 @@ def propose( last_hidden_states, hidden_states = ret_hidden_states hidden_states = hidden_states[:batch_size] - draft_token_ids = self._greedy_sample(last_hidden_states[:batch_size]) + draft_token_ids = self.model.sample_chain(last_hidden_states[:batch_size], use_local_argmax_reduction=self.use_local_argmax_reduction) draft_token_ids_list.append(draft_token_ids) # [batch_size, num_speculative_tokens] diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py index f6b097d47f87..2a6ec5d4354b 100644 --- a/vllm/v1/spec_decode/eagle_mixin.py +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -8,9 +8,9 @@ class EagleMixin: compute_logits: Any - def sample_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: + def sample_chain(self, hidden_states: torch.Tensor, use_local_argmax_reduction: bool = False) -> torch.Tensor: """Greedy-sample draft tokens from hidden states.""" - if self.use_local_argmax_reduction: + if use_local_argmax_reduction: return self.get_top_tokens(hidden_states) return self.compute_logits(hidden_states).argmax(dim=-1) @@ -54,9 +54,9 @@ def compute_logits( logits_new[:, targets] = logits return logits_new - def sample_chain(self, hidden_states: torch.Tensor) -> torch.Tensor: + def sample_chain(self, hidden_states: torch.Tensor, use_local_argmax_reduction: bool = False) -> torch.Tensor: if ( - self.use_local_argmax_reduction + use_local_argmax_reduction and not (hasattr(self, "draft_id_to_target_id") and self.draft_id_to_target_id is not None) ): From 369b9edd65a4537f933d600e3e2fb7cfb3d96b1d Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 27 Feb 2026 20:55:17 +0000 Subject: [PATCH 12/13] lint Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle.py | 15 ++++++++++++--- vllm/v1/spec_decode/eagle_mixin.py | 15 +++++++++------ 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index af3bd8110626..9ce1e25ef848 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -490,7 +490,10 @@ def propose( # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1 or self.parallel_drafting: - draft_token_ids = self.model.sample_chain(sample_hidden_states, use_local_argmax_reduction=self.use_local_argmax_reduction) + draft_token_ids = self.model.sample_chain( + sample_hidden_states, + use_local_argmax_reduction=self.use_local_argmax_reduction, + ) return draft_token_ids.view(-1, self.num_speculative_tokens) if self.uses_mrope: @@ -521,7 +524,10 @@ def propose( # [batch_size, num_tree_tokens] return torch.cat(draft_token_ids_list, dim=1) - draft_token_ids = self.model.sample_chain(sample_hidden_states, use_local_argmax_reduction=self.use_local_argmax_reduction) + draft_token_ids = self.model.sample_chain( + sample_hidden_states, + use_local_argmax_reduction=self.use_local_argmax_reduction, + ) if self.allowed_attn_types is not None and not isinstance( attn_metadata, self.allowed_attn_types @@ -679,7 +685,10 @@ def propose( last_hidden_states, hidden_states = ret_hidden_states hidden_states = hidden_states[:batch_size] - draft_token_ids = self.model.sample_chain(last_hidden_states[:batch_size], use_local_argmax_reduction=self.use_local_argmax_reduction) + draft_token_ids = self.model.sample_chain( + last_hidden_states[:batch_size], + use_local_argmax_reduction=self.use_local_argmax_reduction, + ) draft_token_ids_list.append(draft_token_ids) # [batch_size, num_speculative_tokens] diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py index 2a6ec5d4354b..a803fe4c2392 100644 --- a/vllm/v1/spec_decode/eagle_mixin.py +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -8,7 +8,9 @@ class EagleMixin: compute_logits: Any - def sample_chain(self, hidden_states: torch.Tensor, use_local_argmax_reduction: bool = False) -> torch.Tensor: + def sample_chain( + self, hidden_states: torch.Tensor, use_local_argmax_reduction: bool = False + ) -> torch.Tensor: """Greedy-sample draft tokens from hidden states.""" if use_local_argmax_reduction: return self.get_top_tokens(hidden_states) @@ -54,11 +56,12 @@ def compute_logits( logits_new[:, targets] = logits return logits_new - def sample_chain(self, hidden_states: torch.Tensor, use_local_argmax_reduction: bool = False) -> torch.Tensor: - if ( - use_local_argmax_reduction - and not (hasattr(self, "draft_id_to_target_id") - and self.draft_id_to_target_id is not None) + def sample_chain( + self, hidden_states: torch.Tensor, use_local_argmax_reduction: bool = False + ) -> torch.Tensor: + if use_local_argmax_reduction and not ( + hasattr(self, "draft_id_to_target_id") + and self.draft_id_to_target_id is not None ): return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) else: From 6f92f4dab29a93883ba301d872397d04236e298e Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 27 Feb 2026 20:58:08 +0000 Subject: [PATCH 13/13] lint Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- vllm/v1/spec_decode/eagle_mixin.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/spec_decode/eagle_mixin.py b/vllm/v1/spec_decode/eagle_mixin.py index a803fe4c2392..3748398fa5ff 100644 --- a/vllm/v1/spec_decode/eagle_mixin.py +++ b/vllm/v1/spec_decode/eagle_mixin.py @@ -6,6 +6,8 @@ class EagleMixin: + lm_head: Any + logits_processor: Any compute_logits: Any def sample_chain( @@ -13,7 +15,7 @@ def sample_chain( ) -> torch.Tensor: """Greedy-sample draft tokens from hidden states.""" if use_local_argmax_reduction: - return self.get_top_tokens(hidden_states) + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) return self.compute_logits(hidden_states).argmax(dim=-1)