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, 12:15 UTC
[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE (#17211)
Signed-off-by: Bryan Lu <yuzhelu@amazon.com>
Bryan Lu committed
Apr 29, 2025 at 14:10 UTC
70788bdbdc590e7fbf9bddb3fa9bc92ac3181733
6 files changed
+152
-53
examples/offline_inference/eagle.py
+12
-2
@@ -36,6 +36,10 @@ def parse_args():
36
help="downloaded from the eagle repo " \
37
"https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/"
38
)
39
+ parser.add_argument("--method",
40
+ type=str,
41
+ default='eagle',
42
+ choices=['eagle', 'eagle3'])
43
parser.add_argument("--max_num_seqs", type=int, default=8)
44
parser.add_argument("--num_prompts", type=int, default=80)
45
parser.add_argument("--num_spec_tokens", type=int, default=2)
@@ -53,7 +57,13 @@ def main():
57
args = parse_args()
58
59
model_dir = "meta-llama/Llama-3.1-8B-Instruct"
56
- eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
60
+
61
+ if args.method == 'eagle':
62
+ eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
63
+ elif args.method == 'eagle3':
64
+ eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
65
+ else:
66
+ raise ValueError(f"unknown method: {args.method}")
67
68
max_model_len = 2048
69
@@ -81,7 +91,7 @@ def main():
91
max_num_seqs=args.max_num_seqs,
92
gpu_memory_utilization=0.8,
93
speculative_config={
84
- "method": "eagle3" if "eagle3" in eagle_dir.lower() else "eagle",
94
+ "method": args.method,
95
"model": eagle_dir,
96
"num_speculative_tokens": args.num_spec_tokens,
97
"draft_tensor_parallel_size": args.draft_tp,
vllm/compilation/backends.py
+12
-3
@@ -347,8 +347,12 @@ class VllmBackend:
347
PASS_KEY = "post_grad_custom_post_pass"
348
if PASS_KEY in inductor_config:
349
# Config should automatically wrap all inductor passes
350
- assert isinstance(inductor_config[PASS_KEY], InductorPass)
351
- self.post_grad_pass_manager.add(inductor_config[PASS_KEY])
350
+ if isinstance(inductor_config[PASS_KEY], PostGradPassManager):
351
+ assert (inductor_config[PASS_KEY].uuid() ==
352
+ self.post_grad_pass_manager.uuid())
353
+ else:
354
+ assert isinstance(inductor_config[PASS_KEY], InductorPass)
355
+ self.post_grad_pass_manager.add(inductor_config[PASS_KEY])
356
inductor_config[PASS_KEY] = self.post_grad_pass_manager
357
358
def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:
@@ -408,8 +412,13 @@ class VllmBackend:
412
)
413
self.compilation_config.cache_dir = cache_dir
414
411
- cache_dir = self.compilation_config.cache_dir
415
+ if compilation_counter.num_graphs_seen > 0:
416
+ cache_dir = self.compilation_config.cache_dir + \
417
+ f'-{compilation_counter.num_graphs_seen}'
418
+ else:
419
+ cache_dir = self.compilation_config.cache_dir
420
os.makedirs(cache_dir, exist_ok=True)
421
+ self.compilation_config.cache_dir = cache_dir
422
rank = vllm_config.parallel_config.rank
423
dp_rank = vllm_config.parallel_config.data_parallel_rank
424
local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
vllm/model_executor/models/llama_eagle.py
+14
-11
@@ -6,7 +6,8 @@ import torch
6
import torch.nn as nn
7
from transformers import LlamaConfig
8
9
-from vllm.config import ModelConfig
9
+from vllm.compilation.decorators import support_torch_compile
10
+from vllm.config import VllmConfig
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.logits_processor import LogitsProcessor
13
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -37,17 +38,19 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
38
self.input_layernorm = nn.Identity()
39
40
41
+@support_torch_compile
42
class LlamaModel(nn.Module):
43
44
def __init__(
45
self,
46
*,
45
- model_config: ModelConfig,
46
- start_layer_id: int = 0,
47
+ vllm_config: VllmConfig,
48
prefix: str = "",
49
+ start_layer_id: int = 0,
50
) -> None:
51
super().__init__()
50
- self.config = model_config.hf_config
52
+ self.config = vllm_config. \
53
+ speculative_config.draft_model_config.hf_config
54
self.vocab_size = self.config.vocab_size
55
self.embed_tokens = VocabParallelEmbedding(
56
self.config.vocab_size,
@@ -75,8 +78,7 @@ class LlamaModel(nn.Module):
78
hidden_states = self.fc(
79
torch.cat((input_embeds, hidden_states), dim=-1))
80
residual = None
78
- for i in range(len(self.layers)):
79
- layer = self.layers[i]
81
+ for layer in self.layers:
82
hidden_states, residual = layer(
83
positions,
84
hidden_states,
@@ -117,12 +119,13 @@ class LlamaModel(nn.Module):
119
120
class EagleLlamaForCausalLM(LlamaForCausalLM):
121
120
- def __init__(self, *, model_config: ModelConfig, start_layer_id: int = 0):
122
+ def __init__(self, *, vllm_config: VllmConfig, start_layer_id: int = 0):
123
nn.Module.__init__(self)
122
- self.config = model_config.hf_config
123
- self.model = LlamaModel(model_config=model_config,
124
- start_layer_id=start_layer_id,
125
- prefix="model")
124
+ self.config = vllm_config. \
125
+ speculative_config.draft_model_config.hf_config
126
+ self.model = LlamaModel(vllm_config=vllm_config,
127
+ prefix="model",
128
+ start_layer_id=start_layer_id)
129
130
logit_scale = getattr(self.config, "logit_scale", 1.0)
131
self.logits_processor = LogitsProcessor(self.config.vocab_size,
vllm/model_executor/models/llama_eagle3.py
+3
-2
@@ -6,7 +6,7 @@ import torch
6
import torch.nn as nn
7
from transformers import LlamaConfig
8
9
-from vllm.config import ModelConfig
9
+from vllm.config import ModelConfig, VllmConfig
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
@@ -167,8 +167,9 @@ class LlamaModel(nn.Module):
167
168
class Eagle3LlamaForCausalLM(LlamaForCausalLM):
169
170
- def __init__(self, *, model_config: ModelConfig, start_layer_id: int = 0):
170
+ def __init__(self, *, vllm_config: VllmConfig, start_layer_id: int = 0):
171
nn.Module.__init__(self)
172
+ model_config = vllm_config.speculative_config.draft_model_config
173
self.config = model_config.hf_config
174
self.model = LlamaModel(model_config=model_config,
175
start_layer_id=start_layer_id,
vllm/v1/spec_decode/eagle.py
+100
-22
@@ -4,7 +4,7 @@ import torch.nn as nn
4
import triton
5
import triton.language as tl
6
7
-from vllm.config import VllmConfig, set_current_vllm_config
7
+from vllm.config import CompilationLevel, 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
@@ -26,10 +26,41 @@ class EagleProposer:
26
device: torch.device,
27
):
28
self.vllm_config = vllm_config
29
+ self.method = self.vllm_config.speculative_config.method
30
self.num_speculative_tokens = (
31
vllm_config.speculative_config.num_speculative_tokens)
32
self.max_model_len = vllm_config.model_config.max_model_len
33
self.block_size = vllm_config.cache_config.block_size
34
+
35
+ self.dtype = vllm_config.model_config.dtype
36
+
37
+ self.max_num_tokens = vllm_config.scheduler_config \
38
+ .max_num_batched_tokens
39
+
40
+ self.hidden_size = vllm_config.model_config.get_hidden_size()
41
+
42
+ # TODO: make eagle3 compatible with cudagraph
43
+ self.use_cuda_graph = self.method != 'eagle3' and \
44
+ (self.vllm_config.compilation_config.level
45
+ == CompilationLevel.PIECEWISE and
46
+ not self.vllm_config.model_config.enforce_eager)
47
+
48
+ self.cudagraph_batch_sizes = list(
49
+ reversed(
50
+ self.vllm_config.compilation_config.cudagraph_capture_sizes))
51
+
52
+ # persistent buffers for cuda graph
53
+ self.input_ids = torch.zeros(self.max_num_tokens,
54
+ dtype=torch.int32,
55
+ device=device)
56
+ self.positions = torch.zeros(self.max_num_tokens,
57
+ dtype=torch.int64,
58
+ device=device)
59
+
60
+ self.hidden_states = torch.zeros(
61
+ (self.max_num_tokens, self.hidden_size),
62
+ dtype=self.dtype,
63
+ device=device)
64
# We need +1 here because the arange is used to set query_start_loc,
65
# which has one more element than batch_size.
66
self.arange = torch.arange(vllm_config.scheduler_config.max_num_seqs +
@@ -59,13 +90,12 @@ class EagleProposer:
90
batch_size = next_token_ids.shape[0]
91
last_token_indices = cu_num_tokens[1:] - 1
92
62
- input_ids = torch.empty_like(target_token_ids)
93
# Shift the input ids by one token.
94
# E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
65
- input_ids[:-1] = target_token_ids[1:]
95
+ self.input_ids[:num_tokens - 1] = target_token_ids[1:]
96
# Replace the last token with the next token.
97
# E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
68
- input_ids[last_token_indices] = next_token_ids
98
+ self.input_ids[last_token_indices] = next_token_ids
99
100
# FA requires seq_len to have dtype int32.
101
seq_lens = (target_positions[last_token_indices] + 1).int()
@@ -88,14 +118,30 @@ class EagleProposer:
118
prefix_kv_lens=None,
119
suffix_kv_lens=None,
120
)
121
+ if self.use_cuda_graph and \
122
+ num_tokens <= self.cudagraph_batch_sizes[-1]:
123
+ num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
124
+ else:
125
+ num_input_tokens = num_tokens
126
+ # copy inputs to buffer for cudagraph
127
+ self.positions[:num_tokens] = target_positions
128
92
- with set_forward_context(attn_metadata, self.vllm_config):
93
- hidden_states_logits, hidden_states_fwd = self.model(
94
- input_ids=input_ids,
95
- hidden_states=target_hidden_states,
96
- positions=target_positions,
129
+ if self.method == 'eagle':
130
+ self.hidden_states[:num_tokens] = target_hidden_states
131
+ hidden_states = self.hidden_states
132
+ else:
133
+ # TODO: make eagle3 compatible with cuda graph
134
+ hidden_states = target_hidden_states
135
+
136
+ with set_forward_context(attn_metadata,
137
+ self.vllm_config,
138
+ num_tokens=num_input_tokens):
139
+ last_hidden_states, hidden_states = self.model(
140
+ input_ids=self.input_ids[:num_input_tokens],
141
+ positions=self.positions[:num_input_tokens],
142
+ hidden_states=hidden_states[:num_input_tokens],
143
)
98
- sample_hidden_states = hidden_states_logits[last_token_indices]
144
+ sample_hidden_states = last_hidden_states[last_token_indices]
145
logits = self.model.compute_logits(sample_hidden_states, None)
146
draft_token_ids = logits.argmax(dim=-1)
147
@@ -108,13 +154,20 @@ class EagleProposer:
154
draft_token_ids_list = [draft_token_ids]
155
156
positions = target_positions[last_token_indices]
111
- hidden_states = hidden_states_fwd[last_token_indices]
157
+ hidden_states = hidden_states[last_token_indices]
158
+ if self.use_cuda_graph and \
159
+ batch_size <= self.cudagraph_batch_sizes[-1]:
160
+ input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
161
+ else:
162
+ input_batch_size = batch_size
163
attn_metadata.num_actual_tokens = batch_size
164
attn_metadata.max_query_len = 1
165
attn_metadata.query_start_loc = self.arange[:batch_size + 1]
166
for _ in range(self.num_speculative_tokens - 1):
167
# Update the inputs.
117
- input_ids = draft_token_ids_list[-1]
168
+ # cast to int32 is crucial when eagle model is compiled.
169
+ # tensor.argmax() returns int64 by default.
170
+ input_ids = draft_token_ids_list[-1].int()
171
positions += 1
172
173
# NOTE(woosuk): We should handle the case where the draft model
@@ -152,14 +205,27 @@ class EagleProposer:
205
attn_metadata.slot_mapping.masked_fill_(exceeds_max_model_len,
206
PADDING_SLOT_ID)
207
208
+ # copy inputs to buffer for cudagraph
209
+ self.input_ids[:batch_size] = input_ids
210
+ self.positions[:batch_size] = clamped_positions
211
+
212
+ if self.method == 'eagle':
213
+ # TODO: make eagle3 compatible with cudagraph.
214
+ self.hidden_states[:batch_size] = hidden_states
215
+ hidden_states = self.hidden_states
216
+
217
# Run the model.
156
- with set_forward_context(attn_metadata, self.vllm_config):
157
- hidden_states_logits, hidden_states = self.model(
158
- input_ids=input_ids,
159
- hidden_states=hidden_states,
160
- positions=clamped_positions,
218
+ with set_forward_context(attn_metadata,
219
+ self.vllm_config,
220
+ num_tokens=input_batch_size):
221
+ last_hidden_states, hidden_states = self.model(
222
+ input_ids=self.input_ids[:input_batch_size],
223
+ positions=self.positions[:input_batch_size],
224
+ hidden_states=hidden_states[:input_batch_size],
225
)
162
- logits = self.model.compute_logits(hidden_states_logits, None)
226
+ hidden_states = hidden_states[:batch_size]
227
+ logits = self.model.compute_logits(last_hidden_states[:batch_size],
228
+ None)
229
draft_token_ids = logits.argmax(dim=-1)
230
draft_token_ids_list.append(draft_token_ids)
231
@@ -227,13 +293,11 @@ class EagleProposer:
293
draft_model_cls, arch = ModelRegistry.resolve_model_cls(
294
draft_model_config.architectures)
295
self.model = draft_model_cls(
230
- model_config=draft_model_config,
296
+ vllm_config=self.vllm_config,
297
start_layer_id=target_layer_num).to(target_device)
298
299
loaded_weights = self.model.load_weights(
234
- loader.get_all_weights(
235
- self.vllm_config.speculative_config.draft_model_config,
236
- self.model))
300
+ loader.get_all_weights(draft_model_config, self.model))
301
if self.vllm_config.speculative_config.method == "eagle3":
302
if "model.embed_tokens.weight" not in loaded_weights:
303
logger.info(
@@ -243,6 +307,20 @@ class EagleProposer:
307
logger.info("Loading EAGLE LM head weights from the target model.")
308
self.model.lm_head = target_model.lm_head
309
310
+ @torch.inference_mode()
311
+ def dummy_run(
312
+ self,
313
+ num_tokens: int,
314
+ ) -> None:
315
+ with set_forward_context(None, self.vllm_config,
316
+ num_tokens=num_tokens):
317
+ if self.method == 'eagle':
318
+ self.model(
319
+ input_ids=self.input_ids[:num_tokens],
320
+ positions=self.positions[:num_tokens],
321
+ hidden_states=self.hidden_states[:num_tokens],
322
+ )
323
+
324
325
# NOTE(woosuk): Currently, the below code is not used and we always use argmax
326
# to sample the draft tokens. We will use this after we find a way to manage
vllm/v1/worker/gpu_model_runner.py
+11
-13
@@ -1106,7 +1106,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1106
# For mid-pipeline stages, return the hidden states.
1107
return hidden_states
1108
1109
- hidden_states = hidden_states[:num_scheduled_tokens]
1109
sample_hidden_states = hidden_states[logits_indices]
1110
logits = self.model.compute_logits(sample_hidden_states, None)
1111
@@ -1172,7 +1171,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1171
1172
# Compute prompt logprobs if needed.
1173
prompt_logprobs_dict = self._get_prompt_logprobs_dict(
1175
- hidden_states,
1174
+ hidden_states[:num_scheduled_tokens],
1175
scheduler_output,
1176
)
1177
@@ -1222,15 +1221,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1221
1222
if spec_decode_metadata is None:
1223
# input_ids can be None for multimodal models.
1225
- # We need to slice token_ids, positions, and hidden_states
1226
- # because the eagle head does not use cuda graph and should
1227
- # not include padding.
1224
target_token_ids = self.input_ids[:num_scheduled_tokens]
1225
target_positions = positions[:num_scheduled_tokens]
1226
if self.use_aux_hidden_state_outputs:
1231
- target_hidden_states = [
1232
- h[:num_scheduled_tokens] for h in aux_hidden_states
1233
- ]
1227
+ target_hidden_states = torch.cat(
1228
+ [h[:num_scheduled_tokens] for h in aux_hidden_states],
1229
+ dim=-1)
1230
else:
1231
target_hidden_states = hidden_states[:num_scheduled_tokens]
1232
target_slot_mapping = attn_metadata.slot_mapping
@@ -1254,15 +1250,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1250
target_token_ids = self.input_ids[token_indices]
1251
target_positions = positions[token_indices]
1252
if self.use_aux_hidden_state_outputs:
1257
- target_hidden_states = [
1258
- h[token_indices] for h in aux_hidden_states
1259
- ]
1253
+ target_hidden_states = torch.cat(
1254
+ [h[token_indices] for h in aux_hidden_states], dim=-1)
1255
else:
1256
target_hidden_states = hidden_states[token_indices]
1257
target_slot_mapping = attn_metadata.slot_mapping[token_indices]
1258
1264
- if self.use_aux_hidden_state_outputs:
1265
- target_hidden_states = torch.cat(target_hidden_states, dim=-1)
1259
draft_token_ids = self.drafter.propose(
1260
target_token_ids=target_token_ids,
1261
target_positions=target_positions,
@@ -1506,6 +1499,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1499
else:
1500
hidden_states = outputs
1501
1502
+ if self.use_spec_decode and \
1503
+ self.speculative_config.method in ('eagle', 'eagle3'):
1504
+ assert isinstance(self.drafter, EagleProposer)
1505
+ self.drafter.dummy_run(num_tokens)
1506
+
1507
logit_indices = np.cumsum(num_scheduled_tokens) - 1
1508
return hidden_states[logit_indices]
1509