[Bugfix][Beam Search] Respect skip_special_tokens during decoding (#56211)
Signed-off-by: lxy-alexander <lxy_alexander@outlook.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Alexander Lee committed
Sep 14, 2026 at 01:49 UTC
10e8d5614c50ad180acbe61878c93d98ca7098a9
6 files changed
+46
-3
tests/samplers/test_beam_search_online.py
+37
-1
@@ -11,8 +11,13 @@ from vllm.sampling_params import BeamSearchParams
11
12
class _Tokenizer:
13
eos_token_id = 0
14
+ special_token_id = eos_token_id
15
15
- def decode(self, token_ids: list[int]) -> str:
16
+ def decode(self, token_ids: list[int], skip_special_tokens: bool = False) -> str:
17
+ if skip_special_tokens:
18
+ token_ids = [
19
+ token_id for token_id in token_ids if token_id != self.special_token_id
20
+ ]
21
return " ".join(str(token_id) for token_id in token_ids)
22
23
@@ -72,3 +77,34 @@ async def test_beam_search_handles_extra_logprob_candidates() -> None:
77
assert outputs[0].outputs[0].finish_reason == "stop"
78
assert outputs[0].outputs[0].token_ids == []
79
assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1)
80
+
81
+
82
+@pytest.mark.asyncio
83
+@pytest.mark.parametrize(
84
+ ("skip_special_tokens", "expected_text"),
85
+ [
86
+ pytest.param(True, "", id="skip"),
87
+ pytest.param(False, "0", id="keep"),
88
+ ],
89
+)
90
+async def test_beam_search_respects_skip_special_tokens(
91
+ skip_special_tokens: bool, expected_text: str
92
+) -> None:
93
+ prompt = {
94
+ "type": "token",
95
+ "prompt": "prompt",
96
+ "prompt_token_ids": [1],
97
+ }
98
+ params = BeamSearchParams(
99
+ beam_width=1,
100
+ max_tokens=1,
101
+ ignore_eos=True,
102
+ skip_special_tokens=skip_special_tokens,
103
+ )
104
+
105
+ outputs = [
106
+ output async for output in _Serving().beam_search(prompt, "request", params)
107
+ ]
108
+
109
+ assert outputs[0].outputs[0].text == expected_text
110
+ assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id]
vllm/entrypoints/generate/beam_search/offline.py
+3
-1
@@ -186,7 +186,9 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin):
186
best_beams = sorted_completed[:beam_width]
187
188
for beam in best_beams:
189
- beam.text = tokenizer.decode(beam.tokens)
189
+ beam.text = tokenizer.decode(
190
+ beam.tokens, skip_special_tokens=params.skip_special_tokens
191
+ )
192
193
outputs.append(BeamSearchOutput(sequences=best_beams))
194
vllm/entrypoints/generate/beam_search/online.py
+3
-1
@@ -197,7 +197,9 @@ class BeamSearchOnlineMixin(ABC):
197
tokens = beam.tokens[tokenized_length:-1]
198
else:
199
tokens = beam.tokens[tokenized_length:]
200
- beam.text = tokenizer.decode(tokens)
200
+ beam.text = tokenizer.decode(
201
+ tokens, skip_special_tokens=params.skip_special_tokens
202
+ )
203
204
yield RequestOutput(
205
request_id=request_id,
vllm/entrypoints/openai/chat_completion/protocol.py
+1
@@ -659,6 +659,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
659
temperature=temperature,
660
length_penalty=self.length_penalty,
661
include_stop_str_in_output=self.include_stop_str_in_output,
662
+ skip_special_tokens=self.skip_special_tokens,
663
)
664
665
def extract_structured_outputs(self) -> StructuredOutputsParams | None:
vllm/entrypoints/openai/completion/protocol.py
+1
@@ -303,6 +303,7 @@ class CompletionRequest(OpenAIBaseModel):
303
temperature=temperature,
304
length_penalty=self.length_penalty,
305
include_stop_str_in_output=self.include_stop_str_in_output,
306
+ skip_special_tokens=self.skip_special_tokens,
307
)
308
309
def extract_structured_outputs(self) -> StructuredOutputsParams | None:
vllm/sampling_params.py
+1
@@ -1313,6 +1313,7 @@ class BeamSearchParams(
1313
length_penalty: float = 1.0
1314
include_stop_str_in_output: bool = False
1315
structured_outputs: StructuredOutputsParams | None = None
1316
+ skip_special_tokens: bool = True
1317
1318
def __post_init__(self) -> None:
1319
_verify_num_sequences(self.beam_width, "beam_width")