iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Fix beam search candidate indexing when logprobs count varies (#47126)

Browse Source 
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
This commit is contained in:
Chauncey
2026-07-01 07:07:06 +00:00
committed by GitHub
parent 5b431b905c
commit 697c34b97b
2 changed files with 119 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/samplers/test_beam_search_online.py
+74
View File
@@ -0,0 +1,74 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm import CompletionOutput, RequestOutput
from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin
from vllm.logprobs import Logprob
from vllm.sampling_params import BeamSearchParams
class _Tokenizer:
eos_token_id = 0
def decode(self, token_ids: list[int]) -> str:
return " ".join(str(token_id) for token_id in token_ids)
class _Renderer:
def get_tokenizer(self) -> _Tokenizer:
return _Tokenizer()
class _EngineClient:
async def generate(self, prompt, *args, **kwargs):
yield RequestOutput(
request_id=kwargs.get("request_id", "test-request"),
prompt=prompt.get("prompt"),
prompt_token_ids=prompt["prompt_token_ids"],
prompt_logprobs=None,
outputs=[
CompletionOutput(
index=0,
text="",
token_ids=[],
cumulative_logprob=None,
logprobs=[
{
11: Logprob(logprob=-1.0),
12: Logprob(logprob=-2.0),
13: Logprob(logprob=-3.0),
14: Logprob(logprob=-4.0),
_Tokenizer.eos_token_id: Logprob(logprob=-0.1),
}
],
finish_reason=None,
)
],
finished=True,
)
class _Serving(BeamSearchOnlineMixin):
renderer = _Renderer()
engine_client = _EngineClient()
@pytest.mark.asyncio
async def test_beam_search_handles_extra_logprob_candidates() -> None:
prompt = {
"type": "token",
"prompt": "prompt",
"prompt_token_ids": [1],
}
params = BeamSearchParams(beam_width=2, max_tokens=1)
outputs = [
output async for output in _Serving().beam_search(prompt, "request", params)
]
assert len(outputs) == 1
assert outputs[0].outputs[0].finish_reason == "stop"
assert outputs[0].outputs[0].token_ids == []
assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1)
vllm/entrypoints/generate/beam_search/online.py
+45 -51
View File
@@ -97,12 +97,7 @@ class BeamSearchOnlineMixin(ABC):
output = [x[0] for x in await asyncio.gather(*tasks)]
new_beams = []
# Store all new tokens generated by beam
all_beams_token_id = []
# Store the cumulative probability of all tokens
# generated by beam search
all_beams_logprob = []
candidates = []
# Iterate through all beam inference results
for i, result in enumerate(output):
current_beam = all_beams[i]
@@ -131,65 +126,64 @@ class BeamSearchOnlineMixin(ABC):
if result.outputs[0].logprobs is not None:
logprobs = result.outputs[0].logprobs[0]
all_beams_token_id.extend(list(logprobs.keys()))
all_beams_logprob.extend(
[
current_beam.cum_logprob + obj.logprob
for obj in logprobs.values()
]
)
# Handle the token for the end of sentence (EOS)
all_beams_token_id = np.array(all_beams_token_id)
all_beams_logprob = np.array(all_beams_logprob)
if not ignore_eos:
# Get the index position of eos token in all generated results
eos_idx = np.where(all_beams_token_id == eos_token_id)[0]
for idx in eos_idx:
current_beam = all_beams[idx // logprobs_num]
result = output[idx // logprobs_num]
assert result.outputs[0].logprobs is not None
logprobs_entry = result.outputs[0].logprobs[0]
completed.append(
BeamSearchSequence(
orig_prompt=prompt,
tokens=current_beam.tokens + [eos_token_id]
if include_stop_str_in_output
else current_beam.tokens,
logprobs=current_beam.logprobs + [logprobs_entry],
cum_logprob=float(all_beams_logprob[idx]),
finish_reason="stop",
stop_reason=eos_token_id,
for token_id, logprob_obj in logprobs.items():
candidate_logprob = (
current_beam.cum_logprob + logprob_obj.logprob
)
)
# After processing, set the log probability of the eos condition
# to negative infinity.
all_beams_logprob[eos_idx] = -np.inf
if token_id == eos_token_id and not ignore_eos:
completed.append(
BeamSearchSequence(
orig_prompt=prompt,
tokens=current_beam.tokens + [eos_token_id]
if include_stop_str_in_output
else current_beam.tokens,
logprobs=current_beam.logprobs + [logprobs],
cum_logprob=candidate_logprob,
finish_reason="stop",
stop_reason=eos_token_id,
)
)
else:
candidates.append(
(
candidate_logprob,
int(token_id),
current_beam,
logprobs,
)
)
# Processing non-EOS tokens
# Get indices of the top beam_width probabilities
topn_idx = np.argpartition(np.negative(all_beams_logprob), beam_width)[
:beam_width
]
candidate_logprobs = np.fromiter(
(candidate[0] for candidate in candidates),
dtype=np.float64,
count=len(candidates),
)
if len(candidates) <= beam_width:
topn_idx = np.argsort(-candidate_logprobs)
else:
topn_idx = np.argpartition(
-candidate_logprobs,
beam_width - 1,
)[:beam_width]
topn_idx = topn_idx[np.argsort(-candidate_logprobs[topn_idx])]
new_beams = []
for idx in topn_idx:
current_beam = all_beams[idx // logprobs_num]
result = output[idx // logprobs_num]
token_id = int(all_beams_token_id[idx])
assert result.outputs[0].logprobs is not None
logprobs_entry = result.outputs[0].logprobs[0]
cum_logprob, token_id, current_beam, logprobs = candidates[int(idx)]
new_beams.append(
BeamSearchSequence(
orig_prompt=prompt,
tokens=current_beam.tokens + [token_id],
logprobs=current_beam.logprobs + [logprobs_entry],
logprobs=current_beam.logprobs + [logprobs],
lora_request=current_beam.lora_request,
cum_logprob=float(all_beams_logprob[idx]),
cum_logprob=cum_logprob,
)
)
all_beams = new_beams
if not all_beams:
break
completed.extend(all_beams)
sorted_completed = sorted(completed, key=sort_beams_key, reverse=True)