diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py new file mode 100644 index 0000000000..14481f79c6 --- /dev/null +++ b/tests/samplers/test_beam_search_online.py @@ -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) diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index 4d101e9434..1cd821f9db 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -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)