Expose raw logits for GLM KLD capture - #32
Closed
voipmonitor wants to merge 2 commits into
Closed
Conversation
Keep the B12X DCP global top-k path rank-consistent through graph warmup and MTP draft execution. The target GLM path already allocated score buffers for B12X DCP top-k, but the MTP draft layer only allocated indices. With B12X sparse indexer global-topk enabled, the draft capture either failed or fell back to non-comparable slow paths. Allocate and pass top-k score buffers for MTP when DCP and B12X sparse indexer are active. Also keep empty DCP prefill chunks participating in the global-topk merge and prewarm the actual B12X merge/extend paths so rank-sensitive graph warmup enters the same collectives on every rank.
📝 WalkthroughWalkthroughThis PR adds two new ChangesRaw Logits Output Pipeline
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
vllm/v1/worker/gpu_model_runner.py (3)
3658-3666:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate
_bookkeeping_sync’s return annotation.The function now returns
prompt_logits_dictas an additional tuple element, but the annotation still describes the old 7-element shape.Proposed fix
) -> tuple[ dict[str, int], LogprobsLists | None, list[list[int]], dict[str, LogprobsTensors | None], + dict[str, torch.Tensor | None], list[str], dict[str, int], list[int], ]:Also applies to: 3782-3790
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/v1/worker/gpu_model_runner.py` around lines 3658 - 3666, The return type annotation for the `_bookkeeping_sync` method is outdated and does not reflect the actual return tuple structure. The method now returns an additional `prompt_logits_dict` element but the tuple annotation at lines 3658-3666 still shows the old 7-element shape. Add the missing `dict[str, int]` type annotation to the tuple to account for the `prompt_logits_dict` return value. Also update the matching return annotation at lines 3782-3790 to ensure consistency.
4674-4680:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftThread
sampler_output.sample_logitsinto this runner output.
prompt_logits_dictis now forwarded, but sample logits are not attached toModelRunnerOutputor copied by this file’s async wrapper, soreturn_sample_logitscan be silently dropped on this path.Mirror the sample-logit CPU transfer/output handling used by the updated async output path for both sync and async scheduling.
Also applies to: 4730-4738
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/v1/worker/gpu_model_runner.py` around lines 4674 - 4680, The ModelRunnerOutput being constructed is missing the sample logits from sampler_output.sample_logits, even though prompt_logits_dict is being forwarded. Add the sample_logits field to the ModelRunnerOutput construction at the location where prompt_logits_dict is being set, ensuring it captures sampler_output.sample_logits. Additionally, ensure that the CPU transfer and output handling for sample logits mirrors the approach used in the async output path (referenced at lines 4730-4738) so that sample logits are properly handled consistently across both sync and async scheduling paths and the return_sample_logits flag is not silently dropped.
5588-5612:⚠️ Potential issue | 🟠 Major | ⚡ Quick winFlush buffered logits and don’t suppress requested prompt logprobs.
The new
continuemakesreturn_prompt_logitsskipprompt_logprobsoutput when both are requested, and thenum_logits <= 0completion path can clear already-buffered prompt logits without returning them.Proposed fix
if num_logits <= 0: # This can happen for the final chunk if we prefilled exactly # (num_prompt_tokens - 1) tokens for this request in the prior # step. There are no more prompt logprobs to produce. + if req_id in completed_prefill_reqs: + if ( + return_prompt_logits + and request.in_progress_prompt_logits is not None + ): + prompt_logits_dict[req_id] = torch.cat( + request.in_progress_prompt_logits, dim=0 + ) + if request.in_progress_prompt_logprobs_cpu is not None: + prompt_logprobs_dict[req_id] = ( + request.in_progress_prompt_logprobs_cpu + ) continue @@ if req_id in completed_prefill_reqs: prompt_logits_dict[req_id] = torch.cat( request.in_progress_prompt_logits, dim=0 ) request.in_progress_prompt_logits = None - continue + if ( + request.sampling_params is None + or request.sampling_params.prompt_logprobs is None + ): + continue🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/v1/worker/gpu_model_runner.py` around lines 5588 - 5612, The code has two issues: first, the `continue` statement inside the `return_prompt_logits` block prevents subsequent prompt_logprobs processing from running when both features are requested, and second, the `num_logits <= 0` early continue path exits without flushing any buffered logits that may have been accumulated in `request.in_progress_prompt_logits`. To fix this, remove the `continue` statement in the `return_prompt_logits` block to allow prompt_logprobs processing to proceed, and in the `num_logits <= 0` condition, before continuing, check if `request.in_progress_prompt_logits` has buffered data and flush it to `prompt_logits_dict` with the accumulated logits (using torch.cat like in the completed prefill case).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vllm/v1/engine/output_processor.py`:
- Line 371: In DELTA mode (RequestOutputKind.DELTA), the output should only
include logits for newly generated tokens, but the code at line 371 is including
self.logprobs_processor.prompt_logits which represents cumulative/non-delta
data. Remove the prompt_logits field when RequestOutputKind is DELTA and ensure
only delta sample_logits (logits for newly sampled tokens) are emitted. Apply
this same fix at the second location mentioned at line 409 where the same
pattern appears.
In `@vllm/v1/worker/gpu/sample/prompt_logprob.py`:
- Line 95: The issue is that line 95 uses np.any(returns_prompt_logits) to
convert a per-request boolean array into a single batch-wide flag, which causes
the helper to materialize raw prompt logits for all requests in the batch even
when only some requests need them. Then lines 253-255 copy all prompt-logit
chunks to CPU unnecessarily. To fix this, either pass the per-request
returns_prompt_logits array directly as a mask instead of aggregating it with
np.any(), or move the CPU copy logic at lines 253-255 to only execute for
requests where the corresponding element in returns_prompt_logits is True,
ensuring raw logits are materialized only for flagged requests.
- Around line 145-146: The condition in the if statement that checks
logits_chunk.numel() and the compound boolean expression (req_is_prompt_chunked
or prompt_logits_list) exceeds the 88 character line length limit. Break this
condition across multiple lines by wrapping it after the logical operators,
ensuring each line stays within the character limit while maintaining code
readability. Use parentheses to group related conditions together if needed.
---
Outside diff comments:
In `@vllm/v1/worker/gpu_model_runner.py`:
- Around line 3658-3666: The return type annotation for the `_bookkeeping_sync`
method is outdated and does not reflect the actual return tuple structure. The
method now returns an additional `prompt_logits_dict` element but the tuple
annotation at lines 3658-3666 still shows the old 7-element shape. Add the
missing `dict[str, int]` type annotation to the tuple to account for the
`prompt_logits_dict` return value. Also update the matching return annotation at
lines 3782-3790 to ensure consistency.
- Around line 4674-4680: The ModelRunnerOutput being constructed is missing the
sample logits from sampler_output.sample_logits, even though prompt_logits_dict
is being forwarded. Add the sample_logits field to the ModelRunnerOutput
construction at the location where prompt_logits_dict is being set, ensuring it
captures sampler_output.sample_logits. Additionally, ensure that the CPU
transfer and output handling for sample logits mirrors the approach used in the
async output path (referenced at lines 4730-4738) so that sample logits are
properly handled consistently across both sync and async scheduling paths and
the return_sample_logits flag is not silently dropped.
- Around line 5588-5612: The code has two issues: first, the `continue`
statement inside the `return_prompt_logits` block prevents subsequent
prompt_logprobs processing from running when both features are requested, and
second, the `num_logits <= 0` early continue path exits without flushing any
buffered logits that may have been accumulated in
`request.in_progress_prompt_logits`. To fix this, remove the `continue`
statement in the `return_prompt_logits` block to allow prompt_logprobs
processing to proceed, and in the `num_logits <= 0` condition, before
continuing, check if `request.in_progress_prompt_logits` has buffered data and
flush it to `prompt_logits_dict` with the accumulated logits (using torch.cat
like in the completed prefill case).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ec963c98-238d-495c-afb4-9c1902cd7754
📒 Files selected for processing (16)
vllm/outputs.pyvllm/sampling_params.pyvllm/v1/core/sched/scheduler.pyvllm/v1/engine/__init__.pyvllm/v1/engine/logprobs.pyvllm/v1/engine/output_processor.pyvllm/v1/outputs.pyvllm/v1/worker/gpu/async_utils.pyvllm/v1/worker/gpu/model_runner.pyvllm/v1/worker/gpu/sample/output.pyvllm/v1/worker/gpu/sample/prompt_logprob.pyvllm/v1/worker/gpu/sample/sampler.pyvllm/v1/worker/gpu/sample/states.pyvllm/v1/worker/gpu/spec_decode/rejection_sampler.pyvllm/v1/worker/gpu_input_batch.pyvllm/v1/worker/gpu_model_runner.py
| prompt_logprobs=prompt_logprobs, | ||
| outputs=cast(list[CompletionOutput], outputs), | ||
| finished=finished, | ||
| prompt_logits=self.logprobs_processor.prompt_logits, |
There was a problem hiding this comment.
DELTA mode is emitting non-delta logits payloads.
RequestOutputKind.DELTA sends delta token_ids, but this path currently re-emits
stored prompt_logits and cumulative sample_logits. That breaks delta output
shape/semantics and can duplicate logits during downstream merge.
Proposed fix
@@
- if self.output_kind == RequestOutputKind.DELTA:
+ prompt_logits = self.logprobs_processor.prompt_logits
+ if self.output_kind == RequestOutputKind.DELTA:
# Side effect: logprobs processor forgets prompt logprobs
prompt_logprobs = self.logprobs_processor.pop_prompt_logprobs()
+ # Emit prompt logits once in DELTA mode.
+ self.logprobs_processor.prompt_logits = None
else:
prompt_logprobs = self.logprobs_processor.prompt_logprobs
@@
- prompt_logits=self.logprobs_processor.prompt_logits,
+ prompt_logits=prompt_logits,
@@
- return CompletionOutput(
+ sample_logits = self.logprobs_processor.sample_logits
+ if delta:
+ sample_logits = (
+ sample_logits[-len(token_ids):]
+ if sample_logits is not None and token_ids
+ else None
+ )
+
+ return CompletionOutput(
@@
- sample_logits=self.logprobs_processor.sample_logits,
+ sample_logits=sample_logits,Also applies to: 409-409
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/engine/output_processor.py` at line 371, In DELTA mode
(RequestOutputKind.DELTA), the output should only include logits for newly
generated tokens, but the code at line 371 is including
self.logprobs_processor.prompt_logits which represents cumulative/non-delta
data. Remove the prompt_logits field when RequestOutputKind is DELTA and ensure
only delta sample_logits (logits for newly sampled tokens) are emitted. Apply
this same fix at the second location mentioned at line 409 where the same
pattern appears.
| hidden_states[: input_batch.num_tokens], | ||
| logits_fn, | ||
| max_num_prompt_logprobs, | ||
| return_prompt_logits=np.any(returns_prompt_logits), |
There was a problem hiding this comment.
Only materialize raw prompt logits for flagged requests.
Line 95 turns on helper-level logit capture for the whole batch when any request asks for prompt logits, and Lines 253-255 then copy every prompt-logit chunk to CPU. In mixed batches, one debug request can force full-vocab CPU copies for unrelated prompt-logprob requests. Pass a per-token/request mask or move the CPU copy to flagged request slices so raw logits are materialized only for return_prompt_logits requests.
Also applies to: 253-255
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/worker/gpu/sample/prompt_logprob.py` at line 95, The issue is that
line 95 uses np.any(returns_prompt_logits) to convert a per-request boolean
array into a single batch-wide flag, which causes the helper to materialize raw
prompt logits for all requests in the batch even when only some requests need
them. Then lines 253-255 copy all prompt-logit chunks to CPU unnecessarily. To
fix this, either pass the per-request returns_prompt_logits array directly as a
mask instead of aggregating it with np.any(), or move the CPU copy logic at
lines 253-255 to only execute for requests where the corresponding element in
returns_prompt_logits is True, ensuring raw logits are materialized only for
flagged requests.
Comment on lines
+145
to
+146
| if logits_chunk.numel() and (req_is_prompt_chunked or prompt_logits_list): | ||
| prompt_logits_list.append(logits_chunk) |
There was a problem hiding this comment.
Wrap this condition to stay within the Python line limit.
Suggested formatting fix
- if logits_chunk.numel() and (req_is_prompt_chunked or prompt_logits_list):
+ if (
+ logits_chunk.numel()
+ and (req_is_prompt_chunked or prompt_logits_list)
+ ):
prompt_logits_list.append(logits_chunk)As per coding guidelines, Maintain a line length limit of 88 characters for Python code.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if logits_chunk.numel() and (req_is_prompt_chunked or prompt_logits_list): | |
| prompt_logits_list.append(logits_chunk) | |
| if ( | |
| logits_chunk.numel() | |
| and (req_is_prompt_chunked or prompt_logits_list) | |
| ): | |
| prompt_logits_list.append(logits_chunk) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/worker/gpu/sample/prompt_logprob.py` around lines 145 - 146, The
condition in the if statement that checks logits_chunk.numel() and the compound
boolean expression (req_is_prompt_chunked or prompt_logits_list) exceeds the 88
character line length limit. Break this condition across multiple lines by
wrapping it after the logical operators, ensuring each line stays within the
character limit while maintaining code readability. Use parentheses to group
related conditions together if needed.
Source: Coding guidelines
voipmonitor
force-pushed
the
codex/dark-devotion-dcp4-mtp3-globaltopk-fix-20260621
branch
from
June 21, 2026 15:29
000807e to
79f154c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SamplingParams.return_prompt_logitsandreturn_sample_logitsdebug flags.Why
GLM-5.2 KLD comparisons currently require copying a local vLLM overlay into the container before running prefill/decode captures. This PR makes that capture path reproducible from a branch/image instead of a one-off overlay.
Validation
python3 -m py_compileon all 16 modified Python files.Stack
Stacked on PR #31 (
codex/dark-devotion-dcp4-mtp3-globaltopk-fix-20260621) so this diff only contains the raw-logits capture plumbing.Summary by CodeRabbit
return_prompt_logitsandreturn_sample_logitssampling options to retrieve raw logits tensors for offline numerical comparison and analysis.