← back to Exo
Update exo bench (#1357)
21d477f1cbc50ae2b2370b330728eeb245bc434c · 2026-02-02 15:46:15 +0000 · rltakashige
## Motivation
Make exo bench faster for longer prompts, lengthen default timeouts and
use pairs for pp and tg.
## Changes
- Uses binary search to find correct prompt
- Flag to force all combinations if that is desired
Files touched
Diff
commit 21d477f1cbc50ae2b2370b330728eeb245bc434c
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Feb 2 15:46:15 2026 +0000
Update exo bench (#1357)
## Motivation
Make exo bench faster for longer prompts, lengthen default timeouts and
use pairs for pp and tg.
## Changes
- Uses binary search to find correct prompt
- Flag to force all combinations if that is desired
---
bench/exo_bench.py | 247 +++++++++++++++++++++++++++++++++++++----------------
1 file changed, 174 insertions(+), 73 deletions(-)
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index c8bbba5c..83390b0b 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -5,10 +5,13 @@ from __future__ import annotations
import argparse
import contextlib
import http.client
+import itertools
import json
import os
+import sys
import time
from collections.abc import Callable
+from pathlib import Path
from statistics import mean
from typing import Any
from urllib.parse import urlencode
@@ -16,6 +19,84 @@ from urllib.parse import urlencode
from loguru import logger
from transformers import AutoTokenizer
+# Monkey-patch for transformers 5.x compatibility
+# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
+# which was moved in transformers 5.0.0rc2
+try:
+ import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
+ from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+ if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
+ gpt2_tokenization.bytes_to_unicode = bytes_to_unicode # type: ignore[attr-defined]
+except ImportError:
+ pass # transformers < 5.0 or bytes_to_unicode not available
+
+
+def load_tokenizer_for_bench(model_id: str) -> Any:
+ """
+ Load tokenizer for benchmarking, with special handling for Kimi models.
+
+ Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer.
+ This function replicates the logic from utils_mlx.py for bench compatibility.
+ """
+ model_id_lower = model_id.lower()
+
+ if "kimi-k2" in model_id_lower:
+ import importlib.util
+ import types
+
+ from huggingface_hub import snapshot_download
+
+ # Download/get the model path
+ model_path = Path(
+ snapshot_download(
+ model_id,
+ allow_patterns=["*.json", "*.py", "*.tiktoken"],
+ )
+ )
+
+ sys.path.insert(0, str(model_path))
+
+ # Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
+ tool_decl_path = model_path / "tool_declaration_ts.py"
+ if tool_decl_path.exists():
+ spec = importlib.util.spec_from_file_location(
+ "tool_declaration_ts", tool_decl_path
+ )
+ if spec and spec.loader:
+ tool_decl_module = importlib.util.module_from_spec(spec)
+ sys.modules["tool_declaration_ts"] = tool_decl_module
+ spec.loader.exec_module(tool_decl_module)
+
+ # Load tokenization_kimi with patched source (convert relative to absolute import)
+ tok_path = model_path / "tokenization_kimi.py"
+ source = tok_path.read_text()
+ source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
+ spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
+ if spec:
+ tok_module = types.ModuleType("tokenization_kimi")
+ tok_module.__file__ = str(tok_path)
+ sys.modules["tokenization_kimi"] = tok_module
+ exec(compile(source, tok_path, "exec"), tok_module.__dict__) # noqa: S102
+ TikTokenTokenizer = tok_module.TikTokenTokenizer # noqa: N806
+ else:
+ from tokenization_kimi import TikTokenTokenizer # type: ignore[import-not-found] # noqa: I001
+
+ hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)
+
+ # Patch encode to use internal tiktoken model directly
+ # transformers 5.x has a bug in the encode->pad path for slow tokenizers
+ def _patched_encode(text: str, **kwargs: object) -> list[int]:
+ # Pass allowed_special="all" to handle special tokens like <|im_user|>
+ return list(hf_tokenizer.model.encode(text, allowed_special="all"))
+
+ hf_tokenizer.encode = _patched_encode
+
+ return hf_tokenizer
+
+ # Default: use AutoTokenizer
+ return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+
class ExoHttpError(RuntimeError):
def __init__(self, status: int, reason: str, body_preview: str):
@@ -24,7 +105,7 @@ class ExoHttpError(RuntimeError):
class ExoClient:
- def __init__(self, host: str, port: int, timeout_s: float = 600.0):
+ def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
self.host = host
self.port = port
self.timeout_s = timeout_s
@@ -180,14 +261,7 @@ def parse_int_list(values: list[str]) -> list[int]:
part = part.strip()
if part:
items.append(int(part))
-
- seen: set[int] = set()
- out: list[int] = []
- for x in items:
- if x not in seen:
- out.append(x)
- seen.add(x)
- return out
+ return items
def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
@@ -240,7 +314,11 @@ def run_one_completion(
stats = out.get("generation_stats")
- preview = (out.get("choices") or [{}])[0]["message"]["content"][:200]
+ # Extract preview, handling None content (common for thinking models)
+ choices = out.get("choices") or [{}]
+ message = choices[0].get("message", {}) if choices else {}
+ content = message.get("content") or ""
+ preview = content[:200] if content else ""
return {
"elapsed_s": elapsed,
@@ -277,12 +355,29 @@ class PromptSizer:
f"Target ({target}) is smaller than template overhead ({self.base_tokens})."
)
- content = ""
+ # Estimate tokens per atom using a sample
+ sample_count = 100
+ sample_content = self.atom * sample_count
+ sample_tokens = self.count_fn(sample_content) - self.base_tokens
+ tokens_per_atom = sample_tokens / sample_count
+
+ # Estimate starting point
+ needed_tokens = target - self.base_tokens
+ estimated_atoms = int(needed_tokens / tokens_per_atom)
+
+ # Binary search to find exact atom count
+ low, high = 0, estimated_atoms * 2 + 100
+ while low < high:
+ mid = (low + high) // 2
+ tok = self.count_fn(self.atom * mid)
+ if tok < target:
+ low = mid + 1
+ else:
+ high = mid
+
+ content = self.atom * low
tok = self.count_fn(content)
-
- while tok < target:
- content += self.atom
- tok = self.count_fn(content)
+ logger.info(f"{tok=}")
if tok != target:
raise RuntimeError(
@@ -348,7 +443,7 @@ def main() -> int:
help="Warmup runs per placement (uses first pp/tg).",
)
ap.add_argument(
- "--timeout", type=float, default=600.0, help="HTTP timeout (seconds)."
+ "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
)
ap.add_argument(
"--json-out",
@@ -358,6 +453,11 @@ def main() -> int:
ap.add_argument(
"--dry-run", action="store_true", help="List selected placements and exit."
)
+ ap.add_argument(
+ "--all-combinations",
+ action="store_true",
+ help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
+ )
args = ap.parse_args()
pp_list = parse_int_list(args.pp)
@@ -369,6 +469,15 @@ def main() -> int:
logger.error("--repeat must be >= 1")
return 2
+ # Log pairing mode
+ use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
+ if use_combinations:
+ logger.info(
+ f"pp/tg mode: combinations (product) - {len(pp_list) * len(tg_list)} pairs"
+ )
+ else:
+ logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs")
+
client = ExoClient(args.host, args.port, timeout_s=args.timeout)
short_id, full_model_id = resolve_model_short_id(client, args.model)
@@ -377,10 +486,7 @@ def main() -> int:
)
previews = previews_resp.get("previews") or []
- tokenizer = AutoTokenizer.from_pretrained(
- full_model_id,
- trust_remote_code=True,
- )
+ tokenizer = load_tokenizer_for_bench(full_model_id)
if tokenizer is None:
raise RuntimeError("[exo-bench] tokenizer load failed")
@@ -486,60 +592,55 @@ def main() -> int:
)
logger.debug(f" warmup {i + 1}/{args.warmup} done")
- for pp in pp_list:
- # if (
- # pp * n_nodes > 2048
- # and "ring" in instance_meta.lower()
- # and "tensor" in sharding.lower()
- # ):
- # model_card = MODEL_CARDS[short_id]
- # if model_card.metadata.storage_size > Memory.from_gb(10):
- # logger.info(
- # f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
- # )
- # continue
- for tg in tg_list:
- runs: list[dict[str, Any]] = []
- for r in range(args.repeat):
- time.sleep(3)
- try:
- row, actual_pp_tokens = run_one_completion(
- client, full_model_id, pp, tg, prompt_sizer
- )
- except Exception as e:
- logger.error(e)
- continue
- row.update(
- {
- "model_short_id": short_id,
- "model_id": full_model_id,
- "placement_sharding": sharding,
- "placement_instance_meta": instance_meta,
- "placement_nodes": n_nodes,
- "instance_id": instance_id,
- "pp_tokens": actual_pp_tokens,
- "tg": tg,
- "repeat_index": r,
- }
- )
- runs.append(row)
- all_rows.append(row)
-
- if runs:
- prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
- gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
- ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
- gtok = mean(x["stats"]["generation_tokens"] for x in runs)
- peak = mean(
- x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
- )
-
- logger.info(
- f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
- f"prompt_tokens={ptok} gen_tokens={gtok} "
- f"peak_memory={format_peak_memory(peak)}\n"
+ # If pp and tg lists have same length, run in tandem (zip)
+ # Otherwise (or if --all-combinations), run all combinations (cartesian product)
+ if use_combinations:
+ pp_tg_pairs = list(itertools.product(pp_list, tg_list))
+ else:
+ pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
+
+ for pp, tg in pp_tg_pairs:
+ runs: list[dict[str, Any]] = []
+ for r in range(args.repeat):
+ time.sleep(3)
+ try:
+ row, actual_pp_tokens = run_one_completion(
+ client, full_model_id, pp, tg, prompt_sizer
)
- time.sleep(2)
+ except Exception as e:
+ logger.error(e)
+ continue
+ row.update(
+ {
+ "model_short_id": short_id,
+ "model_id": full_model_id,
+ "placement_sharding": sharding,
+ "placement_instance_meta": instance_meta,
+ "placement_nodes": n_nodes,
+ "instance_id": instance_id,
+ "pp_tokens": actual_pp_tokens,
+ "tg": tg,
+ "repeat_index": r,
+ }
+ )
+ runs.append(row)
+ all_rows.append(row)
+
+ if runs:
+ prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
+ gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+ ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
+ gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+ peak = mean(
+ x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
+ )
+
+ logger.info(
+ f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
+ f"prompt_tokens={ptok} gen_tokens={gtok} "
+ f"peak_memory={format_peak_memory(peak)}\n"
+ )
+ time.sleep(2)
finally:
try:
client.request_json("DELETE", f"/instance/{instance_id}")
← b2579c78 nix: add macmon to PATH in wrapper scripts on Darwin
·
back to Exo
·
feat: add Claude Messages API and OpenAI Responses API suppo c3537980 →