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, 16:00 UTC
[V1][Spec Decode] Implement Eagle Proposer [1/N] (#15729)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Apr 1, 2025 at 12:33 UTC
e75a6301bda30a5a8453eb95804aedc2104936f7
6 files changed
+378
-21
vllm/config.py
+3
-2
@@ -2152,9 +2152,10 @@ class SpeculativeConfig:
2152
2153
# Replace hf_config for EAGLE draft_model
2154
if self.method == "eagle":
2155
- if self.enable_chunked_prefill:
2155
+ if self.enable_chunked_prefill and not envs.VLLM_USE_V1:
2156
raise ValueError(
2157
- "Chunked prefill and EAGLE are not compatible.")
2157
+ "Chunked prefill and EAGLE are not compatible "
2158
+ "when using V0.")
2159
2160
from vllm.transformers_utils.configs.eagle import (
2161
EAGLEConfig)
vllm/engine/arg_utils.py
+16
-6
@@ -1468,15 +1468,21 @@ class EngineArgs:
1468
1469
# Only Ngram speculative decoding so far.
1470
is_ngram_enabled = False
1471
+ is_eagle_enabled = False
1472
if self.speculative_config is not None:
1473
# This is supported but experimental (handled below).
1473
- if (("method" in self.speculative_config
1474
- and self.speculative_config["method"] in ("ngram", "[ngram]"))
1475
- or
1476
- ("model" in self.speculative_config and
1477
- self.speculative_config["model"] in ("ngram", "[ngram]"))):
1478
- is_ngram_enabled = True
1474
+ speculative_method = self.speculative_config.get("method")
1475
+ if speculative_method:
1476
+ if speculative_method in ("ngram", "[ngram]"):
1477
+ is_ngram_enabled = True
1478
+ elif speculative_method == "eagle":
1479
+ is_eagle_enabled = True
1480
else:
1481
+ speculative_model = self.speculative_config.get("model")
1482
+ if speculative_model in ("ngram", "[ngram]"):
1483
+ is_ngram_enabled = True
1484
+ if not (is_ngram_enabled or is_eagle_enabled):
1485
+ # Other speculative decoding methods are not supported yet.
1486
_raise_or_fallback(feature_name="Speculative Decoding",
1487
recommend_to_remove=False)
1488
return False
@@ -1523,6 +1529,10 @@ class EngineArgs:
1529
if is_ngram_enabled and _warn_or_fallback("ngram"):
1530
return False
1531
1532
+ # Eagle is under development, so we don't support it yet.
1533
+ if is_eagle_enabled and _warn_or_fallback("Eagle"):
1534
+ return False
1535
+
1536
# Non-CUDA is supported on V1, but off by default for now.
1537
not_cuda = not current_platform.is_cuda()
1538
if not_cuda and _warn_or_fallback( # noqa: SIM103
vllm/v1/spec_decode/eagle.py
new
+262
@@ -0,0 +1,262 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import torch
3
+import torch.nn as nn
4
+import triton
5
+import triton.language as tl
6
+
7
+from vllm.config import VllmConfig
8
+from vllm.forward_context import set_forward_context
9
+from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
10
+from vllm.v1.sample.metadata import SamplingMetadata
11
+
12
+
13
+class EagleProposer:
14
+
15
+ def __init__(
16
+ self,
17
+ vllm_config: VllmConfig,
18
+ device: torch.device,
19
+ ):
20
+ self.vllm_config = vllm_config
21
+ self.num_speculative_tokens = (
22
+ vllm_config.speculative_config.num_speculative_tokens)
23
+ self.block_size = vllm_config.cache_config.block_size
24
+ self.arange = torch.arange(vllm_config.scheduler_config.max_num_seqs,
25
+ device=device)
26
+
27
+ def propose(
28
+ self,
29
+ # [num_tokens]
30
+ target_token_ids: torch.Tensor,
31
+ # [num_tokens]
32
+ target_positions: torch.Tensor,
33
+ # [num_tokens, hidden_size]
34
+ target_hidden_states: torch.Tensor,
35
+ # [num_tokens]
36
+ target_slot_mapping: torch.Tensor,
37
+ # [batch_size]
38
+ next_token_ids: torch.Tensor,
39
+ # [batch_size + 1] starting with 0
40
+ cu_num_tokens: torch.Tensor,
41
+ # [batch_size, max_num_blocks_per_req]
42
+ block_table: torch.Tensor,
43
+ sampling_metadata: SamplingMetadata,
44
+ ) -> tuple[torch.Tensor, torch.Tensor]:
45
+ num_tokens = target_token_ids.shape[0]
46
+ batch_size = next_token_ids.shape[0]
47
+ last_token_indices = cu_num_tokens[1:] - 1
48
+
49
+ input_ids = torch.empty_like(target_token_ids)
50
+ # Shift the input ids by one token.
51
+ # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
52
+ input_ids[:-1] = target_token_ids[1:]
53
+ # Replace the last token with the next token.
54
+ # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
55
+ input_ids[last_token_indices] = next_token_ids
56
+
57
+ seq_lens = target_positions[last_token_indices] + 1
58
+ # FIXME(woosuk): The below two ops cause synchronization. Optimize.
59
+ max_seq_len = seq_lens.max().item()
60
+ max_num_tokens = (cu_num_tokens[1:] - cu_num_tokens[:-1]).max().item()
61
+ attn_metadata = FlashAttentionMetadata(
62
+ num_actual_tokens=num_tokens,
63
+ max_query_len=max_num_tokens,
64
+ query_start_loc=cu_num_tokens,
65
+ max_seq_len=max_seq_len,
66
+ seq_lens=seq_lens,
67
+ block_table=block_table,
68
+ slot_mapping=target_slot_mapping,
69
+ # TODO(woosuk): Support cascade attention.
70
+ use_cascade=False,
71
+ common_prefix_len=0,
72
+ cu_prefix_query_lens=None,
73
+ prefix_kv_lens=None,
74
+ suffix_kv_lens=None,
75
+ )
76
+
77
+ with set_forward_context(attn_metadata, self.vllm_config):
78
+ hidden_states = self.model(
79
+ input_ids=input_ids,
80
+ hidden_states=target_hidden_states,
81
+ positions=target_positions,
82
+ )
83
+ sample_hidden_states = hidden_states[last_token_indices]
84
+ logits = self.model.compute_logits(sample_hidden_states, None)
85
+ draft_token_ids, draft_probs = compute_probs_and_sample_next_token(
86
+ logits, sampling_metadata)
87
+
88
+ # Early exit if there is only one draft token to be generated.
89
+ if self.num_speculative_tokens == 1:
90
+ # [batch_size, 1] and [batch_size, 1, vocab_size]
91
+ return draft_token_ids.view(-1, 1), draft_probs.unsqueeze(dim=1)
92
+
93
+ # Generate the remaining draft tokens.
94
+ draft_token_ids_list = [draft_token_ids]
95
+ draft_probs_list = [draft_probs]
96
+
97
+ positions = target_positions[last_token_indices]
98
+ hidden_states = sample_hidden_states
99
+ attn_metadata.num_actual_tokens = batch_size
100
+ attn_metadata.max_query_len = 1
101
+ attn_metadata.query_start_loc = self.arange[:batch_size]
102
+ for _ in range(self.num_speculative_tokens - 1):
103
+ # Update the inputs.
104
+ input_ids = draft_token_ids_list[-1]
105
+ positions += 1
106
+ attn_metadata.max_seq_len += 1
107
+ attn_metadata.seq_lens += 1
108
+ # Compute the slot mapping.
109
+ block_numbers = positions // self.block_size
110
+ block_ids = block_table.gather(dim=1,
111
+ index=block_numbers.view(-1, 1))
112
+ block_ids = block_ids.view(-1)
113
+ attn_metadata.slot_mapping = (block_ids * self.block_size +
114
+ positions % self.block_size)
115
+
116
+ # Run the model.
117
+ with set_forward_context(attn_metadata, self.vllm_config):
118
+ hidden_states = self.model(
119
+ input_ids=input_ids,
120
+ hidden_states=hidden_states,
121
+ positions=positions,
122
+ )
123
+ logits = self.model.compute_logits(hidden_states, None)
124
+ draft_token_ids, probs = compute_probs_and_sample_next_token(
125
+ logits, sampling_metadata)
126
+ draft_token_ids_list.append(draft_token_ids)
127
+ draft_probs_list.append(probs)
128
+
129
+ # [batch_size, num_speculative_tokens]
130
+ draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
131
+ # [batch_size, num_speculative_tokens, vocab_size]
132
+ draft_probs = torch.stack(draft_probs_list, dim=1)
133
+ return draft_token_ids, draft_probs
134
+
135
+ @staticmethod
136
+ def prepare_inputs(
137
+ # [batch_size + 1]
138
+ cu_target_query_lens: torch.Tensor,
139
+ # [batch_size]
140
+ num_rejected_tokens: torch.Tensor,
141
+ ) -> tuple[torch.Tensor, torch.Tensor]:
142
+ # cu_target_query_lens: [0, a, a + b, a + b + c]
143
+ # num_rejected_tokens: [n1, n2, n3]
144
+ # num_tokens_per_req: [a - n1, b - n2, c - n3]
145
+ # cu_num_tokens: [0, a - n1, a + b - n1 - n2, a + b + c - n1 - n2 - n3]
146
+ # token_indices: [0, 1, ..., a - n1 - 1,
147
+ # a, a + 1, ..., a + b - n2 - 1,
148
+ # a + b, a + b + 1, ..., a + b + c - n3 - 1]
149
+
150
+ # [0, a, a + b, a + b + c] -> [a, b, c]
151
+ query_len_per_req = (cu_target_query_lens[1:] -
152
+ cu_target_query_lens[:-1])
153
+ # [a, b, c] -> [a - n1, b - n2, c - n3]
154
+ num_tokens_per_req = query_len_per_req - num_rejected_tokens
155
+
156
+ cu_num_tokens = torch.empty_like(cu_target_query_lens)
157
+ torch.cumsum(num_tokens_per_req, dim=0, out=cu_num_tokens[1:])
158
+ cu_num_tokens[0] = 0
159
+
160
+ # FIXME(woosuk): Avoid synchronization.
161
+ num_tokens = cu_num_tokens[-1].item()
162
+ token_indices = torch.empty(
163
+ num_tokens,
164
+ dtype=torch.int32,
165
+ device=cu_num_tokens.device,
166
+ )
167
+
168
+ batch_size = num_rejected_tokens.shape[0]
169
+ BLOCK_SIZE = 1024
170
+ prepare_input_kernel[(batch_size, )](
171
+ token_indices,
172
+ cu_target_query_lens,
173
+ cu_num_tokens,
174
+ BLOCK_SIZE=BLOCK_SIZE,
175
+ )
176
+ return cu_num_tokens, token_indices
177
+
178
+ def load_model(self, target_model: nn.Module) -> None:
179
+ self.model = DummyEagleModel()
180
+ self.model.get_input_embeddings = target_model.get_input_embeddings
181
+ self.model.compute_logits = target_model.compute_logits
182
+
183
+
184
+# FIXME(woosuk): This is a dummy model for testing.
185
+# Remove this once we have a real model.
186
+class DummyEagleModel(nn.Module):
187
+
188
+ def __init__(self):
189
+ super().__init__()
190
+
191
+ def forward(
192
+ self,
193
+ input_ids: torch.Tensor,
194
+ hidden_states: torch.Tensor,
195
+ positions: torch.Tensor,
196
+ ) -> torch.Tensor:
197
+ input_embeddings = self.get_input_embeddings(input_ids)
198
+ return hidden_states + input_embeddings # Dummy return.
199
+
200
+
201
+# FIXME(woosuk): The logic here is duplicated with the main sampling code.
202
+# We should refactor this to reuse the same sampling implementation.
203
+def compute_probs_and_sample_next_token(
204
+ logits: torch.Tensor,
205
+ sampling_metadata: SamplingMetadata,
206
+) -> tuple[torch.Tensor, torch.Tensor]:
207
+ if sampling_metadata.all_greedy:
208
+ # For greedy requests, draft_probs is not used in rejection sampling.
209
+ # Therefore, we can just return the logits.
210
+ probs = logits
211
+ next_token_ids = logits.argmax(dim=-1)
212
+ return next_token_ids, probs
213
+
214
+ is_greedy = sampling_metadata.temperature == -1
215
+ temperature = torch.where(is_greedy, 1.0, sampling_metadata.temperature)
216
+ logits.div_(temperature.view(-1, 1))
217
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
218
+
219
+ # NOTE(woosuk): Currently, we ignore most of the sampling parameters in
220
+ # generating the draft tokens. We only use the temperature. While this
221
+ # could degrade the acceptance rate, it does not affect the distribution
222
+ # of the generated tokens after rejection sampling.
223
+
224
+ # TODO(woosuk): Consider seeds.
225
+ q = torch.empty_like(probs)
226
+ q.exponential_()
227
+ next_token_ids = probs.div_(q).argmax(dim=-1).view(-1)
228
+ if not sampling_metadata.all_random:
229
+ greedy_token_ids = probs.argmax(dim=-1)
230
+ next_token_ids = torch.where(
231
+ is_greedy,
232
+ greedy_token_ids,
233
+ next_token_ids,
234
+ )
235
+ return next_token_ids, probs
236
+
237
+
238
+@triton.jit
239
+def prepare_input_kernel(
240
+ out_ptr,
241
+ cu_query_lens_ptr,
242
+ cu_num_tokens_ptr,
243
+ BLOCK_SIZE: tl.constexpr,
244
+):
245
+ pid = tl.program_id(0)
246
+
247
+ # [start_pos, end_pos)
248
+ start_pos = tl.load(cu_num_tokens_ptr + pid)
249
+ end_pos = tl.load(cu_num_tokens_ptr + pid + 1)
250
+ num_tokens = end_pos - start_pos
251
+
252
+ index_start = tl.load(cu_query_lens_ptr + pid)
253
+ indices = index_start + tl.arange(0, BLOCK_SIZE)
254
+
255
+ num_blocks = tl.cdiv(num_tokens, BLOCK_SIZE)
256
+ for i in tl.range(num_blocks):
257
+ offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
258
+ tl.store(
259
+ out_ptr + start_pos + offset,
260
+ indices,
261
+ mask=offset < num_tokens,
262
+ )
vllm/v1/spec_decode/ngram_proposer.py
+9
@@ -4,9 +4,14 @@ from typing import Optional
4
import numpy as np
5
from numba import jit
6
7
+from vllm.config import VllmConfig
8
+
9
10
class NgramProposer:
11
12
+ def __init__(self, vllm_config: VllmConfig):
13
+ self.vllm_config = vllm_config
14
+
15
def propose(
16
self,
17
context_token_ids: np.ndarray,
@@ -50,6 +55,10 @@ class NgramProposer:
55
return result
56
return None
57
58
+ def load_model(self, *args, **kwargs):
59
+ # No model to load.
60
+ pass
61
+
62
63
@jit(nopython=True)
64
def _kmp_lps_array(pattern: np.ndarray) -> np.ndarray:
vllm/v1/worker/gpu_input_batch.py
+10
-1
@@ -39,9 +39,18 @@ class CachedRequestState:
39
40
lora_request: Optional[LoRARequest] = None
41
42
+ def __post_init__(self):
43
+ self.num_prompt_tokens = len(self.prompt_token_ids)
44
+
45
@property
46
def num_tokens(self) -> int:
44
- return len(self.prompt_token_ids) + len(self.output_token_ids)
47
+ return self.num_prompt_tokens + len(self.output_token_ids)
48
+
49
+ def get_token_id(self, idx: int) -> int:
50
+ if idx < self.num_prompt_tokens:
51
+ return self.prompt_token_ids[idx]
52
+ else:
53
+ return self.output_token_ids[idx - self.num_prompt_tokens]
54
55
56
class InputBatch:
vllm/v1/worker/gpu_model_runner.py
+78
-12
@@ -35,6 +35,7 @@ from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
35
ModelRunnerOutput)
36
from vllm.v1.sample.metadata import SamplingMetadata
37
from vllm.v1.sample.rejection_sampler import RejectionSampler
38
+from vllm.v1.spec_decode.eagle import EagleProposer
39
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
40
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
41
from vllm.v1.spec_decode.utils import is_spec_decode_supported
@@ -157,18 +158,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
158
self.use_spec_decode = False
159
if self.speculative_config:
160
self.use_spec_decode = True
160
- assert self.speculative_config.method == "ngram", \
161
- "Currently, only ngram spec decode is supported in V1."
161
if get_pp_group().is_last_rank:
163
- self.drafter = NgramProposer()
164
- # Trigger Numba JIT compilation for N-gram proposer.
165
- # This usually takes less than 1 second.
166
- self.drafter.propose(
167
- np.zeros(1024, dtype=np.int32),
168
- self.speculative_config.prompt_lookup_min,
169
- self.speculative_config.prompt_lookup_max,
170
- self.speculative_config.num_speculative_tokens,
171
- )
162
+ if self.speculative_config.method == "ngram":
163
+ self.drafter = NgramProposer(self.vllm_config)
164
+ elif self.speculative_config.method == "eagle":
165
+ self.drafter = EagleProposer(self.vllm_config,
166
+ self.device) # type: ignore
167
+ else:
168
+ raise ValueError("Unknown speculative decoding method: "
169
+ f"{self.speculative_config.method}")
170
self.rejection_sampler = RejectionSampler()
171
172
# Request states.
@@ -1144,10 +1142,75 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1142
valid_sampled_token_ids[i].clear()
1143
1144
if not self.use_spec_decode:
1145
+ # Speculative decoding is not enabled.
1146
spec_token_ids = None
1148
- else:
1147
+ elif self.speculative_config.method == "ngram":
1148
+ assert isinstance(self.drafter, NgramProposer)
1149
spec_token_ids = self.generate_draft_token_ids(
1150
valid_sampled_token_ids, sampling_metadata)
1151
+ elif self.speculative_config.method == "eagle":
1152
+ assert isinstance(self.drafter, EagleProposer)
1153
+ # TODO(woosuk): Refactor the loop.
1154
+ next_token_ids: list[int] = []
1155
+ for i, token_ids in enumerate(valid_sampled_token_ids):
1156
+ if token_ids:
1157
+ # Common case.
1158
+ next_token_id = token_ids[-1]
1159
+ else:
1160
+ # Partial prefill (rare case).
1161
+ # Get the next token id from the request state.
1162
+ req_id = self.input_batch.req_ids[i]
1163
+ req_state = self.requests[req_id]
1164
+ seq_len = (req_state.num_computed_tokens +
1165
+ scheduler_output.num_scheduled_tokens[req_id])
1166
+ next_token_id = req_state.get_token_id(seq_len)
1167
+ next_token_ids.append(next_token_id)
1168
+ next_token_ids = torch.tensor(next_token_ids,
1169
+ dtype=torch.int32,
1170
+ device=self.device)
1171
+
1172
+ if spec_decode_metadata is None:
1173
+ # input_ids can be None for multimodal models.
1174
+ target_token_ids = self.input_ids[:num_scheduled_tokens]
1175
+ target_positions = positions
1176
+ target_hidden_states = hidden_states
1177
+ target_slot_mapping = attn_metadata.slot_mapping
1178
+ cu_num_tokens = attn_metadata.query_start_loc
1179
+ else:
1180
+ # TODO(woosuk): Refactor this.
1181
+ num_draft_tokens = spec_decode_metadata.num_draft_tokens
1182
+ num_rejected_tokens = [
1183
+ n + 1 - len(valid_sampled_token_ids[i]) if n > 0 else 0
1184
+ for i, n in enumerate(num_draft_tokens)
1185
+ ]
1186
+ num_rejected_tokens = torch.tensor(
1187
+ num_rejected_tokens,
1188
+ dtype=torch.int32,
1189
+ device=self.device,
1190
+ )
1191
+ cu_num_tokens, token_indices = self.drafter.prepare_inputs(
1192
+ attn_metadata.query_start_loc,
1193
+ num_rejected_tokens,
1194
+ )
1195
+ target_token_ids = self.input_ids[token_indices]
1196
+ target_positions = positions[token_indices]
1197
+ target_hidden_states = hidden_states[token_indices]
1198
+ target_slot_mapping = attn_metadata.slot_mapping[token_indices]
1199
+
1200
+ draft_token_ids, draft_probs = self.drafter.propose(
1201
+ target_token_ids=target_token_ids,
1202
+ target_positions=target_positions,
1203
+ target_hidden_states=target_hidden_states,
1204
+ target_slot_mapping=target_slot_mapping,
1205
+ next_token_ids=next_token_ids,
1206
+ cu_num_tokens=cu_num_tokens,
1207
+ block_table=attn_metadata.block_table,
1208
+ sampling_metadata=sampling_metadata,
1209
+ )
1210
+ spec_token_ids = draft_token_ids.tolist()
1211
+ # TODO(woosuk): Cache draft_probs and use it for rejection sampling
1212
+ # in the next step.
1213
+ del draft_probs
1214
1215
return ModelRunnerOutput(
1216
req_ids=self.input_batch.req_ids,
@@ -1205,6 +1268,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1268
self.scheduler_config,
1269
self.lora_config,
1270
self.device)
1271
+ if hasattr(self, "drafter"):
1272
+ logger.info("Loading drafter model...")
1273
+ self.drafter.load_model(self.model)
1274
time_after_load = time.perf_counter()
1275
self.model_memory_usage = m.consumed_memory
1276
logger.info("Model loading took %.4f GiB and %.6f seconds",