From 75932cbcca333f27f4e7daecd9b613f3ba09d8dd Mon Sep 17 00:00:00 2001
From: dmcc73 <david_correia1998@hotmail.com>
Date: Mon, 30 Mar 2026 18:44:01 +0100
Subject: [PATCH] Fix stop token dropping valid tokens before it
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 .../mlx/speculative/mtp_batch_generator.py    | 37 ++++++++++++++-----
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py b/src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py
index 2e5dc8d4..fa1fe412 100644
--- a/src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py
+++ b/src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py
@@ -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: