[Bugfix][Spec Decode] Trim token_ids/logprobs left past a stop string under speculative decoding (#47616)

Signed-off-by: Ting Sun <suntcrick@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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

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]
 

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)