[Bugfix][Frontend] Handle aborted requests in beam search (#56249)
Signed-off-by: 子华 <huaxi.shx@alibaba-inc.com> Co-authored-by: Codex <noreply@openai.com>
shaohuaxi committed
Sep 15, 2026 at 01:17 UTC
ca67438c088d5ada7da901b77f9aeaebc1b9542c
5 files changed
+162
-9
tests/samplers/test_beam_search.py
+85
@@ -6,13 +6,17 @@ 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.logprobs import Logprob, SampleLogprobs
20
from vllm.platforms import current_platform
21
from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams
22
@@ -45,6 +49,87 @@ MM_BEAM_WIDTHS = [2]
49
MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
50
51
52
+@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)])
53
+@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]])
54
+def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts(
55
+ monkeypatch,
56
+ abort_after: int,
57
+ prompt_token: int,
58
+ terminal_logprobs: SampleLogprobs | None,
59
+) -> None:
60
+ """Return tokens and scores from the last completed step for aborted prompts.
61
+
62
+ Other prompts in the batch continue generating.
63
+ """
64
+
65
+ def run_requests(prompts, **kwargs):
66
+ results = []
67
+ for prompt in prompts:
68
+ tokens = prompt["prompt_token_ids"]
69
+ if tokens[0] == prompt_token:
70
+ assert len(tokens) <= abort_after + 1, "Continued after abort"
71
+ aborted = (
72
+ tokens[0] == prompt_token
73
+ and len(tokens) == abort_after + 1
74
+ and tokens[-1] != 12
75
+ )
76
+ results.append(
77
+ RequestOutput(
78
+ request_id="inner",
79
+ prompt=None,
80
+ prompt_token_ids=tokens,
81
+ prompt_logprobs=None,
82
+ finished=True,
83
+ outputs=[
84
+ CompletionOutput(
85
+ index=0,
86
+ text="",
87
+ token_ids=[] if aborted else [11],
88
+ cumulative_logprob=None,
89
+ logprobs=terminal_logprobs
90
+ if aborted
91
+ else [{11: Logprob(-0.1), 12: Logprob(-0.2)}],
92
+ finish_reason="abort" if aborted else "length",
93
+ )
94
+ ],
95
+ )
96
+ )
97
+ return results
98
+
99
+ llm = LLM.__new__(LLM)
100
+ llm.llm_engine = Mock()
101
+ tokenizer = SimpleNamespace(
102
+ eos_token_id=0,
103
+ decode=lambda tokens, skip_special_tokens=False: " ".join(
104
+ str(token) for token in tokens if not skip_special_tokens or token != 0
105
+ ),
106
+ )
107
+ llm.renderer = Mock(get_tokenizer=Mock(return_value=tokenizer))
108
+ monkeypatch.setattr(llm, "_preprocess_cmpl", lambda prompts: prompts)
109
+ monkeypatch.setattr(llm, "_render_and_run_requests", run_requests)
110
+ prompts = [
111
+ {"type": "token", "prompt_token_ids": [token]} for token in [prompt_token, 2]
112
+ ]
113
+ params = BeamSearchParams(beam_width=2, max_tokens=3)
114
+ outputs = llm.beam_search(prompts, params)
115
+
116
+ aborted, normal = outputs
117
+ expected_tokens = (
118
+ [[prompt_token, 11], [prompt_token, 12]] if abort_after else [[prompt_token]]
119
+ )
120
+ assert [beam.tokens for beam in aborted.sequences] == expected_tokens
121
+ for beam in aborted.sequences:
122
+ assert beam.finish_reason == "abort"
123
+ assert beam.text == tokenizer.decode(
124
+ beam.tokens, skip_special_tokens=params.skip_special_tokens
125
+ )
126
+ assert len(beam.logprobs) == abort_after
127
+ assert aborted.sequences[0].cum_logprob == pytest.approx(-0.1 * abort_after)
128
+ assert len(normal.sequences) == 2
129
+ assert all(len(beam.tokens) == 4 for beam in normal.sequences)
130
+ assert all(beam.finish_reason != "abort" for beam in normal.sequences)
131
+
132
+
133
@pytest.mark.parametrize("model", MODELS)
134
@pytest.mark.parametrize("dtype", ["half"])
135
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
tests/samplers/test_beam_search_online.py
+53
-1
@@ -5,7 +5,7 @@ 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
8
+from vllm.logprobs import Logprob, SampleLogprobs
9
from vllm.sampling_params import BeamSearchParams
10
11
@@ -108,3 +108,55 @@ async def test_beam_search_respects_skip_special_tokens(
108
109
assert outputs[0].outputs[0].text == expected_text
110
assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id]
111
+
112
+
113
+@pytest.mark.asyncio
114
+@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)])
115
+@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]])
116
+async def test_beam_search_abort_returns_partial_outputs(
117
+ monkeypatch,
118
+ abort_after: int,
119
+ prompt_token: int,
120
+ terminal_logprobs: SampleLogprobs | None,
121
+) -> None:
122
+ """Abort returns each beam's tokens and scores from the last completed step."""
123
+ calls = 0
124
+
125
+ async def generate(prompt, *args, **kwargs):
126
+ nonlocal calls
127
+ calls += 1
128
+ depth = len(prompt["prompt_token_ids"]) - 1
129
+ assert depth <= abort_after, "Beam search continued after abort"
130
+ result = await anext(_EngineClient().generate(prompt, *args, **kwargs))
131
+ output = result.outputs[0]
132
+ output.token_ids = [11]
133
+ output.logprobs = [{11: Logprob(-0.1), 12: Logprob(-0.2)}]
134
+ output.finish_reason = "length"
135
+ if depth == abort_after and prompt["prompt_token_ids"][-1] != 12:
136
+ output.token_ids = []
137
+ output.logprobs = terminal_logprobs
138
+ output.finish_reason = "abort"
139
+ yield result
140
+
141
+ serving = _Serving()
142
+ monkeypatch.setattr(serving.engine_client, "generate", generate)
143
+ prompt = {"type": "token", "prompt": "prompt", "prompt_token_ids": [prompt_token]}
144
+ outputs = [
145
+ output
146
+ async for output in serving.beam_search(
147
+ prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4)
148
+ )
149
+ ]
150
+
151
+ assert calls == 1 + 2 * abort_after
152
+ assert len(outputs) == 1
153
+ assert outputs[0].finished
154
+ assert outputs[0].request_id == "request"
155
+ expected_tokens = [[11], [12]] if abort_after else [[]]
156
+ assert [output.token_ids for output in outputs[0].outputs] == expected_tokens
157
+ for output in outputs[0].outputs:
158
+ assert output.finish_reason == "abort"
159
+ assert output.text == _Tokenizer().decode(list(output.token_ids))
160
+ assert output.logprobs is not None
161
+ assert len(output.logprobs) == abort_after
162
+ assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after)
vllm/entrypoints/generate/beam_search/offline.py
+12
-2
@@ -293,6 +293,18 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
293
allowed_sets[i] = set(entry[1])
294
295
for (start, end), instance in zip(instance_start_and_end, instances_batch):
296
+ instance_output = output[start:end]
297
+ if any(
298
+ result is not None and result.outputs[0].finish_reason == "abort"
299
+ for result in instance_output
300
+ ):
301
+ for beam, result in zip(all_beams[start:end], instance_output):
302
+ if result is not None:
303
+ beam.finish_reason = "abort"
304
+ instance.completed.append(beam)
305
+ instance.beams = []
306
+ continue
307
+
308
instance_new_beams = []
309
for i in range(start, end):
310
current_beam = all_beams[i]
@@ -302,8 +314,6 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
314
continue
315
316
if result.outputs[0].logprobs is not None:
305
- # if logprobs is None, the sequence completed
306
- # due to max-model-len or abortion.
317
logprobs = result.outputs[0].logprobs[0]
318
allowed = allowed_sets[i]
319
for token_id, logprob_obj in logprobs.items():
vllm/entrypoints/generate/beam_search/online.py
+10
-5
@@ -99,11 +99,7 @@ class BeamSearchOnlineMixin(ABC):
99
100
output = [x[0] for x in await asyncio.gather(*tasks)]
101
102
- candidates = []
103
- # Iterate through all beam inference results
104
- for i, result in enumerate(output):
105
- current_beam = all_beams[i]
106
-
102
+ for result in output:
103
# check for error finish reason and abort beam search
104
if result.outputs[0].finish_reason == "error":
105
# yield error output and terminate beam search
@@ -126,6 +122,15 @@ class BeamSearchOnlineMixin(ABC):
122
)
123
return
124
125
+ if any(result.outputs[0].finish_reason == "abort" for result in output):
126
+ for beam in all_beams:
127
+ beam.finish_reason = "abort"
128
+ break
129
+
130
+ candidates = []
131
+ for i, result in enumerate(output):
132
+ current_beam = all_beams[i]
133
+
134
if result.outputs[0].logprobs is not None:
135
logprobs = result.outputs[0].logprobs[0]
136
for token_id, logprob_obj in logprobs.items():
vllm/entrypoints/generate/beam_search/utils.py
+2
-1
@@ -150,7 +150,8 @@ def get_beam_search_score(
150
if tokens[-1] == eos_token_id:
151
seq_len -= 1
152
153
- return cumulative_logprob / (seq_len**length_penalty)
153
+ # An aborted beam may contain only an EOS prompt token.
154
+ return cumulative_logprob / (max(seq_len, 1) ** length_penalty)
155
156
157
def create_sort_beams_key_function(eos_token_id: int, length_penalty: float):