← back to Exo
Prepend <think> tag to stream for thinking models like GLM-4.7 (#1186)
39f0ed6018d09682763caf7039818e54cafc73e3 · 2026-01-19 19:44:51 +0000 · Alex Cheema
## Motivation
For thinking models like GLM-4.7, the `<think>` tag is inserted by the
tokenizer's `apply_chat_template()` into the **prompt** (input). The
model generates tokens starting *after* this tag, so `<think>` never
appears in the streamed output. The frontend expects
`<think>...</think>` tags to extract and display thinking content.
**Log evidence:**
```
[gMASK]<sop><|system|>...<|user|>...<|assistant|><think>
```
The prompt ends with `<think>`, so the model generates content after it,
never returning the opening tag.
## Changes
- Added `detect_thinking_prompt_suffix()` helper function in
`utils_mlx.py` to detect if a prompt ends with `<think>` tag
- Added `parse_thinking_models()` generator wrapper in `runner.py` that
prepends the thinking tag to the output stream
- Modified the main generation loop to use the thinking wrapper for
non-GptOssModel models when a thinking prefix is detected
- Updated test mocks to handle the new `apply_chat_template` call
## Why It Works
The solution follows the same pattern as `parse_gpt_oss()` - a generator
wrapper that transforms the output stream. When the chat template ends
with `<think>`, we prepend this tag to the first generated token so the
frontend receives the complete `<think>...</think>` structure it
expects.
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
- Run exo: `uv run exo`
- Send a chat request to GLM-4.7:
```bash
curl http://localhost:52415/v1/chat/completions -H "Content-Type:
application/json" -d '{
"model": "mlx-community/GLM-4.7-8bit-gs32",
"messages": [{"role": "user", "content": "What is 2+2?"}],
"stream": true
}'
```
- Verify the streamed response starts with `<think>` tag
- Verify the frontend dashboard correctly shows the thinking section
collapsed
### Automated Testing
- All 72 worker tests pass: `uv run pytest src/exo/worker/`
- Type checker passes: `uv run basedpyright`
- Linter passes: `uv run ruff check`
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
Files touched
M src/exo/worker/engines/mlx/__init__.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/runner.pyM src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
Diff
commit 39f0ed6018d09682763caf7039818e54cafc73e3
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Mon Jan 19 19:44:51 2026 +0000
Prepend <think> tag to stream for thinking models like GLM-4.7 (#1186)
## Motivation
For thinking models like GLM-4.7, the `<think>` tag is inserted by the
tokenizer's `apply_chat_template()` into the **prompt** (input). The
model generates tokens starting *after* this tag, so `<think>` never
appears in the streamed output. The frontend expects
`<think>...</think>` tags to extract and display thinking content.
**Log evidence:**
```
[gMASK]<sop><|system|>...<|user|>...<|assistant|><think>
```
The prompt ends with `<think>`, so the model generates content after it,
never returning the opening tag.
## Changes
- Added `detect_thinking_prompt_suffix()` helper function in
`utils_mlx.py` to detect if a prompt ends with `<think>` tag
- Added `parse_thinking_models()` generator wrapper in `runner.py` that
prepends the thinking tag to the output stream
- Modified the main generation loop to use the thinking wrapper for
non-GptOssModel models when a thinking prefix is detected
- Updated test mocks to handle the new `apply_chat_template` call
## Why It Works
The solution follows the same pattern as `parse_gpt_oss()` - a generator
wrapper that transforms the output stream. When the chat template ends
with `<think>`, we prepend this tag to the first generated token so the
frontend receives the complete `<think>...</think>` structure it
expects.
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
- Run exo: `uv run exo`
- Send a chat request to GLM-4.7:
```bash
curl http://localhost:52415/v1/chat/completions -H "Content-Type:
application/json" -d '{
"model": "mlx-community/GLM-4.7-8bit-gs32",
"messages": [{"role": "user", "content": "What is 2+2?"}],
"stream": true
}'
```
- Verify the streamed response starts with `<think>` tag
- Verify the frontend dashboard correctly shows the thinking section
collapsed
### Automated Testing
- All 72 worker tests pass: `uv run pytest src/exo/worker/`
- Type checker passes: `uv run basedpyright`
- Linter passes: `uv run ruff check`
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
---
src/exo/worker/engines/mlx/__init__.py | 26 ----------------
src/exo/worker/engines/mlx/generator/generate.py | 6 +---
src/exo/worker/engines/mlx/utils_mlx.py | 10 ++++++
src/exo/worker/runner/runner.py | 36 ++++++++++++++++++++++
.../unittests/test_runner/test_event_ordering.py | 4 +++
5 files changed, 51 insertions(+), 31 deletions(-)
diff --git a/src/exo/worker/engines/mlx/__init__.py b/src/exo/worker/engines/mlx/__init__.py
index d6f0b6b3..740b0eff 100644
--- a/src/exo/worker/engines/mlx/__init__.py
+++ b/src/exo/worker/engines/mlx/__init__.py
@@ -1,5 +1,3 @@
-from typing import Any
-
import mlx.core as mx
import mlx.nn as nn
from mlx_lm.models.cache import KVCache
@@ -17,27 +15,3 @@ class Model(nn.Module):
cache: list[KVCache] | None,
input_embeddings: mx.array | None = None,
) -> mx.array: ...
-
-
-class Detokenizer:
- def reset(self) -> None: ...
- def add_token(self, token: int) -> None: ...
- def finalize(self) -> None: ...
-
- @property
- def last_segment(self) -> str: ...
-
-
-class TokenizerWrapper:
- bos_token: str | None
- eos_token_ids: list[int]
- detokenizer: Detokenizer
-
- def encode(self, text: str, add_special_tokens: bool = True) -> list[int]: ...
-
- def apply_chat_template(
- self,
- messages_dicts: list[dict[str, Any]],
- tokenize: bool = False,
- add_generation_prompt: bool = True,
- ) -> str: ...
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index c42bf5d4..2482cc0c 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -119,6 +119,7 @@ def mlx_generate(
model: Model,
tokenizer: TokenizerWrapper,
task: ChatCompletionTaskParams,
+ prompt: str,
) -> Generator[GenerationResponse]:
# Ensure that generation stats only contains peak memory for this generation
mx.reset_peak_memory()
@@ -130,11 +131,6 @@ def mlx_generate(
if task.seed is not None:
mx.random.seed(task.seed)
- prompt = apply_chat_template(
- tokenizer=tokenizer,
- chat_task_data=task,
- )
-
caches = make_kv_cache(model=model)
logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 258f7e0d..c7b6be9f 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -431,6 +431,16 @@ def apply_chat_template(
return prompt
+def detect_thinking_prompt_suffix(prompt: str, tokenizer: TokenizerWrapper) -> bool:
+ """
+ Detect if prompt ends with a thinking opening tag that should be
+ prepended to the output stream.
+ """
+ think_token = tokenizer.think_start
+
+ return think_token is not None and prompt.rstrip().endswith(think_token)
+
+
class NullKVCache(KVCache):
"""
A KVCache that pretends to exist but holds zero tokens.
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 2b4e5b49..792b9f30 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -4,6 +4,7 @@ from functools import cache
import mlx.core as mx
from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm.tokenizer_utils import TokenizerWrapper
from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
HarmonyEncodingName,
Role,
@@ -50,6 +51,8 @@ from exo.shared.types.worker.runners import (
from exo.utils.channels import MpReceiver, MpSender
from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
from exo.worker.engines.mlx.utils_mlx import (
+ apply_chat_template,
+ detect_thinking_prompt_suffix,
initialize_mlx,
load_mlx_items,
mlx_force_oom,
@@ -177,17 +180,28 @@ def main(
try:
_check_for_debug_prompts(task_params.messages[0].content)
+ # Build prompt once - used for both generation and thinking detection
+ prompt = apply_chat_template(tokenizer, task_params)
+
# Generate responses using the actual MLX generation
mlx_generator = mlx_generate(
model=model,
tokenizer=tokenizer,
task=task_params,
+ prompt=prompt,
)
# GPT-OSS specific parsing to match other model formats.
if isinstance(model, GptOssModel):
mlx_generator = parse_gpt_oss(mlx_generator)
+ # For other thinking models (GLM, etc.), check if we need to
+ # prepend the thinking tag that was consumed by the chat template
+ if detect_thinking_prompt_suffix(prompt, tokenizer):
+ mlx_generator = parse_thinking_models(
+ mlx_generator, tokenizer
+ )
+
# TODO: Add tool call parser here
for response in mlx_generator:
@@ -293,6 +307,28 @@ def parse_gpt_oss(
break
+def parse_thinking_models(
+ responses: Generator[GenerationResponse],
+ tokenizer: TokenizerWrapper,
+) -> Generator[GenerationResponse]:
+ """
+ For models that inject thinking tags in the prompt (like GLM-4.7),
+ prepend the thinking tag to the output stream so the frontend
+ can properly parse thinking content.
+ """
+ first = True
+ for response in responses:
+ if first:
+ first = False
+ yield response.model_copy(
+ update={
+ "text": tokenizer.think_start,
+ "token": tokenizer.think_start_id, # type: ignore
+ }
+ )
+ yield response
+
+
EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
index a8b4b1a5..9bf8ce4d 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
@@ -114,6 +114,10 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, 1)))
monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
+ # Mock apply_chat_template since we're using a fake tokenizer (integer 1).
+ # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None.
+ monkeypatch.setattr(mlx_runner, "apply_chat_template", make_nothin("test prompt"))
+ monkeypatch.setattr(mlx_runner, "detect_thinking_prompt_suffix", make_nothin(False))
def fake_generate(*_1: object, **_2: object):
yield GenerationResponse(token=0, text="hi", finish_reason="stop")
← ee43b598 Split NodePerformanceProfile into granular state mappings (#
·
back to Exo
·
Load layers individually (#1211) f5e6aa82 →