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 3, 21:30 UTC
[Model Runner V2] Bug fix: logprob dtype int64/int32 issue (#41761)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
May 11, 2026 at 17:55 UTC
d7af6b34d83fc691ad69347ee4d066231e5678ab
3 files changed
+83
-7
tests/v1/engine/test_logprobs_processor.py
new
+66
@@ -0,0 +1,66 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for LogprobsProcessor.
4
+
5
+These tests exercise the truncation invariant that the MRV2 sampler relies
6
+on: when the sampler returns a row wider than a request's own
7
+`num_logprobs + 1` (because another request in the batch needed a wider
8
+row), the trailing positions are populated with sentinel values
9
+(`token_id=0`, `logprob=-inf`). LogprobsProcessor must read only the first
10
+`num_logprobs + 1` entries so those sentinels never reach the user.
11
+"""
12
+
13
+import numpy as np
14
+
15
+from vllm.logprobs import create_sample_logprobs
16
+from vllm.v1.engine.logprobs import LogprobsProcessor
17
+from vllm.v1.outputs import LogprobsLists
18
+
19
+
20
+def _make_processor(num_logprobs: int) -> LogprobsProcessor:
21
+ return LogprobsProcessor(
22
+ tokenizer=None,
23
+ logprobs=create_sample_logprobs(flat_logprobs=False),
24
+ prompt_logprobs=None,
25
+ cumulative_logprob=0.0,
26
+ num_logprobs=num_logprobs,
27
+ num_prompt_logprobs=None,
28
+ )
29
+
30
+
31
+def test_drops_trailing_sentinel_columns():
32
+ """A request that asked for 3 custom token logprobs but ended up in a
33
+ batch padded to width 5 must not surface the trailing -inf entries."""
34
+ processor = _make_processor(num_logprobs=3)
35
+
36
+ sampled = 42
37
+ # Layout: [sampled, custom_1, custom_2, custom_3, SENTINEL, SENTINEL]
38
+ # Use float32-exact values so cumulative_logprob compares cleanly.
39
+ token_ids = np.array([[sampled, 100, 200, 300, 0, 0]], dtype=np.int32)
40
+ logprobs = np.array([[-0.5, -1.0, -2.0, -3.0, -np.inf, -np.inf]], dtype=np.float32)
41
+ ranks = np.array([1], dtype=np.int32)
42
+
43
+ processor._update_sample_logprobs(LogprobsLists(token_ids, logprobs, ranks))
44
+
45
+ assert len(processor.logprobs) == 1
46
+ pos = processor.logprobs[0]
47
+ # Exactly sampled + 3 requested tokens; trailing sentinels dropped.
48
+ assert set(pos.keys()) == {sampled, 100, 200, 300}
49
+ assert 0 not in pos
50
+ assert all(np.isfinite(lp.logprob) for lp in pos.values())
51
+ # cumulative_logprob comes from the sampled token's logprob only.
52
+ assert processor.cumulative_logprob == -0.5
53
+
54
+
55
+def test_accepts_exactly_sized_row():
56
+ """When the row is exactly num_logprobs+1, no truncation needed."""
57
+ processor = _make_processor(num_logprobs=2)
58
+
59
+ token_ids = np.array([[7, 11, 13]], dtype=np.int32)
60
+ logprobs = np.array([[-0.5, -1.5, -2.5]], dtype=np.float32)
61
+ ranks = np.array([1], dtype=np.int32)
62
+
63
+ processor._update_sample_logprobs(LogprobsLists(token_ids, logprobs, ranks))
64
+
65
+ pos = processor.logprobs[0]
66
+ assert set(pos.keys()) == {7, 11, 13}
vllm/sampling_params.py
+8
@@ -690,6 +690,14 @@ class SamplingParams(
690
parameter="logprob_token_ids",
691
value=n,
692
)
693
+ if self.logprobs is not None and self.logprobs != n:
694
+ raise VLLMValidationError(
695
+ f"When both logprobs and logprob_token_ids are set, "
696
+ f"logprobs must equal len(logprob_token_ids). Got "
697
+ f"logprobs={self.logprobs}, len(logprob_token_ids)={n}.",
698
+ parameter="logprob_token_ids",
699
+ value=n,
700
+ )
701
702
# Validate prompt logprobs.
703
if num_prompt_logprobs := self.prompt_logprobs:
vllm/v1/worker/gpu/sample/logprob.py
+9
-7
@@ -124,9 +124,13 @@ def compute_topk_logprobs(
124
# tokens where applicable.
125
assert logprob_token_ids_state is not None
126
assert expanded_idx_mapping is not None
127
- topk_indices = None
127
+
128
if num_logprobs > 0:
129
- topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
129
+ topk_token_ids = torch.topk(logits, num_logprobs, dim=-1).indices
130
+ topk_token_ids = topk_token_ids.to(torch.int32)
131
+ else:
132
+ # This tensor just used as an int32 pointer, data not accessed.
133
+ topk_token_ids = logprob_token_ids_state.token_ids.gpu
134
135
num_cols = max(num_logprobs, max_per_req_token_ids)
136
logprob_token_ids = sampled_token_ids.new_zeros((batch_size, 1 + num_cols))
@@ -137,8 +141,8 @@ def compute_topk_logprobs(
141
valid_mask,
142
valid_mask.stride(0),
143
sampled_token_ids,
140
- topk_indices if topk_indices is not None else logprob_token_ids,
141
- topk_indices.stride(0) if topk_indices is not None else 0,
144
+ topk_token_ids,
145
+ topk_token_ids.stride(0),
146
expanded_idx_mapping,
147
logprob_token_ids_state.num_token_ids.gpu,
148
logprob_token_ids_state.token_ids.gpu,
@@ -202,14 +206,12 @@ def _fill_logprob_token_ids_kernel(
206
# Override topk with per-request custom tokens.
207
src = per_req_token_ids_ptr + req_state_idx * per_req_token_ids_stride
208
valid = col < num_custom
205
- # per_req_token_ids is int32; output is int64.
206
- tokens = tl.load(src + col, mask=valid, other=0).to(tl.int64)
209
else:
210
# Fill with topk indices (no-op when NUM_TOPK == 0).
211
src = topk_indices_ptr + batch_idx * topk_indices_stride
212
valid = col < NUM_TOPK
211
- tokens = tl.load(src + col, mask=valid, other=0)
213
214
+ tokens = tl.load(src + col, mask=valid, other=0).to(tl.int64)
215
tl.store(tid_base + col, tokens, mask=valid)
216
tl.store(mask_base + col, tl.full([PADDED_COLS], 1, tl.int1), mask=valid)
217