@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 301d986473
Fork 0

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, 08:00 UTC

[Frontend] Consolidate beam search by BeamSearchMixin. (#42946)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

wang.yuqi committed May 19, 2026 at 15:37 UTC 301d986473a0ffc1df563422e01eac4a1efd59e0
9 files changed +459 -375
.buildkite/test-amd.yaml
+1 -1
@@ -2003,7 +2003,7 @@ steps:
2003 - vllm/model_executor/layers
2004 - vllm/sampling_metadata.py
2005 - vllm/v1/sample/
2006 - - vllm/beam_search.py
2006 + - vllm/entrypoints/generate/beam_search/
2007 - tests/samplers
2008 - tests/conftest.py
2009 - vllm/_aiter_ops.py
.buildkite/test_areas/samplers.yaml
+1
@@ -11,6 +11,7 @@ steps:
11 - vllm/sampling_metadata.py
12 - tests/samplers
13 - tests/conftest.py
14 + - vllm/entrypoints/generate/beam_search
15 commands:
16 # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
17 # values explicitly to still cover the PyTorch-native (Triton) path.
vllm/entrypoints/generate/__init__.py
vllm/entrypoints/generate/beam_search/__init__.py
vllm/entrypoints/generate/beam_search/offline.py new
+226
@@ -0,0 +1,226 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import itertools
5 +from abc import ABC, abstractmethod
6 +from collections.abc import Callable, Iterable, Sequence
7 +from typing import Any
8 +
9 +from tqdm import tqdm
10 +
11 +from vllm import PromptType, RequestOutput, TextPrompt, TokensPrompt
12 +from vllm.inputs import EngineInput
13 +from vllm.logger import init_logger
14 +from vllm.lora.request import LoRARequest
15 +from vllm.renderers import BaseRenderer
16 +from vllm.sampling_params import BeamSearchParams, SamplingParams
17 +
18 +from .utils import (
19 + BeamSearchInstance,
20 + BeamSearchOutput,
21 + BeamSearchSequence,
22 + create_sort_beams_key_function,
23 +)
24 +
25 +logger = init_logger(__name__)
26 +
27 +
28 +class BeamSearchOfflineMixin(ABC):
29 + """Offline inference for beam search"""
30 +
31 + renderer: BaseRenderer
32 +
33 + def beam_search(
34 + self,
35 + prompts: list[TokensPrompt | TextPrompt],
36 + params: BeamSearchParams,
37 + lora_request: list[LoRARequest] | LoRARequest | None = None,
38 + use_tqdm: bool = False,
39 + concurrency_limit: int | None = None,
40 + ) -> list[BeamSearchOutput]:
41 + """
42 + Generate sequences using beam search.
43 +
44 + Args:
45 + prompts: A list of prompts. Each prompt can be a string or a list
46 + of token IDs.
47 + params: The beam search parameters.
48 + lora_request: LoRA request to use for generation, if any.
49 + use_tqdm: Whether to use tqdm to display the progress bar.
50 + concurrency_limit: The maximum number of concurrent requests.
51 + If None, the number of concurrent requests is unlimited.
52 + """
53 + # TODO: how does beam search work together with length penalty,
54 + # frequency, penalty, and stopping criteria, etc.?
55 + beam_width = params.beam_width
56 + max_tokens = params.max_tokens
57 + temperature = params.temperature
58 + ignore_eos = params.ignore_eos
59 + length_penalty = params.length_penalty
60 +
61 + tokenizer = self.renderer.get_tokenizer()
62 + eos_token_id = tokenizer.eos_token_id
63 + sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty)
64 +
65 + engine_inputs = self._preprocess_cmpl(prompts)
66 + lora_requests = self._lora_request_to_seq(lora_request, len(engine_inputs))
67 +
68 + if use_tqdm and concurrency_limit is not None:
69 + logger.warning(
70 + "Progress bar is not supported when using concurrency_limit. "
71 + "Disabling progress bar."
72 + )
73 + use_tqdm = False
74 +
75 + if concurrency_limit is None:
76 + concurrency_limit = len(engine_inputs)
77 +
78 + # generate 2 * beam_width candidates at each step
79 + # following the huggingface transformers implementation
80 + # at https://github.com/huggingface/transformers/blob/e15687fffe5c9d20598a19aeab721ae0a7580f8a/src/transformers/generation/beam_search.py#L534 # noqa
81 + sampling_params = SamplingParams(
82 + logprobs=2 * beam_width,
83 + max_tokens=1,
84 + temperature=temperature,
85 + skip_clone=True, # Internal beam search, safe to skip clone
86 + )
87 + instances: list[BeamSearchInstance] = []
88 +
89 + for lora_req, prompt in zip(lora_requests, engine_inputs):
90 + if prompt["type"] == "embeds":
91 + raise NotImplementedError(
92 + "Embedding prompt not supported for beam search"
93 + )
94 +
95 + instances.append(
96 + BeamSearchInstance(
97 + prompt,
98 + lora_request=lora_req,
99 + logprobs=None,
100 + ),
101 + )
102 +
103 + for prompt_start in range(0, len(instances), concurrency_limit):
104 + instances_batch = instances[prompt_start : prompt_start + concurrency_limit]
105 +
106 + token_iter = range(max_tokens)
107 + if use_tqdm:
108 + token_iter = tqdm(
109 + token_iter, desc="Beam search", unit="token", unit_scale=False
110 + )
111 + logger.warning(
112 + "The progress bar shows the upper bound on token steps and "
113 + "may finish early due to stopping conditions. It does not "
114 + "reflect instance-level progress."
115 + )
116 + for _ in token_iter:
117 + all_beams: list[BeamSearchSequence] = list(
118 + sum((instance.beams for instance in instances_batch), [])
119 + )
120 + pos = [0] + list(
121 + itertools.accumulate(
122 + len(instance.beams) for instance in instances_batch
123 + )
124 + )
125 + instance_start_and_end: list[tuple[int, int]] = list(
126 + zip(pos[:-1], pos[1:])
127 + )
128 +
129 + if len(all_beams) == 0:
130 + break
131 +
132 + # only runs for one step
133 + # we don't need to use tqdm here
134 + output = self._render_and_run_requests(
135 + prompts=(beam.get_prompt() for beam in all_beams),
136 + params=self._params_to_seq(sampling_params, len(all_beams)),
137 + output_type=RequestOutput,
138 + lora_requests=[beam.lora_request for beam in all_beams],
139 + use_tqdm=False,
140 + )
141 +
142 + for (start, end), instance in zip(
143 + instance_start_and_end, instances_batch
144 + ):
145 + instance_new_beams = []
146 + for i in range(start, end):
147 + current_beam = all_beams[i]
148 + result = output[i]
149 +
150 + if result.outputs[0].logprobs is not None:
151 + # if `result.outputs[0].logprobs` is None, it means
152 + # the sequence is completed because of the
153 + # max-model-len or abortion. we don't need to add
154 + # it to the new beams.
155 + logprobs = result.outputs[0].logprobs[0]
156 + for token_id, logprob_obj in logprobs.items():
157 + new_beam = BeamSearchSequence(
158 + current_beam.orig_prompt,
159 + tokens=current_beam.tokens + [token_id],
160 + logprobs=current_beam.logprobs + [logprobs],
161 + lora_request=current_beam.lora_request,
162 + cum_logprob=current_beam.cum_logprob
163 + + logprob_obj.logprob,
164 + )
165 +
166 + if token_id == eos_token_id and not ignore_eos:
167 + instance.completed.append(new_beam)
168 + else:
169 + instance_new_beams.append(new_beam)
170 + sorted_beams = sorted(
171 + instance_new_beams, key=sort_beams_key, reverse=True
172 + )
173 + instance.beams = sorted_beams[:beam_width]
174 +
175 + outputs = []
176 + for instance in instances:
177 + instance.completed.extend(instance.beams)
178 + sorted_completed = sorted(
179 + instance.completed, key=sort_beams_key, reverse=True
180 + )
181 + best_beams = sorted_completed[:beam_width]
182 +
183 + for beam in best_beams:
184 + beam.text = tokenizer.decode(beam.tokens)
185 +
186 + outputs.append(BeamSearchOutput(sequences=best_beams))
187 +
188 + return outputs
189 +
190 + @abstractmethod
191 + def _preprocess_cmpl(
192 + self,
193 + prompts: Sequence[PromptType],
194 + tokenization_kwargs: dict[str, Any] | None = None,
195 + mm_processor_kwargs: dict[str, Any] | None = None,
196 + ) -> Sequence[EngineInput]:
197 + raise NotImplementedError
198 +
199 + @abstractmethod
200 + def _lora_request_to_seq(
201 + self,
202 + lora_request: LoRARequest | None | Sequence[LoRARequest | None],
203 + num_requests: int,
204 + ) -> Sequence[LoRARequest | None]:
205 + raise NotImplementedError
206 +
207 + @abstractmethod
208 + def _params_to_seq(
209 + self,
210 + params: SamplingParams,
211 + num_requests: int,
212 + ) -> Sequence[SamplingParams]:
213 + raise NotImplementedError
214 +
215 + @abstractmethod
216 + def _render_and_run_requests(
217 + self,
218 + prompts: Iterable[EngineInput],
219 + params: Sequence[SamplingParams],
220 + output_type: type[RequestOutput],
221 + *,
222 + lora_requests: Sequence[LoRARequest | None] | None = None,
223 + priorities: Sequence[int] | None = None,
224 + use_tqdm: bool | Callable[..., tqdm] = True,
225 + ):
226 + raise NotImplementedError
vllm/entrypoints/generate/beam_search/online.py new
+225
@@ -0,0 +1,225 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import asyncio
5 +from abc import ABC
6 +from collections.abc import AsyncGenerator, Mapping
7 +
8 +import numpy as np
9 +
10 +from vllm import CompletionOutput, RequestOutput
11 +from vllm.engine.protocol import EngineClient
12 +from vllm.inputs import EngineInput
13 +from vllm.lora.request import LoRARequest
14 +from vllm.renderers import BaseRenderer
15 +from vllm.sampling_params import BeamSearchParams, SamplingParams
16 +from vllm.utils import random_uuid
17 +from vllm.utils.async_utils import collect_from_async_generator
18 +
19 +from .utils import BeamSearchSequence, create_sort_beams_key_function
20 +
21 +
22 +class BeamSearchOnlineMixin(ABC):
23 + """online serving for beam search"""
24 +
25 + renderer: BaseRenderer
26 + engine_client: EngineClient
27 +
28 + async def beam_search(
29 + self,
30 + prompt: EngineInput,
31 + request_id: str,
32 + params: BeamSearchParams,
33 + lora_request: LoRARequest | None = None,
34 + trace_headers: Mapping[str, str] | None = None,
35 + ) -> AsyncGenerator[RequestOutput, None]:
36 + beam_width = params.beam_width
37 + max_tokens = params.max_tokens
38 + ignore_eos = params.ignore_eos
39 + temperature = params.temperature
40 + length_penalty = params.length_penalty
41 + include_stop_str_in_output = params.include_stop_str_in_output
42 +
43 + tokenizer = self.renderer.get_tokenizer()
44 + eos_token_id = tokenizer.eos_token_id
45 + sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty)
46 +
47 + if prompt["type"] == "embeds":
48 + raise NotImplementedError("Embedding prompt not supported for beam search")
49 +
50 + # Extract prompt tokens and text based on model type
51 + decoder_prompt = (
52 + prompt if prompt["type"] != "enc_dec" else prompt["decoder_prompt"]
53 + )
54 + prompt_text = decoder_prompt.get("prompt")
55 + prompt_token_ids = decoder_prompt["prompt_token_ids"]
56 +
57 + tokenized_length = len(prompt_token_ids)
58 +
59 + logprobs_num = 2 * beam_width
60 + sampling_params = SamplingParams(
61 + logprobs=logprobs_num,
62 + max_tokens=1,
63 + temperature=temperature,
64 + )
65 + all_beams = [
66 + BeamSearchSequence(
67 + orig_prompt=prompt,
68 + tokens=prompt_token_ids,
69 + cum_logprob=0,
70 + logprobs=[],
71 + lora_request=lora_request,
72 + )
73 + ]
74 + completed = []
75 +
76 + for _ in range(max_tokens):
77 + tasks = []
78 + request_id_batch = f"{request_id}-{random_uuid()}"
79 +
80 + for i, beam in enumerate(all_beams):
81 + prompt_item = beam.get_prompt()
82 + lora_request_item = beam.lora_request
83 + request_id_item = f"{request_id_batch}-beam-{i}"
84 + task = asyncio.create_task(
85 + collect_from_async_generator(
86 + self.engine_client.generate(
87 + prompt_item,
88 + sampling_params,
89 + request_id_item,
90 + lora_request=lora_request_item,
91 + trace_headers=trace_headers,
92 + )
93 + )
94 + )
95 + tasks.append(task)
96 +
97 + output = [x[0] for x in await asyncio.gather(*tasks)]
98 +
99 + new_beams = []
100 + # Store all new tokens generated by beam
101 + all_beams_token_id = []
102 + # Store the cumulative probability of all tokens
103 + # generated by beam search
104 + all_beams_logprob = []
105 + # Iterate through all beam inference results
106 + for i, result in enumerate(output):
107 + current_beam = all_beams[i]
108 +
109 + # check for error finish reason and abort beam search
110 + if result.outputs[0].finish_reason == "error":
111 + # yield error output and terminate beam search
112 + yield RequestOutput(
113 + request_id=request_id,
114 + prompt=prompt_text,
115 + outputs=[
116 + CompletionOutput(
117 + index=0,
118 + text="",
119 + token_ids=[],
120 + cumulative_logprob=None,
121 + logprobs=None,
122 + finish_reason="error",
123 + )
124 + ],
125 + finished=True,
126 + prompt_token_ids=prompt_token_ids,
127 + prompt_logprobs=None,
128 + )
129 + return
130 +
131 + if result.outputs[0].logprobs is not None:
132 + logprobs = result.outputs[0].logprobs[0]
133 + all_beams_token_id.extend(list(logprobs.keys()))
134 + all_beams_logprob.extend(
135 + [
136 + current_beam.cum_logprob + obj.logprob
137 + for obj in logprobs.values()
138 + ]
139 + )
140 +
141 + # Handle the token for the end of sentence (EOS)
142 + all_beams_token_id = np.array(all_beams_token_id)
143 + all_beams_logprob = np.array(all_beams_logprob)
144 +
145 + if not ignore_eos:
146 + # Get the index position of eos token in all generated results
147 + eos_idx = np.where(all_beams_token_id == eos_token_id)[0]
148 + for idx in eos_idx:
149 + current_beam = all_beams[idx // logprobs_num]
150 + result = output[idx // logprobs_num]
151 + assert result.outputs[0].logprobs is not None
152 + logprobs_entry = result.outputs[0].logprobs[0]
153 + completed.append(
154 + BeamSearchSequence(
155 + orig_prompt=prompt,
156 + tokens=current_beam.tokens + [eos_token_id]
157 + if include_stop_str_in_output
158 + else current_beam.tokens,
159 + logprobs=current_beam.logprobs + [logprobs_entry],
160 + cum_logprob=float(all_beams_logprob[idx]),
161 + finish_reason="stop",
162 + stop_reason=eos_token_id,
163 + )
164 + )
165 + # After processing, set the log probability of the eos condition
166 + # to negative infinity.
167 + all_beams_logprob[eos_idx] = -np.inf
168 +
169 + # Processing non-EOS tokens
170 + # Get indices of the top beam_width probabilities
171 + topn_idx = np.argpartition(np.negative(all_beams_logprob), beam_width)[
172 + :beam_width
173 + ]
174 +
175 + for idx in topn_idx:
176 + current_beam = all_beams[idx // logprobs_num]
177 + result = output[idx // logprobs_num]
178 + token_id = int(all_beams_token_id[idx])
179 + assert result.outputs[0].logprobs is not None
180 + logprobs_entry = result.outputs[0].logprobs[0]
181 + new_beams.append(
182 + BeamSearchSequence(
183 + orig_prompt=prompt,
184 + tokens=current_beam.tokens + [token_id],
185 + logprobs=current_beam.logprobs + [logprobs_entry],
186 + lora_request=current_beam.lora_request,
187 + cum_logprob=float(all_beams_logprob[idx]),
188 + )
189 + )
190 +
191 + all_beams = new_beams
192 +
193 + completed.extend(all_beams)
194 + sorted_completed = sorted(completed, key=sort_beams_key, reverse=True)
195 + best_beams = sorted_completed[:beam_width]
196 +
197 + for beam in best_beams:
198 + if beam.tokens[-1] == eos_token_id and not ignore_eos:
199 + # Skip the eos token in the text.
200 + tokens = beam.tokens[tokenized_length:-1]
201 + else:
202 + tokens = beam.tokens[tokenized_length:]
203 + beam.text = tokenizer.decode(tokens)
204 +
205 + yield RequestOutput(
206 + request_id=request_id,
207 + prompt=prompt_text,
208 + outputs=[
209 + CompletionOutput(
210 + text=beam.text, # type: ignore
211 + cumulative_logprob=beam.cum_logprob,
212 + token_ids=beam.tokens[tokenized_length:],
213 + index=i,
214 + logprobs=beam.logprobs,
215 + finish_reason=beam.finish_reason
216 + if beam.finish_reason is not None
217 + else "length",
218 + stop_reason=beam.stop_reason,
219 + )
220 + for (i, beam) in enumerate(best_beams)
221 + ],
222 + finished=True,
223 + prompt_token_ids=prompt_token_ids,
224 + prompt_logprobs=None,
225 + )
vllm/entrypoints/generate/beam_search/utils.py renamed
vllm/entrypoints/llm.py
+3 -168
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import itertools
4 from collections.abc import Callable, Iterable, Sequence
5 from pathlib import Path
6 from typing import TYPE_CHECKING, Any
@@ -12,12 +11,6 @@ from pydantic import ValidationError
11 from tqdm.auto import tqdm
12 from typing_extensions import TypeVar, overload
13
15 -from vllm.beam_search import (
16 - BeamSearchInstance,
17 - BeamSearchOutput,
18 - BeamSearchSequence,
19 - create_sort_beams_key_function,
20 -)
14 from vllm.config import (
15 AttentionConfig,
16 CompilationConfig,
@@ -45,13 +38,12 @@ from vllm.entrypoints.chat_utils import (
38 ChatTemplateContentFormatOption,
39 load_chat_template,
40 )
41 +from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin
42 from vllm.entrypoints.pooling.offline import PoolingOfflineMixin
43 from vllm.entrypoints.utils import log_non_default_args
44 from vllm.inputs import (
45 EngineInput,
46 PromptType,
53 - TextPrompt,
54 - TokensPrompt,
47 )
48 from vllm.logger import init_logger
49 from vllm.lora.request import LoRARequest
@@ -65,7 +57,7 @@ from vllm.renderers.inputs.preprocess import (
57 parse_model_prompt,
58 prompt_to_seq,
59 )
68 -from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
60 +from vllm.sampling_params import RequestOutputKind, SamplingParams
61 from vllm.tokenizers import TokenizerLike
62 from vllm.usage.usage_lib import UsageContext
63 from vllm.utils.counter import Counter
@@ -89,7 +81,7 @@ _P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
81 _R = TypeVar("_R", default=Any)
82
83
92 -class LLM(PoolingOfflineMixin):
84 +class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin):
85 """An LLM for generating texts from given prompts and sampling parameters.
86
87 This class includes a tokenizer, a language model (possibly distributed
@@ -675,163 +667,6 @@ class LLM(PoolingOfflineMixin):
667 """
668 return self.llm_engine.apply_model(func)
669
678 - def beam_search(
679 - self,
680 - prompts: list[TokensPrompt | TextPrompt],
681 - params: BeamSearchParams,
682 - lora_request: list[LoRARequest] | LoRARequest | None = None,
683 - use_tqdm: bool = False,
684 - concurrency_limit: int | None = None,
685 - ) -> list[BeamSearchOutput]:
686 - """
687 - Generate sequences using beam search.
688 -
689 - Args:
690 - prompts: A list of prompts. Each prompt can be a string or a list
691 - of token IDs.
692 - params: The beam search parameters.
693 - lora_request: LoRA request to use for generation, if any.
694 - use_tqdm: Whether to use tqdm to display the progress bar.
695 - concurrency_limit: The maximum number of concurrent requests.
696 - If None, the number of concurrent requests is unlimited.
697 - """
698 - # TODO: how does beam search work together with length penalty,
699 - # frequency, penalty, and stopping criteria, etc.?
700 - beam_width = params.beam_width
701 - max_tokens = params.max_tokens
702 - temperature = params.temperature
703 - ignore_eos = params.ignore_eos
704 - length_penalty = params.length_penalty
705 -
706 - tokenizer = self.renderer.get_tokenizer()
707 - eos_token_id = tokenizer.eos_token_id
708 - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty)
709 -
710 - engine_inputs = self._preprocess_cmpl(prompts)
711 - lora_requests = self._lora_request_to_seq(lora_request, len(engine_inputs))
712 -
713 - if use_tqdm and concurrency_limit is not None:
714 - logger.warning(
715 - "Progress bar is not supported when using concurrency_limit. "
716 - "Disabling progress bar."
717 - )
718 - use_tqdm = False
719 -
720 - if concurrency_limit is None:
721 - concurrency_limit = len(engine_inputs)
722 -
723 - # generate 2 * beam_width candidates at each step
724 - # following the huggingface transformers implementation
725 - # at https://github.com/huggingface/transformers/blob/e15687fffe5c9d20598a19aeab721ae0a7580f8a/src/transformers/generation/beam_search.py#L534 # noqa
726 - sampling_params = SamplingParams(
727 - logprobs=2 * beam_width,
728 - max_tokens=1,
729 - temperature=temperature,
730 - skip_clone=True, # Internal beam search, safe to skip clone
731 - )
732 - instances: list[BeamSearchInstance] = []
733 -
734 - for lora_req, prompt in zip(lora_requests, engine_inputs):
735 - if prompt["type"] == "embeds":
736 - raise NotImplementedError(
737 - "Embedding prompt not supported for beam search"
738 - )
739 -
740 - instances.append(
741 - BeamSearchInstance(
742 - prompt,
743 - lora_request=lora_req,
744 - logprobs=None,
745 - ),
746 - )
747 -
748 - for prompt_start in range(0, len(instances), concurrency_limit):
749 - instances_batch = instances[prompt_start : prompt_start + concurrency_limit]
750 -
751 - token_iter = range(max_tokens)
752 - if use_tqdm:
753 - token_iter = tqdm(
754 - token_iter, desc="Beam search", unit="token", unit_scale=False
755 - )
756 - logger.warning(
757 - "The progress bar shows the upper bound on token steps and "
758 - "may finish early due to stopping conditions. It does not "
759 - "reflect instance-level progress."
760 - )
761 - for _ in token_iter:
762 - all_beams: list[BeamSearchSequence] = list(
763 - sum((instance.beams for instance in instances_batch), [])
764 - )
765 - pos = [0] + list(
766 - itertools.accumulate(
767 - len(instance.beams) for instance in instances_batch
768 - )
769 - )
770 - instance_start_and_end: list[tuple[int, int]] = list(
771 - zip(pos[:-1], pos[1:])
772 - )
773 -
774 - if len(all_beams) == 0:
775 - break
776 -
777 - # only runs for one step
778 - # we don't need to use tqdm here
779 - output = self._render_and_run_requests(
780 - prompts=(beam.get_prompt() for beam in all_beams),
781 - params=self._params_to_seq(sampling_params, len(all_beams)),
782 - output_type=RequestOutput,
783 - lora_requests=[beam.lora_request for beam in all_beams],
784 - use_tqdm=False,
785 - )
786 -
787 - for (start, end), instance in zip(
788 - instance_start_and_end, instances_batch
789 - ):
790 - instance_new_beams = []
791 - for i in range(start, end):
792 - current_beam = all_beams[i]
793 - result = output[i]
794 -
795 - if result.outputs[0].logprobs is not None:
796 - # if `result.outputs[0].logprobs` is None, it means
797 - # the sequence is completed because of the
798 - # max-model-len or abortion. we don't need to add
799 - # it to the new beams.
800 - logprobs = result.outputs[0].logprobs[0]
801 - for token_id, logprob_obj in logprobs.items():
802 - new_beam = BeamSearchSequence(
803 - current_beam.orig_prompt,
804 - tokens=current_beam.tokens + [token_id],
805 - logprobs=current_beam.logprobs + [logprobs],
806 - lora_request=current_beam.lora_request,
807 - cum_logprob=current_beam.cum_logprob
808 - + logprob_obj.logprob,
809 - )
810 -
811 - if token_id == eos_token_id and not ignore_eos:
812 - instance.completed.append(new_beam)
813 - else:
814 - instance_new_beams.append(new_beam)
815 - sorted_beams = sorted(
816 - instance_new_beams, key=sort_beams_key, reverse=True
817 - )
818 - instance.beams = sorted_beams[:beam_width]
819 -
820 - outputs = []
821 - for instance in instances:
822 - instance.completed.extend(instance.beams)
823 - sorted_completed = sorted(
824 - instance.completed, key=sort_beams_key, reverse=True
825 - )
826 - best_beams = sorted_completed[:beam_width]
827 -
828 - for beam in best_beams:
829 - beam.text = tokenizer.decode(beam.tokens)
830 -
831 - outputs.append(BeamSearchOutput(sequences=best_beams))
832 -
833 - return outputs
834 -
670 def _preprocess_cmpl(
671 self,
672 prompts: Sequence[PromptType],
vllm/entrypoints/openai/engine/serving.py
+3 -206
@@ -1,25 +1,23 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import asyncio
3 import contextlib
4 import json
5 import time
7 -from collections.abc import AsyncGenerator, Awaitable, Mapping
6 +from collections.abc import Awaitable, Mapping
7 from dataclasses import dataclass, field
8 from http import HTTPStatus
9 from typing import Any, ClassVar, Generic, Protocol, TypeAlias, TypeVar
10
12 -import numpy as np
11 from fastapi import Request
12 from openai.types.responses import ToolChoiceFunction
13 from pydantic import ConfigDict, TypeAdapter, ValidationError
14 from starlette.datastructures import Headers
15
16 import vllm.envs as envs
19 -from vllm.beam_search import BeamSearchSequence, create_sort_beams_key_function
17 from vllm.config import ModelConfig
18 from vllm.engine.protocol import EngineClient
19 from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
20 +from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin
21 from vllm.entrypoints.logger import RequestLogger
22 from vllm.entrypoints.openai.chat_completion.protocol import (
23 BatchChatCompletionRequest,
@@ -56,7 +54,6 @@ from vllm.inputs import EngineInput, PromptType
54 from vllm.logger import init_logger
55 from vllm.logprobs import Logprob, PromptLogprobs
56 from vllm.lora.request import LoRARequest
59 -from vllm.outputs import CompletionOutput, RequestOutput
57 from vllm.renderers import ChatParams, TokenizeParams
58 from vllm.renderers.inputs.preprocess import (
59 extract_prompt_components,
@@ -71,7 +68,6 @@ from vllm.tracing import (
68 log_tracing_disabled_warning,
69 )
70 from vllm.utils import random_uuid
74 -from vllm.utils.async_utils import collect_from_async_generator
71 from vllm.utils.mistral import is_mistral_tool_parser
72
73 logger = init_logger(__name__)
@@ -133,7 +129,7 @@ class ServeContext(Generic[RequestT]):
129 model_config = ConfigDict(arbitrary_types_allowed=True)
130
131
136 -class OpenAIServing:
132 +class OpenAIServing(BeamSearchOnlineMixin):
133 request_id_prefix: ClassVar[str] = """
134 A short string prepended to every request’s ID.
135 """
@@ -174,205 +170,6 @@ class OpenAIServing:
170 # Never fail server startup over the fingerprint.
171 self.system_fingerprint = None
172
177 - async def beam_search(
178 - self,
179 - prompt: EngineInput,
180 - request_id: str,
181 - params: BeamSearchParams,
182 - lora_request: LoRARequest | None = None,
183 - trace_headers: Mapping[str, str] | None = None,
184 - ) -> AsyncGenerator[RequestOutput, None]:
185 - beam_width = params.beam_width
186 - max_tokens = params.max_tokens
187 - ignore_eos = params.ignore_eos
188 - temperature = params.temperature
189 - length_penalty = params.length_penalty
190 - include_stop_str_in_output = params.include_stop_str_in_output
191 -
192 - tokenizer = self.renderer.get_tokenizer()
193 - eos_token_id = tokenizer.eos_token_id
194 - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty)
195 -
196 - if prompt["type"] == "embeds":
197 - raise NotImplementedError("Embedding prompt not supported for beam search")
198 -
199 - # Extract prompt tokens and text based on model type
200 - decoder_prompt = (
201 - prompt if prompt["type"] != "enc_dec" else prompt["decoder_prompt"]
202 - )
203 - prompt_text = decoder_prompt.get("prompt")
204 - prompt_token_ids = decoder_prompt["prompt_token_ids"]
205 -
206 - tokenized_length = len(prompt_token_ids)
207 -
208 - logprobs_num = 2 * beam_width
209 - sampling_params = SamplingParams(
210 - logprobs=logprobs_num,
211 - max_tokens=1,
212 - temperature=temperature,
213 - )
214 - all_beams = [
215 - BeamSearchSequence(
216 - orig_prompt=prompt,
217 - tokens=prompt_token_ids,
218 - cum_logprob=0,
219 - logprobs=[],
220 - lora_request=lora_request,
221 - )
222 - ]
223 - completed = []
224 -
225 - for _ in range(max_tokens):
226 - tasks = []
227 - request_id_batch = f"{request_id}-{random_uuid()}"
228 -
229 - for i, beam in enumerate(all_beams):
230 - prompt_item = beam.get_prompt()
231 - lora_request_item = beam.lora_request
232 - request_id_item = f"{request_id_batch}-beam-{i}"
233 - task = asyncio.create_task(
234 - collect_from_async_generator(
235 - self.engine_client.generate(
236 - prompt_item,
237 - sampling_params,
238 - request_id_item,
239 - lora_request=lora_request_item,
240 - trace_headers=trace_headers,
241 - )
242 - )
243 - )
244 - tasks.append(task)
245 -
246 - output = [x[0] for x in await asyncio.gather(*tasks)]
247 -
248 - new_beams = []
249 - # Store all new tokens generated by beam
250 - all_beams_token_id = []
251 - # Store the cumulative probability of all tokens
252 - # generated by beam search
253 - all_beams_logprob = []
254 - # Iterate through all beam inference results
255 - for i, result in enumerate(output):
256 - current_beam = all_beams[i]
257 -
258 - # check for error finish reason and abort beam search
259 - if result.outputs[0].finish_reason == "error":
260 - # yield error output and terminate beam search
261 - yield RequestOutput(
262 - request_id=request_id,
263 - prompt=prompt_text,
264 - outputs=[
265 - CompletionOutput(
266 - index=0,
267 - text="",
268 - token_ids=[],
269 - cumulative_logprob=None,
270 - logprobs=None,
271 - finish_reason="error",
272 - )
273 - ],
274 - finished=True,
275 - prompt_token_ids=prompt_token_ids,
276 - prompt_logprobs=None,
277 - )
278 - return
279 -
280 - if result.outputs[0].logprobs is not None:
281 - logprobs = result.outputs[0].logprobs[0]
282 - all_beams_token_id.extend(list(logprobs.keys()))
283 - all_beams_logprob.extend(
284 - [
285 - current_beam.cum_logprob + obj.logprob
286 - for obj in logprobs.values()
287 - ]
288 - )
289 -
290 - # Handle the token for the end of sentence (EOS)
291 - all_beams_token_id = np.array(all_beams_token_id)
292 - all_beams_logprob = np.array(all_beams_logprob)
293 -
294 - if not ignore_eos:
295 - # Get the index position of eos token in all generated results
296 - eos_idx = np.where(all_beams_token_id == eos_token_id)[0]
297 - for idx in eos_idx:
298 - current_beam = all_beams[idx // logprobs_num]
299 - result = output[idx // logprobs_num]
300 - assert result.outputs[0].logprobs is not None
301 - logprobs_entry = result.outputs[0].logprobs[0]
302 - completed.append(
303 - BeamSearchSequence(
304 - orig_prompt=prompt,
305 - tokens=current_beam.tokens + [eos_token_id]
306 - if include_stop_str_in_output
307 - else current_beam.tokens,
308 - logprobs=current_beam.logprobs + [logprobs_entry],
309 - cum_logprob=float(all_beams_logprob[idx]),
310 - finish_reason="stop",
311 - stop_reason=eos_token_id,
312 - )
313 - )
314 - # After processing, set the log probability of the eos condition
315 - # to negative infinity.
316 - all_beams_logprob[eos_idx] = -np.inf
317 -
318 - # Processing non-EOS tokens
319 - # Get indices of the top beam_width probabilities
320 - topn_idx = np.argpartition(np.negative(all_beams_logprob), beam_width)[
321 - :beam_width
322 - ]
323 -
324 - for idx in topn_idx:
325 - current_beam = all_beams[idx // logprobs_num]
326 - result = output[idx // logprobs_num]
327 - token_id = int(all_beams_token_id[idx])
328 - assert result.outputs[0].logprobs is not None
329 - logprobs_entry = result.outputs[0].logprobs[0]
330 - new_beams.append(
331 - BeamSearchSequence(
332 - orig_prompt=prompt,
333 - tokens=current_beam.tokens + [token_id],
334 - logprobs=current_beam.logprobs + [logprobs_entry],
335 - lora_request=current_beam.lora_request,
336 - cum_logprob=float(all_beams_logprob[idx]),
337 - )
338 - )
339 -
340 - all_beams = new_beams
341 -
342 - completed.extend(all_beams)
343 - sorted_completed = sorted(completed, key=sort_beams_key, reverse=True)
344 - best_beams = sorted_completed[:beam_width]
345 -
346 - for beam in best_beams:
347 - if beam.tokens[-1] == eos_token_id and not ignore_eos:
348 - # Skip the eos token in the text.
349 - tokens = beam.tokens[tokenized_length:-1]
350 - else:
351 - tokens = beam.tokens[tokenized_length:]
352 - beam.text = tokenizer.decode(tokens)
353 -
354 - yield RequestOutput(
355 - request_id=request_id,
356 - prompt=prompt_text,
357 - outputs=[
358 - CompletionOutput(
359 - text=beam.text, # type: ignore
360 - cumulative_logprob=beam.cum_logprob,
361 - token_ids=beam.tokens[tokenized_length:],
362 - index=i,
363 - logprobs=beam.logprobs,
364 - finish_reason=beam.finish_reason
365 - if beam.finish_reason is not None
366 - else "length",
367 - stop_reason=beam.stop_reason,
368 - )
369 - for (i, beam) in enumerate(best_beams)
370 - ],
371 - finished=True,
372 - prompt_token_ids=prompt_token_ids,
373 - prompt_logprobs=None,
374 - )
375 -
173 @staticmethod
174 def create_error_response(
175 message: str | Exception,