diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 145a0275dd0a..d9fe12404c5a 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -1023,9 +1023,19 @@ def _handle_response(self, response: "GenerationExecutor.Response"): is not None) and self.tokenizer is not None: for beam_output in self.outputs: beam_output._last_text_len = len(beam_output.text) - if hasattr( - self.tokenizer, 'decode_incrementally' - ) and self._streaming and not self.sampling_params.use_beam_search: + # A stop word matched by token id is trimmed from token_ids in + # _handle_sequence. If some of its ids came in earlier streamed + # responses, the incremental decoder has already turned them + # into text, or holds the lead bytes of a character that the + # final flush would emit as U+FFFD. Decode the kept ids afresh + # instead, as the non-streaming path does. + trimmed_streamed_ids = (len(beam_output.token_ids) + < beam_output._last_token_ids_len) + use_incremental = (self._streaming + and not self.sampling_params.use_beam_search + and not trimmed_streamed_ids and hasattr( + self.tokenizer, 'decode_incrementally')) + if use_incremental: beam_output.text, beam_output._incremental_states = self.tokenizer.decode_incrementally( beam_output.token_ids_diff, prev_text=beam_output.text, @@ -1036,6 +1046,13 @@ def _handle_response(self, response: "GenerationExecutor.Response"): else: beam_output.text = self.tokenizer.decode( beam_output.token_ids, **kwargs) + if trimmed_streamed_ids: + # Drop the pending flush, and move the watermark to the + # kept ids so that later responses for sibling + # candidates (n > 1) do not decode this output again. + beam_output._incremental_states = None + beam_output._last_token_ids_len = len( + beam_output.token_ids) is_generating = not self._done is_finished_with_stop_or_length = ( diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index edf634c48fcf..bd5c4ee2d27a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -206,6 +206,7 @@ l0_cpu: - unittest/executor/test_stats_serializer.py - unittest/executor/test_spec_dec_perf_metrics.py - unittest/executor/test_stop_word_side_output_trim.py + - unittest/executor/test_stop_word_streaming_text.py - unittest/executor/test_ray_stub.py - unittest/executor/test_generation_request_array_prompt.py - unittest/grpc/openengine/test_coordinator.py diff --git a/tests/unittest/executor/test_stop_word_streaming_text.py b/tests/unittest/executor/test_stop_word_streaming_text.py new file mode 100644 index 000000000000..d4104ac456fd --- /dev/null +++ b/tests/unittest/executor/test_stop_word_streaming_text.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A streamed request that stops on a multi-token stop word must end with the non-streamed text. + +When the sampler matches a stop word by token id, ``_handle_sequence`` trims +the stop ids from ``token_ids``. In a streamed request some of those ids may +have arrived in earlier responses and already been decoded incrementally, so +the final ``.text`` has to be decoded again from the kept ids: otherwise it +keeps the start of the stop string (``"42<0xB9> until the last byte completes the character. +VOCAB = { + "": 0, + "<0xE0>": 1, + "<0xB9>": 2, + "<0x91>": 3, + "<0x92>": 4, + "ข้อ": 5, + " ": 6, + "": 7, + "42": 8, + "": 11, +} +THAI_DIGIT_TWO = [1, 2, 4] +# "ข้อ ๑ ข้อ ๒" +THAI_GENERATED = [5, 6, 1, 2, 3, 6, 5, 6, *THAI_DIGIT_TWO] +ANSWER_CLOSE = [9, 10, 11] +# "42" +ANSWER_GENERATED = [7, 8, *ANSWER_CLOSE] + + +@pytest.fixture(scope="module") +def tokenizer() -> TransformersTokenizer: + backend = Tokenizer(WordLevel(VOCAB, unk_token="")) + backend.decoder = Sequence([ByteFallback(), Fuse()]) + return TransformersTokenizer(PreTrainedTokenizerFast(tokenizer_object=backend)) + + +@pytest.fixture(params=["HF", "TRTLLM"], autouse=True) +def detokenization_backend(request, monkeypatch: pytest.MonkeyPatch) -> str: + from tensorrt_llm.tokenizer import tokenizer as tokenizer_module + + monkeypatch.setattr(tokenizer_module, "TLLM_INCREMENTAL_DETOKENIZATION_BACKEND", request.param) + return request.param + + +def _response(token_ids, is_final, finish_reason, sequence_index=0): + result = SimpleNamespace( + is_final=is_final, + context_phase_params=None, + decoding_iter=1, + avg_decoded_tokens_per_iter=None, + finish_reasons=[finish_reason], + output_token_ids=[list(token_ids)], + sequence_index=sequence_index, + cum_log_probs=None, + log_probs=None, + generation_logits=None, + per_pos_drafted=None, + per_pos_accepted=None, + spec_dec_totals=None, + time_breakdown_metrics=None, + request_perf_metrics=None, + additional_context_outputs=None, + additional_generation_outputs=None, + context_logits=None, + ) + return SimpleNamespace(has_error=lambda: False, result=result) + + +def _sampling_params(stop, stop_ids, **kwargs): + sampling_params = SamplingParams(stop=stop, **kwargs) + # Stands in for SamplingParams._setup tokenizing the stop string. + sampling_params._stop_word_ids = [stop_ids] + sampling_params._stream_interval = 1 + return sampling_params + + +def _run(tokenizer, sampling_params, chunks, finish_reason, streaming=True): + """Feed one response per chunk; the last one finishes the request with finish_reason.""" + result = DetokenizedGenerationResultBase( + 0, sampling_params, tokenizer=tokenizer, streaming=streaming + ) + deltas = [] + for index, chunk in enumerate(chunks): + is_final = index == len(chunks) - 1 + result._handle_response( + _response( + chunk, is_final, finish_reason if is_final else tllm.FinishReason.NOT_FINISHED + ) + ) + deltas.append(result.outputs[0].text_diff) + return result.outputs[0], "".join(deltas) + + +def _chunked(token_ids, size): + return [token_ids[i : i + size] for i in range(0, len(token_ids), size)] + + +# The last field is the joined text_diff values. For "" they still end in +# "", + ANSWER_CLOSE, + ANSWER_GENERATED, + "42", + "42 1, a candidate that stopped on a trimmed stop word is not decoded again. + + Candidate 0 streams to its stop word first; candidate 1 then streams on and + each of its responses walks every output. Candidate 0 must keep the + non-streamed text and an empty text_diff without another full decode. + """ + sampling_params = _sampling_params( + "๒", THAI_DIGIT_TWO, n=2, best_of=2, temperature=0.8, top_p=0.9 + ) + result = DetokenizedGenerationResultBase( + 0, sampling_params, tokenizer=tokenizer, streaming=True + ) + full_decodes = [] + decode = tokenizer.decode + + def counting_decode(token_ids, *args, **kwargs): + full_decodes.append(list(token_ids)) + return decode(token_ids, *args, **kwargs) + + monkeypatch.setattr(tokenizer, "decode", counting_decode) + + sibling = [5, 6, 7, 8] + steps = [ + (0, [token_id], tllm.FinishReason.STOP_WORDS if index == len(THAI_GENERATED) - 1 else None) + for index, token_id in enumerate(THAI_GENERATED) + ] + [ + (1, [token_id], tllm.FinishReason.LENGTH if index == len(sibling) - 1 else None) + for index, token_id in enumerate(sibling) + ] + first, second = result._outputs + first_diffs_after_stop = [] + for step, (sequence_index, chunk, finish_reason) in enumerate(steps): + is_final = step == len(steps) - 1 + result._handle_response( + _response( + chunk, + is_final, + finish_reason or tllm.FinishReason.NOT_FINISHED, + sequence_index=sequence_index, + ) + ) + if sequence_index == 1: + first_diffs_after_stop.append(first.text_diff) + + assert first.text == "ข้อ ๑ ข้อ " + assert first.finish_reason == "stop" + assert first_diffs_after_stop == [""] * len(sibling) + assert full_decodes == [THAI_GENERATED[:-3]] + assert second.text == "ข้อ 42" + assert second.finish_reason == "length"