electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Tighten up timings a bit more

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-03-27 23:46:00 +00:00
parent e9d911ccfc
commit 9b03b561d3
2 changed files with 27 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bench/exo_bench.py
+9 -12
View File
@@ -494,7 +494,7 @@ def main() -> int:
except Exception as e:
logger.error(f"Concurrent request failed: {e}")
batch_errors += 1
batch_wall_s = time.perf_counter() - batch_t0
batch_wall_s = max(x["elapsed_s"] for x, _ in batch_results) if batch_results else time.perf_counter() - batch_t0
for idx, (row, actual_pp_tokens) in enumerate(
batch_results
@@ -523,28 +523,25 @@ def main() -> int:
all_rows.append(row)
if batch_results:
total_gen_tokens = sum(
x["stats"]["generation_tokens"]
for x, _ in batch_results
)
wall_agg_tps = total_gen_tokens / batch_wall_s if batch_wall_s > 0 else 0.0
server_per_req_tps = mean(
valid_gen_tps = [
x["stats"]["generation_tps"]
for x, _ in batch_results
if x["stats"]["generation_tps"] > 0
)
server_agg_tps = server_per_req_tps * concurrency
]
per_req_tps = max(valid_gen_tps) if valid_gen_tps else 0.0
agg_gen_tps = per_req_tps * concurrency
logger.info(
f"[concurrent {concurrency}x] "
f"wall_agg_tps={wall_agg_tps:.2f} "
f"server_agg_tps={server_agg_tps:.2f} "
f"agg_gen_tps={agg_gen_tps:.2f} "
f"per_req_tps={per_req_tps:.2f} "
f"wall_s={batch_wall_s:.2f} "
f"errors={batch_errors}"
)
if runs:
prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
per_req_tps = mean(x["stats"]["generation_tps"] for x in runs)
valid_gen = [x["stats"]["generation_tps"] for x in runs if x["stats"]["generation_tps"] > 0]
per_req_tps = max(valid_gen) if valid_gen else 0.0
gen_tps = per_req_tps * concurrency
ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
src/exo/worker/engines/mlx/generator/batch_generate.py
+18 -9
View File
@@ -70,6 +70,8 @@ class _EngineTask:
in_thinking: bool = False
reasoning_tokens: int = 0
prefill_tps: float = 0.0
first_gen_token_time: float | None = None
last_gen_token_time: float | None = None
@dataclass(eq=False)
@@ -242,6 +244,10 @@ class ExoBatchGenerator:
continue
state = self._active_tasks[response.uid]
now = time.perf_counter()
if state.first_gen_token_time is None:
state.first_gen_token_time = now
state.last_gen_token_time = now
if state.on_generation_token is not None:
state.on_generation_token()
if response.finish_reason != "stop":
@@ -250,6 +256,9 @@ class ExoBatchGenerator:
state.detokenizer.finalize()
text = state.detokenizer.last_segment
state.completion_tokens += 1
if state.task_params.bench:
delta = now - state.first_gen_token_time
logger.debug(f"[bench] uid={response.uid} tok#{state.completion_tokens} {text!r} t={delta:.4f}s")
state.generated_text_parts.append(text)
state.potential_stop_sequence_text += text
@@ -304,15 +313,15 @@ class ExoBatchGenerator:
stats: GenerationStats | None = None
usage: Usage | None = None
if is_done:
gen_time_delta = (
self._mlx_gen._stats.generation_time
- state.generation_time_at_start
)
generation_tps = (
state.completion_tokens / gen_time_delta
if gen_time_delta > 0
else 0.0
)
if (
state.first_gen_token_time is not None
and state.last_gen_token_time is not None
and state.completion_tokens > 1
):
gen_span = state.last_gen_token_time - state.first_gen_token_time
generation_tps = (state.completion_tokens - 1) / gen_span if gen_span > 0 else 0.0
else:
generation_tps = 0.0
stats = GenerationStats(
prompt_tps=state.prefill_tps,