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] Add minor clarification comments for Eagle (#29332)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 24, 2025 at 11:28 UTC
97588c4d1231287eb380cf9fb95ec77f88479b85
1 file changed
+11
vllm/v1/worker/gpu/spec_decode/eagle.py
+11
@@ -65,6 +65,12 @@ class EagleSpeculator:
65
# [num_reqs]
66
next_prefill_tokens: torch.Tensor,
67
) -> torch.Tensor:
68
+ # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
69
+ # number of rejected tokens, we maintain the size of eagle's input_ids and
70
+ # hidden_states the same as the target model's. This means, we pad each
71
+ # request's query length to include any rejected positions. By doing so,
72
+ # we can also reuse the attention metadata (e.g., query_start_loc,
73
+ # seq_lens) of the target model.
74
if aux_hidden_states:
75
assert self.method == "eagle3"
76
hidden_states = self.model.combine_hidden_states(
@@ -110,6 +116,11 @@ class EagleSpeculator:
116
# NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
117
# used for draft and target sampling.
118
pos = input_batch.positions[last_token_indices] + 1
119
+ # NOTE(woosuk): For draft sampling, we only consider the temperature
120
+ # and ignore the other sampling parameters such as top_k and top_p,
121
+ # for simplicity and performance.
122
+ # While this may slightly degrade the acceptance rate, it does not
123
+ # affect the output distribution after rejection sampling.
124
draft_tokens = gumbel_sample(
125
logits, temperature, seed, pos, apply_temperature=True
126
)