← back to Exo
Fix bench after recent updates (#1331)
9dabde7e5784a93a8b64938ca75e9756632a7100 · 2026-01-29 19:14:40 +0000 · rltakashige
## Motivation
A lot of changes happened without much attention to the state of exo
bench.
## Changes
Use TaggedModel for BenchChatCompletion so it serialises properly.
Don't break after gpt oss tool call to preserve parity with the rest of
the codebase.
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
<img width="2856" height="678" alt="image"
src="https://github.com/user-attachments/assets/2e18cf0d-c0f8-467c-9763-1a6a59c8a327"
/>
Also tested GPT OSS tool calling in OpenCode
Files touched
M src/exo/shared/types/api.pyM src/exo/shared/types/commands.pyM src/exo/shared/types/tasks.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/runner/runner.pyA tmp/config_examples/opencode.json
Diff
commit 9dabde7e5784a93a8b64938ca75e9756632a7100
Author: rltakashige <rl.takashige@gmail.com>
Date: Thu Jan 29 19:14:40 2026 +0000
Fix bench after recent updates (#1331)
## Motivation
A lot of changes happened without much attention to the state of exo
bench.
## Changes
Use TaggedModel for BenchChatCompletion so it serialises properly.
Don't break after gpt oss tool call to preserve parity with the rest of
the codebase.
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
<img width="2856" height="678" alt="image"
src="https://github.com/user-attachments/assets/2e18cf0d-c0f8-467c-9763-1a6a59c8a327"
/>
Also tested GPT OSS tool calling in OpenCode
---
src/exo/shared/types/api.py | 8 +++++---
src/exo/shared/types/commands.py | 3 ++-
src/exo/shared/types/tasks.py | 3 ++-
src/exo/worker/engines/mlx/generator/generate.py | 14 +++++++++-----
src/exo/worker/runner/runner.py | 1 -
tmp/config_examples/opencode.json | 18 ++++++++++++++++++
6 files changed, 36 insertions(+), 11 deletions(-)
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 4dd3fff6..1a758a3c 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -3,7 +3,7 @@ from collections.abc import Generator
from typing import Annotated, Any, Literal
from fastapi import UploadFile
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic_core import PydanticUseDefault
from exo.shared.models.model_cards import ModelCard, ModelId
@@ -11,7 +11,7 @@ from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel
+from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
FinishReason = Literal[
"stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -170,7 +170,9 @@ class BenchChatCompletionResponse(ChatCompletionResponse):
generation_stats: GenerationStats | None = None
-class ChatCompletionTaskParams(BaseModel):
+class ChatCompletionTaskParams(TaggedModel):
+ model_config = ConfigDict(extra="ignore")
+
model: str
frequency_penalty: float | None = None
messages: list[ChatCompletionMessage]
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index ed086f35..a6d4d5c1 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -2,6 +2,7 @@ from pydantic import Field
from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.api import (
+ BenchChatCompletionTaskParams,
ChatCompletionTaskParams,
ImageEditsInternalParams,
ImageGenerationTaskParams,
@@ -22,7 +23,7 @@ class TestCommand(BaseCommand):
class ChatCompletion(BaseCommand):
- request_params: ChatCompletionTaskParams
+ request_params: ChatCompletionTaskParams | BenchChatCompletionTaskParams
class ImageGeneration(BaseCommand):
diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py
index 2c58087e..4f61a59a 100644
--- a/src/exo/shared/types/tasks.py
+++ b/src/exo/shared/types/tasks.py
@@ -3,6 +3,7 @@ from enum import Enum
from pydantic import Field
from exo.shared.types.api import (
+ BenchChatCompletionTaskParams,
ChatCompletionTaskParams,
ImageEditsInternalParams,
ImageGenerationTaskParams,
@@ -54,7 +55,7 @@ class StartWarmup(BaseTask): # emitted by Worker
class ChatCompletion(BaseTask): # emitted by Master
command_id: CommandId
- task_params: ChatCompletionTaskParams
+ task_params: ChatCompletionTaskParams | BenchChatCompletionTaskParams
error_type: str | None = Field(default=None)
error_message: str | None = Field(default=None)
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index b09b384f..56327a0f 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -39,7 +39,7 @@ def prefill(
sampler: Callable[[mx.array], mx.array],
prompt_tokens: mx.array,
cache: KVCacheType,
-) -> float:
+) -> tuple[float, int]:
"""Prefill the KV cache with prompt tokens.
This runs the model over the prompt tokens to populate the cache,
@@ -50,7 +50,7 @@ def prefill(
"""
num_tokens = len(prompt_tokens)
if num_tokens == 0:
- return 0.0
+ return 0.0, 0
logger.debug(f"Prefilling {num_tokens} tokens...")
start_time = time.perf_counter()
@@ -85,7 +85,7 @@ def prefill(
f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s "
f"({tokens_per_sec:.1f} tok/s)"
)
- return tokens_per_sec
+ return tokens_per_sec, num_tokens
def warmup_inference(
@@ -169,6 +169,8 @@ def mlx_generate(
mx.reset_peak_memory()
is_bench: bool = isinstance(task, BenchChatCompletionTaskParams)
+ logger.info(f"{is_bench=}")
+
# Currently we support chat-completion tasks only.
logger.debug(f"task_params: {task}")
@@ -204,7 +206,9 @@ def mlx_generate(
)
# Prefill cache with all tokens except the last one
- prefill_tps = prefill(model, tokenizer, sampler, prompt_tokens[:-1], caches)
+ prefill_tps, prefill_tokens = prefill(
+ model, tokenizer, sampler, prompt_tokens[:-1], caches
+ )
# stream_generate starts from the last token
last_token = prompt_tokens[-1:]
@@ -233,7 +237,7 @@ def mlx_generate(
stats = GenerationStats(
prompt_tps=float(prefill_tps or out.prompt_tps),
generation_tps=float(out.generation_tps),
- prompt_tokens=int(out.prompt_tokens),
+ prompt_tokens=int(prefill_tokens + out.prompt_tokens),
generation_tokens=int(out.generation_tokens),
peak_memory_usage=Memory.from_gb(out.peak_memory),
)
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 764cdaa8..937aa7a9 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -538,7 +538,6 @@ def parse_gpt_oss(
]
)
tool_arg_parts = []
- break
current_tool_name = recipient
# If inside a tool call, accumulate arguments
diff --git a/tmp/config_examples/opencode.json b/tmp/config_examples/opencode.json
new file mode 100644
index 00000000..9553934b
--- /dev/null
+++ b/tmp/config_examples/opencode.json
@@ -0,0 +1,18 @@
+{
+ "$schema": "https://opencode.ai/config.json",
+ "model": "exo/mlx-community/gpt-oss-120b-MXFP4-Q8",
+ "provider": {
+ "exo": {
+ "api": "http://localhost:52415/v1",
+ "models": {
+ "mlx-community/gpt-oss-120b-MXFP4-Q8": {
+ "name": "GPT OSS 120B",
+ "limit": {
+ "context": 32768,
+ "output": 8192
+ }
+ }
+ }
+ }
+ }
+}
← a31942ce Ciaran/image non streaming (#1328)
·
back to Exo
·
Add usage stats (#1333) d9eca758 →