A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Add EAGLE-3 Speculative Decoding Support for Qwen3 MoE (#26485)
Signed-off-by: Rahul Tuli <rtuli@redhat.com>
Rahul Tuli committed
Oct 11, 2025 at 15:44 UTC
d2a71530c159c361f991a1ed986e64209651cc92
1 file changed
+33
-4
vllm/model_executor/models/qwen3_moe.py
+33
-4
@@ -64,7 +64,7 @@ from vllm.model_executor.model_loader.weight_utils import (
64
from vllm.model_executor.models.utils import sequence_parallel_chunk
65
from vllm.sequence import IntermediateTensors
66
67
-from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
67
+from .interfaces import MixtureOfExperts, SupportsEagle3, SupportsLoRA, SupportsPP
68
from .utils import (
69
AutoWeightsLoader,
70
PPMissingLayer,
@@ -422,6 +422,8 @@ class Qwen3MoeModel(nn.Module):
422
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
423
["hidden_states", "residual"], config.hidden_size
424
)
425
+ # Track layers for auxiliary hidden state outputs (EAGLE3)
426
+ self.aux_hidden_state_layers: tuple[int, ...] = ()
427
428
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
429
return self.embed_tokens(input_ids)
@@ -432,7 +434,9 @@ class Qwen3MoeModel(nn.Module):
434
positions: torch.Tensor,
435
intermediate_tensors: Optional[IntermediateTensors] = None,
436
inputs_embeds: Optional[torch.Tensor] = None,
435
- ) -> Union[torch.Tensor, IntermediateTensors]:
437
+ ) -> Union[
438
+ torch.Tensor, IntermediateTensors, tuple[torch.Tensor, list[torch.Tensor]]
439
+ ]:
440
if get_pp_group().is_first_rank:
441
if inputs_embeds is not None:
442
hidden_states = inputs_embeds
@@ -443,13 +447,29 @@ class Qwen3MoeModel(nn.Module):
447
assert intermediate_tensors is not None
448
hidden_states = intermediate_tensors["hidden_states"]
449
residual = intermediate_tensors["residual"]
446
- for layer in islice(self.layers, self.start_layer, self.end_layer):
450
+
451
+ aux_hidden_states = []
452
+ for layer_idx, layer in enumerate(
453
+ islice(self.layers, self.start_layer, self.end_layer),
454
+ start=self.start_layer,
455
+ ):
456
+ # Collect auxiliary hidden states if specified
457
+ if layer_idx in self.aux_hidden_state_layers:
458
+ aux_hidden_state = (
459
+ hidden_states + residual if residual is not None else hidden_states
460
+ )
461
+ aux_hidden_states.append(aux_hidden_state)
462
hidden_states, residual = layer(positions, hidden_states, residual)
463
+
464
if not get_pp_group().is_last_rank:
465
return IntermediateTensors(
466
{"hidden_states": hidden_states, "residual": residual}
467
)
468
hidden_states, _ = self.norm(hidden_states, residual)
469
+
470
+ # Return auxiliary hidden states if collected
471
+ if len(aux_hidden_states) > 0:
472
+ return hidden_states, aux_hidden_states
473
return hidden_states
474
475
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
@@ -606,7 +626,9 @@ class Qwen3MoeModel(nn.Module):
626
return loaded_params
627
628
609
-class Qwen3MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExperts):
629
+class Qwen3MoeForCausalLM(
630
+ nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3, MixtureOfExperts
631
+):
632
packed_modules_mapping = {
633
"qkv_proj": [
634
"q_proj",
@@ -702,6 +724,13 @@ class Qwen3MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExperts)
724
moe.n_redundant_experts = self.num_redundant_experts
725
moe.experts.update_expert_map()
726
727
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
728
+ self.model.aux_hidden_state_layers = layers
729
+
730
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
731
+ num_layers = len(self.model.layers)
732
+ return (2, num_layers // 2, num_layers - 3)
733
+
734
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
735
return self.model.get_input_embeddings(input_ids)
736