[object Object]

← back to Exo

Fix reasoning_tokens always reported as 0 for thinking models (#1836)

24420eb10a5dd4271410bd02c8b05f5592af2e1f · 2026-04-05 03:05:18 +0300 · mlpy0

When `enable_thinking` is set, chat templates (Qwen3, DeepSeek, etc.)
append `<think>` to the prompt. The model starts generating thinking
content directly without emitting a `<think>` token in the output
stream.

Both generators initialized `in_thinking = False` and only set it to
`True` on seeing a `<think>` token in output. Since that token was part
of the prompt, the flag never flipped and `reasoning_tokens` stayed at 0
in the usage response.

Fix: initialize `in_thinking` from `detect_thinking_prompt_suffix()`,
which already exists and is used by `model_output_parsers` for routing
thinking content correctly.

Files touched

Diff

commit 24420eb10a5dd4271410bd02c8b05f5592af2e1f
Author: mlpy0 <96022931+mlpy0@users.noreply.github.com>
Date:   Sun Apr 5 03:05:18 2026 +0300

    Fix reasoning_tokens always reported as 0 for thinking models (#1836)
    
    When `enable_thinking` is set, chat templates (Qwen3, DeepSeek, etc.)
    append `<think>` to the prompt. The model starts generating thinking
    content directly without emitting a `<think>` token in the output
    stream.
    
    Both generators initialized `in_thinking = False` and only set it to
    `True` on seeing a `<think>` token in output. Since that token was part
    of the prompt, the flag never flipped and `reasoning_tokens` stayed at 0
    in the usage response.
    
    Fix: initialize `in_thinking` from `detect_thinking_prompt_suffix()`,
    which already exists and is used by `model_output_parsers` for routing
    thinking content correctly.
---
 src/exo/worker/engines/mlx/generator/batch_generate.py | 2 ++
 src/exo/worker/engines/mlx/generator/generate.py       | 3 ++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index dba462aa..e06cab59 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -41,6 +41,7 @@ from exo.worker.engines.mlx.generator.generate import (
     prefill,
 )
 from exo.worker.engines.mlx.utils_mlx import (
+    detect_thinking_prompt_suffix,
     fix_unmatched_think_end_tokens,
     system_prompt_token_count,
 )
@@ -269,6 +270,7 @@ class ExoBatchGenerator:
             prefill_tps=_prefill_tps,
             generation_time_at_start=self._mlx_gen._stats.generation_time,
             media_regions=media_regions,
+            in_thinking=detect_thinking_prompt_suffix(prompt, self.tokenizer),
         )
 
         return uid
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index bb89dc7e..eef11101 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -53,6 +53,7 @@ from exo.worker.engines.mlx.constants import (
 )
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
+    detect_thinking_prompt_suffix,
     fix_unmatched_think_end_tokens,
     mx_barrier,
     system_prompt_token_count,
@@ -596,7 +597,7 @@ def mlx_generate(
     generated_text_parts: list[str] = []
     generation_start_time = time.perf_counter()
     usage: Usage | None = None
-    in_thinking = False
+    in_thinking = detect_thinking_prompt_suffix(prompt, tokenizer)
     reasoning_tokens = 0
     think_start = tokenizer.think_start
     think_end = tokenizer.think_end

← 59669c11 Tighten EXO bench concurrency numbers and explain methodolog  ·  back to Exo  ·  Fix BatchGenerator in line with upstream refactor (and preve 43b3df45 →