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, 10:30 UTC
[Model Runner V2] Fix prompt logprobs calculation `Sizes of tensors must match` error (#42778)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 18, 2026 at 11:27 UTC
e26736973a1981dbb4054dc1ac430e78d8006ef2
1 file changed
+8
-2
vllm/v1/worker/gpu/sample/prompt_logprob.py
+8
-2
@@ -99,6 +99,7 @@ class PromptLogprobsWorker:
99
continue
100
101
req_is_prompt_chunked = is_prompt_chunked[i]
102
+ req_num_prompt_logprobs = int(num_prompt_logprobs[i])
103
start_idx = query_start_loc_np[i]
104
end_idx = query_start_loc_np[i + 1]
105
assert start_idx < end_idx, (
@@ -107,13 +108,18 @@ class PromptLogprobsWorker:
108
if not req_is_prompt_chunked:
109
end_idx -= 1
110
111
+ width = (
112
+ prompt_logprobs.shape[1]
113
+ if req_num_prompt_logprobs == -1
114
+ else req_num_prompt_logprobs + 1
115
+ )
116
# no logprobs if start_idx >= end_idx
117
logprobs = (
118
None
119
if start_idx >= end_idx
120
else LogprobsTensors(
115
- logprob_token_ids=prompt_token_ids[start_idx:end_idx],
116
- logprobs=prompt_logprobs[start_idx:end_idx],
121
+ logprob_token_ids=prompt_token_ids[start_idx:end_idx, :width],
122
+ logprobs=prompt_logprobs[start_idx:end_idx, :width],
123
selected_token_ranks=prompt_ranks[start_idx:end_idx],
124
)
125
)