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 15, 13:30 UTC
[V1][Spec Decode] EAGLE-3 Support (#16937)
Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai> Co-authored-by: Bryan Lu <yuzhelu@amazon.com>
Benjamin Chislett committed
Apr 25, 2025 at 18:43 UTC
a0e619e62a778073d53fcab2429b5f2ac67d687c
12 files changed
+359
-35
examples/offline_inference/eagle.py
+11
-3
@@ -52,8 +52,8 @@ def main():
52
53
args = parse_args()
54
55
- model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
56
- eagle_dir = "abhigoyal/EAGLE-LLaMA3-Instruct-8B-vllm"
55
+ model_dir = "meta-llama/Llama-3.1-8B-Instruct"
56
+ eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
57
58
max_model_len = 2048
59
@@ -81,7 +81,7 @@ def main():
81
max_num_seqs=args.max_num_seqs,
82
gpu_memory_utilization=0.8,
83
speculative_config={
84
- "method": "eagle",
84
+ "method": "eagle3" if "eagle3" in eagle_dir.lower() else "eagle",
85
"model": eagle_dir,
86
"num_speculative_tokens": args.num_spec_tokens,
87
"draft_tensor_parallel_size": args.draft_tp,
@@ -95,6 +95,9 @@ def main():
95
outputs = llm.generate(prompt_token_ids=prompt_ids,
96
sampling_params=sampling_params)
97
98
+ if not hasattr(outputs, "metrics") or outputs.metrics is None:
99
+ return
100
+
101
# calculate the average number of accepted tokens per forward pass, +1 is
102
# to account for the token from the target model that's always going to be
103
# accepted
@@ -109,6 +112,11 @@ def main():
112
{sum(acceptance_counts) / acceptance_counts[0]:.2f}")
113
print("-" * 50)
114
115
+ # print acceptance at each token position
116
+ for i in range(len(acceptance_counts)):
117
+ print(f"acceptance at token {i}:"
118
+ f"{acceptance_counts[i] / (acceptance_counts[0]):.2f}")
119
+
120
121
if __name__ == "__main__":
122
main()
tests/models/registry.py
+4
@@ -393,6 +393,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
393
trust_remote_code=True,
394
speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
395
tokenizer="meta-llama/Meta-Llama-3-8B-Instruct"), # noqa: E501
396
+ "Eagle3LlamaForCausalLM": _HfExamplesInfo("yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", # noqa: E501
397
+ trust_remote_code=True,
398
+ speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
399
+ tokenizer="meta-llama/Llama-3.1-8B-Instruct"),
400
}
401
402
_TRANSFORMERS_MODELS = {
tests/v1/e2e/test_spec_decode.py
+16
-8
@@ -50,12 +50,15 @@ def sampling_config():
50
51
@pytest.fixture
52
def model_name():
53
- return "meta-llama/Meta-Llama-3-8B-Instruct"
53
+ return "meta-llama/Llama-3.1-8B-Instruct"
54
55
56
-@pytest.fixture
56
def eagle_model_name():
58
- return "yuhuili/EAGLE-LLaMA3-Instruct-8B"
57
+ return "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
58
+
59
+
60
+def eagle3_model_name():
61
+ return "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
62
63
64
def test_ngram_correctness(
@@ -102,12 +105,13 @@ def test_ngram_correctness(
105
del spec_llm
106
107
108
+@pytest.mark.parametrize("use_eagle3", [False, True], ids=["eagle", "eagle3"])
109
def test_eagle_correctness(
110
monkeypatch: pytest.MonkeyPatch,
111
test_prompts: list[list[dict[str, Any]]],
112
sampling_config: SamplingParams,
113
model_name: str,
110
- eagle_model_name: str,
114
+ use_eagle3: bool,
115
):
116
'''
117
Compare the outputs of a original LLM and a speculative LLM
@@ -116,18 +120,22 @@ def test_eagle_correctness(
120
with monkeypatch.context() as m:
121
m.setenv("VLLM_USE_V1", "1")
122
119
- ref_llm = LLM(model=model_name, max_model_len=1024)
123
+ ref_llm = LLM(model=model_name, max_model_len=2048)
124
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
125
del ref_llm
126
127
+ spec_model_name = eagle3_model_name(
128
+ ) if use_eagle3 else eagle_model_name()
129
spec_llm = LLM(
130
model=model_name,
131
+ trust_remote_code=True,
132
speculative_config={
126
- "method": "eagle",
127
- "model": eagle_model_name,
133
+ "method": "eagle3" if use_eagle3 else "eagle",
134
+ "model": spec_model_name,
135
"num_speculative_tokens": 3,
136
+ "max_model_len": 2048,
137
},
130
- max_model_len=1024,
138
+ max_model_len=2048,
139
)
140
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
141
matches = 0
vllm/config.py
+10
-3
@@ -2339,9 +2339,10 @@ class SpeculativeConfig:
2339
)
2340
2341
# Automatically detect the method
2342
- if self.method == 'eagle':
2342
+ if self.method in ('eagle', 'eagle3'):
2343
pass
2344
- elif "eagle-" in self.draft_model_config.model.lower():
2344
+ elif "eagle-" in self.draft_model_config.model.lower() or \
2345
+ "eagle3-" in self.draft_model_config.model.lower():
2346
self.method = "eagle"
2347
elif self.draft_model_config.hf_config.model_type == "medusa":
2348
self.method = "medusa"
@@ -2352,7 +2353,7 @@ class SpeculativeConfig:
2353
self.method = "draft_model"
2354
2355
# Replace hf_config for EAGLE draft_model
2355
- if self.method == "eagle":
2356
+ if self.method in ("eagle", "eagle3"):
2357
if self.enable_chunked_prefill and not envs.VLLM_USE_V1:
2358
raise ValueError(
2359
"Chunked prefill and EAGLE are not compatible "
@@ -2549,6 +2550,12 @@ class SpeculativeConfig:
2550
"speculative decoding is > 1, but got "
2551
f"{self.disable_by_batch_size=}")
2552
2553
+ if self.method == "eagle3" and self.target_model_config and \
2554
+ "llama" not in self.target_model_config.hf_text_config.model_type:
2555
+ raise ValueError(
2556
+ "Eagle3 is only supported for Llama models. "
2557
+ f"Got {self.target_model_config.hf_text_config.model_type=}")
2558
+
2559
@property
2560
def num_lookahead_slots(self) -> int:
2561
"""The number of additional slots the scheduler should allocate per
vllm/engine/arg_utils.py
+1
-1
@@ -1459,7 +1459,7 @@ class EngineArgs:
1459
if speculative_method:
1460
if speculative_method in ("ngram", "[ngram]"):
1461
is_ngram_enabled = True
1462
- elif speculative_method == "eagle":
1462
+ elif speculative_method in ("eagle", "eagle3"):
1463
is_eagle_enabled = True
1464
else:
1465
speculative_model = self.speculative_config.get("model")
vllm/model_executor/models/llama.py
+17
-1
@@ -330,6 +330,8 @@ class LlamaModel(nn.Module):
330
else:
331
self.norm = PPMissingLayer()
332
333
+ self.aux_hidden_state_layers: tuple[int] = tuple()
334
+
335
self.make_empty_intermediate_tensors = (
336
make_empty_intermediate_tensors_factory(
337
["hidden_states", "residual"], config.hidden_size))
@@ -355,7 +357,11 @@ class LlamaModel(nn.Module):
357
hidden_states = intermediate_tensors["hidden_states"]
358
residual = intermediate_tensors["residual"]
359
358
- for layer in self.layers[self.start_layer:self.end_layer]:
360
+ aux_hidden_states = []
361
+ for idx, layer in enumerate(
362
+ self.layers[self.start_layer:self.end_layer]):
363
+ if idx in self.aux_hidden_state_layers:
364
+ aux_hidden_states.append(hidden_states + residual)
365
hidden_states, residual = layer(positions, hidden_states, residual)
366
367
if not get_pp_group().is_last_rank:
@@ -365,6 +371,9 @@ class LlamaModel(nn.Module):
371
})
372
373
hidden_states, _ = self.norm(hidden_states, residual)
374
+
375
+ if len(aux_hidden_states) > 0:
376
+ return hidden_states, aux_hidden_states
377
return hidden_states
378
379
def load_weights(self, weights: Iterable[Tuple[str,
@@ -517,6 +526,13 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
526
self.make_empty_intermediate_tensors = (
527
self.model.make_empty_intermediate_tensors)
528
529
+ def set_aux_hidden_state_layers(self, layers: tuple[int]) -> None:
530
+ self.model.aux_hidden_state_layers = layers
531
+
532
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int]:
533
+ num_layers = len(self.model.layers)
534
+ return (2, num_layers // 2, num_layers - 3)
535
+
536
def _init_model(self,
537
vllm_config: VllmConfig,
538
prefix: str = "",
vllm/model_executor/models/llama_eagle.py
+2
-1
@@ -82,7 +82,8 @@ class LlamaModel(nn.Module):
82
hidden_states,
83
residual,
84
)
85
- return hidden_states + residual
85
+ hidden_states = hidden_states + residual
86
+ return hidden_states, hidden_states
87
88
def load_weights(self, weights: Iterable[Tuple[str,
89
torch.Tensor]]) -> Set[str]:
vllm/model_executor/models/llama_eagle3.py
new
+232
@@ -0,0 +1,232 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from typing import Iterable, Optional, Set, Tuple
4
+
5
+import torch
6
+import torch.nn as nn
7
+from transformers import LlamaConfig
8
+
9
+from vllm.config import ModelConfig
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.layers.layernorm import RMSNorm
12
+from vllm.model_executor.layers.linear import QKVParallelLinear
13
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
14
+from vllm.model_executor.layers.quantization.base_config import (
15
+ QuantizationConfig)
16
+from vllm.model_executor.layers.vocab_parallel_embedding import (
17
+ DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
18
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
19
+from vllm.model_executor.models.llama import (LlamaDecoderLayer,
20
+ LlamaForCausalLM)
21
+from vllm.v1.sample.metadata import SamplingMetadata
22
+
23
+from .utils import AutoWeightsLoader, maybe_prefix
24
+
25
+logger = init_logger(__name__)
26
+
27
+
28
+class LlamaDecoderLayer(LlamaDecoderLayer):
29
+
30
+ def __init__(
31
+ self,
32
+ config: LlamaConfig,
33
+ quant_config: Optional[QuantizationConfig] = None,
34
+ prefix: str = "",
35
+ ) -> None:
36
+ super().__init__(config, quant_config=quant_config, prefix=prefix)
37
+
38
+ # override qkv
39
+ self.self_attn.qkv_proj = QKVParallelLinear(
40
+ 2 * self.hidden_size,
41
+ self.self_attn.head_dim,
42
+ self.self_attn.total_num_heads,
43
+ self.self_attn.total_num_kv_heads,
44
+ bias=False,
45
+ quant_config=quant_config,
46
+ prefix=maybe_prefix(prefix, "qkv_proj"),
47
+ )
48
+
49
+ self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
50
+
51
+ def forward(
52
+ self,
53
+ positions: torch.Tensor,
54
+ embeds: torch.Tensor,
55
+ hidden_states: torch.Tensor,
56
+ residual: Optional[torch.Tensor],
57
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
58
+
59
+ residual = hidden_states
60
+ embeds = self.input_layernorm(embeds)
61
+ hidden_states = self.hidden_norm(hidden_states)
62
+
63
+ hidden_states = torch.cat([embeds, hidden_states], dim=-1)
64
+ # Self Attention
65
+ hidden_states = self.self_attn(
66
+ positions=positions,
67
+ hidden_states=hidden_states,
68
+ )
69
+
70
+ hidden_states, residual = self.post_attention_layernorm(
71
+ hidden_states, residual)
72
+
73
+ # Fully Connected
74
+ hidden_states = self.mlp(hidden_states)
75
+
76
+ return hidden_states, residual
77
+
78
+
79
+class LlamaModel(nn.Module):
80
+
81
+ def __init__(
82
+ self,
83
+ *,
84
+ model_config: ModelConfig,
85
+ start_layer_id: int = 0,
86
+ prefix: str = "",
87
+ ) -> None:
88
+ super().__init__()
89
+ self.config = model_config.hf_config
90
+ self.vocab_size = self.config.vocab_size
91
+ self.embed_tokens = VocabParallelEmbedding(
92
+ self.config.vocab_size,
93
+ self.config.hidden_size,
94
+ prefix=maybe_prefix(prefix, "embed_tokens"),
95
+ )
96
+ self.layers = nn.ModuleList([
97
+ LlamaDecoderLayer(
98
+ self.config,
99
+ prefix=maybe_prefix(prefix, f"layers.{start_layer_id}"),
100
+ )
101
+ ])
102
+ if hasattr(self.config, "target_hidden_size"):
103
+ self.fc = torch.nn.Linear(self.config.target_hidden_size * 3,
104
+ self.config.hidden_size,
105
+ bias=False)
106
+ else:
107
+ self.fc = torch.nn.Linear(self.config.hidden_size * 3,
108
+ self.config.hidden_size,
109
+ bias=False)
110
+ self.norm = RMSNorm(
111
+ self.config.hidden_size,
112
+ eps=self.config.rms_norm_eps,
113
+ )
114
+
115
+ def forward(
116
+ self,
117
+ input_ids: torch.Tensor,
118
+ positions: torch.Tensor,
119
+ hidden_states: torch.Tensor,
120
+ ) -> torch.Tensor:
121
+ input_embeds = self.embed_tokens(input_ids)
122
+ if (hidden_states.shape[-1] != input_embeds.shape[-1]):
123
+ hidden_states = self.fc(hidden_states)
124
+
125
+ residual = None
126
+ hidden_states, residual = self.layers[0](
127
+ positions,
128
+ input_embeds,
129
+ hidden_states,
130
+ residual,
131
+ )
132
+
133
+ hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
134
+ return hidden_states, hidden_prenorm
135
+
136
+ def load_weights(self, weights: Iterable[Tuple[str,
137
+ torch.Tensor]]) -> Set[str]:
138
+ stacked_params_mapping = [
139
+ # (param_name, shard_name, shard_id)
140
+ (".qkv_proj", ".q_proj", "q"),
141
+ (".qkv_proj", ".k_proj", "k"),
142
+ (".qkv_proj", ".v_proj", "v"),
143
+ (".gate_up_proj", ".gate_proj", 0),
144
+ (".gate_up_proj", ".up_proj", 1),
145
+ ]
146
+ params_dict = dict(self.named_parameters())
147
+ loaded_params: Set[str] = set()
148
+ for name, loaded_weight in weights:
149
+ if 'midlayer.' in name:
150
+ name = name.replace('midlayer.', 'layers.0.')
151
+ for param_name, weight_name, shard_id in stacked_params_mapping:
152
+ if weight_name not in name:
153
+ continue
154
+ name = name.replace(weight_name, param_name)
155
+ param = params_dict[name]
156
+ weight_loader = param.weight_loader
157
+ weight_loader(param, loaded_weight, shard_id)
158
+ break
159
+ else:
160
+ param = params_dict[name]
161
+ weight_loader = getattr(param, "weight_loader",
162
+ default_weight_loader)
163
+ weight_loader(param, loaded_weight)
164
+ loaded_params.add(name)
165
+ return loaded_params
166
+
167
+
168
+class Eagle3LlamaForCausalLM(LlamaForCausalLM):
169
+
170
+ def __init__(self, *, model_config: ModelConfig, start_layer_id: int = 0):
171
+ nn.Module.__init__(self)
172
+ self.config = model_config.hf_config
173
+ self.model = LlamaModel(model_config=model_config,
174
+ start_layer_id=start_layer_id,
175
+ prefix="model")
176
+
177
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
178
+ self.lm_head = ParallelLMHead(
179
+ self.config.draft_vocab_size,
180
+ self.config.hidden_size,
181
+ org_num_embeddings=self.config.draft_vocab_size,
182
+ padding_size=(DEFAULT_VOCAB_PADDING_SIZE),
183
+ prefix="")
184
+ self.logits_processor = LogitsProcessor(self.config.draft_vocab_size,
185
+ scale=logit_scale)
186
+ self.draft_id_to_target_id = nn.Parameter(
187
+ torch.zeros((self.config.draft_vocab_size),
188
+ dtype=torch.long).type(torch.LongTensor),
189
+ requires_grad=False,
190
+ )
191
+
192
+ def forward(
193
+ self,
194
+ input_ids: torch.Tensor,
195
+ positions: torch.Tensor,
196
+ hidden_states: torch.Tensor,
197
+ ) -> torch.Tensor:
198
+ return self.model(input_ids, positions, hidden_states)
199
+
200
+ def compute_logits(
201
+ self,
202
+ hidden_states: torch.Tensor,
203
+ sampling_metadata: SamplingMetadata,
204
+ ) -> Optional[torch.Tensor]:
205
+ logits = self.logits_processor(self.lm_head, hidden_states,
206
+ sampling_metadata)
207
+ base = torch.arange(self.config.draft_vocab_size, device=logits.device)
208
+ targets = base + self.draft_id_to_target_id
209
+ logits_new = logits.new_full((
210
+ logits.shape[0],
211
+ self.config.vocab_size,
212
+ ), float('-inf'))
213
+ logits_new[:, targets] = logits
214
+ return logits_new
215
+
216
+ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
217
+ loader = AutoWeightsLoader(
218
+ self,
219
+ skip_prefixes=None,
220
+ )
221
+
222
+ model_weights = {}
223
+ for name, loaded_weight in weights:
224
+ if "t2d" in name:
225
+ continue
226
+ if "d2t" in name:
227
+ name = name.replace("d2t", "draft_id_to_target_id")
228
+ elif "lm_head" not in name:
229
+ name = "model." + name
230
+ model_weights[name] = loaded_weight
231
+
232
+ return loader.load_weights(model_weights.items())
vllm/model_executor/models/registry.py
+1
@@ -214,6 +214,7 @@ _MULTIMODAL_MODELS = {
214
_SPECULATIVE_DECODING_MODELS = {
215
"EAGLEModel": ("eagle", "EAGLE"),
216
"EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
217
+ "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
218
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
219
"MedusaModel": ("medusa", "Medusa"),
220
"MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
vllm/v1/core/sched/scheduler.py
+1
-1
@@ -126,7 +126,7 @@ class Scheduler(SchedulerInterface):
126
self.num_spec_tokens = self.num_lookahead_tokens = 0
127
if speculative_config:
128
self.num_spec_tokens = speculative_config.num_speculative_tokens
129
- if speculative_config.method == "eagle":
129
+ if speculative_config.method in ("eagle", "eagle3"):
130
self.num_lookahead_tokens = self.num_spec_tokens
131
132
def schedule(self) -> SchedulerOutput:
vllm/v1/spec_decode/eagle.py
+28
-11
@@ -6,12 +6,16 @@ import triton.language as tl
6
7
from vllm.config import VllmConfig, set_current_vllm_config
8
from vllm.forward_context import set_forward_context
9
+from vllm.logger import init_logger
10
from vllm.model_executor.model_loader.loader import get_model_loader
11
from vllm.model_executor.model_loader.utils import set_default_torch_dtype
12
from vllm.model_executor.models.llama_eagle import EagleLlamaForCausalLM
13
+from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
14
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
15
from vllm.v1.sample.metadata import SamplingMetadata
16
17
+logger = init_logger(__name__)
18
+
19
PADDING_SLOT_ID = -1
20
21
@@ -87,12 +91,12 @@ class EagleProposer:
91
)
92
93
with set_forward_context(attn_metadata, self.vllm_config):
90
- hidden_states = self.model(
94
+ hidden_states_logits, hidden_states_fwd = self.model(
95
input_ids=input_ids,
96
hidden_states=target_hidden_states,
97
positions=target_positions,
98
)
95
- sample_hidden_states = hidden_states[last_token_indices]
99
+ sample_hidden_states = hidden_states_logits[last_token_indices]
100
logits = self.model.compute_logits(sample_hidden_states, None)
101
draft_token_ids = logits.argmax(dim=-1)
102
@@ -105,7 +109,7 @@ class EagleProposer:
109
draft_token_ids_list = [draft_token_ids]
110
111
positions = target_positions[last_token_indices]
108
- hidden_states = sample_hidden_states
112
+ hidden_states = hidden_states_fwd[last_token_indices]
113
attn_metadata.num_actual_tokens = batch_size
114
attn_metadata.max_query_len = 1
115
attn_metadata.query_start_loc = self.arange[:batch_size + 1]
@@ -151,12 +155,12 @@ class EagleProposer:
155
156
# Run the model.
157
with set_forward_context(attn_metadata, self.vllm_config):
154
- hidden_states = self.model(
158
+ hidden_states_logits, hidden_states = self.model(
159
input_ids=input_ids,
160
hidden_states=hidden_states,
161
positions=clamped_positions,
162
)
159
- logits = self.model.compute_logits(hidden_states, None)
163
+ logits = self.model.compute_logits(hidden_states_logits, None)
164
draft_token_ids = logits.argmax(dim=-1)
165
draft_token_ids_list.append(draft_token_ids)
166
@@ -221,15 +225,28 @@ class EagleProposer:
225
with set_default_torch_dtype(
226
draft_model_config.dtype), set_current_vllm_config(
227
self.vllm_config):
224
- self.model = EagleLlamaForCausalLM(
225
- model_config=draft_model_config,
226
- start_layer_id=target_layer_num).to(target_device)
227
-
228
- self.model.load_weights(
228
+ if self.vllm_config.speculative_config.method == "eagle":
229
+ self.model = EagleLlamaForCausalLM(
230
+ model_config=draft_model_config,
231
+ start_layer_id=target_layer_num).to(target_device)
232
+ else:
233
+ assert self.vllm_config.speculative_config.method == "eagle3"
234
+ self.model = Eagle3LlamaForCausalLM(
235
+ model_config=draft_model_config,
236
+ start_layer_id=target_layer_num).to(target_device)
237
+
238
+ loaded_weights = self.model.load_weights(
239
loader.get_all_weights(
240
self.vllm_config.speculative_config.draft_model_config,
241
self.model))
232
- self.model.lm_head = target_model.lm_head
242
+ if self.vllm_config.speculative_config.method == "eagle3":
243
+ if "model.embed_tokens.weight" not in loaded_weights:
244
+ logger.info(
245
+ "Loading EAGLE embedding weights from the target model.")
246
+ self.model.model.embed_tokens = target_model.model.embed_tokens
247
+ else:
248
+ logger.info("Loading EAGLE LM head weights from the target model.")
249
+ self.model.lm_head = target_model.lm_head
250
251
252
# NOTE(woosuk): Currently, the below code is not used and we always use argmax
vllm/v1/worker/gpu_model_runner.py
+36
-6
@@ -165,14 +165,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
165
166
# Set up speculative decoding.
167
self.use_spec_decode = False
168
+ self.use_aux_hidden_state_outputs = False
169
if self.speculative_config:
170
self.use_spec_decode = True
171
if get_pp_group().is_last_rank:
172
if self.speculative_config.method == "ngram":
173
self.drafter = NgramProposer(self.vllm_config)
173
- elif self.speculative_config.method == "eagle":
174
+ elif self.speculative_config.method == "eagle" or \
175
+ self.speculative_config.method == "eagle3":
176
self.drafter = EagleProposer(self.vllm_config,
177
self.device) # type: ignore
178
+ if self.speculative_config.method == "eagle3":
179
+ self.use_aux_hidden_state_outputs = True
180
else:
181
raise ValueError("Unknown speculative decoding method: "
182
f"{self.speculative_config.method}")
@@ -1079,12 +1083,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1083
# Run the decoder.
1084
# Use persistent buffers for CUDA graphs.
1085
with set_forward_context(attn_metadata, self.vllm_config):
1082
- hidden_states = self.model(
1086
+ output = self.model(
1087
input_ids=input_ids,
1088
positions=positions,
1089
intermediate_tensors=intermediate_tensors,
1090
inputs_embeds=inputs_embeds,
1091
)
1092
+
1093
+ if self.use_aux_hidden_state_outputs:
1094
+ hidden_states, aux_hidden_states = output
1095
+ else:
1096
+ hidden_states = output
1097
+
1098
if not get_pp_group().is_last_rank:
1099
# For mid-pipeline stages, return the hidden states.
1100
return hidden_states
@@ -1182,7 +1192,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1192
assert isinstance(self.drafter, NgramProposer)
1193
spec_token_ids = self.generate_draft_token_ids(
1194
valid_sampled_token_ids, sampling_metadata)
1185
- elif self.speculative_config.method == "eagle":
1195
+ elif self.speculative_config.method == "eagle" or \
1196
+ self.speculative_config.method == "eagle3":
1197
assert isinstance(self.drafter, EagleProposer)
1198
# TODO(woosuk): Refactor the loop.
1199
next_token_ids: list[int] = []
@@ -1210,7 +1221,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1221
# not include padding.
1222
target_token_ids = self.input_ids[:num_scheduled_tokens]
1223
target_positions = positions[:num_scheduled_tokens]
1213
- target_hidden_states = hidden_states[:num_scheduled_tokens]
1224
+ if self.use_aux_hidden_state_outputs:
1225
+ target_hidden_states = [
1226
+ h[:num_scheduled_tokens] for h in aux_hidden_states
1227
+ ]
1228
+ else:
1229
+ target_hidden_states = hidden_states[:num_scheduled_tokens]
1230
target_slot_mapping = attn_metadata.slot_mapping
1231
cu_num_tokens = attn_metadata.query_start_loc
1232
else:
@@ -1231,9 +1247,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1247
)
1248
target_token_ids = self.input_ids[token_indices]
1249
target_positions = positions[token_indices]
1234
- target_hidden_states = hidden_states[token_indices]
1250
+ if self.use_aux_hidden_state_outputs:
1251
+ target_hidden_states = [
1252
+ h[token_indices] for h in aux_hidden_states
1253
+ ]
1254
+ else:
1255
+ target_hidden_states = hidden_states[token_indices]
1256
target_slot_mapping = attn_metadata.slot_mapping[token_indices]
1257
1258
+ if self.use_aux_hidden_state_outputs:
1259
+ target_hidden_states = torch.cat(target_hidden_states, dim=-1)
1260
draft_token_ids = self.drafter.propose(
1261
target_token_ids=target_token_ids,
1262
target_positions=target_positions,
@@ -1311,6 +1334,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1334
if hasattr(self, "drafter"):
1335
logger.info("Loading drafter model...")
1336
self.drafter.load_model(self.model)
1337
+ if self.use_aux_hidden_state_outputs:
1338
+ self.model.set_aux_hidden_state_layers(
1339
+ self.model.get_eagle3_aux_hidden_state_layers())
1340
time_after_load = time.perf_counter()
1341
self.model_memory_usage = m.consumed_memory
1342
logger.info("Model loading took %.4f GiB and %.6f seconds",
@@ -1463,12 +1489,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1489
with set_forward_context(None,
1490
self.vllm_config,
1491
num_tokens=num_tokens):
1466
- hidden_states = model(
1492
+ outputs = model(
1493
input_ids=input_ids,
1494
positions=positions,
1495
intermediate_tensors=intermediate_tensors,
1496
inputs_embeds=inputs_embeds,
1497
)
1498
+ if self.use_aux_hidden_state_outputs:
1499
+ hidden_states, _ = outputs
1500
+ else:
1501
+ hidden_states = outputs
1502
1503
logit_indices = np.cumsum(num_scheduled_tokens) - 1
1504
return hidden_states[logit_indices]