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, 13:30 UTC
[BugFix][Spec Decode] No in-place update to draft probs (#16952)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Apr 21, 2025 at 19:54 UTC
1311913f5537b36a7b12f481ebd15f7ad775db58
1 file changed
+3
-1
vllm/v1/spec_decode/eagle.py
+3
-1
@@ -264,7 +264,9 @@ def compute_probs_and_sample_next_token(
264
# TODO(woosuk): Consider seeds.
265
q = torch.empty_like(probs)
266
q.exponential_()
267
- next_token_ids = probs.div_(q).argmax(dim=-1).view(-1)
267
+ # NOTE(woosuk): We shouldn't use `probs.div_(q)` because the draft_probs
268
+ # will be used later for rejection sampling.
269
+ next_token_ids = probs.div(q).argmax(dim=-1).view(-1)
270
if not sampling_metadata.all_random:
271
greedy_token_ids = probs.argmax(dim=-1)
272
next_token_ids = torch.where(