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 6, 21:45 UTC
Add: Eagle3 support for Qwen3.5 (#36658)
Signed-off-by: Rahul-Tuli <rtuli@redhat.com>
Rahul Tuli committed
Mar 11, 2026 at 15:37 UTC
9d07a3d6e472c8e5a231a34ec9c38084605b037d
2 files changed
+25
-2
vllm/model_executor/models/qwen3_5.py
+11
@@ -75,6 +75,7 @@ from .interfaces import (
75
IsHybrid,
76
MixtureOfExperts,
77
MultiModalEmbeddings,
78
+ SupportsEagle3,
79
SupportsLoRA,
80
SupportsPP,
81
_require_is_multimodal,
@@ -353,6 +354,8 @@ class Qwen3_5Model(Qwen3NextModel):
354
else:
355
self.norm = PPMissingLayer()
356
357
+ self.aux_hidden_state_layers: tuple[int, ...] = ()
358
+
359
def load_fused_expert_weights(
360
self,
361
name: str,
@@ -536,6 +539,7 @@ class Qwen3_5Model(Qwen3NextModel):
539
class Qwen3_5ForCausalLMBase(
540
nn.Module,
541
HasInnerState,
542
+ SupportsEagle3,
543
SupportsLoRA,
544
SupportsPP,
545
):
@@ -592,6 +596,13 @@ class Qwen3_5ForCausalLMBase(
596
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
597
return self.model.embed_input_ids(input_ids)
598
599
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
600
+ self.model.aux_hidden_state_layers = layers
601
+
602
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
603
+ num_layers = len(self.model.layers)
604
+ return (2, num_layers // 2, num_layers - 3)
605
+
606
def forward(
607
self,
608
input_ids: torch.Tensor,
vllm/model_executor/models/qwen3_next.py
+14
-2
@@ -1148,6 +1148,8 @@ class Qwen3NextModel(nn.Module):
1148
else:
1149
self.norm = PPMissingLayer()
1150
1151
+ self.aux_hidden_state_layers: tuple[int, ...] = ()
1152
+
1153
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1154
return self.embed_tokens(input_ids)
1155
@@ -1157,7 +1159,7 @@ class Qwen3NextModel(nn.Module):
1159
positions: torch.Tensor,
1160
intermediate_tensors: IntermediateTensors | None = None,
1161
inputs_embeds: torch.Tensor | None = None,
1160
- ) -> torch.Tensor:
1162
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1163
if get_pp_group().is_first_rank:
1164
if inputs_embeds is not None:
1165
hidden_states = inputs_embeds
@@ -1169,7 +1171,15 @@ class Qwen3NextModel(nn.Module):
1171
hidden_states = intermediate_tensors["hidden_states"]
1172
residual = intermediate_tensors["residual"]
1173
1172
- for layer in islice(self.layers, self.start_layer, self.end_layer):
1174
+ aux_hidden_states = []
1175
+ for layer_idx, layer in enumerate(
1176
+ islice(self.layers, self.start_layer, self.end_layer),
1177
+ start=self.start_layer,
1178
+ ):
1179
+ if layer_idx in self.aux_hidden_state_layers:
1180
+ aux_hidden_states.append(
1181
+ hidden_states + residual if residual is not None else hidden_states
1182
+ )
1183
hidden_states, residual = layer(
1184
positions=positions,
1185
hidden_states=hidden_states,
@@ -1181,6 +1191,8 @@ class Qwen3NextModel(nn.Module):
1191
{"hidden_states": hidden_states, "residual": residual}
1192
)
1193
hidden_states, _ = self.norm(hidden_states, residual)
1194
+ if aux_hidden_states:
1195
+ return hidden_states, aux_hidden_states
1196
return hidden_states
1197
1198
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: