Fix stop token dropping valid tokens before it

When <|im_end|> appeared in accepted drafts, all preceding tokens in
the cycle were returned with finish_reason="stop", causing exo to
drop them (exo skips adding tokens with finish_reason="stop").

Symptom: γ=0 outputs "20", γ=1 outputs "2", γ=2 outputs nothing —
losing exactly γ tokens at the end.

Fix: yield tokens before the stop normally (no finish_reason), buffer
the stop token, let _yield_buffered return it with finish_reason="stop".

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
dmcc73
2026-03-30 18:44:01 +01:00
parent 199a4ab7e0
commit 75932cbcca
@@ -245,27 +245,44 @@ class MTPBatchGenerator(BatchGenerator):
(batch.tokens[0], mx.array([t for t, _ in all_tokens[1:]])))
batch.num_tokens[0] += len(all_tokens)
# 10. Check stop conditions
# 10. Check stop conditions — truncate at stop token
toc = time.perf_counter()
self._stats.generation_time += toc - tic
self._stats.generation_tokens += len(all_tokens)
finish_reason = None
for tok, _ in all_tokens:
# Find first stop token or length limit in all_tokens
stop_idx = None
for idx, (tok, _) in enumerate(all_tokens):
if tok in self.stop_tokens:
finish_reason = "stop"
stop_idx = idx
break
if batch.num_tokens[0] >= batch.max_tokens[0]:
finish_reason = "length"
stop_idx = idx
break
first_tok, first_lp = all_tokens[0]
if finish_reason:
cache = batch.extract_cache(0)
self.active_batch = None
self._cleanup_uid(uid)
return [self.Response(uid, first_tok, first_lp, finish_reason, cache)]
if stop_idx is not None:
# Tokens before the stop are valid output — buffer them
# The stop token itself triggers finish_reason
valid_tokens = all_tokens[:stop_idx]
if valid_tokens:
# Yield first, buffer rest + a final stop entry
if len(valid_tokens) > 1:
self._token_buffer[uid] = valid_tokens[1:]
# Append stop marker as last buffered token
stop_tok, stop_lp = all_tokens[stop_idx]
if uid not in self._token_buffer:
self._token_buffer[uid] = []
self._token_buffer[uid].append((stop_tok, stop_lp))
mx.async_eval(batch.y)
return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
else:
# Stop token is the first token — finish immediately
cache = batch.extract_cache(0)
self.active_batch = None
self._cleanup_uid(uid)
return [self.Response(uid, first_tok, first_lp, "stop", cache)]
# Buffer remaining tokens
if len(all_tokens) > 1: