From 51144dd83ec557afa6fc6032bbf03d7f50a55b89 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sat, 10 Oct 2026 05:22:31 +0800 Subject: [PATCH] [Bugfix][Spec Decode] Trim token_ids/logprobs left past a stop string under speculative decoding (#47616) Signed-off-by: Ting Sun Signed-off-by: Nick Hill Co-authored-by: Nick Hill Co-authored-by: Claude Opus 5.5 --- ...stop_string_while_stop_model_terminates.py | 113 +++++++++++++++++- vllm/v1/engine/detokenizer.py | 24 +++- vllm/v1/engine/output_processor.py | 34 +++++- 3 files changed, 158 insertions(+), 13 deletions(-) diff --git a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py index 3b72d612dd..38d860969f 100644 --- a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py +++ b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest from vllm.sampling_params import SamplingParams -from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason from vllm.v1.engine.detokenizer import BaseIncrementalDetokenizer +from vllm.v1.engine.output_processor import OutputProcessor +from vllm.v1.outputs import LogprobsLists, SamplingMaskLists @pytest.fixture(params=[True, False]) @@ -82,7 +85,7 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo # Simulate that the last token ('Z') is a stop token (stop_terminated=True). result = detok.update(new_token_ids=token_ids, stop_terminated=True) - # The update should not report a stop string + # The update should report the matched stop string. assert result == stop_string # Output text should reflect stop-string handling: @@ -91,9 +94,111 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo expected_text = "abcd" if include_stop_str_in_output else "ab" assert detok.output_text == expected_text - # The skipped final token should still be recorded in token_ids. - assert detok.output_token_ids == token_ids + # Tokens after the stop string, including the stop token, are dropped. + assert detok.output_token_ids == token_ids[:4] # get_next_output_text should return the full text when finished=True. # (Buffering only applies during streaming when finished=False.) assert detok.get_next_output_text(finished=True, delta=False) == expected_text + + +@pytest.mark.parametrize( + "min_tokens,num_kept", + [(0, 2), (1, 2), (2, 4), (3, 4), (4, None)], +) +@pytest.mark.parametrize("num_prev_tokens", [0, 1]) +def test_stop_string_min_tokens_within_step( + include_stop_str_in_output: bool, min_tokens, num_kept, num_prev_tokens +): + """Stop strings completing within min_tokens are ignored, even when the + min_tokens boundary falls in the middle of a multi-token step.""" + token_ids = [ord(c) for c in "cdcdx"] + req = _make_request( + stop=["cd"], + include_stop_str_in_output=include_stop_str_in_output, + min_tokens=min_tokens, + ) + detok = _DummyDetokenizer(req) + + if num_prev_tokens: + assert detok.update(token_ids[:num_prev_tokens], False) is None + result = detok.update(token_ids[num_prev_tokens:], False) + + if num_kept is None: + assert result is None + assert detok.output_text == "cdcdx" + assert detok.output_token_ids == token_ids + return + assert result == "cd" + text_len = num_kept if include_stop_str_in_output else num_kept - 2 + assert detok.output_text == "cdcdx"[:text_len] + assert detok.output_token_ids == token_ids[:num_kept] + + +@pytest.mark.parametrize( + "decoded_tokens,keep", + [ + (["a", "b", "c", "d", "e", "f"], 4), + (["ab", "cdx", "ef"], 2), + (["ab", "", "c", "", "d", "", "ef"], 5), + ], +) +@pytest.mark.parametrize("stop_terminated", [False, True]) +def test_stop_string_trims_speculative_overflow( + include_stop_str_in_output: bool, + decoded_tokens, + keep, + stop_terminated: bool, + monkeypatch, +): + if stop_terminated: + # The engine also stopped on a stop token at the end of the batch. + decoded_tokens = decoded_tokens + [""] + token_ids = list(range(len(decoded_tokens))) + stop_string = "cd" + expected_token_ids = token_ids[:keep] + + req = _make_request( + stop=[stop_string], include_stop_str_in_output=include_stop_str_in_output + ) + req.external_req_id = req.request_id + assert req.sampling_params is not None + req.sampling_params.logprobs = 0 + detok = _DummyDetokenizer(req) + monkeypatch.setattr(detok, "decode_next", decoded_tokens.__getitem__) + processor = OutputProcessor(tokenizer=None, log_stats=False) + processor.add_request(req, prompt=None) + processor.request_states[req.request_id].detokenizer = detok + + outputs = processor.process_outputs( + [ + EngineCoreOutput( + request_id=req.request_id, + new_token_ids=token_ids, + finish_reason=FinishReason.STOP if stop_terminated else None, + new_logprobs=LogprobsLists( + np.array(token_ids).reshape(-1, 1), + np.zeros((len(token_ids), 1)), + np.ones(len(token_ids), dtype=int), + None, + ), + new_sampling_mask=SamplingMaskLists( + np.array(token_ids), np.arange(len(token_ids) + 1) + ), + routed_experts=np.array(token_ids).reshape(-1, 1, 1), + ) + ] + ) + result = outputs.request_outputs[0].outputs[0] + + assert result.stop_reason == stop_string + expected_text = "abcd" if include_stop_str_in_output else "ab" + assert detok.output_text == expected_text + assert detok.output_token_ids == expected_token_ids + assert result.token_ids == expected_token_ids + assert result.logprobs is not None + assert len(result.logprobs) == len(expected_token_ids) + assert result.sampling_mask is not None + assert result.sampling_mask.token_ids == [[token] for token in expected_token_ids] + assert result.routed_experts is not None + assert result.routed_experts.ravel().tolist() == expected_token_ids diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index c27e01895c..adc0be3158 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import sys from abc import ABC, abstractmethod +from bisect import bisect_left import tokenizers.decoders from tokenizers import Tokenizer @@ -112,21 +113,27 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): skipped_stop_token_id = None # 1) Detokenize the new token ids incrementally. - stop_check_offset = len(self.output_text) + num_prev_tokens = self.num_output_tokens() + first_new_token_index = len(self.token_ids) + # Output text offsets at the boundaries of the new tokens. + token_offsets = [len(self.output_text)] if self.stop else None for new_token_id in new_token_ids: self.token_ids.append(new_token_id) self.output_text += self.decode_next(new_token_id) - # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014 - if self.min_tokens and self.num_output_tokens() <= self.min_tokens: - stop_check_offset = len(self.output_text) + if token_offsets is not None: + token_offsets.append(len(self.output_text)) if skipped_stop_token_id is not None: # Cleanup after skipping detokenization. self.token_ids.append(skipped_stop_token_id) + if token_offsets is not None: + token_offsets.append(len(self.output_text)) # 2) Evaluate stop strings. stop_string = None - if self.stop and self.num_output_tokens() > self.min_tokens: + if token_offsets is not None and self.num_output_tokens() > self.min_tokens: + # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014 + stop_check_offset = token_offsets[max(self.min_tokens - num_prev_tokens, 0)] stop = check_stop_strings( output_text=self.output_text, new_char_count=len(self.output_text) - stop_check_offset, @@ -135,6 +142,13 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): ) if stop is not None: stop_string, truncate_to = stop + stop_end = len(self.output_text) if truncate_to == -1 else truncate_to + if not self.include_stop_str_in_output: + stop_end += len(stop_string) + # Drop tokens generated after the one completing the stop string + # (possible when multiple tokens are appended in one step). + num_kept = bisect_left(token_offsets, stop_end, lo=1) + del self.token_ids[first_new_token_index + num_kept :] if truncate_to != -1: self.output_text = self.output_text[:truncate_to] diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index c22c5474c7..abb45f7239 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -736,17 +736,22 @@ class OutputProcessor: if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None - if engine_core_output.new_sampling_mask is not None: - req_state.sampling_mask_chunks.append( - engine_core_output.new_sampling_mask - ) # 2) Detokenize the token ids into text and perform stop checks. + num_prev_tokens = req_state.detokenizer.num_output_tokens() stop_string = req_state.detokenizer.update( new_token_ids, finish_reason == FinishReason.STOP ) if stop_string: finish_reason = FinishReason.STOP stop_reason = stop_string + self._trim_surplus_tokens( + req_state, engine_core_output, num_prev_tokens + ) + + if engine_core_output.new_sampling_mask is not None: + req_state.sampling_mask_chunks.append( + engine_core_output.new_sampling_mask + ) # 3) Compute sample and prompt logprobs for request, # if required. @@ -799,6 +804,27 @@ class OutputProcessor: reqs_to_abort=reqs_to_abort, ) + @staticmethod + def _trim_surplus_tokens( + req_state: RequestState, output: EngineCoreOutput, num_prev_tokens: int + ) -> None: + """Drop new tokens generated past a stop string matched by the + detokenizer, e.g. when multiple tokens are accepted in one step.""" + assert req_state.detokenizer is not None + num_kept = req_state.detokenizer.num_output_tokens() - num_prev_tokens + num_dropped = len(output.new_token_ids) - num_kept + if num_dropped <= 0: + return + del output.new_token_ids[num_kept:] + if output.new_logprobs is not None: + output.new_logprobs = output.new_logprobs.slice_request(0, num_kept) + if output.new_sampling_mask is not None: + output.new_sampling_mask = output.new_sampling_mask.slice_request( + 0, num_kept + ) + if output.routed_experts is not None: + req_state.routed_experts_chunks[-1] = output.routed_experts[:-num_dropped] + def _finish_request(self, req_state: RequestState) -> None: req_id = req_state.request_id self.request_states.pop(req_id)