From 9b03b561d36de8ac92da14f34b32ba95381f0081 Mon Sep 17 00:00:00 2001 From: Ryuichi Leo Takashige Date: Fri, 27 Mar 2026 23:46:00 +0000 Subject: [PATCH] Tighten up timings a bit more --- bench/exo_bench.py | 21 +++++++-------- .../engines/mlx/generator/batch_generate.py | 27 ++++++++++++------- 2 files changed, 27 insertions(+), 21 deletions(-) diff --git a/bench/exo_bench.py b/bench/exo_bench.py index bc9f2ed5..40b8eb9f 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -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) diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 3b026051..0f955365 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -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,