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: Support for multiple hidden layers in Eagle3 (#26164)
Signed-off-by: Rahul Tuli <rtuli@redhat.com>
Rahul Tuli committed
Oct 9, 2025 at 13:00 UTC
cf4cd6c24f48e5efdcfa1a94915c77e1437a94b4
2 files changed
+29
-13
tests/speculative_decoding/speculators/test_eagle3.py
+4
@@ -22,6 +22,10 @@ from vllm.model_executor.models.interfaces import supports_eagle3
22
"nm-testing/Speculator-Qwen3-8B-Eagle3-converted-071-quantized-w4a16",
23
id="qwen3-eagle3-speculator-w4a16-verifier",
24
),
25
+ pytest.param(
26
+ "nm-testing/random-weights-llama3.1.8b-2layer-eagle3",
27
+ id="llama3-eagl3-multiple-layers",
28
+ ),
29
],
30
)
31
def test_eagle3_speculators_model(
vllm/model_executor/models/llama_eagle3.py
+25
-13
@@ -34,15 +34,20 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
34
vllm_config: VllmConfig,
35
prefix: str = "",
36
config: Optional[LlamaConfig] = None,
37
+ layer_idx: int = 0,
38
) -> None:
39
super().__init__(vllm_config, prefix=prefix, config=config)
40
41
config = config or vllm_config.model_config.hf_config
42
quant_config = self.get_quant_config(vllm_config)
43
44
+ # First layer uses 2*hidden_size (embeds + hidden_states concatenated)
45
+ # Subsequent layers use hidden_size (only hidden_states, no embeds)
46
+ qkv_input_size = 2 * self.hidden_size if layer_idx == 0 else self.hidden_size
47
+
48
# override qkv
49
self.self_attn.qkv_proj = QKVParallelLinear(
45
- 2 * self.hidden_size,
50
+ qkv_input_size,
51
self.self_attn.head_dim,
52
self.self_attn.total_num_heads,
53
self.self_attn.total_num_kv_heads,
@@ -52,6 +57,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
57
)
58
59
self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
60
+ self.layer_idx = layer_idx
61
62
if getattr(config, "norm_before_residual", False):
63
self._residual_norm = self._norm_before_residual
@@ -90,11 +96,15 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
96
hidden_states: torch.Tensor,
97
residual: Optional[torch.Tensor],
98
) -> tuple[torch.Tensor, torch.Tensor]:
93
- embeds = self.input_layernorm(embeds)
94
-
95
- hidden_states, residual = self._residual_norm(hidden_states=hidden_states)
99
+ if self.layer_idx == 0:
100
+ # First layer: concatenate embeds with hidden_states
101
+ embeds = self.input_layernorm(embeds)
102
+ hidden_states, residual = self._residual_norm(hidden_states=hidden_states)
103
+ hidden_states = torch.cat([embeds, hidden_states], dim=-1)
104
+ else:
105
+ # Subsequent layers: process hidden_states and residuals only
106
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
107
97
- hidden_states = torch.cat([embeds, hidden_states], dim=-1)
108
# Self Attention
109
hidden_states = self.self_attn(
110
positions=positions,
@@ -133,9 +143,11 @@ class LlamaModel(nn.Module):
143
[
144
LlamaDecoderLayer(
145
current_vllm_config,
136
- prefix=maybe_prefix(prefix, f"layers.{start_layer_id}"),
146
+ prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
147
config=self.config,
148
+ layer_idx=layer_idx,
149
)
150
+ for layer_idx in range(self.config.num_hidden_layers)
151
]
152
)
153
if hasattr(self.config, "target_hidden_size"):
@@ -166,13 +178,13 @@ class LlamaModel(nn.Module):
178
assert hidden_states.shape[-1] == input_embeds.shape[-1]
179
180
residual = None
169
- hidden_states, residual = self.layers[0](
170
- positions,
171
- input_embeds,
172
- hidden_states,
173
- residual,
174
- )
175
-
181
+ for layer in self.layers:
182
+ hidden_states, residual = layer(
183
+ positions=positions,
184
+ embeds=input_embeds,
185
+ hidden_states=hidden_states,
186
+ residual=residual,
187
+ )
188
hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
189
return hidden_states, hidden_prenorm
190