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, 11:00 UTC
[BugFix][Spec Decode] Fix hidden size mismatch between target and eagle head (#17740)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
May 6, 2025 at 19:51 UTC
8d84d836d17bdabe4c640021bc6f8bd11a546a44
1 file changed
+13
-14
vllm/v1/spec_decode/eagle.py
+13
-14
@@ -28,23 +28,25 @@ class EagleProposer:
28
device: torch.device,
29
):
30
self.vllm_config = vllm_config
31
- self.method = self.vllm_config.speculative_config.method
32
- self.num_speculative_tokens = (
33
- vllm_config.speculative_config.num_speculative_tokens)
34
- self.max_model_len = vllm_config.model_config.max_model_len
35
- self.block_size = vllm_config.cache_config.block_size
31
+ self.speculative_config = vllm_config.speculative_config
32
+ self.draft_model_config = self.speculative_config.draft_model_config
33
+ self.method = self.speculative_config.method
34
35
self.dtype = vllm_config.model_config.dtype
38
-
39
- self.max_num_tokens = vllm_config.scheduler_config \
40
- .max_num_batched_tokens
41
-
42
- self.hidden_size = vllm_config.model_config.get_hidden_size()
36
+ self.max_model_len = vllm_config.model_config.max_model_len
37
+ self.block_size = vllm_config.cache_config.block_size
38
+ self.num_speculative_tokens = (
39
+ self.speculative_config.num_speculative_tokens)
40
+ self.max_num_tokens = (
41
+ vllm_config.scheduler_config.max_num_batched_tokens)
42
+ # We need to get the hidden size from the draft model config because
43
+ # the draft model's hidden size can be different from the target model's
44
+ # hidden size (e.g., Llama 3.3 70B).
45
+ self.hidden_size = self.draft_model_config.get_hidden_size()
46
47
self.use_cuda_graph = (self.vllm_config.compilation_config.level
48
== CompilationLevel.PIECEWISE and
49
not self.vllm_config.model_config.enforce_eager)
47
-
50
self.cudagraph_batch_sizes = list(
51
reversed(
52
self.vllm_config.compilation_config.cudagraph_capture_sizes))
@@ -56,7 +58,6 @@ class EagleProposer:
58
self.positions = torch.zeros(self.max_num_tokens,
59
dtype=torch.int64,
60
device=device)
59
-
61
self.hidden_states = torch.zeros(
62
(self.max_num_tokens, self.hidden_size),
63
dtype=self.dtype,
@@ -131,7 +132,6 @@ class EagleProposer:
132
num_input_tokens = num_tokens
133
# copy inputs to buffer for cudagraph
134
self.positions[:num_tokens] = target_positions
134
-
135
self.hidden_states[:num_tokens] = target_hidden_states
136
137
with set_forward_context(attn_metadata,
@@ -209,7 +209,6 @@ class EagleProposer:
209
# copy inputs to buffer for cudagraph
210
self.input_ids[:batch_size] = input_ids
211
self.positions[:batch_size] = clamped_positions
212
-
212
self.hidden_states[:batch_size] = hidden_states
213
214
# Run the model.