From 38f0c09175419bc2864affeeb77cfed44973ee65 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mustafa=20Alp=20Y=C4=B1lmaz?=
 <96022931+mustafalpyilmaz@users.noreply.github.com>
Date: Tue, 10 Mar 2026 19:10:22 +0300
Subject: [PATCH] fix: use StreamingDetokenizer in batch generator to fix
 emoji/UTF-8 corruption (#1691)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Problem

Emojis and other multi-byte UTF-8 characters are rendered as `\ufffd`
(Unicode Replacement Character) in batch streaming responses.

Byte-level BPE tokenizers (like Qwen's) can split multi-byte UTF-8
characters (e.g. a 4-byte emoji) across multiple tokens. The batch
generator decodes each token independently with
`tokenizer.decode([token_id])`, which produces U+FFFD for partial byte
sequences.

The sequential generator doesn't have this problem — it uses
`stream_generate()` from mlx_lm, which internally uses
`StreamingDetokenizer` to buffer incomplete bytes.

## Fix

Use `StreamingDetokenizer` in the batch generator, matching the
sequential path:

- Each `_EngineTask` gets its own `StreamingDetokenizer` instance
- `add_token()` buffers tokens, holding back incomplete UTF-8 byte
sequences until they form valid characters
- `last_segment` returns only complete, valid text
- `finalize()` flushes any remaining buffered bytes when generation
completes

Empty text segments during buffering are harmless — they're already
handled correctly by the downstream streaming pipeline.

---------

Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
---
 .mlx_typings/mlx_lm/tokenizer_utils.pyi            |  8 ++++----
 .../worker/engines/mlx/generator/batch_generate.py | 14 ++++++++------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/.mlx_typings/mlx_lm/tokenizer_utils.pyi b/.mlx_typings/mlx_lm/tokenizer_utils.pyi
index dd12af3b..1326f059 100644
--- a/.mlx_typings/mlx_lm/tokenizer_utils.pyi
+++ b/.mlx_typings/mlx_lm/tokenizer_utils.pyi
@@ -39,11 +39,11 @@ class StreamingDetokenizer:
     """
 
     __slots__ = ...
-    def reset(self): ...
-    def add_token(self, token): ...
-    def finalize(self): ...
+    def reset(self) -> None: ...
+    def add_token(self, token: int) -> None: ...
+    def finalize(self) -> None: ...
     @property
-    def last_segment(self):
+    def last_segment(self) -> str:
         """Return the last segment of readable text since last time this property was accessed."""
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index 3b72e0a7..1911a97c 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -8,7 +8,7 @@ from mlx_lm.generate import (
 )
 from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
-from mlx_lm.tokenizer_utils import TokenizerWrapper
+from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
 
 from exo.shared.types.api import (
     CompletionTokensDetails,
@@ -57,6 +57,7 @@ class _EngineTask:
     prefix_hit_length: int
     matched_index: int | None
     cache_snapshots: list[CacheSnapshot] | None
+    detokenizer: StreamingDetokenizer
     on_generation_token: Callable[[], None] | None = None
     generated_text_parts: list[str] = field(default_factory=list)
     potential_stop_sequence_text: str = ""
@@ -205,6 +206,7 @@ class ExoBatchGenerator:
             prefix_hit_length=prefix_hit_length,
             matched_index=matched_index,
             cache_snapshots=cache_snapshots or None,
+            detokenizer=self.tokenizer.detokenizer,
             on_generation_token=on_generation_token,
             generation_start_time=time.perf_counter(),
         )
@@ -229,11 +231,11 @@ class ExoBatchGenerator:
             state = self._active_tasks[response.uid]
             if state.on_generation_token is not None:
                 state.on_generation_token()
-            text = (
-                ""
-                if response.finish_reason == "stop"
-                else self.tokenizer.decode([response.token])
-            )
+            if response.finish_reason != "stop":
+                state.detokenizer.add_token(response.token)
+            if response.finish_reason is not None:
+                state.detokenizer.finalize()
+            text = state.detokenizer.last_segment
             state.completion_tokens += 1
             state.generated_text_parts.append(text)
             state.potential_stop_sequence_text += text