Stop sequence matching mishandles overlapping prefixes and buffered tokens
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
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.
Source: Lightning-AI/litgpt