#2302·litgpt

Stop sequence matching mishandles overlapping prefixes and buffered tokens

Author: tandedeCreated Aug 19, 2026Updated Aug 19, 2026

Bug description

generate_fn() and batched_generate_fn() keep one prefix length for each stop sequence and reset that length to zero after any mismatch. This misses a valid match when the mismatching token is also the beginning of a new occurrence of the same sequence.

For example, with stop sequence [1, 2] and generated tokens [1, 1, 2], the second 1 should restart the match. Instead, the current matcher resets to zero and returns all three tokens rather than stopping before the final [1, 2].

There is a related buffering problem when multiple stop sequences are active. With stop sequences [1, 3] and [2], generated tokens [1, 2] should return [1]: the first token is temporarily buffered as a possible prefix, then becomes safe when [2] completes. The current implementation returns immediately on the completed sequence and drops the buffered 1 as well.

Minimal reproduction

python
from unittest.mock import MagicMock, patch

import torch
import litgpt.generate.base as generation
from litgpt.chat.base import generate


def run(generated, stop_tokens):
    model = MagicMock()
    model.max_seq_length = 100
    samples = iter(generated)
    with patch.object(
        generation,
        "multinomial_num_samples_1",
        lambda *args, **kwargs: torch.tensor([next(samples)]),
    ):
        output = list(
            generate(
                model,
                torch.tensor([5, 3]),
                len(generated) + 2,
                stop_tokens=stop_tokens,
            )
        )
    return torch.cat(output).tolist() if output else []


print(run([1, 1, 2], ([1, 2],)))       # actual: [1, 1, 2], expected: [1]
print(run([1, 2], ([1, 3], [2])))      # actual: [], expected: [1]

Expected behavior

Stop matching should retain the longest suffix that is also a prefix of each stop sequence. When a sequence completes, tokens buffered only for other partial matches should be emitted before generation stops. The result should not depend on the ordering of stop_tokens.

This affects both the single-prompt and batched streaming generators because they use the same reset-on-mismatch logic.