[Bugfix] Fix beam search candidate indexing when logprobs count varies (#47126)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Jul 1, 2026 at 15:07 UTC
697c34b97b4d7a5377ab223e846f545111a9a352
2 files changed
+119
-51
tests/samplers/test_beam_search_online.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm import CompletionOutput, RequestOutput
7
+from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin
8
+from vllm.logprobs import Logprob
9
+from vllm.sampling_params import BeamSearchParams
10
+
11
+
12
+class _Tokenizer:
13
+ eos_token_id = 0
14
+
15
+ def decode(self, token_ids: list[int]) -> str:
16
+ return " ".join(str(token_id) for token_id in token_ids)
17
+
18
+
19
+class _Renderer:
20
+ def get_tokenizer(self) -> _Tokenizer:
21
+ return _Tokenizer()
22
+
23
+
24
+class _EngineClient:
25
+ async def generate(self, prompt, *args, **kwargs):
26
+ yield RequestOutput(
27
+ request_id=kwargs.get("request_id", "test-request"),
28
+ prompt=prompt.get("prompt"),
29
+ prompt_token_ids=prompt["prompt_token_ids"],
30
+ prompt_logprobs=None,
31
+ outputs=[
32
+ CompletionOutput(
33
+ index=0,
34
+ text="",
35
+ token_ids=[],
36
+ cumulative_logprob=None,
37
+ logprobs=[
38
+ {
39
+ 11: Logprob(logprob=-1.0),
40
+ 12: Logprob(logprob=-2.0),
41
+ 13: Logprob(logprob=-3.0),
42
+ 14: Logprob(logprob=-4.0),
43
+ _Tokenizer.eos_token_id: Logprob(logprob=-0.1),
44
+ }
45
+ ],
46
+ finish_reason=None,
47
+ )
48
+ ],
49
+ finished=True,
50
+ )
51
+
52
+
53
+class _Serving(BeamSearchOnlineMixin):
54
+ renderer = _Renderer()
55
+ engine_client = _EngineClient()
56
+
57
+
58
+@pytest.mark.asyncio
59
+async def test_beam_search_handles_extra_logprob_candidates() -> None:
60
+ prompt = {
61
+ "type": "token",
62
+ "prompt": "prompt",
63
+ "prompt_token_ids": [1],
64
+ }
65
+ params = BeamSearchParams(beam_width=2, max_tokens=1)
66
+
67
+ outputs = [
68
+ output async for output in _Serving().beam_search(prompt, "request", params)
69
+ ]
70
+
71
+ assert len(outputs) == 1
72
+ assert outputs[0].outputs[0].finish_reason == "stop"
73
+ assert outputs[0].outputs[0].token_ids == []
74
+ assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1)
vllm/entrypoints/generate/beam_search/online.py
+45
-51
@@ -97,12 +97,7 @@ class BeamSearchOnlineMixin(ABC):
97
98
output = [x[0] for x in await asyncio.gather(*tasks)]
99
100
- new_beams = []
101
- # Store all new tokens generated by beam
102
- all_beams_token_id = []
103
- # Store the cumulative probability of all tokens
104
- # generated by beam search
105
- all_beams_logprob = []
100
+ candidates = []
101
# Iterate through all beam inference results
102
for i, result in enumerate(output):
103
current_beam = all_beams[i]
@@ -131,65 +126,64 @@ class BeamSearchOnlineMixin(ABC):
126
127
if result.outputs[0].logprobs is not None:
128
logprobs = result.outputs[0].logprobs[0]
134
- all_beams_token_id.extend(list(logprobs.keys()))
135
- all_beams_logprob.extend(
136
- [
137
- current_beam.cum_logprob + obj.logprob
138
- for obj in logprobs.values()
139
- ]
140
- )
141
-
142
- # Handle the token for the end of sentence (EOS)
143
- all_beams_token_id = np.array(all_beams_token_id)
144
- all_beams_logprob = np.array(all_beams_logprob)
145
-
146
- if not ignore_eos:
147
- # Get the index position of eos token in all generated results
148
- eos_idx = np.where(all_beams_token_id == eos_token_id)[0]
149
- for idx in eos_idx:
150
- current_beam = all_beams[idx // logprobs_num]
151
- result = output[idx // logprobs_num]
152
- assert result.outputs[0].logprobs is not None
153
- logprobs_entry = result.outputs[0].logprobs[0]
154
- completed.append(
155
- BeamSearchSequence(
156
- orig_prompt=prompt,
157
- tokens=current_beam.tokens + [eos_token_id]
158
- if include_stop_str_in_output
159
- else current_beam.tokens,
160
- logprobs=current_beam.logprobs + [logprobs_entry],
161
- cum_logprob=float(all_beams_logprob[idx]),
162
- finish_reason="stop",
163
- stop_reason=eos_token_id,
129
+ for token_id, logprob_obj in logprobs.items():
130
+ candidate_logprob = (
131
+ current_beam.cum_logprob + logprob_obj.logprob
132
)
165
- )
166
- # After processing, set the log probability of the eos condition
167
- # to negative infinity.
168
- all_beams_logprob[eos_idx] = -np.inf
133
+ if token_id == eos_token_id and not ignore_eos:
134
+ completed.append(
135
+ BeamSearchSequence(
136
+ orig_prompt=prompt,
137
+ tokens=current_beam.tokens + [eos_token_id]
138
+ if include_stop_str_in_output
139
+ else current_beam.tokens,
140
+ logprobs=current_beam.logprobs + [logprobs],
141
+ cum_logprob=candidate_logprob,
142
+ finish_reason="stop",
143
+ stop_reason=eos_token_id,
144
+ )
145
+ )
146
+ else:
147
+ candidates.append(
148
+ (
149
+ candidate_logprob,
150
+ int(token_id),
151
+ current_beam,
152
+ logprobs,
153
+ )
154
+ )
155
156
# Processing non-EOS tokens
171
- # Get indices of the top beam_width probabilities
172
- topn_idx = np.argpartition(np.negative(all_beams_logprob), beam_width)[
173
- :beam_width
174
- ]
157
+ candidate_logprobs = np.fromiter(
158
+ (candidate[0] for candidate in candidates),
159
+ dtype=np.float64,
160
+ count=len(candidates),
161
+ )
162
+ if len(candidates) <= beam_width:
163
+ topn_idx = np.argsort(-candidate_logprobs)
164
+ else:
165
+ topn_idx = np.argpartition(
166
+ -candidate_logprobs,
167
+ beam_width - 1,
168
+ )[:beam_width]
169
+ topn_idx = topn_idx[np.argsort(-candidate_logprobs[topn_idx])]
170
171
+ new_beams = []
172
for idx in topn_idx:
177
- current_beam = all_beams[idx // logprobs_num]
178
- result = output[idx // logprobs_num]
179
- token_id = int(all_beams_token_id[idx])
180
- assert result.outputs[0].logprobs is not None
181
- logprobs_entry = result.outputs[0].logprobs[0]
173
+ cum_logprob, token_id, current_beam, logprobs = candidates[int(idx)]
174
new_beams.append(
175
BeamSearchSequence(
176
orig_prompt=prompt,
177
tokens=current_beam.tokens + [token_id],
186
- logprobs=current_beam.logprobs + [logprobs_entry],
178
+ logprobs=current_beam.logprobs + [logprobs],
179
lora_request=current_beam.lora_request,
188
- cum_logprob=float(all_beams_logprob[idx]),
180
+ cum_logprob=cum_logprob,
181
)
182
)
183
184
all_beams = new_beams
185
+ if not all_beams:
186
+ break
187
188
completed.extend(all_beams)
189
sorted_completed = sorted(completed, key=sort_beams_key, reverse=True)