[bugfix] Fix prompt logprobs on request eviction during chunked prefill

[joa-stdn]

Summary

• Fix prompt logprobs computation for chunked prefill: the computed_prefill < prompt_lens - 1 check incorrectly skipped the last prompt token, causing prompt logprobs to not be computed when they should be. Changed to computed_prefill < prompt_lens.
• Move in_progress_prompt_logprobs_cpu state from InputBatch dict to CachedRequestState, ensuring prompt logprobs accumulation is tied to the request lifecycle rather than the batch.
• Include prompt_logprobs in VllmRunner test helper output and add None handling to _logprobs_match for prompt logprob entries.
• Add prompt_logprobs=2 test cases to async scheduling e2e tests.

Test Plan

• Added dict(prompt_logprobs=2) and dict(prompt_logprobs=2, logprobs=2) to both test_without_spec_decoding and test_with_eagle3_spec_decoding
• pytest tests/v1/e2e/general/test_async_scheduling.py -v

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request refactors the management of prompt logprobs by moving the in_progress_prompt_logprobs_cpu state from a dictionary within InputBatch to the CachedRequestState object. This change streamlines how logprob tensors are tracked across prefill steps and ensures proper cleanup during request removal or updates. I have no feedback to provide.

[njhill]

Thanks @joa-stdn, very clean fix!

Hopefully this same bug isn't in model runner v2, maybe you could check that too? (gpu/model_runner.py)

It would also be great to add or extend a test to cover this. You could look for example at https://github.com/vllm-project/vllm/blob/main/tests/v1/e2e/general/test_async_scheduling.py which forces preemptions.

[njhill]

Thanks again @joa-stdn!

[njhill]

The test changes are failing I think because we need to modify the test to also take prompt logprobs into account when comparing outputs.

[joa-stdn]

Thanks a lot for your review!

Hopefully this same bug isn't in model runner v2

I just checked my repro in model runner v2 and everything is fine there!

The test changes are failing I think because we need to modify the test to also take prompt logprobs into account when comparing outputs.

Yeah thanks I'll look into it!

[mergify]

Hi @joa-stdn, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Documentation preview: https://vllm--41411.org.readthedocs.build/en/41411/

[mergify]

Documentation preview: https://vllm--41411.org.readthedocs.build/en/41411/

[njhill]

Thanks again @joa-stdn, especially for also exposing and fixing the MRV2 bug!