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 4, 07:30 UTC
[bugfix] Fix prompt logprobs on request eviction during chunked prefill (#41411)
Signed-off-by: Joachim Studnia <joachim@mistral.ai> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Joachim Studnia committed
May 4, 2026 at 11:46 UTC
422dd0259853e5ce2752341109b38293b9bbb5b8
5 files changed
+19
-13
tests/conftest.py
+2
@@ -996,6 +996,8 @@ class VllmRunner:
996
req_sample_output_ids: list[list[int]] = []
997
req_sample_output_strs: list[str] = []
998
req_logprobs = []
999
+ if req_output.prompt_logprobs:
1000
+ req_logprobs.extend(req_output.prompt_logprobs)
1001
for sample in req_output.outputs:
1002
output_str = sample.text
1003
output_ids = list(sample.token_ids)
tests/v1/e2e/general/test_async_scheduling.py
+10
-1
@@ -57,6 +57,8 @@ def test_without_spec_decoding(
57
dict(bad_words=["the", " the"]),
58
dict(logprobs=2),
59
dict(logprobs=2, frequency_penalty=-1.0),
60
+ dict(prompt_logprobs=2),
61
+ dict(prompt_logprobs=2, logprobs=2),
62
dict(structured_outputs=struct_outputs),
63
dict(
64
structured_outputs=struct_outputs,
@@ -126,6 +128,8 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
128
dict(bad_words=["the", " the"]),
129
dict(logprobs=2),
130
dict(logprobs=2, frequency_penalty=-1.0),
131
+ dict(prompt_logprobs=2),
132
+ dict(prompt_logprobs=2, logprobs=2),
133
dict(structured_outputs=struct_outputs),
134
dict(
135
structured_outputs=struct_outputs,
@@ -413,7 +417,12 @@ def _all_logprobs_match(req_a, req_b) -> bool:
417
)
418
419
416
-def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool:
420
+def _logprobs_match(
421
+ lps_a: dict[int, Logprob] | None,
422
+ lps_b: dict[int, Logprob] | None,
423
+) -> bool:
424
+ if lps_a is None or lps_b is None:
425
+ return lps_a is lps_b
426
rel_tol, abs_tol = 1e-3, 1e-6
427
return (
428
len(lps_a) == len(lps_b)
vllm/v1/worker/gpu/sample/prompt_logprob.py
+1
-3
@@ -55,10 +55,8 @@ class PromptLogprobsWorker:
55
56
num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np]
57
prompt_lens = prompt_lens[idx_mapping_np]
58
- # NOTE(woosuk): -1 because the last prompt token's hidden state is not
59
- # needed for prompt logprobs.
58
computed_prefill = num_computed_prefill_tokens[idx_mapping_np]
61
- includes_prompt = computed_prefill < prompt_lens - 1
59
+ includes_prompt = computed_prefill < prompt_lens
60
# NOTE(woosuk): If the request was resumed after preemption, its prompt
61
# logprobs must have been computed before preemption. Skip.
62
resumed_after_prompt = prompt_lens < prefill_lens[idx_mapping_np]
vllm/v1/worker/gpu_input_batch.py
+2
-4
@@ -49,6 +49,8 @@ class CachedRequestState:
49
50
lora_request: LoRARequest | None = None
51
prompt_embeds: torch.Tensor | None = None
52
+ # To accumulate prompt logprobs tensor chunks across prefill steps.
53
+ in_progress_prompt_logprobs_cpu: LogprobsTensors | None = None
54
55
# Per-position mask for mixed-mode inputs (e.g chat completion with
56
# prompt_embeds content parts). See `Request.prompt_is_token_ids`.
@@ -255,9 +257,6 @@ class InputBatch:
257
# More efficient than num_logprobs=-1 when only a few tokens are needed
258
self.logprob_token_ids: dict[str, list[int]] = {}
259
258
- # To accumulate prompt logprobs tensor chunks across prefill steps.
259
- self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
260
-
260
# Internal representation of per-step batch state changes, used for
261
# reordering persistent batch and generating logitsprocs batch state
262
# updates. Should reset each step.
@@ -552,7 +551,6 @@ class InputBatch:
551
self.generators.pop(req_index, None)
552
self.num_logprobs.pop(req_id, None)
553
self.logprob_token_ids.pop(req_id, None)
555
- self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
554
if self.prev_req_id_to_index is not None:
555
self.prev_req_id_to_index.pop(req_id, None)
556
vllm/v1/worker/gpu_model_runner.py
+4
-5
@@ -5094,7 +5094,6 @@ class GPUModelRunner(
5094
if not num_prompt_logprobs_dict:
5095
return {}
5096
5097
- in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu
5097
prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
5098
5099
# Since prompt logprobs are a rare feature, prioritize simple,
@@ -5118,14 +5117,14 @@ class GPUModelRunner(
5117
)
5118
5119
# Set up target LogprobsTensors object.
5121
- logprobs_tensors = in_progress_dict.get(req_id)
5122
- if not logprobs_tensors:
5120
+ logprobs_tensors = request.in_progress_prompt_logprobs_cpu
5121
+ if logprobs_tensors is None:
5122
# Create empty logprobs CPU tensors for the entire prompt.
5123
# If chunked, we'll copy in slice by slice.
5124
logprobs_tensors = LogprobsTensors.empty_cpu(
5125
num_prompt_tokens - 1, num_prompt_logprobs + 1
5126
)
5128
- in_progress_dict[req_id] = logprobs_tensors
5127
+ request.in_progress_prompt_logprobs_cpu = logprobs_tensors
5128
5129
# Determine number of logits to retrieve.
5130
start_idx = request.num_computed_tokens
@@ -5182,7 +5181,7 @@ class GPUModelRunner(
5181
# num_prompt_logprobs_dict.
5182
for req_id in completed_prefill_reqs:
5183
del num_prompt_logprobs_dict[req_id]
5185
- del in_progress_dict[req_id]
5184
+ self.requests[req_id].in_progress_prompt_logprobs_cpu = None
5185
5186
# Must synchronize the non-blocking GPU->CPU transfers.
5187
if prompt_logprobs_dict: