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 11, 05:00 UTC
[BugFix] Fix logprobs with spec decode and modified logits (#30846)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Dec 18, 2025 at 19:58 UTC
2ac85a4544cf9488037e61bf9ed7a87d0c3696bb
2 files changed
+30
-11
tests/v1/sample/test_logprobs.py
+23
-10
@@ -547,6 +547,13 @@ def test_spec_decode_logprobs(
547
sampling_params = SamplingParams(
548
temperature=0, logprobs=top_logprobs, max_tokens=10, ignore_eos=False
549
)
550
+ penalty_sampling_params = SamplingParams(
551
+ temperature=0,
552
+ logprobs=top_logprobs,
553
+ max_tokens=10,
554
+ ignore_eos=False,
555
+ presence_penalty=-1.0,
556
+ )
557
method, model_name, spec_model_name = model_setup
558
max_model_len = 256
559
@@ -558,14 +565,17 @@ def test_spec_decode_logprobs(
565
seed=42,
566
logprobs_mode=logprobs_mode,
567
gpu_memory_utilization=0.4,
568
+ enable_prefix_caching=False,
569
+ )
570
+ ref_results = ref_llm.generate(
571
+ [prompt, prompt], [sampling_params, penalty_sampling_params]
572
)
562
- ref_results = ref_llm.generate([prompt], sampling_params)
573
# Collect logprobs outputs from reference LLM.
574
ref_logprobs = []
565
- for output in ref_results[0].outputs:
566
- for logprobs in output.logprobs:
567
- for token_id in logprobs:
568
- ref_logprobs.append(logprobs[token_id])
575
+ for results in ref_results:
576
+ for output in results.outputs:
577
+ for logprobs in output.logprobs:
578
+ ref_logprobs.extend(logprobs.values())
579
del ref_llm
580
torch.cuda.empty_cache()
581
cleanup_dist_env_and_memory()
@@ -587,14 +597,17 @@ def test_spec_decode_logprobs(
597
# Force prefill chunking
598
enable_chunked_prefill=True,
599
max_num_batched_tokens=32,
600
+ enable_prefix_caching=False,
601
+ )
602
+ spec_results = spec_llm.generate(
603
+ [prompt, prompt], [sampling_params, penalty_sampling_params]
604
)
591
- spec_results = spec_llm.generate([prompt], sampling_params)
605
# Collect logprobs outputs from spec decode LLM.
606
spec_logprobs = []
594
- for output in spec_results[0].outputs:
595
- for logprobs in output.logprobs:
596
- for token_id in logprobs:
597
- spec_logprobs.append(logprobs[token_id])
607
+ for results in spec_results:
608
+ for output in results.outputs:
609
+ for logprobs in output.logprobs:
610
+ spec_logprobs.extend(logprobs.values())
611
del spec_llm
612
torch.cuda.empty_cache()
613
cleanup_dist_env_and_memory()
vllm/v1/sample/rejection_sampler.py
+7
-1
@@ -119,8 +119,14 @@ class RejectionSampler(nn.Module):
119
raw_target_logits = logits[target_logits_indices]
120
# Use float32 for the target_logits.
121
raw_target_logits = raw_target_logits.to(torch.float32)
122
+ target_logits = raw_target_logits
123
+ if not self.is_processed_logprobs_mode:
124
+ # Clone raw_target_logits before applying processors to preserve
125
+ # the original raw logits for logprobs computation, since
126
+ # apply_logits_processors modifies the tensor in-place.
127
+ target_logits = target_logits.clone()
128
target_logits = self.apply_logits_processors(
123
- raw_target_logits, sampling_metadata, metadata
129
+ target_logits, sampling_metadata, metadata
130
)
131
# [num_tokens, vocab_size]
132
# NOTE(woosuk): `target_logits` can be updated in place inside the