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 12, 14:45 UTC
[Model Runner V2] Implement Single-step Eagle 1 (#29300)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 24, 2025 at 09:32 UTC
cc313cb73d75cd5ac2715fc45bfadb89888cf8cd
5 files changed
+300
-2
vllm/v1/worker/gpu/input_batch.py
+3
@@ -37,6 +37,9 @@ class InputBuffers:
37
self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
38
self.cu_num_logits = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
39
40
+ # Spec decoding.
41
+ self.next_prefill_tokens = self._make_buffer(max_num_reqs, dtype=torch.int32)
42
+
43
# Structured outputs.
44
self.bitmask_indices = self._make_buffer(max_num_reqs, dtype=torch.int32)
45
self.grammar_bitmask = self._make_buffer(
vllm/v1/worker/gpu/model_runner.py
+79
@@ -45,6 +45,7 @@ from vllm.v1.worker.gpu.input_batch import (
45
prepare_prefill_inputs,
46
)
47
from vllm.v1.worker.gpu.sampler import Sampler, compute_prompt_logprobs
48
+from vllm.v1.worker.gpu.spec_decode import init_speculator
49
from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
50
from vllm.v1.worker.gpu.states import RequestState, SamplingMetadata
51
from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
@@ -97,16 +98,20 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
98
if self.use_async_scheduling:
99
self.input_prep_event = torch.cuda.Event()
100
self.structured_outputs_event = torch.cuda.Event()
101
+ self.spec_decode_event = torch.cuda.Event()
102
else:
103
self.input_prep_event = None
104
self.structured_outputs_event = None
105
+ self.spec_decode_event = None
106
107
if self.speculative_config is not None:
108
self.do_spec_decode = True
109
self.num_speculative_steps = self.speculative_config.num_speculative_tokens
110
+ self.speculator = init_speculator(self.vllm_config, self.device)
111
else:
112
self.do_spec_decode = False
113
self.num_speculative_steps = 0
114
+ self.speculator = None
115
116
self.req_states = RequestState(
117
max_num_reqs=self.max_num_reqs,
@@ -153,6 +158,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
158
self.vllm_config,
159
self.device,
160
)
161
+ if self.do_spec_decode:
162
+ self.speculator.load_model(self.model)
163
time_after_load = time.perf_counter()
164
165
self.model_memory_usage = m.consumed_memory
@@ -285,6 +292,33 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
292
logits = self.model.compute_logits(hidden_states)
293
self.sampler(logits, sampling_metadata)
294
295
+ @torch.inference_mode()
296
+ def _dummy_speculator_run(
297
+ self,
298
+ hidden_states: torch.Tensor,
299
+ aux_hidden_states: list[torch.Tensor] | None,
300
+ ) -> None:
301
+ num_tokens = hidden_states.shape[0]
302
+ num_reqs = min(num_tokens, self.max_num_reqs)
303
+ input_batch = InputBatch.make_dummy(
304
+ num_reqs=num_reqs,
305
+ num_tokens=num_tokens,
306
+ input_buffers=self.input_buffers,
307
+ device=self.device,
308
+ )
309
+ sampling_metadata = SamplingMetadata.make_dummy(
310
+ num_reqs=num_reqs,
311
+ device=self.device,
312
+ )
313
+ num_sampled = torch.ones(num_reqs, dtype=torch.int32, device=self.device)
314
+ self.propose_draft(
315
+ input_batch=input_batch,
316
+ sampling_metadata=sampling_metadata,
317
+ last_hidden_states=hidden_states,
318
+ aux_hidden_states=aux_hidden_states,
319
+ num_sampled=num_sampled,
320
+ )
321
+
322
@torch.inference_mode()
323
def profile_run(self) -> None:
324
hidden_states, sample_hidden_states = self._dummy_run(
@@ -292,6 +326,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
326
skip_attn=True,
327
)
328
self._dummy_sampler_run(sample_hidden_states)
329
+ if self.do_spec_decode:
330
+ self._dummy_speculator_run(hidden_states, None)
331
torch.cuda.synchronize()
332
del hidden_states, sample_hidden_states
333
gc.collect()
@@ -727,6 +763,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
763
self.req_states.prefill_len.np[idx_mapping_np],
764
)
765
766
+ @torch.inference_mode()
767
+ def propose_draft(
768
+ self,
769
+ input_batch: InputBatch,
770
+ sampling_metadata: SamplingMetadata,
771
+ last_hidden_states: torch.Tensor,
772
+ aux_hidden_states: list[torch.Tensor] | None,
773
+ num_sampled: torch.Tensor,
774
+ ) -> torch.Tensor:
775
+ num_reqs = input_batch.num_reqs
776
+ idx_mapping_np = input_batch.idx_mapping_np
777
+ with async_barrier(self.spec_decode_event):
778
+ self.input_buffers.next_prefill_tokens.np[:num_reqs] = (
779
+ self.req_states.prefill_token_ids[
780
+ idx_mapping_np,
781
+ self.req_states.num_computed_prefill_tokens[idx_mapping_np],
782
+ ]
783
+ )
784
+ next_prefill_tokens = self.input_buffers.next_prefill_tokens.copy_to_gpu(
785
+ num_reqs
786
+ )
787
+
788
+ assert self.speculator is not None
789
+ draft_tokens = self.speculator.propose(
790
+ input_batch,
791
+ sampling_metadata,
792
+ last_hidden_states,
793
+ aux_hidden_states,
794
+ num_sampled,
795
+ self.req_states.last_sampled_tokens,
796
+ next_prefill_tokens,
797
+ )
798
+ self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
799
+ return draft_tokens
800
+
801
def get_cudagraph_and_dp_padding(
802
self,
803
scheduler_output: SchedulerOutput,
@@ -913,6 +984,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
984
self.postprocess(
985
input_batch, sampler_output.sampled_token_ids, num_sampled_tokens
986
)
987
+ if self.do_spec_decode:
988
+ _ = self.propose_draft(
989
+ input_batch,
990
+ sampling_metadata,
991
+ hidden_states,
992
+ None, # aux_hidden_states
993
+ num_sampled_tokens,
994
+ )
995
996
if self.use_async_scheduling:
997
return async_output
vllm/v1/worker/gpu/sampler.py
+3
-2
@@ -100,8 +100,9 @@ def _gumbel_sample_kernel(
100
mask=mask,
101
other=float("-inf"),
102
)
103
+ logits = logits.to(tl.float32)
104
104
- temp = tl.load(temp_ptr + req_idx)
105
+ temp = tl.load(temp_ptr + req_idx).to(tl.float32)
106
if temp != 0.0:
107
# Calculate the seed for gumbel noise.
108
seed = tl.load(seeds_ptr + req_idx)
@@ -116,7 +117,7 @@ def _gumbel_sample_kernel(
117
# Apply temperature.
118
if APPLY_TEMPERATURE:
119
# NOTE(woosuk): Use div_rn to match the behavior of torch.
119
- logits = tl.div_rn(logits, temp.to(tl.float32))
120
+ logits = tl.div_rn(logits, temp)
121
122
# Apply gumbel noise.
123
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
vllm/v1/worker/gpu/spec_decode/__init__.py
+18
@@ -0,0 +1,18 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.config import VllmConfig
6
+
7
+
8
+def init_speculator(
9
+ vllm_config: VllmConfig,
10
+ device: torch.device,
11
+):
12
+ speculative_config = vllm_config.speculative_config
13
+ assert speculative_config is not None
14
+ if speculative_config.use_eagle():
15
+ from vllm.v1.worker.gpu.spec_decode.eagle import EagleSpeculator
16
+
17
+ return EagleSpeculator(vllm_config, device)
18
+ raise NotImplementedError(f"{speculative_config.method} is not supported yet.")
vllm/v1/worker/gpu/spec_decode/eagle.py
new
+197
@@ -0,0 +1,197 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+import torch.nn as nn
5
+
6
+from vllm.config import VllmConfig
7
+from vllm.config.compilation import CUDAGraphMode
8
+from vllm.forward_context import set_forward_context
9
+from vllm.model_executor.model_loader import get_model
10
+from vllm.triton_utils import tl, triton
11
+from vllm.v1.worker.gpu.input_batch import InputBatch
12
+from vllm.v1.worker.gpu.sampler import gumbel_sample
13
+from vllm.v1.worker.gpu.states import SamplingMetadata
14
+
15
+
16
+class EagleSpeculator:
17
+ def __init__(self, vllm_config: VllmConfig, device: torch.device):
18
+ self.vllm_config = vllm_config
19
+ self.device = device
20
+
21
+ self.speculative_config = vllm_config.speculative_config
22
+ assert self.speculative_config is not None
23
+ self.method = self.speculative_config.method
24
+ self.num_speculative_steps = self.speculative_config.num_speculative_tokens
25
+ self.draft_model_config = self.speculative_config.draft_model_config
26
+
27
+ self.scheduler_config = vllm_config.scheduler_config
28
+ self.max_num_reqs = self.scheduler_config.max_num_seqs
29
+ self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
30
+
31
+ self.input_ids = torch.zeros(
32
+ self.max_num_tokens, dtype=torch.int32, device=device
33
+ )
34
+ self.positions = torch.zeros(
35
+ self.max_num_tokens, dtype=torch.int64, device=device
36
+ )
37
+
38
+ def load_model(self, target_model: nn.Module) -> None:
39
+ from vllm.compilation.backends import set_model_tag
40
+
41
+ with set_model_tag("eagle_head"):
42
+ self.model = get_model(
43
+ vllm_config=self.vllm_config, model_config=self.draft_model_config
44
+ )
45
+
46
+ share_lm_head = True
47
+ if share_lm_head and hasattr(target_model, "lm_head"):
48
+ if hasattr(self.model, "lm_head"):
49
+ del self.model.lm_head
50
+ self.model.lm_head = target_model.lm_head
51
+
52
+ @torch.inference_mode()
53
+ def propose(
54
+ self,
55
+ input_batch: InputBatch,
56
+ sampling_metadata: SamplingMetadata,
57
+ # [num_tokens, hidden_size]
58
+ last_hidden_states: torch.Tensor,
59
+ # num_layers x [num_tokens, hidden_size]
60
+ aux_hidden_states: list[torch.Tensor] | None,
61
+ # [num_reqs]
62
+ num_sampled: torch.Tensor,
63
+ # [max_num_reqs, 1]
64
+ last_sampled: torch.Tensor,
65
+ # [num_reqs]
66
+ next_prefill_tokens: torch.Tensor,
67
+ ) -> torch.Tensor:
68
+ if aux_hidden_states:
69
+ assert self.method == "eagle3"
70
+ hidden_states = self.model.combine_hidden_states(
71
+ torch.cat(aux_hidden_states, dim=-1)
72
+ )
73
+ else:
74
+ hidden_states = last_hidden_states
75
+
76
+ # Get the input ids and last token indices for the speculator.
77
+ last_token_indices = prepare_eagle_inputs(
78
+ self.input_ids,
79
+ input_batch,
80
+ num_sampled,
81
+ last_sampled,
82
+ next_prefill_tokens,
83
+ )
84
+ input_ids = self.input_ids[: input_batch.num_tokens_after_padding]
85
+
86
+ # Prefill: Run the eagle speculator with eager mode.
87
+ with set_forward_context(
88
+ input_batch.attn_metadata,
89
+ self.vllm_config,
90
+ num_tokens=input_batch.num_tokens_after_padding,
91
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
92
+ ):
93
+ ret_hidden_states = self.model(
94
+ input_ids=input_ids,
95
+ positions=input_batch.positions,
96
+ hidden_states=hidden_states,
97
+ )
98
+ if self.method == "mtp":
99
+ last_hidden_states = ret_hidden_states
100
+ hidden_states = ret_hidden_states
101
+ else:
102
+ last_hidden_states, hidden_states = ret_hidden_states
103
+ sample_hidden_states = last_hidden_states[last_token_indices]
104
+ logits = self.model.compute_logits(sample_hidden_states)
105
+
106
+ num_reqs = input_batch.num_reqs
107
+ cu_num_logits = input_batch.cu_num_logits[:num_reqs]
108
+ temperature = sampling_metadata.temperature[cu_num_logits]
109
+ seed = sampling_metadata.seeds[cu_num_logits]
110
+ # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
111
+ # used for draft and target sampling.
112
+ pos = input_batch.positions[last_token_indices] + 1
113
+ draft_tokens = gumbel_sample(
114
+ logits, temperature, seed, pos, apply_temperature=True
115
+ )
116
+ if self.num_speculative_steps == 1:
117
+ # Early exit.
118
+ return draft_tokens.view(-1, 1)
119
+ raise NotImplementedError("num_speculative_steps > 1 is not supported yet.")
120
+
121
+
122
+@triton.jit
123
+def _prepare_eagle_inputs_kernel(
124
+ last_token_indices_ptr,
125
+ eagle_input_ids_ptr,
126
+ target_input_ids_ptr,
127
+ idx_mapping_ptr,
128
+ last_sampled_ptr,
129
+ next_prefill_tokens_ptr,
130
+ num_sampled_ptr,
131
+ query_start_loc_ptr,
132
+ cu_num_logits_ptr,
133
+ BLOCK_SIZE: tl.constexpr,
134
+):
135
+ batch_idx = tl.program_id(0)
136
+ query_start = tl.load(query_start_loc_ptr + batch_idx)
137
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
138
+ query_len = query_end - query_start
139
+
140
+ # Get the true query length and next token after accounting for rejected tokens.
141
+ num_sampled = tl.load(num_sampled_ptr + batch_idx)
142
+ if num_sampled > 0:
143
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
144
+ next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
145
+
146
+ logits_start = tl.load(cu_num_logits_ptr + batch_idx)
147
+ logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
148
+ num_logits = logits_end - logits_start
149
+
150
+ num_rejected = num_logits - num_sampled
151
+ query_len -= num_rejected
152
+ else:
153
+ # Chunked prefilling.
154
+ # Get the next prefill token.
155
+ next_token = tl.load(next_prefill_tokens_ptr + batch_idx)
156
+
157
+ # Shift target_input_ids by one.
158
+ for i in range(1, query_len, BLOCK_SIZE):
159
+ block = i + tl.arange(0, BLOCK_SIZE)
160
+ mask = block < query_len
161
+ input_ids = tl.load(target_input_ids_ptr + query_start + block, mask=mask)
162
+ tl.store(eagle_input_ids_ptr + query_start + block - 1, input_ids, mask=mask)
163
+
164
+ last_token_index = query_start + query_len - 1
165
+ tl.store(last_token_indices_ptr + batch_idx, last_token_index)
166
+ tl.store(eagle_input_ids_ptr + last_token_index, next_token)
167
+
168
+
169
+def prepare_eagle_inputs(
170
+ eagle_input_ids: torch.Tensor,
171
+ input_batch: InputBatch,
172
+ # [num_reqs]
173
+ num_sampled: torch.Tensor,
174
+ # [max_num_reqs, 1]
175
+ last_sampled: torch.Tensor,
176
+ # [max_num_reqs]
177
+ next_prefill_tokens: torch.Tensor,
178
+) -> torch.Tensor:
179
+ num_reqs = input_batch.num_reqs
180
+ last_token_indices = torch.empty(
181
+ num_reqs,
182
+ dtype=torch.int64,
183
+ device=eagle_input_ids.device,
184
+ )
185
+ _prepare_eagle_inputs_kernel[(num_reqs,)](
186
+ last_token_indices,
187
+ eagle_input_ids,
188
+ input_batch.input_ids,
189
+ input_batch.idx_mapping,
190
+ last_sampled,
191
+ next_prefill_tokens,
192
+ num_sampled,
193
+ input_batch.query_start_loc,
194
+ input_batch.cu_num_logits,
195
+ BLOCK_SIZE=1024,
196
+ )
197
+ return last_token_indices