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 7, 22:45 UTC
[Model Runner V2] Support Eagle3 (no CUDA graph) (#35029)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 21, 2026 at 12:55 UTC
a4047d4ea993fd52038433d87c16e603bee4f214
7 files changed
+169
-49
vllm/v1/worker/gpu/model_runner.py
+67
-34
@@ -66,6 +66,9 @@ from vllm.v1.worker.gpu.sample.output import SamplerOutput
66
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
67
from vllm.v1.worker.gpu.sample.sampler import Sampler
68
from vllm.v1.worker.gpu.spec_decode import init_speculator
69
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
70
+ set_eagle3_aux_hidden_state_layers,
71
+)
72
from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
73
from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
74
from vllm.v1.worker.gpu.states import RequestState
@@ -133,14 +136,42 @@ class GPUModelRunner(LoRAModelRunnerMixin):
136
self.output_copy_stream = torch.cuda.Stream(self.device)
137
self.output_copy_event = torch.cuda.Event()
138
139
+ # Pipeline parallelism.
140
+ self.pp_size = self.parallel_config.pipeline_parallel_size
141
+ self.use_pp = self.pp_size > 1
142
+ if self.use_pp:
143
+ self.is_first_pp_rank = get_pp_group().is_first_rank
144
+ self.is_last_pp_rank = get_pp_group().is_last_rank
145
+ else:
146
+ self.is_first_pp_rank = True
147
+ self.is_last_pp_rank = True
148
+
149
+ # Decode context parallelism.
150
+ self.dcp_size = self.parallel_config.decode_context_parallel_size
151
+ self.use_dcp = self.dcp_size > 1
152
+ self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
153
+ self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
154
+
155
+ self.speculator = None
156
+ self.use_aux_hidden_state_outputs = False
157
if self.speculative_config is not None:
158
self.do_spec_decode = True
159
self.num_speculative_steps = self.speculative_config.num_speculative_tokens
139
- self.speculator = init_speculator(self.vllm_config, self.device)
160
+ if self.is_last_pp_rank:
161
+ self.speculator = init_speculator(self.vllm_config, self.device)
162
+
163
+ if self.speculative_config.method == "eagle3":
164
+ # EAGLE3 may require auxiliary hidden states from target model outputs.
165
+ self.use_aux_hidden_state_outputs = True
166
+ if self.pp_size > 1:
167
+ raise ValueError("EAGLE3 with pipeline parallel is not supported.")
168
else:
169
self.do_spec_decode = False
170
self.num_speculative_steps = 0
143
- self.speculator = None
171
+
172
+ # Draft tokens propagation - for spec-dec + struct outputs.
173
+ self.draft_tokens_handler = DraftTokensHandler(self.device)
174
+
175
self.req_states = RequestState(
176
max_num_reqs=self.max_num_reqs,
177
max_model_len=self.max_model_len,
@@ -176,28 +207,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
207
)
208
# LoRA-related workers.
209
self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
179
-
180
- # Draft tokens propagation - for spec-dec + struct outputs.
181
- self.draft_tokens_handler = DraftTokensHandler(self.device)
182
-
210
# KV Connector if configured.
211
self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
212
186
- # Pipeline parallelism.
187
- self.use_pp = self.parallel_config.pipeline_parallel_size > 1
188
- if self.use_pp:
189
- self.is_first_pp_rank = get_pp_group().is_first_rank
190
- self.is_last_pp_rank = get_pp_group().is_last_rank
191
- else:
192
- self.is_first_pp_rank = True
193
- self.is_last_pp_rank = True
194
-
195
- # Decode context parallelism.
196
- self.dcp_size = self.parallel_config.decode_context_parallel_size
197
- self.use_dcp = self.dcp_size > 1
198
- self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
199
- self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
200
-
213
def update_max_model_len(self, max_model_len: int) -> None:
214
self.max_model_len = max_model_len
215
self.req_states.max_model_len = max_model_len
@@ -220,7 +232,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
232
self.model = self.load_lora_model(
233
self.model, self.vllm_config, self.device
234
)
223
- if self.do_spec_decode:
235
+
236
+ if self.use_aux_hidden_state_outputs:
237
+ assert self.speculative_config is not None
238
+ set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
239
+ if self.speculator is not None:
240
self.speculator.load_model(self.model)
241
time_after_load = time.perf_counter()
242
@@ -271,7 +287,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
287
self.kv_cache_config, self.vllm_config, self.device
288
)
289
check_attention_cp_compatibility(self.vllm_config)
274
- if self.do_spec_decode:
290
+ if self.speculator is not None:
291
# HACK(woosuk)
292
self.speculator.set_attn(
293
self.kv_cache_config,
@@ -359,7 +375,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
375
return None, None
376
377
assert self.execute_model_state is not None
362
- hidden_states, input_batch, _ = self.execute_model_state
378
+ hidden_states, _, input_batch, _ = self.execute_model_state
379
assert hidden_states is not None # Last PP rank always has hidden_states
380
sample_hidden_states = hidden_states[input_batch.logits_indices]
381
return hidden_states, sample_hidden_states
@@ -399,7 +415,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
415
assert sample_hidden_states is not None
416
self._dummy_sampler_run(sample_hidden_states)
417
402
- if self.do_spec_decode:
418
+ if self.speculator is not None:
419
num_tokens_across_dp = make_num_tokens_across_dp(
420
self.parallel_config.data_parallel_size, self.max_num_tokens
421
)
@@ -465,7 +481,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
481
kv_cache_config=self.kv_cache_config,
482
has_lora=self.lora_config is not None,
483
)
468
- if self.do_spec_decode:
484
+ if self.speculator is not None:
485
self.speculator.capture_model()
486
487
end_time = time.perf_counter()
@@ -964,9 +980,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
980
# NOTE(woosuk): Here, we don't need to pass the input tensors,
981
# because they are already copied to the CUDA graph input buffers.
982
self.kv_connector.pre_forward(scheduler_output)
967
- hidden_states = self.cudagraph_manager.run_fullgraph(
983
+ model_output = self.cudagraph_manager.run_fullgraph(
984
input_batch.num_tokens_after_padding
985
)
986
+ if self.use_aux_hidden_state_outputs:
987
+ hidden_states, aux_hidden_states = model_output
988
+ else:
989
+ hidden_states = model_output
990
+ aux_hidden_states = None
991
else:
992
# For piecewise and eager mode, just call model().
993
positions = input_batch.positions
@@ -998,12 +1019,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1019
slot_mapping=input_batch.slot_mappings,
1020
):
1021
self.kv_connector.pre_forward(scheduler_output)
1001
- hidden_states = self.model(
1022
+ model_output = self.model(
1023
input_ids=input_ids,
1024
positions=positions,
1025
inputs_embeds=inputs_embeds,
1026
intermediate_tensors=intermediate_tensors,
1027
)
1028
+ if self.use_aux_hidden_state_outputs:
1029
+ hidden_states, aux_hidden_states = model_output
1030
+ else:
1031
+ hidden_states = model_output
1032
+ aux_hidden_states = None
1033
1034
kv_connector_output = self.kv_connector.post_forward(scheduler_output)
1035
@@ -1011,12 +1037,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1037
# Non-last PP rank: return IntermediateTensors for sending.
1038
assert isinstance(hidden_states, IntermediateTensors)
1039
hidden_states.kv_connector_output = kv_connector_output
1014
- self.execute_model_state = (None, input_batch, kv_connector_output)
1040
+ self.execute_model_state = (None, None, input_batch, kv_connector_output)
1041
return hidden_states
1042
1017
- assert isinstance(hidden_states, torch.Tensor)
1043
# Last rank (or no PP): hidden_states is a tensor for sampling.
1019
- self.execute_model_state = (hidden_states, input_batch, kv_connector_output)
1044
+ assert isinstance(hidden_states, torch.Tensor)
1045
+ self.execute_model_state = (
1046
+ hidden_states,
1047
+ aux_hidden_states,
1048
+ input_batch,
1049
+ kv_connector_output,
1050
+ )
1051
return None
1052
1053
@torch.inference_mode()
@@ -1024,7 +1055,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1055
self, grammar_output: GrammarOutput | None
1056
) -> AsyncOutput | ModelRunnerOutput | None:
1057
assert self.execute_model_state is not None
1027
- hidden_states, input_batch, kv_connector_output = self.execute_model_state
1058
+ hidden_states, aux_hidden_states, input_batch, kv_connector_output = (
1059
+ self.execute_model_state
1060
+ )
1061
self.execute_model_state = None # type: ignore
1062
1063
if not self.is_last_pp_rank:
@@ -1084,11 +1117,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1117
self.postprocess(
1118
input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
1119
)
1087
- if self.do_spec_decode:
1120
+ if self.speculator is not None:
1121
draft_tokens = self.propose_draft(
1122
input_batch,
1123
hidden_states,
1091
- None, # aux_hidden_states
1124
+ aux_hidden_states,
1125
num_sampled,
1126
num_rejected,
1127
)
vllm/v1/worker/gpu/spec_decode/__init__.py
+1
-1
@@ -9,7 +9,7 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
9
speculative_config = vllm_config.speculative_config
10
assert speculative_config is not None
11
if speculative_config.use_eagle():
12
- from vllm.v1.worker.gpu.spec_decode.eagle import EagleSpeculator
12
+ from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator
13
14
return EagleSpeculator(vllm_config, device)
15
raise NotImplementedError(f"{speculative_config.method} is not supported yet.")
vllm/v1/worker/gpu/spec_decode/eagle/__init__.py
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
renamed
vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import cast
4
+
5
+import torch.nn as nn
6
+
7
+from vllm.config import SpeculativeConfig
8
+from vllm.logger import init_logger
9
+from vllm.model_executor.models.interfaces import SupportsEagle3, supports_eagle3
10
+
11
+logger = init_logger(__name__)
12
+
13
+
14
+def set_eagle3_aux_hidden_state_layers(
15
+ model: nn.Module,
16
+ spec_config: SpeculativeConfig,
17
+) -> None:
18
+ if not supports_eagle3(model):
19
+ raise RuntimeError("Model does not support EAGLE3 interface")
20
+ # mypy may infer the class-level overload for supports_eagle3.
21
+ # Narrow explicitly to the runtime protocol instance.
22
+ if isinstance(model, type):
23
+ raise RuntimeError("Expected model instance for EAGLE3 configuration")
24
+ eagle3_model = cast(SupportsEagle3, model)
25
+
26
+ aux_layers = get_eagle3_aux_layers_from_config(spec_config)
27
+ if aux_layers:
28
+ logger.info("Using Eagle3 auxiliary layers from config: %s", aux_layers)
29
+ else:
30
+ aux_layers = eagle3_model.get_eagle3_aux_hidden_state_layers()
31
+ logger.info("Using Eagle3 auxiliary layers from model: %s", aux_layers)
32
+ eagle3_model.set_aux_hidden_state_layers(aux_layers)
33
+
34
+
35
+def get_eagle3_aux_layers_from_config(
36
+ spec_config: SpeculativeConfig,
37
+) -> tuple[int, ...] | None:
38
+ if not (spec_config and spec_config.draft_model_config):
39
+ return None
40
+ hf_config = spec_config.draft_model_config.hf_config
41
+ if not hasattr(hf_config, "eagle_aux_hidden_state_layer_ids"):
42
+ return None
43
+ layer_ids = hf_config.eagle_aux_hidden_state_layer_ids
44
+ if layer_ids and isinstance(layer_ids, (list, tuple)):
45
+ return tuple(layer_ids)
46
+ return None
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
renamed
+3
-14
@@ -9,7 +9,6 @@ from vllm.config import VllmConfig
9
from vllm.config.compilation import CUDAGraphMode
10
from vllm.forward_context import BatchDescriptor, set_forward_context
11
from vllm.logger import init_logger
12
-from vllm.model_executor.model_loader import get_model
12
from vllm.triton_utils import tl, triton
13
from vllm.v1.attention.backend import AttentionMetadataBuilder
14
from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -20,7 +19,8 @@ from vllm.v1.worker.gpu.attn_utils import (
19
from vllm.v1.worker.gpu.block_table import BlockTables
20
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
23
-from vllm.v1.worker.gpu.spec_decode.eagle_cudagraph import EagleCudaGraphManager
22
+from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import EagleCudaGraphManager
23
+from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
24
25
logger = init_logger(__name__)
26
@@ -73,18 +73,7 @@ class EagleSpeculator:
73
self.cudagraph_manager = EagleCudaGraphManager(vllm_config, device)
74
75
def load_model(self, target_model: nn.Module) -> None:
76
- from vllm.compilation.backends import set_model_tag
77
-
78
- with set_model_tag("eagle_head"):
79
- self.model = get_model(
80
- vllm_config=self.vllm_config, model_config=self.draft_model_config
81
- )
82
-
83
- share_lm_head = True
84
- if share_lm_head and hasattr(target_model, "lm_head"):
85
- if hasattr(self.model, "lm_head"):
86
- del self.model.lm_head
87
- self.model.lm_head = target_model.lm_head
76
+ self.model = load_eagle_model(target_model, self.vllm_config)
77
78
def set_attn(
79
self,
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
new
+52
@@ -0,0 +1,52 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch.nn as nn
4
+
5
+from vllm.config import VllmConfig
6
+from vllm.model_executor.model_loader import get_model
7
+
8
+
9
+def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
10
+ from vllm.compilation.backends import set_model_tag
11
+
12
+ speculative_config = vllm_config.speculative_config
13
+ assert speculative_config is not None
14
+ draft_model_config = speculative_config.draft_model_config
15
+ with set_model_tag("eagle_head"):
16
+ eagle_model = get_model(
17
+ vllm_config=vllm_config, model_config=draft_model_config
18
+ )
19
+
20
+ # Share target embeddings when the draft checkpoint does not include
21
+ # its own vocab embedding table.
22
+ share_embeddings = True
23
+ if hasattr(eagle_model, "has_own_embed_tokens"):
24
+ share_embeddings = not eagle_model.has_own_embed_tokens
25
+ if share_embeddings:
26
+ target_language_model = (
27
+ target_model.get_language_model()
28
+ if hasattr(target_model, "get_language_model")
29
+ else target_model
30
+ )
31
+ inner_model = getattr(target_language_model, "model", None)
32
+ target_embed_tokens = None
33
+ if inner_model is not None:
34
+ if hasattr(inner_model, "embed_tokens"):
35
+ target_embed_tokens = inner_model.embed_tokens
36
+ elif hasattr(inner_model, "embedding"):
37
+ target_embed_tokens = inner_model.embedding
38
+ if target_embed_tokens is not None and hasattr(eagle_model, "model"):
39
+ if hasattr(eagle_model.model, "embed_tokens"):
40
+ del eagle_model.model.embed_tokens
41
+ eagle_model.model.embed_tokens = target_embed_tokens
42
+
43
+ # Only share target lm_head when the draft model does not own one.
44
+ share_lm_head = True
45
+ if hasattr(eagle_model, "has_own_lm_head"):
46
+ share_lm_head = not eagle_model.has_own_lm_head
47
+ if share_lm_head and hasattr(target_model, "lm_head"):
48
+ if hasattr(eagle_model, "lm_head"):
49
+ del eagle_model.lm_head
50
+ eagle_model.lm_head = target_model.lm_head
51
+
52
+ return eagle_model