[Bugfix][V1/V2] Fix prompt_logprobs to respect logprobs_mode#47680
[Bugfix][V1/V2] Fix prompt_logprobs to respect logprobs_mode#47680aoshen02 wants to merge 3 commits into
Conversation
aoshen02
requested review from
ProExpertProg,
WoosukKwon,
houseroad,
mgoin,
njhill,
robertgshaw2-redhat,
tlrmchlsmth,
yewentao256 and
youkaichao
as code owners
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
aoshen02
force-pushed
the
optimized-36539
branch
2 times, most recently
from
3615999 to
bc02352
Compare
fede-kamel
left a comment
fede-kamel
left a comment
There was a problem hiding this comment.
Thanks @aoshen02 for consolidating this — V2 coverage plus the shared compute_topk_logits utility makes this strictly better than either of the prior PRs, and I'm happy to be listed as co-author. Reviewed the full diff; one blocking concern and two smaller ones.
1. Removing the config gate regresses V2 + raw_logits from "works via V1 fallback" to "crash at model load".
The gate entry in _get_v2_model_runner_unsupported_features is what makes default-V2 models silently fall back to the V1 runner for logits modes (use_v2_model_runner warns and returns False — vllm/config/vllm.py, ~L545). With the entry removed, those configs now stay on V2 — but the V2 sampler still hard-rejects logits modes at construction:
# vllm/v1/worker/gpu/sample/sampler.py
if logprobs_mode not in ("processed_logprobs", "raw_logprobs"):
raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")and Sampler is constructed unconditionally for generate models in gpu/model_runner.py::load_model, so logprobs_mode="raw_logits" now dies during engine startup with NotImplementedError in the worker instead of transparently running on V1. It also means sampled-token logprobs still don't respect logits modes on V2 — this PR threads the mode through the prompt path only.
Two clean ways out:
- Keep the gate entry and scope this PR to prompt logprobs (the V2 threading becomes groundwork that activates once the V2 sampler learns logits modes), or
- Extend the V2
Samplerto supportraw_logits/processed_logitsfor sampled tokens in this PR — then removing the gate is correct and the whole feature works end-to-end on V2.
I'd lean to the first for reviewability; either works.
2. Pre-commit will likely fail on prompt_logprob.py — there's trailing whitespace in the compute_fn(...) call args, and the compute_fn = (...) ternary line exceeds the 88-char limit. pre-commit run --all-files should catch both.
3. Possible GPU OOM flake in the new e2e test. test_sampling_params_e2e.py uses a module-scoped llm fixture, so its engine (default gpu_memory_utilization=0.9) can still be alive when this test constructs two more engines. Suggest moving the test to its own file, or calling cleanup_dist_env_and_memory() between engines and passing a lower gpu_memory_utilization.
Everything else LGTM: compute_topk_logits rank semantics match the existing _ranks_kernel convention ((logits >= target).sum()), log_softmax→logits is rank-preserving so ranks stay consistent across modes, and the inline V1 branch is equivalent to what #36539 did.
cc @WoosukKwon @njhill @22quinn @houseroad — this consolidates #35885 + #36539 (fixes #35832); maintainer review would be appreciated once the gate question above is settled.
|
One more thought on sequencing, following up on my review above. Since the gate question means the V2 half needs another iteration either way (keep the gate, or extend This PR would then become the V2 completion on top of it: @aoshen02 does that sequencing work for you? Happy to review the V2 follow-up promptly. If a maintainer agrees, #36539 just needs the |
|
@codex review |
done |
I don't have any problems. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fa0ba744bf
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| logprobs_tensors = compute_topk_logits( | ||
| logits, | ||
| num_logprobs, | ||
| sampled, |
There was a problem hiding this comment.
When a request uses SamplingParams(logprob_token_ids=[...]) with raw_logits or processed_logits, this branch ignores the max_per_req_token_ids/logprob_token_ids_state computed above. If logprobs is unset, num_logprobs becomes 0 and the response contains only the sampled token; if both are set, it returns top-k logits instead of the explicitly requested token IDs. The logprobs branch below passes the custom-token state, so the newly enabled logits modes regress the generative-scoring/logprob_token_ids path.
Useful? React with 👍 / 👎.
| if self.sampler.logprobs_mode.endswith("logits"): | ||
| return compute_topk_logits(logits, max_num_logprobs, flat_sampled) |
There was a problem hiding this comment.
In V2 speculative decoding with raw_logits or processed_logits and logprobs enabled, this early return drops the input_batch.cu_num_logits_np offsets that the logprobs branch passes below. The resulting LogprobsTensors has cu_num_generated_tokens=None, so LogprobsLists.slice_request treats the request index as the row offset; when any earlier request has multiple speculative logits, later requests get logprobs sliced from the wrong rows.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8f0a95c8bf
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self._pending_logprobs[slot.item()] = compute_topk_scores( | ||
| flat_logits[start : start + k_i], | ||
| max_num_logprobs, | ||
| argmax_tokens[local_idx][:k_i], |
There was a problem hiding this comment.
When DiffusionGemma runs with logprobs_mode="raw_logits" or "processed_logits" and the request asks for logprobs, this custom sampler now calls compute_topk_scores without the new logits_mode flag, so it defaults to False and takes the log-softmax path. Since this commit also removes the V2/worker guard for logits modes, those requests no longer fail fast but return log probabilities in the logprob field instead of logits. Store/pass the sampler's logprobs_mode.endswith("logits") here before stashing _pending_logprobs.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3c4dc291a4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| result = compute_topk_scores( | ||
| prompt_logits, | ||
| requested_num_prompt_logprobs, | ||
| prompt_token_ids[start_idx:end_idx], | ||
| requested_num, | ||
| chunk_tgt, | ||
| logits_mode=logits_mode, |
There was a problem hiding this comment.
When logprobs_mode is processed_logits or processed_logprobs, this path still feeds the unprocessed prompt_logits directly into compute_topk_scores; the new flag only switches between returning logits and log_softmax. Since ModelConfig.logprobs_mode documents processed modes for prompt_logprobs as after processors such as logit bias, bad words, temperature, and top_k/top_p, V2 requests with prompt_logprobs plus any such sampling processor will now be accepted but return the same raw prompt scores as the raw modes.
Useful? React with 👍 / 👎.
|
@codex review, processed logprobs is not apply to prompt logprobs, prompt will not process by sampler |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3c4dc291a4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| result = compute_topk_scores( | ||
| prompt_logits, | ||
| requested_num_prompt_logprobs, | ||
| prompt_token_ids[start_idx:end_idx], | ||
| requested_num, | ||
| chunk_tgt, | ||
| logits_mode=logits_mode, |
There was a problem hiding this comment.
When a request uses logprobs_mode='processed_logits' or 'processed_logprobs' with prompt_logprobs, this path only switches between gathering logits vs log-softmax from the unmodified prompt_logits; it never applies the sampler processors that define the processed modes, such as temperature/top_k/top_p, logit bias, bad words, or penalties. Since this commit also unblocks processed_logits for the V2 runner, those requests now succeed but return raw prompt scores while generated-token logprobs are processed. Please either run prompt logits through the same processing path or keep processed modes unsupported for prompt logprobs.
Useful? React with 👍 / 👎.
njhill
left a comment
njhill
left a comment
There was a problem hiding this comment.
Thanks @aoshen02 @wojciech-wais this looks generally good to me.
@WoosukKwon might have some opinion on renaming logprobs to scores... though they might now technically not be logprobs that naming might still be clearer.
We could even consider logprobs_or_logits. Or maybe scores is best afer all. I'm really not sure about this just thinking aloud.
| else num_prompt_logprobs | ||
| ) | ||
| prompt_logprobs = compute_topk_logprobs( | ||
| chunk_tgt = prompt_token_ids[start_idx:end_idx] |
|
|
||
| if return_logprobs: | ||
| if self.logprobs_mode == "processed_logprobs": | ||
| if self.logprobs_mode.startswith("processed"): |
There was a problem hiding this comment.
It feels nasty to use string operations for these checks. Though I know it kind of predates this PR.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 617d132c6e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1) | ||
| logprobs = compute_token_logprobs(logits, logprob_token_ids) | ||
| if logits_mode: | ||
| scores = logits.to(torch.float32).gather(-1, logprob_token_ids) |
There was a problem hiding this comment.
In raw_logits/processed_logits modes this casts the entire [batch, vocab] logits tensor to FP32 before gathering the few requested columns. On common bf16/fp16 models with large vocabs or prompt-logprob chunks, that adds a full extra logits-sized allocation and can OOM requests; the logprob path avoids materializing a full score matrix. Gather from logits first and cast the small gathered result instead; the same pattern appears in the custom-token branch below.
Useful? React with 👍 / 👎.
| logprobs = [] | ||
| scores = [] | ||
| ranks = [] | ||
| logits_mode = logprobs_mode.endswith("logits") |
There was a problem hiding this comment.
Now that the V2 unsupported-feature guard for processed_logits is removed, that mode reaches this path, but this suffix check treats raw_logits and processed_logits identically and scores the unprocessed prompt_logits. For requests that use logit processors or sampling transforms, prompt_logprobs in processed_logits mode are therefore the same as raw logits instead of processed scores as the model config promises. Either apply the per-request processing before scoring prompt logits or keep processed_logits gated for this path.
Useful? React with 👍 / 👎.
|
|
||
| class PromptLogprobsWorker: | ||
| def __init__(self, max_num_reqs: int): | ||
| def __init__(self, max_num_reqs: int, logprobs_mode: str = "raw_logprobs"): |
There was a problem hiding this comment.
Both PromptLogprobsWorker.init and compute_prompt_logprobs_with_chunking declare logprobs_mode: str.
Every other caller in the diff uses LogprobsMode.
With str, a typo like "raw_logrobs" silently falls through to the log-softmax path because "raw_logrobs".endswith("logits") is False - no error, wrong behavior.
|
|
||
|
|
||
| def test_prompt_logprobs_respects_logprobs_mode(): | ||
| """raw_logits must return positive logits, raw_logprobs must return |
There was a problem hiding this comment.
The assertion only checks values["raw_logprobs"] <= 0 and values["raw_logits"] != values["raw_logprobs"].
There is no assertion for the following condition
raw_logits > 0
and there shouldn't be, because raw logits can be negative for any token.
The docstring is technically false and would mislead anyone reading the test to understand the sense of it.
| assert out_1[0].outputs[0].text != out_3[0].outputs[0].text | ||
|
|
||
|
|
||
| def test_prompt_logprobs_respects_logprobs_mode(): |
There was a problem hiding this comment.
The e2e test only iterates over ("raw_logits", "raw_logprobs").
Those two processed_* modes have no test coverage at all. IMO this is missing unit test.
| logprobs = self.sampler.compute_logprobs(logits) | ||
| # Compute prompt scores respecting logprobs_mode. | ||
| if self.model_config.logprobs_mode.endswith("logits"): | ||
| scores = logits.to(torch.float32) |
There was a problem hiding this comment.
Since both runners return raw unprocessed logits for processed_logits maybe worth to flag it somehow in the code/docstring or in the comments.
prompt_logprobs ignored the logprobs_mode config — all modes returned log_softmax results. This fix threads the mode through both V1 and V2 GPU model runners so that *_logits modes return raw logits and *_logprobs modes return log_softmax values. V1 runner (gpu_model_runner.py): inline mode branch before gather. V2 runner: add logits_mode flag to compute_topk_scores, thread logprobs_mode through Sampler, PromptLogprobsWorker, and RejectionSampler. Config (vllm.py): remove V2 runner unsupported-feature gate for raw_logits/processed_logits now that V2 handles them. Fixes: vllm-project#35832 Co-authored-by: Federico Kamelhar <209537060+fede-kamel@users.noreply.github.com> Co-authored-by: Allen Shen <aoshen@inferact.ai> Signed-off-by: Wojciech Wais <wojciech.wais@gmail.com> Signed-off-by: Federico Kamelhar <209537060+fede-kamel@users.noreply.github.com> Signed-off-by: Allen Shen <aoshen@inferact.ai>
|
Re-reviewed at the current head — the blocker from my earlier review is resolved: the V2 sampler no longer rejects logits modes at construction, and the mode is threaded through the sampled-token, spec-decode, and DiffusionGemma paths, so dropping the config gate is safe now. The four-mode e2e test and the LogprobsMode typing cover the rest of what I'd flagged. LGTM from my side — this supersedes #36539, and I'll close that one once this merges. Thanks for carrying the co-author trailers through; please keep them in the squash message. |
Thank you! |
|
@njhill on the naming question: scores seems right to me — the tensor holds either logits or logprobs depending on mode, and logprobs_or_logits bakes today's two options into the name. The LogprobsTensors.logprobs field has the same ambiguity, but renaming that touches a much wider surface, so probably better as a follow-up if at all. |
|
Heads-up from #36539's CI run (build 76426): test_prompt_logprobs_respects_logprobs_mode here will likely hit the same OOM once full CI runs — the four LLMs are created with the default gpu_memory_utilization while the module-scoped llm fixture in test_sampling_params_e2e.py still holds its memory, and on the 16 GiB CI GPU engine init then fails with ~0.7 GiB free (my variant of this test failed exactly that way on both the NVIDIA and AMD runners). The cleanup_dist_env_and_memory() between iterations helps but doesn't release the fixture's engine. What fixed it for me: pass gpu_memory_utilization=0.05 to the test's LLMs, same idiom as test_logprobs.py. |
|
Correcting my earlier suggestion: CI on #36539 (build 76618) showed that even gpu_memory_utilization=0.05 doesn't fit next to the module-scoped llm fixture — only 0.72 GiB stays free on the 16 GiB runner, below even that reduced request. What worked is moving the test into its own module (test_prompt_logprobs_mode_e2e.py), collected before the fixture module so the engines never coexist. Worth doing the same for the four-mode test here. |
4 participants
Summary
Fix
prompt_logprobsto respect thelogprobs_modeconfiguration setting across both V1 and V2 GPU model runners.Previously,
prompt_logprobsalways returnedlog_softmaxresults regardless oflogprobs_mode. When a user setlogprobs_mode="raw_logits", output logprobs correctly returned raw logits, but prompt logprobs silently ignored the mode and returned logprobs instead.Changes
gpu_model_runner.py): inline mode branch —*_logitsreturnslogits.float32(),*_logprobscallscompute_logprobs()prompt_logprob.py): threadlogprobs_modethroughPromptLogprobsWorker→compute_prompt_logprobs_with_chunking(), dispatch tocompute_topk_logitsorcompute_topk_logprobslogprob.py): addcompute_topk_logits()— sibling ofcompute_topk_logprobs()that gathers raw logits instead of applyinglog_softmaxvia Triton kernelvllm.py): remove the V2 runner unsupported-feature gate forraw_logits/processed_logits, since V2 now handles them correctlytest_sampling_params_e2e.py, unit test intest_prompt_logprobs_mode.pyAcknowledgements
This PR consolidates and extends the work from two prior PRs that independently identified and fixed this bug:
prompt_logprob.pyThis PR takes the best of both: V1+V2 coverage,
compute_topk_logitsas a reusable sibling inlogprob.py(not a private helper), config gate removal to unblock V2, and tests.Fixes: #35832
Test plan
pytest tests/v1/sample/test_prompt_logprobs_mode.py -v— unit test for mode branch logicpytest tests/v1/sample/test_sampling_params_e2e.py::test_prompt_logprobs_respects_logprobs_mode -v— e2e:raw_logitsvsraw_logprobsreturn different valueslogprobs_mode=raw_logitsAI-assisted: Claude was used for code generation and review. All changes reviewed and understood by the submitter.