Report the matched stop sequence on GenerationBatch.Response

[mloiterman]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: The code change and the test were written with Claude Code under my direction. The first draft of this description was also written by Claude Code; I reviewed and edited every line, evaluated the test results on both sides of the fix, and am responsible for the result.

Summary

StopSequenceMatcher builds an Aho-Corasick trie whose terminal nodes
store the sequence that matched (_build_trie:
node["__match__"] = (tuple(seq), idx)), but match() reduces that to
a boolean and GenerationBatch.next reports only finish_reason = "stop":

@staticmethod
def match(state, trie, x):
    """Advance by one token. Returns (new_state, matched)."""
    node = _step_trie(state, trie, x)
    return node, node.get("__match__") is not None

A multi-token stop sequence matches only on its final token. The earlier
tokens were already yielded on previous steps as ordinary, non-terminal
responses.

A BatchGenerator consumer that streams therefore cannot tell how many
of its already-emitted tokens belong to the stop, or which stop sequence fired.
The existing upstream text-level path keeps this information for itself —
TextStateMachine.step uses len(match[0]) to find where a match
starts in its buffer. However, the token-level path discards it before any
caller can see it. There is no reason the two paths should not expose the
same information.

Response.match_sequence existed before #1501 (on the old
SequenceStateMachine) and was dropped along with that class.

Impact (downstream)

This was encountered in a production server built directly on
BatchGenerator (not mlx_lm.server) that works on tokens. Our server
holds back a fixed number of the most recent tokens and trims the
client's stop string at the stop terminal. mlx_lm.server does not see
this because it works only on decoded text. It routes every chunk
through TextStateMachine with the stop words as transitions. The state
machine holds back anything that might be the start of a stop word, and
the server calls discard on stop.

A token-level consumer has no exact trim length without the field. One
alternative is to re-decode the held-back text and search it for each
stop string, which is only a guess when the same string tokenizes
differently.

The other is to run its own copy of the trie in lockstep with this one,
which duplicates library internals in application code and remains
exact only while it stays a faithful duplicate of upstream's code.

Diff

Additive. Matcher.matched exposes the sequence completed at the
current position; Response.match_sequence defaults to None, so no
existing constructor call or consumer changes, and advance()'s
signature is untouched.

--- a/mlx_lm/generate.py
+++ b/mlx_lm/generate.py
@@ -918,6 +918,12 @@ class StopSequences:
             self._node = node
             return node.get("__match__") is not None
 
+        @property
+        def matched(self) -> Optional[Tuple[int, ...]]:
+            """The stop sequence completed at the current position, if any."""
+            match = self._node.get("__match__")
+            return match[0] if match is not None else None
+
     def __init__(self, stop_sequences: Optional[Sequence[Sequence[int]]] = None):

@@ -1281,6 +1287,8 @@ class GenerationBatch:
         finish_reason: Optional[str]
         prompt_cache: Optional[List[Any]]
         all_tokens: Optional[List[int]]
+        # The stop sequence that matched, when finish_reason == "stop".
+        match_sequence: Optional[Tuple[int, ...]] = None

@@ -1457,8 +1465,10 @@ class GenerationBatch:
             if self._num_tokens[i] >= self.max_tokens[i]:
                 finish_reason = "length"
 
+            match_sequence = None
             if self._matchers[i].advance(tokens[i]):
                 finish_reason = "stop"
+                match_sequence = self._matchers[i].matched

@@ -1469,6 +1479,7 @@ class GenerationBatch:
                         finish_reason=finish_reason,
                         prompt_cache=self.extract_cache(i),
                         all_tokens=self.tokens[i],
+                        match_sequence=match_sequence,
                     )

__match__[0] is already a tuple, so it is passed through as is. A
"length" finish and every non-terminal response carry None.

Test

test_batch_generate_stop_match_sequence, next to
test_batch_generate_with_stop_sequences and using the same logit-bias
device: three sequences forced to emit token 0, with stops [0],
[0, 0] and [1]. Step 1: (0,) on the single-token stop, None on
the in-progress and non-matching slots. Step 2: (0, 0) on the
two-token stop, None on the slot still running. Fails on main with
AttributeError: 'Response' object has no attribute 'match_sequence';
passes with the fix, as does the existing stop-matcher test.