| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Compare the outputs of HF and vLLM when using beam search. |
| 4 | |
| 5 | Run `pytest tests/samplers/test_beam_search.py`. |
| 6 | """ |
| 7 | |
| 8 | import json |
| 9 | from types import SimpleNamespace |
| 10 | from unittest.mock import Mock |
| 11 | |
| 12 | import jsonschema |
| 13 | import pytest |
| 14 | from transformers import AutoModelForSeq2SeqLM |
| 15 | |
| 16 | from vllm import CompletionOutput, RequestOutput |
| 17 | from vllm.assets.audio import AudioAsset |
| 18 | from vllm.entrypoints.llm import LLM |
| 19 | from vllm.inputs import TokensInput |
| 20 | from vllm.logprobs import Logprob, SampleLogprobs |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams |
| 23 | |
| 24 | # Extra engine kwargs needed for numerically deterministic beam search. |
| 25 | # On ROCm, floating-point reductions in attention and GEMM kernels are |
| 26 | # non-associative and sensitive to batch geometry, so we: |
| 27 | # async_scheduling=False – deterministic batch composition |
| 28 | # enforce_eager=True – no CUDA-graph padding changing effective size |
| 29 | # enable_prefix_caching=False – avoid prefix-sharing side effects |
| 30 | # max_num_seqs=1 – fixed batch size across runs |
| 31 | # On other platforms these are not needed and the dict is empty. |
| 32 | EXTRA_ENGINE_KWARGS: dict = ( |
| 33 | dict( |
| 34 | async_scheduling=False, |
| 35 | enforce_eager=True, |
| 36 | enable_prefix_caching=False, |
| 37 | max_num_seqs=1, |
| 38 | ) |
| 39 | if current_platform.is_rocm() |
| 40 | else dict(async_scheduling=False, max_num_seqs=1) |
| 41 | ) |
| 42 | |
| 43 | # FIXME(zhuohan): The test can not pass if we: |
| 44 | # 1. Increase max_tokens to 256. |
| 45 | # 2. Increase beam_width to 8. |
| 46 | # 3. Use the model "huggyllama/llama-7b". |
| 47 | MAX_TOKENS = [64] |
| 48 | BEAM_WIDTHS = [4] |
| 49 | MM_BEAM_WIDTHS = [2] |
| 50 | MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"] |
| 51 | |
| 52 | |
| 53 | @pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)]) |
| 54 | @pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]]) |
| 55 | def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts( |
| 56 | monkeypatch, |
| 57 | abort_after: int, |
| 58 | prompt_token: int, |
| 59 | terminal_logprobs: SampleLogprobs | None, |
| 60 | ) -> None: |
| 61 | """Return tokens and scores from the last completed step for aborted prompts. |
| 62 | |
| 63 | Other prompts in the batch continue generating. |
| 64 | """ |
| 65 | |
| 66 | def run_requests(prompts, **kwargs): |
| 67 | results = [] |
| 68 | for prompt in prompts: |
| 69 | tokens = prompt["prompt_token_ids"] |
| 70 | if tokens[0] == prompt_token: |
| 71 | assert len(tokens) <= abort_after + 1, "Continued after abort" |
| 72 | aborted = ( |
| 73 | tokens[0] == prompt_token |
| 74 | and len(tokens) == abort_after + 1 |
| 75 | and tokens[-1] != 12 |
| 76 | ) |
| 77 | results.append( |
| 78 | RequestOutput( |
| 79 | request_id="inner", |
| 80 | prompt=None, |
| 81 | prompt_token_ids=tokens, |
| 82 | prompt_logprobs=None, |
| 83 | finished=True, |
| 84 | outputs=[ |
| 85 | CompletionOutput( |
| 86 | index=0, |
| 87 | text="", |
| 88 | token_ids=[] if aborted else [11], |
| 89 | cumulative_logprob=None, |
| 90 | logprobs=terminal_logprobs |
| 91 | if aborted |
| 92 | else [{11: Logprob(-0.1), 12: Logprob(-0.2)}], |
| 93 | finish_reason="abort" if aborted else "length", |
| 94 | ) |
| 95 | ], |
| 96 | ) |
| 97 | ) |
| 98 | return results |
| 99 | |
| 100 | llm = LLM.__new__(LLM) |
| 101 | llm.llm_engine = Mock() |
| 102 | tokenizer = SimpleNamespace( |
| 103 | eos_token_id=0, |
| 104 | decode=lambda tokens, skip_special_tokens=False: " ".join( |
| 105 | str(token) for token in tokens if not skip_special_tokens or token != 0 |
| 106 | ), |
| 107 | ) |
| 108 | llm.renderer = Mock(get_tokenizer=Mock(return_value=tokenizer)) |
| 109 | monkeypatch.setattr(llm, "_preprocess_cmpl", lambda prompts: prompts) |
| 110 | monkeypatch.setattr(llm, "_render_and_run_requests", run_requests) |
| 111 | # Engine inputs, since _preprocess_cmpl is patched to the identity |
| 112 | prompts: list[TokensInput] = [ |
| 113 | {"type": "token", "prompt_token_ids": [token]} for token in [prompt_token, 2] |
| 114 | ] |
| 115 | params = BeamSearchParams(beam_width=2, max_tokens=3) |
| 116 | outputs = llm.beam_search(prompts, params) |
| 117 | |
| 118 | aborted, normal = outputs |
| 119 | expected_tokens = ( |
| 120 | [[prompt_token, 11], [prompt_token, 12]] if abort_after else [[prompt_token]] |
| 121 | ) |
| 122 | assert [beam.tokens for beam in aborted.sequences] == expected_tokens |
| 123 | for beam in aborted.sequences: |
| 124 | assert beam.finish_reason == "abort" |
| 125 | assert beam.text == tokenizer.decode( |
| 126 | beam.tokens, skip_special_tokens=params.skip_special_tokens |
| 127 | ) |
| 128 | assert len(beam.logprobs) == abort_after |
| 129 | assert aborted.sequences[0].cum_logprob == pytest.approx(-0.1 * abort_after) |
| 130 | assert len(normal.sequences) == 2 |
| 131 | assert all(len(beam.tokens) == 4 for beam in normal.sequences) |
| 132 | assert all(beam.finish_reason != "abort" for beam in normal.sequences) |
| 133 | |
| 134 | |
| 135 | @pytest.mark.parametrize("model", MODELS) |
| 136 | @pytest.mark.parametrize("dtype", ["half"]) |
| 137 | @pytest.mark.parametrize("max_tokens", MAX_TOKENS) |
| 138 | @pytest.mark.parametrize("beam_width", BEAM_WIDTHS) |
| 139 | def test_beam_search_single_input( |
| 140 | hf_runner, |
| 141 | vllm_runner, |
| 142 | example_prompts, |
| 143 | model: str, |
| 144 | dtype: str, |
| 145 | max_tokens: int, |
| 146 | beam_width: int, |
| 147 | ) -> None: |
| 148 | example_prompts = example_prompts[:1] |
| 149 | with hf_runner(model, dtype=dtype) as hf_model: |
| 150 | hf_outputs = hf_model.generate_beam_search( |
| 151 | example_prompts, beam_width, max_tokens |
| 152 | ) |
| 153 | |
| 154 | with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model: |
| 155 | vllm_outputs = vllm_model.generate_beam_search( |
| 156 | example_prompts, beam_width, max_tokens |
| 157 | ) |
| 158 | |
| 159 | for i in range(len(example_prompts)): |
| 160 | hf_output_ids, hf_output_texts = hf_outputs[i] |
| 161 | vllm_output_ids, vllm_output_texts = vllm_outputs[i] |
| 162 | for j, (hf_text, vllm_text) in enumerate( |
| 163 | zip(hf_output_texts, vllm_output_texts) |
| 164 | ): |
| 165 | print(f">>>{j}-th hf output:") |
| 166 | print(hf_text) |
| 167 | print(f">>>{j}-th vllm output:") |
| 168 | print(vllm_text) |
| 169 | assert len(hf_output_ids) == len(vllm_output_ids) |
| 170 | for j in range(len(hf_output_ids)): |
| 171 | assert hf_output_ids[j] == vllm_output_ids[j], ( |
| 172 | f"Test{i} output{j}:\nHF: {hf_output_ids}\nvLLM: {vllm_output_ids}" |
| 173 | ) |
| 174 | |
| 175 | |
| 176 | @pytest.mark.parametrize("model", MODELS) |
| 177 | @pytest.mark.parametrize("dtype", ["half"]) |
| 178 | @pytest.mark.parametrize("max_tokens", MAX_TOKENS) |
| 179 | @pytest.mark.parametrize("beam_width", BEAM_WIDTHS) |
| 180 | def test_beam_search_with_concurrency_limit( |
| 181 | hf_runner, |
| 182 | vllm_runner, |
| 183 | example_prompts, |
| 184 | model: str, |
| 185 | dtype: str, |
| 186 | max_tokens: int, |
| 187 | beam_width: int, |
| 188 | ) -> None: |
| 189 | # example_prompts[1]&[3]&[7] fails due to unknown reason even without |
| 190 | # concurrency limit. skip them for now. |
| 191 | example_prompts = example_prompts[:8] |
| 192 | concurrency_limit = 2 |
| 193 | assert len(example_prompts) > concurrency_limit |
| 194 | with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model: |
| 195 | outputs_with_limit = vllm_model.generate_beam_search( |
| 196 | example_prompts, |
| 197 | beam_width, |
| 198 | max_tokens, |
| 199 | concurrency_limit=concurrency_limit, |
| 200 | ) |
| 201 | outputs_without_limit = [] |
| 202 | |
| 203 | for i in range(0, len(example_prompts), concurrency_limit): |
| 204 | outputs_without_limit.extend( |
| 205 | vllm_model.generate_beam_search( |
| 206 | example_prompts[i : i + concurrency_limit], |
| 207 | beam_width, |
| 208 | max_tokens, |
| 209 | ) |
| 210 | ) |
| 211 | |
| 212 | correct = True |
| 213 | for i in range(len(example_prompts)): |
| 214 | output_ids_with_limit, output_texts_with_limit = outputs_with_limit[i] |
| 215 | output_ids_without_limit, output_texts_without_limit = outputs_without_limit[i] |
| 216 | for j, (text_with_limit, text_without_limit) in enumerate( |
| 217 | zip(output_texts_with_limit, output_texts_without_limit) |
| 218 | ): |
| 219 | print(f">>>{j}-th with limit output:") |
| 220 | print(text_with_limit) |
| 221 | print(f">>>{j}-th without limit output:") |
| 222 | print(text_without_limit) |
| 223 | assert len(output_ids_with_limit) == len(output_ids_without_limit) |
| 224 | for j in range(len(output_ids_with_limit)): |
| 225 | if output_ids_with_limit[j] != output_ids_without_limit[j]: |
| 226 | print( |
| 227 | f"Test{i} output{j}:\n+limit: {output_ids_with_limit}\n" |
| 228 | f"-limit: {output_ids_without_limit}" |
| 229 | ) |
| 230 | correct = False |
| 231 | assert correct |
| 232 | |
| 233 | |
| 234 | @pytest.mark.parametrize("dtype", ["half"]) |
| 235 | @pytest.mark.parametrize("max_tokens", MAX_TOKENS) |
| 236 | @pytest.mark.parametrize("beam_width", MM_BEAM_WIDTHS) |
| 237 | def test_beam_search_passes_multimodal_data( |
| 238 | hf_runner, |
| 239 | vllm_runner, |
| 240 | dtype: str, |
| 241 | max_tokens: int, |
| 242 | beam_width: int, |
| 243 | ) -> None: |
| 244 | """Ensure that beam search passes multimodal data through correctly.""" |
| 245 | # NOTE - this test is primarily to check that mm data is passed to beams |
| 246 | # correctly. As such, we just need to check one extra modality to make |
| 247 | # sure things pass through properly. |
| 248 | audios = [AudioAsset("mary_had_lamb").audio_and_sample_rate] |
| 249 | model = "Qwen/Qwen2-Audio-7B-Instruct" |
| 250 | audio_seq = "<|audio_bos|><|AUDIO|><|audio_eos|>" |
| 251 | prompts = [ |
| 252 | f"<|im_start|>user\n{audio_seq}Can you transcribe this?<|im_end|>\n<|im_start|>assistant\n" # noqa: E501 |
| 253 | ] |
| 254 | |
| 255 | with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSeq2SeqLM) as hf_model: |
| 256 | audio_token_id = hf_model.config.audio_token_index |
| 257 | eos_token_id = hf_model.tokenizer.eos_token_id # <|im_end|> |
| 258 | hf_outputs = hf_model.generate_beam_search( |
| 259 | prompts, |
| 260 | beam_width=beam_width, |
| 261 | max_tokens=max_tokens, |
| 262 | audios=audios, |
| 263 | ) |
| 264 | |
| 265 | with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model: |
| 266 | vllm_outputs = vllm_model.generate_beam_search( |
| 267 | prompts, |
| 268 | beam_width=beam_width, |
| 269 | max_tokens=max_tokens, |
| 270 | audios=audios, |
| 271 | ) |
| 272 | |
| 273 | seq_with_no_audio_toks = lambda seq: [tok for tok in seq if tok != audio_token_id] |
| 274 | |
| 275 | for i in range(len(prompts)): |
| 276 | hf_output_ids, hf_output_texts = hf_outputs[i] |
| 277 | vllm_output_ids, vllm_output_texts = vllm_outputs[i] |
| 278 | |
| 279 | for j, (hf_text, vllm_text) in enumerate( |
| 280 | zip(hf_output_texts, vllm_output_texts) |
| 281 | ): |
| 282 | print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:") |
| 283 | print(hf_text) |
| 284 | print(f">>>{j}-th vllm output:") |
| 285 | print(vllm_text) |
| 286 | assert len(hf_output_ids) == len(vllm_output_ids) |
| 287 | |
| 288 | for j in range(len(hf_output_ids)): |
| 289 | # Compare everything except for the audio tokens; we do this since |
| 290 | # the IDs returned from the transformers helper expands the audio |
| 291 | # token to match features, while the vLLM helper maintains the |
| 292 | # single audio token in the input text |
| 293 | filtered_hf_output_ids = seq_with_no_audio_toks(hf_output_ids[j]) |
| 294 | filtered_vllm_output_ids = seq_with_no_audio_toks(vllm_output_ids[j]) |
| 295 | |
| 296 | # HF output IDs may contain the end of sequence |
| 297 | if len(filtered_hf_output_ids) == len(filtered_vllm_output_ids) + 1: |
| 298 | assert filtered_hf_output_ids[-1] == eos_token_id |
| 299 | filtered_hf_output_ids = filtered_hf_output_ids[:-1] |
| 300 | |
| 301 | assert filtered_hf_output_ids == filtered_vllm_output_ids |
| 302 | |
| 303 | |
| 304 | # NOTE: encoder/decoder tests are currently located under |
| 305 | # tests/models/multimodal/generation/test_whisper.py |
| 306 | |
| 307 | |
| 308 | @pytest.mark.parametrize("model", MODELS) |
| 309 | @pytest.mark.parametrize("dtype", ["half"]) |
| 310 | @pytest.mark.parametrize("beam_width", BEAM_WIDTHS) |
| 311 | def test_beam_search_structured_output( |
| 312 | model: str, |
| 313 | dtype: str, |
| 314 | beam_width: int, |
| 315 | ) -> None: |
| 316 | """Ensure beam search with structured output produces valid JSON.""" |
| 317 | json_schema = { |
| 318 | "type": "object", |
| 319 | "properties": { |
| 320 | "name": {"type": "string"}, |
| 321 | "age": {"type": "integer"}, |
| 322 | }, |
| 323 | "required": ["name", "age"], |
| 324 | "additionalProperties": False, |
| 325 | } |
| 326 | |
| 327 | llm = LLM( |
| 328 | model=model, |
| 329 | dtype=dtype, |
| 330 | max_model_len=512, |
| 331 | structured_outputs_config=dict( |
| 332 | backend="xgrammar", |
| 333 | disable_any_whitespace=True, |
| 334 | ), |
| 335 | **(dict(enforce_eager=True) | EXTRA_ENGINE_KWARGS), |
| 336 | ) |
| 337 | |
| 338 | params = BeamSearchParams( |
| 339 | beam_width=beam_width, |
| 340 | max_tokens=64, |
| 341 | structured_outputs=StructuredOutputsParams(json=json_schema), |
| 342 | ) |
| 343 | |
| 344 | prompts = [ |
| 345 | "Generate a JSON object for a person with name and age:", |
| 346 | ] |
| 347 | |
| 348 | outputs = llm.beam_search(prompts, params) |
| 349 | |
| 350 | assert len(outputs) == len(prompts) |
| 351 | for output in outputs: |
| 352 | assert len(output.sequences) > 0 |
| 353 | for seq in output.sequences: |
| 354 | assert seq.text is not None |
| 355 | print(f"Full text: {seq.text!r}") |
| 356 | # seq.text includes the prompt, extract generated JSON. |
| 357 | gen_start = seq.text.find("{") |
| 358 | assert gen_start != -1, f"No JSON found in output: {seq.text!r}" |
| 359 | generated = seq.text[gen_start:] |
| 360 | generated = generated.replace("</s>", "").strip() |
| 361 | print(f"Generated JSON: {generated!r}") |
| 362 | parsed = json.loads(generated) |
| 363 | jsonschema.validate(instance=parsed, schema=json_schema) |