← back to Exo
Fix reasoning_tokens counting for multi-token thinking tag models (#1848)
e2e17eafb73f64cfe8088c2007d272a72fd0835e · 2026-04-09 12:29:35 +0100 · ciaranbor
## Motivation
`reasoning_tokens` is always 0 in usage stats, even when thinking
content streams correctly via `reasoning_content` SSE deltas. The MLX
generators had their own thinking detection comparing individual
detokenized tokens against think tags — this never fires for models
where tags span multiple tokens (e.g. gpt-oss-120b) or are already in
the prompt.
## Changes
- Removed broken per-token thinking detection from `batch_generate.py`
and `generate.py`
- Added `_count_reasoning_tokens` wrapper in `model_output_parsers.py`
that counts `is_thinking=True` responses and patches the total into
Usage on the final response
- Wired it as the outermost stage of `apply_all_parsers`, so it works
regardless of which parser sets `is_thinking`
- Added 3 tests covering `parse_thinking_models` and `parse_gpt_oss`
paths
## Why It Works
The parser pipeline already correctly sets `is_thinking` on each
response. Counting at the output of `apply_all_parsers` means one
counting point that works for all model types, replacing the duplicate
broken logic in two generators.
## Test Plan
### Manual Testing
- 4-node cluster, `mlx-community/gpt-oss-120b-MXFP4-Q8`
- Main branch: `reasoning_tokens: 0` — fix branch: `reasoning_tokens:
25`
### Automated Testing
- 3 new tests: explicit think tags, `starts_in_thinking=True`, and
gpt-oss Harmony analysis channel
Files touched
M src/exo/worker/engines/mlx/generator/batch_generate.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/runner/llm_inference/model_output_parsers.pyM src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.pyM src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
Diff
commit e2e17eafb73f64cfe8088c2007d272a72fd0835e
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Thu Apr 9 12:29:35 2026 +0100
Fix reasoning_tokens counting for multi-token thinking tag models (#1848)
## Motivation
`reasoning_tokens` is always 0 in usage stats, even when thinking
content streams correctly via `reasoning_content` SSE deltas. The MLX
generators had their own thinking detection comparing individual
detokenized tokens against think tags — this never fires for models
where tags span multiple tokens (e.g. gpt-oss-120b) or are already in
the prompt.
## Changes
- Removed broken per-token thinking detection from `batch_generate.py`
and `generate.py`
- Added `_count_reasoning_tokens` wrapper in `model_output_parsers.py`
that counts `is_thinking=True` responses and patches the total into
Usage on the final response
- Wired it as the outermost stage of `apply_all_parsers`, so it works
regardless of which parser sets `is_thinking`
- Added 3 tests covering `parse_thinking_models` and `parse_gpt_oss`
paths
## Why It Works
The parser pipeline already correctly sets `is_thinking` on each
response. Counting at the output of `apply_all_parsers` means one
counting point that works for all model types, replacing the duplicate
broken logic in two generators.
## Test Plan
### Manual Testing
- 4-node cluster, `mlx-community/gpt-oss-120b-MXFP4-Q8`
- Main branch: `reasoning_tokens: 0` — fix branch: `reasoning_tokens:
25`
### Automated Testing
- 3 new tests: explicit think tags, `starts_in_thinking=True`, and
gpt-oss Harmony analysis channel
---
.../worker/engines/mlx/generator/batch_generate.py | 15 +----
src/exo/worker/engines/mlx/generator/generate.py | 17 +----
.../runner/llm_inference/model_output_parsers.py | 60 ++++++++++++------
.../test_runner/test_finish_reason_sse.py | 73 ++++++++++++++++++++++
.../unittests/test_runner/test_parse_gpt_oss.py | 55 +++++++++++++++-
5 files changed, 168 insertions(+), 52 deletions(-)
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index 42e02eea..1d47ad40 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -45,7 +45,6 @@ from exo.worker.engines.mlx.patches.opt_batch_gen import (
take_ready_topk,
)
from exo.worker.engines.mlx.utils_mlx import (
- detect_thinking_prompt_suffix,
fix_unmatched_think_end_tokens,
system_prompt_token_count,
)
@@ -82,8 +81,6 @@ class _EngineTask:
potential_stop_sequence_text: str = ""
completion_tokens: int = 0
generation_start_time: float = 0.0
- in_thinking: bool = False
- reasoning_tokens: int = 0
prefill_tps: float = 0.0
media_regions: list[MediaRegion] = field(default_factory=list)
first_gen_token_time: float | None = None
@@ -273,7 +270,6 @@ class ExoBatchGenerator:
generation_start_time=time.perf_counter(),
prefill_tps=_prefill_tps,
media_regions=media_regions,
- in_thinking=detect_thinking_prompt_suffix(prompt, self.tokenizer),
)
return uid
@@ -323,15 +319,6 @@ class ExoBatchGenerator:
state.generated_text_parts.append(text)
state.potential_stop_sequence_text += text
- think_start = self.tokenizer.think_start
- think_end = self.tokenizer.think_end
- if think_start is not None and text == think_start:
- state.in_thinking = True
- elif think_end is not None and text == think_end:
- state.in_thinking = False
- if state.in_thinking:
- state.reasoning_tokens += 1
-
finish_reason: FinishReason | None = cast(
FinishReason | None, response.finish_reason
)
@@ -402,7 +389,7 @@ class ExoBatchGenerator:
cached_tokens=state.prefix_hit_length
),
completion_tokens_details=CompletionTokensDetails(
- reasoning_tokens=state.reasoning_tokens
+ reasoning_tokens=0
),
)
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index eef11101..43a38bca 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -53,7 +53,6 @@ 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,
@@ -597,11 +596,6 @@ def mlx_generate(
generated_text_parts: list[str] = []
generation_start_time = time.perf_counter()
usage: Usage | None = None
- in_thinking = detect_thinking_prompt_suffix(prompt, tokenizer)
- reasoning_tokens = 0
- think_start = tokenizer.think_start
- think_end = tokenizer.think_end
-
logger.info("Starting decode")
mx_barrier(group)
@@ -623,13 +617,6 @@ def mlx_generate(
generated_text_parts.append(out.text)
accumulated_text += out.text
- if think_start is not None and out.text == think_start:
- in_thinking = True
- elif think_end is not None and out.text == think_end:
- in_thinking = False
- if in_thinking:
- reasoning_tokens += 1
-
# Check for stop sequences
text = out.text
finish_reason: FinishReason | None = cast(
@@ -673,9 +660,7 @@ def mlx_generate(
prompt_tokens_details=PromptTokensDetails(
cached_tokens=prefix_hit_length
),
- completion_tokens_details=CompletionTokensDetails(
- reasoning_tokens=reasoning_tokens
- ),
+ completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
)
# Extract logprobs from the full vocabulary logprobs array
diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py
index 1242909a..9f380a77 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py
@@ -30,6 +30,32 @@ def get_gpt_oss_encoding():
return encoding
+def count_reasoning_tokens(
+ responses: Generator[GenerationResponse | ToolCallResponse | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+ """Count tokens with is_thinking=True and patch the total into Usage on the final response."""
+ reasoning_tokens = 0
+ for response in responses:
+ if response is None:
+ yield None
+ continue
+ if isinstance(response, GenerationResponse) and response.is_thinking:
+ reasoning_tokens += 1
+ if response.usage is not None and reasoning_tokens > 0:
+ response = response.model_copy(
+ update={
+ "usage": response.usage.model_copy(
+ update={
+ "completion_tokens_details": response.usage.completion_tokens_details.model_copy(
+ update={"reasoning_tokens": reasoning_tokens}
+ )
+ }
+ )
+ }
+ )
+ yield response
+
+
def apply_all_parsers(
receiver: Generator[GenerationResponse | None],
prompt: str,
@@ -41,14 +67,6 @@ def apply_all_parsers(
) -> Generator[GenerationResponse | ToolCallResponse | None]:
mlx_generator = receiver
- if tokenizer.has_thinking:
- mlx_generator = parse_thinking_models(
- mlx_generator,
- tokenizer.think_start,
- tokenizer.think_end,
- starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
- )
-
if issubclass(model_type, GptOssModel):
mlx_generator = parse_gpt_oss(mlx_generator)
elif (
@@ -56,10 +74,19 @@ def apply_all_parsers(
and "deepseek" in model_id.normalize().lower()
):
mlx_generator = parse_deepseek_v32(mlx_generator)
- elif tool_parser:
- mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
+ else:
+ if tokenizer.has_thinking:
+ mlx_generator = parse_thinking_models(
+ mlx_generator,
+ tokenizer.think_start,
+ tokenizer.think_end,
+ starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
+ )
+
+ if tool_parser:
+ mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
- return mlx_generator
+ return count_reasoning_tokens(mlx_generator)
def parse_gpt_oss(
@@ -67,7 +94,6 @@ def parse_gpt_oss(
) -> Generator[GenerationResponse | ToolCallResponse | None]:
encoding = get_gpt_oss_encoding()
stream = StreamableParser(encoding, role=Role.ASSISTANT)
- thinking = False
current_tool_name: str | None = None
tool_arg_parts: list[str] = []
@@ -121,14 +147,10 @@ def parse_gpt_oss(
tool_arg_parts = []
continue
- if ch == "analysis" and not thinking:
- thinking = True
-
- if ch != "analysis" and thinking:
- thinking = False
-
if delta:
- yield response.model_copy(update={"text": delta, "is_thinking": thinking})
+ yield response.model_copy(
+ update={"text": delta, "is_thinking": ch == "analysis"}
+ )
if response.finish_reason is not None:
yield response
diff --git a/src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py b/src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py
index 907ccddb..0aee03ba 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py
@@ -1,6 +1,7 @@
from collections.abc import Generator
from typing import Any
+from exo.api.types import CompletionTokensDetails, PromptTokensDetails, Usage
from exo.shared.types.worker.runner_response import (
FinishReason,
GenerationResponse,
@@ -14,6 +15,7 @@ from exo.worker.engines.mlx.dsml_encoding import (
TOOL_CALLS_START,
)
from exo.worker.runner.llm_inference.model_output_parsers import (
+ count_reasoning_tokens,
parse_deepseek_v32,
parse_thinking_models,
parse_tool_calls,
@@ -243,6 +245,77 @@ class TestThinkingModelsFinishReason:
)
assert _got_finish(results)
+ def test_reasoning_tokens_counted(self):
+ """reasoning_tokens in Usage reflects the number of thinking tokens."""
+ usage = Usage(
+ prompt_tokens=10,
+ completion_tokens=4,
+ total_tokens=14,
+ prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+ completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+ )
+ tokens = [
+ _make_response("<think>", 0),
+ _make_response("let me", 1),
+ _make_response(" think", 2),
+ _make_response("</think>", 3),
+ GenerationResponse(text="42", token=4, finish_reason="stop", usage=usage),
+ ]
+ results = _step_until_finish(
+ count_reasoning_tokens(
+ parse_thinking_models(
+ _queue_source(tokens),
+ think_start="<think>",
+ think_end="</think>",
+ starts_in_thinking=False,
+ )
+ )
+ )
+ final = [
+ r
+ for r in results
+ if isinstance(r, GenerationResponse) and r.finish_reason is not None
+ ]
+ assert len(final) == 1
+ assert final[0].usage is not None
+ assert final[0].usage.completion_tokens_details.reasoning_tokens == 2
+
+ def test_reasoning_tokens_starts_in_thinking(self):
+ """reasoning_tokens counts correctly when starts_in_thinking=True."""
+ usage = Usage(
+ prompt_tokens=10,
+ completion_tokens=3,
+ total_tokens=13,
+ prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+ completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+ )
+ tokens = [
+ _make_response("hmm", 0),
+ _make_response("ok", 1),
+ _make_response("</think>", 2),
+ GenerationResponse(
+ text="answer", token=3, finish_reason="stop", usage=usage
+ ),
+ ]
+ results = _step_until_finish(
+ count_reasoning_tokens(
+ parse_thinking_models(
+ _queue_source(tokens),
+ think_start="<think>",
+ think_end="</think>",
+ starts_in_thinking=True,
+ )
+ )
+ )
+ final = [
+ r
+ for r in results
+ if isinstance(r, GenerationResponse) and r.finish_reason is not None
+ ]
+ assert len(final) == 1
+ assert final[0].usage is not None
+ assert final[0].usage.completion_tokens_details.reasoning_tokens == 2
+
# ── parse_tool_calls (generic) ──────────────────────────────────
diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
index ba331336..28bbb7c8 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
@@ -1,11 +1,19 @@
from collections.abc import Generator
-from exo.api.types import FinishReason
+from exo.api.types import (
+ CompletionTokensDetails,
+ FinishReason,
+ PromptTokensDetails,
+ Usage,
+)
from exo.shared.types.worker.runner_response import (
GenerationResponse,
ToolCallResponse,
)
-from exo.worker.runner.llm_inference.model_output_parsers import parse_gpt_oss
+from exo.worker.runner.llm_inference.model_output_parsers import (
+ count_reasoning_tokens,
+ parse_gpt_oss,
+)
# Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer.
# These are stable since they come from the model's vocabulary.
@@ -85,6 +93,7 @@ THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [
def _make_gen_responses(
tokens: list[tuple[int, str]],
last_finish_reason: FinishReason = "stop",
+ last_usage: Usage | None = None,
) -> list[GenerationResponse]:
"""Build GenerationResponse list from (token_id, text) pairs."""
responses: list[GenerationResponse] = []
@@ -95,7 +104,7 @@ def _make_gen_responses(
text=text,
token=tid,
finish_reason=last_finish_reason if is_last else None,
- usage=None,
+ usage=last_usage if is_last else None,
)
)
return responses
@@ -251,3 +260,43 @@ class TestParseGptOssMaxTokensTruncation:
# due to Harmony encoding, so we just check something was emitted)
all_text = "".join(r.text for r in gen_responses)
assert len(all_text) > 0
+
+
+class TestGptOssReasoningTokensCounted:
+ """count_reasoning_tokens must patch Usage when parse_gpt_oss emits thinking tokens."""
+
+ def test_thinking_then_text_counts_reasoning_tokens(self):
+ usage = Usage(
+ prompt_tokens=10,
+ completion_tokens=len(PLAIN_TEXT_TOKENS),
+ total_tokens=10 + len(PLAIN_TEXT_TOKENS),
+ prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+ completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+ )
+ responses = _make_gen_responses(PLAIN_TEXT_TOKENS, last_usage=usage)
+
+ def _gen() -> Generator[GenerationResponse, None, None]:
+ yield from responses
+
+ results = list(
+ x for x in count_reasoning_tokens(parse_gpt_oss(_gen())) if x is not None
+ )
+
+ # Verify thinking tokens were detected
+ thinking = [
+ r for r in results if isinstance(r, GenerationResponse) and r.is_thinking
+ ]
+ assert len(thinking) > 0
+
+ # Verify reasoning_tokens is patched on responses that carry Usage
+ with_usage = [
+ r
+ for r in results
+ if isinstance(r, GenerationResponse) and r.usage is not None
+ ]
+ assert len(with_usage) > 0
+ assert all(
+ r.usage is not None
+ and r.usage.completion_tokens_details.reasoning_tokens == len(thinking)
+ for r in with_usage
+ )
← b12cd1b1 Cancel SSE keep-alive when instance is deleted (#1828)
·
back to Exo
·
prevent some crash loops (#1827) f2e6b1ef →