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 9, 06:15 UTC
Add EAGLE3 support for AFMoE (#33111)
Signed-off-by: AutumnAurelium <88015631+AutumnAurelium@users.noreply.github.com>
AutumnAurelium committed
Jan 30, 2026 at 22:53 UTC
f3888aca83edf5125c37d0983f1df94df27d9920
2 files changed
+30
-5
vllm/config/speculative.py
+1
@@ -682,6 +682,7 @@ class SpeculativeConfig:
682
"gpt_oss",
683
"hunyuan_vl",
684
"hunyuan_v1_dense",
685
+ "afmoe",
686
]
687
if (
688
self.method == "eagle3"
vllm/model_executor/models/afmoe.py
+29
-5
@@ -36,7 +36,11 @@ from vllm.model_executor.model_loader.weight_utils import (
36
default_weight_loader,
37
maybe_remap_kv_scale_name,
38
)
39
-from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
39
+from vllm.model_executor.models.interfaces import (
40
+ SupportsEagle3,
41
+ SupportsLoRA,
42
+ SupportsPP,
43
+)
44
from vllm.model_executor.models.llama import LlamaMLP as AfmoeMLP
45
from vllm.model_executor.models.utils import (
46
AutoWeightsLoader,
@@ -416,6 +420,8 @@ class AfmoeModel(nn.Module):
420
else:
421
self.norm = PPMissingLayer()
422
423
+ self.aux_hidden_state_layers = tuple[int, ...]()
424
+
425
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
426
["hidden_states", "residual"], config.hidden_size
427
)
@@ -429,7 +435,7 @@ class AfmoeModel(nn.Module):
435
positions: torch.Tensor,
436
intermediate_tensors: IntermediateTensors | None = None,
437
inputs_embeds: torch.Tensor | None = None,
432
- ) -> torch.Tensor | IntermediateTensors:
438
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
439
if get_pp_group().is_first_rank:
440
if inputs_embeds is not None:
441
hidden_states = inputs_embeds
@@ -446,7 +452,14 @@ class AfmoeModel(nn.Module):
452
hidden_states = intermediate_tensors["hidden_states"]
453
residual = intermediate_tensors["residual"]
454
449
- for layer in islice(self.layers, self.start_layer, self.end_layer):
455
+ aux_hidden_states = []
456
+ for idx, layer in enumerate(
457
+ islice(self.layers, self.start_layer, self.end_layer)
458
+ ):
459
+ if idx in self.aux_hidden_state_layers:
460
+ aux_hidden_states.append(
461
+ hidden_states + residual if residual is not None else hidden_states
462
+ )
463
hidden_states, residual = layer(positions, hidden_states, residual)
464
465
if not get_pp_group().is_last_rank:
@@ -455,6 +468,10 @@ class AfmoeModel(nn.Module):
468
)
469
470
hidden_states, _ = self.norm(hidden_states, residual)
471
+
472
+ if len(aux_hidden_states) > 0:
473
+ return hidden_states, aux_hidden_states
474
+
475
return hidden_states
476
477
def make_empty_intermediate_tensors(
@@ -586,7 +603,7 @@ class AfmoeModel(nn.Module):
603
return loaded_params
604
605
589
-class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
606
+class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
607
packed_modules_mapping = {
608
"qkv_proj": [
609
"q_proj",
@@ -673,13 +690,20 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
690
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
691
return self.model.embed_input_ids(input_ids)
692
693
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
694
+ self.model.aux_hidden_state_layers = layers
695
+
696
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
697
+ num_layers = len(self.model.layers)
698
+ return (2, num_layers // 2, num_layers - 3)
699
+
700
def forward(
701
self,
702
input_ids: torch.Tensor | None,
703
positions: torch.Tensor,
704
intermediate_tensors: IntermediateTensors | None = None,
705
inputs_embeds: torch.Tensor | None = None,
682
- ) -> torch.Tensor | IntermediateTensors:
706
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
707
hidden_states = self.model(
708
input_ids, positions, intermediate_tensors, inputs_embeds
709
)