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 EAGLE3 (#17504)
Signed-off-by: qizixi <qizixi@meta.com>
qizixi committed
May 1, 2025 at 16:19 UTC
39c0813a7f1d0923adb828ff8319a068e6855c64
2 files changed
+36
-31
vllm/model_executor/models/llama_eagle3.py
+17
-8
@@ -6,7 +6,8 @@ import torch
6
import torch.nn as nn
7
from transformers import LlamaConfig
8
9
-from vllm.config import ModelConfig, VllmConfig
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.layernorm import RMSNorm
13
from vllm.model_executor.layers.linear import QKVParallelLinear
@@ -76,17 +77,19 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
77
return hidden_states, residual
78
79
80
+@support_torch_compile
81
class LlamaModel(nn.Module):
82
83
def __init__(
84
self,
85
*,
84
- model_config: ModelConfig,
86
+ vllm_config: VllmConfig,
87
start_layer_id: int = 0,
88
prefix: str = "",
89
) -> None:
90
super().__init__()
89
- self.config = model_config.hf_config
91
+ self.config = vllm_config. \
92
+ speculative_config.draft_model_config.hf_config
93
self.vocab_size = self.config.vocab_size
94
self.embed_tokens = VocabParallelEmbedding(
95
self.config.vocab_size,
@@ -119,8 +122,7 @@ class LlamaModel(nn.Module):
122
hidden_states: torch.Tensor,
123
) -> tuple[torch.Tensor, torch.Tensor]:
124
input_embeds = self.embed_tokens(input_ids)
122
- if (hidden_states.shape[-1] != input_embeds.shape[-1]):
123
- hidden_states = self.fc(hidden_states)
125
+ assert hidden_states.shape[-1] == input_embeds.shape[-1]
126
127
residual = None
128
hidden_states, residual = self.layers[0](
@@ -169,9 +171,9 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
171
172
def __init__(self, *, vllm_config: VllmConfig, start_layer_id: int = 0):
173
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,
174
+ self.config = vllm_config. \
175
+ speculative_config.draft_model_config.hf_config
176
+ self.model = LlamaModel(vllm_config=vllm_config,
177
start_layer_id=start_layer_id,
178
prefix="model")
179
@@ -214,6 +216,13 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
216
logits_new[:, targets] = logits
217
return logits_new
218
219
+ def combine_hidden_states(
220
+ self,
221
+ hidden_states: torch.Tensor,
222
+ ) -> torch.Tensor:
223
+ # combine multiple auxiliary hidden states returned by eagle3
224
+ return self.model.fc(hidden_states)
225
+
226
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
227
loader = AutoWeightsLoader(
228
self,
vllm/v1/spec_decode/eagle.py
+19
-23
@@ -10,6 +10,7 @@ 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 import ModelRegistry
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
@@ -39,11 +40,9 @@ class EagleProposer:
40
41
self.hidden_size = vllm_config.model_config.get_hidden_size()
42
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)
43
+ self.use_cuda_graph = (self.vllm_config.compilation_config.level
44
+ == CompilationLevel.PIECEWISE and
45
+ not self.vllm_config.model_config.enforce_eager)
46
47
self.cudagraph_batch_sizes = list(
48
reversed(
@@ -90,6 +89,12 @@ class EagleProposer:
89
batch_size = next_token_ids.shape[0]
90
last_token_indices = cu_num_tokens[1:] - 1
91
92
+ if self.method == "eagle3":
93
+ assert isinstance(self.model, Eagle3LlamaForCausalLM)
94
+ target_hidden_states = self.model.combine_hidden_states(
95
+ target_hidden_states)
96
+ assert target_hidden_states.shape[-1] == self.hidden_size
97
+
98
# Shift the input ids by one token.
99
# E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
100
self.input_ids[:num_tokens - 1] = target_token_ids[1:]
@@ -126,12 +131,7 @@ class EagleProposer:
131
# copy inputs to buffer for cudagraph
132
self.positions[:num_tokens] = target_positions
133
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
134
+ self.hidden_states[:num_tokens] = target_hidden_states
135
136
with set_forward_context(attn_metadata,
137
self.vllm_config,
@@ -139,7 +139,7 @@ class EagleProposer:
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],
142
+ hidden_states=self.hidden_states[:num_input_tokens],
143
)
144
sample_hidden_states = last_hidden_states[last_token_indices]
145
logits = self.model.compute_logits(sample_hidden_states, None)
@@ -209,10 +209,7 @@ class EagleProposer:
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
212
+ self.hidden_states[:batch_size] = hidden_states
213
214
# Run the model.
215
with set_forward_context(attn_metadata,
@@ -221,7 +218,7 @@ class EagleProposer:
218
last_hidden_states, hidden_states = self.model(
219
input_ids=self.input_ids[:input_batch_size],
220
positions=self.positions[:input_batch_size],
224
- hidden_states=hidden_states[:input_batch_size],
221
+ hidden_states=self.hidden_states[:input_batch_size],
222
)
223
hidden_states = hidden_states[:batch_size]
224
logits = self.model.compute_logits(last_hidden_states[:batch_size],
@@ -314,12 +311,11 @@ class EagleProposer:
311
) -> None:
312
with set_forward_context(None, self.vllm_config,
313
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
- )
314
+ self.model(
315
+ input_ids=self.input_ids[:num_tokens],
316
+ positions=self.positions[:num_tokens],
317
+ hidden_states=self.hidden_states[:num_tokens],
318
+ )
319
320
321
# NOTE(woosuk): Currently, the below code is not used and we always use argmax