diff --git a/tests/conftest.py b/tests/conftest.py index 40adeda2bd..779bd475f3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -996,6 +996,8 @@ class VllmRunner: req_sample_output_ids: list[list[int]] = [] req_sample_output_strs: list[str] = [] req_logprobs = [] + if req_output.prompt_logprobs: + req_logprobs.extend(req_output.prompt_logprobs) for sample in req_output.outputs: output_str = sample.text output_ids = list(sample.token_ids) diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 28a1bedbe0..c3c4970de3 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -57,6 +57,8 @@ def test_without_spec_decoding( dict(bad_words=["the", " the"]), dict(logprobs=2), dict(logprobs=2, frequency_penalty=-1.0), + dict(prompt_logprobs=2), + dict(prompt_logprobs=2, logprobs=2), dict(structured_outputs=struct_outputs), dict( structured_outputs=struct_outputs, @@ -126,6 +128,8 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke dict(bad_words=["the", " the"]), dict(logprobs=2), dict(logprobs=2, frequency_penalty=-1.0), + dict(prompt_logprobs=2), + dict(prompt_logprobs=2, logprobs=2), dict(structured_outputs=struct_outputs), dict( structured_outputs=struct_outputs, @@ -413,7 +417,12 @@ def _all_logprobs_match(req_a, req_b) -> bool: ) -def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool: +def _logprobs_match( + lps_a: dict[int, Logprob] | None, + lps_b: dict[int, Logprob] | None, +) -> bool: + if lps_a is None or lps_b is None: + return lps_a is lps_b rel_tol, abs_tol = 1e-3, 1e-6 return ( len(lps_a) == len(lps_b) diff --git a/vllm/v1/worker/gpu/sample/prompt_logprob.py b/vllm/v1/worker/gpu/sample/prompt_logprob.py index 11dbf69852..baa48ebf90 100644 --- a/vllm/v1/worker/gpu/sample/prompt_logprob.py +++ b/vllm/v1/worker/gpu/sample/prompt_logprob.py @@ -55,10 +55,8 @@ class PromptLogprobsWorker: num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np] prompt_lens = prompt_lens[idx_mapping_np] - # NOTE(woosuk): -1 because the last prompt token's hidden state is not - # needed for prompt logprobs. computed_prefill = num_computed_prefill_tokens[idx_mapping_np] - includes_prompt = computed_prefill < prompt_lens - 1 + includes_prompt = computed_prefill < prompt_lens # NOTE(woosuk): If the request was resumed after preemption, its prompt # logprobs must have been computed before preemption. Skip. resumed_after_prompt = prompt_lens < prefill_lens[idx_mapping_np] diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 75898f4632..44e0efaaa2 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -49,6 +49,8 @@ class CachedRequestState: lora_request: LoRARequest | None = None prompt_embeds: torch.Tensor | None = None + # To accumulate prompt logprobs tensor chunks across prefill steps. + in_progress_prompt_logprobs_cpu: LogprobsTensors | None = None # Per-position mask for mixed-mode inputs (e.g chat completion with # prompt_embeds content parts). See `Request.prompt_is_token_ids`. @@ -255,9 +257,6 @@ class InputBatch: # More efficient than num_logprobs=-1 when only a few tokens are needed self.logprob_token_ids: dict[str, list[int]] = {} - # To accumulate prompt logprobs tensor chunks across prefill steps. - self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {} - # Internal representation of per-step batch state changes, used for # reordering persistent batch and generating logitsprocs batch state # updates. Should reset each step. @@ -552,7 +551,6 @@ class InputBatch: self.generators.pop(req_index, None) self.num_logprobs.pop(req_id, None) self.logprob_token_ids.pop(req_id, None) - self.in_progress_prompt_logprobs_cpu.pop(req_id, None) if self.prev_req_id_to_index is not None: self.prev_req_id_to_index.pop(req_id, None) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index bcab2ca2d4..0ca530c15b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5094,7 +5094,6 @@ class GPUModelRunner( if not num_prompt_logprobs_dict: return {} - in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {} # Since prompt logprobs are a rare feature, prioritize simple, @@ -5118,14 +5117,14 @@ class GPUModelRunner( ) # Set up target LogprobsTensors object. - logprobs_tensors = in_progress_dict.get(req_id) - if not logprobs_tensors: + logprobs_tensors = request.in_progress_prompt_logprobs_cpu + if logprobs_tensors is None: # Create empty logprobs CPU tensors for the entire prompt. # If chunked, we'll copy in slice by slice. logprobs_tensors = LogprobsTensors.empty_cpu( num_prompt_tokens - 1, num_prompt_logprobs + 1 ) - in_progress_dict[req_id] = logprobs_tensors + request.in_progress_prompt_logprobs_cpu = logprobs_tensors # Determine number of logits to retrieve. start_idx = request.num_computed_tokens @@ -5182,7 +5181,7 @@ class GPUModelRunner( # num_prompt_logprobs_dict. for req_id in completed_prefill_reqs: del num_prompt_logprobs_dict[req_id] - del in_progress_dict[req_id] + self.requests[req_id].in_progress_prompt_logprobs_cpu = None # Must synchronize the non-blocking GPU->CPU transfers. if prompt_logprobs_dict: