Tighten up timings a bit more

bench/exo_bench.py

@@ -22,6 +22,7 @@ import contextlib
 import itertools
 import json
 import sys
+import threading
 import time
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -443,19 +444,44 @@ def main() -> int:
                             all_rows.append(row)
                         else:
                             # Concurrent: fire N requests in parallel
-                            # Each thread gets its own ExoClient (separate HTTP connection)
+                            # Pre-build prompt once, barrier ensures simultaneous dispatch
+                            content, actual_pp = prompt_sizer.build(pp)
+                            pre_built_payload: dict[str, Any] = {
+                                "model": full_model_id,
+                                "messages": [{"role": "user", "content": content}],
+                                "stream": False,
+                                "max_tokens": tg,
+                            }
+                            barrier = threading.Barrier(concurrency)
+                            batch_start = threading.Event()
+                            batch_t0: float = 0.0
                             batch_results: list[tuple[dict[str, Any], int]] = []
                             batch_errors = 0
 
                             def _run_concurrent(
-                                idx: int, *, _pp: int = pp, _tg: int = tg
+                                idx: int,
                             ) -> tuple[dict[str, Any], int]:
+                                nonlocal batch_t0
                                 c = ExoClient(
                                     args.host, args.port, timeout_s=args.timeout
                                 )
-                                return run_one_completion(
-                                    c, full_model_id, _pp, _tg, prompt_sizer
-                                )
+                                if barrier.wait() == 0:
+                                    batch_t0 = time.perf_counter()
+                                    batch_start.set()
+                                else:
+                                    batch_start.wait()
+                                t0 = batch_t0
+                                out = c.post_bench_chat_completions(pre_built_payload)
+                                elapsed = time.perf_counter() - t0
+                                stats = out.get("generation_stats")
+                                choices = out.get("choices") or [{}]
+                                message = choices[0].get("message", {}) if choices else {}
+                                text = message.get("content") or ""
+                                return {
+                                    "elapsed_s": elapsed,
+                                    "output_text_preview": text[:200],
+                                    "stats": stats,
+                                }, actual_pp
 
                             with ThreadPoolExecutor(max_workers=concurrency) as pool:
                                 futures = {
@@ -468,6 +494,7 @@ def main() -> int:
                                     except Exception as e:
                                         logger.error(f"Concurrent request failed: {e}")
                                         batch_errors += 1
+                            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
@@ -501,20 +528,20 @@ def main() -> int:
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
                                 ]
-                                per_req_tps = (
-                                    mean(valid_gen_tps) if valid_gen_tps else 0.0
-                                )
+                                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"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

@@ -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,