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 4, 21:15 UTC
[Model Runner V2] Multiple prompt logprobs support (#39937)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Wentao Ye committed
Apr 21, 2026 at 11:49 UTC
66cc3fa559d71bb4ae2335c26c01115791968099
1 file changed
+48
-18
vllm/v1/worker/gpu/sample/prompt_logprob.py
+48
-18
@@ -17,13 +17,14 @@ class PromptLogprobsWorker:
17
self.max_num_reqs = max_num_reqs
18
19
self.uses_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
20
+ self.num_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=np.int32)
21
# req_idx -> list of in-progress LogprobsTensors
22
self.in_progress_prompt_logprobs: dict[str, list[LogprobsTensors]] = {}
23
24
def add_request(self, req_id: str, req_idx: int, sampling_params: SamplingParams):
24
- # For now, only support prompt logprobs for the prompt tokens (not top-k).
25
uses_prompt_logprobs = sampling_params.prompt_logprobs is not None
26
self.uses_prompt_logprobs[req_idx] = uses_prompt_logprobs
27
+ self.num_prompt_logprobs[req_idx] = sampling_params.prompt_logprobs or 0
28
if uses_prompt_logprobs:
29
self.in_progress_prompt_logprobs[req_id] = []
30
@@ -52,6 +53,7 @@ class PromptLogprobsWorker:
53
# Common case: No request asks for prompt logprobs.
54
return {}
55
56
+ num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np]
57
prompt_lens = prompt_lens[idx_mapping_np]
58
# NOTE(woosuk): -1 because the last prompt token's hidden state is not
59
# needed for prompt logprobs.
@@ -64,6 +66,14 @@ class PromptLogprobsWorker:
66
if not np.any(needs_prompt_logprobs):
67
return {}
68
69
+ # get the maximum number in this batch
70
+ requested_num_prompt_logprobs = num_prompt_logprobs[needs_prompt_logprobs]
71
+ max_num_prompt_logprobs = (
72
+ -1
73
+ if np.any(requested_num_prompt_logprobs == -1)
74
+ else int(requested_num_prompt_logprobs.max())
75
+ )
76
+
77
# Get the prompt logprobs token_ids.
78
prompt_logprobs_token_ids = get_prompt_logprobs_token_ids(
79
input_batch.num_tokens,
@@ -72,45 +82,53 @@ class PromptLogprobsWorker:
82
num_computed_tokens,
83
all_token_ids,
84
)
75
- # Compute the prompt logprobs.
76
- prompt_logprobs, prompt_ranks = compute_prompt_logprobs_with_chunking(
77
- prompt_logprobs_token_ids,
78
- hidden_states[: input_batch.num_tokens],
79
- logits_fn,
85
+ prompt_token_ids, prompt_logprobs, prompt_ranks = (
86
+ compute_prompt_logprobs_with_chunking(
87
+ prompt_logprobs_token_ids,
88
+ hidden_states[: input_batch.num_tokens],
89
+ logits_fn,
90
+ max_num_prompt_logprobs,
91
+ )
92
)
93
94
pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
95
is_prompt_chunked = pos_after_step < prompt_lens
96
97
query_start_loc_np = input_batch.query_start_loc_np
86
- prompt_token_ids = prompt_logprobs_token_ids.unsqueeze(-1)
98
prompt_logprobs_dict: dict[str, LogprobsTensors] = {}
99
for i, req_id in enumerate(input_batch.req_ids):
100
if not needs_prompt_logprobs[i]:
101
continue
102
103
+ req_is_prompt_chunked = is_prompt_chunked[i]
104
start_idx = query_start_loc_np[i]
105
end_idx = query_start_loc_np[i + 1]
106
assert start_idx < end_idx, (
107
f"start_idx ({start_idx}) >= end_idx ({end_idx})"
108
)
97
- if not is_prompt_chunked[i]:
109
+ if not req_is_prompt_chunked:
110
end_idx -= 1
99
- logprobs = LogprobsTensors(
100
- logprob_token_ids=prompt_token_ids[start_idx:end_idx],
101
- logprobs=prompt_logprobs[start_idx:end_idx],
102
- selected_token_ranks=prompt_ranks[start_idx:end_idx],
111
+
112
+ # no logprobs if start_idx >= end_idx
113
+ logprobs = (
114
+ None
115
+ if start_idx >= end_idx
116
+ else LogprobsTensors(
117
+ logprob_token_ids=prompt_token_ids[start_idx:end_idx],
118
+ logprobs=prompt_logprobs[start_idx:end_idx],
119
+ selected_token_ranks=prompt_ranks[start_idx:end_idx],
120
+ )
121
)
122
123
prompt_logprobs_list = self.in_progress_prompt_logprobs[req_id]
106
- if is_prompt_chunked[i]:
107
- # Prompt is chunked. Do not return the logprobs yet.
124
+ if logprobs is not None and (req_is_prompt_chunked or prompt_logprobs_list):
125
prompt_logprobs_list.append(logprobs)
126
+ if req_is_prompt_chunked:
127
+ # Prompt is chunked. Do not return the logprobs yet.
128
continue
129
130
if prompt_logprobs_list:
131
# Merge the in-progress logprobs.
113
- prompt_logprobs_list.append(logprobs)
132
logprobs = LogprobsTensors(
133
logprob_token_ids=torch.cat(
134
[x.logprob_token_ids for x in prompt_logprobs_list]
@@ -122,6 +140,9 @@ class PromptLogprobsWorker:
140
)
141
prompt_logprobs_list.clear()
142
143
+ if logprobs is None:
144
+ continue
145
+
146
prompt_logprobs_dict[req_id] = logprobs
147
return prompt_logprobs_dict
148
@@ -184,10 +205,12 @@ def compute_prompt_logprobs_with_chunking(
205
prompt_token_ids: torch.Tensor,
206
prompt_hidden_states: torch.Tensor,
207
logits_fn: Callable[[torch.Tensor], torch.Tensor],
187
-) -> tuple[torch.Tensor, torch.Tensor]:
208
+ num_prompt_logprobs: int,
209
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
210
# Since materializing the full prompt logits can take too much memory,
211
# we compute it in chunks.
212
CHUNK_SIZE = 1024
213
+ token_ids = []
214
logprobs = []
215
ranks = []
216
prompt_token_ids = prompt_token_ids.to(torch.int64)
@@ -195,14 +218,21 @@ def compute_prompt_logprobs_with_chunking(
218
end_idx = start_idx + CHUNK_SIZE
219
# NOTE(woosuk): logits_fn can be slow because it involves all-gather.
220
prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
221
+ requested_num_prompt_logprobs = (
222
+ prompt_logits.shape[-1]
223
+ if num_prompt_logprobs == -1
224
+ else num_prompt_logprobs
225
+ )
226
prompt_logprobs = compute_topk_logprobs(
227
prompt_logits,
200
- 0, # num_logprobs
228
+ requested_num_prompt_logprobs,
229
prompt_token_ids[start_idx:end_idx],
230
)
231
+ token_ids.append(prompt_logprobs.logprob_token_ids)
232
logprobs.append(prompt_logprobs.logprobs)
233
ranks.append(prompt_logprobs.selected_token_ranks)
234
235
+ token_ids = torch.cat(token_ids, dim=0) if len(token_ids) > 1 else token_ids[0]
236
logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
237
ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
208
- return logprobs, ranks
238
+ return token_ids, logprobs, ranks