[object Object]

← back to Exo

Fix exo bench prefill and decode tps (#1749)

6ee673147db47e5a372a8684f383804c52a89214 · 2026-03-17 13:36:44 +0000 · rltakashige

Files touched

Diff

commit 6ee673147db47e5a372a8684f383804c52a89214
Author: rltakashige <rl.takashige@gmail.com>
Date:   Tue Mar 17 13:36:44 2026 +0000

    Fix exo bench prefill and decode tps (#1749)
---
 .../worker/engines/mlx/generator/batch_generate.py | 25 ++++++++++++----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index 3462a873..cb3ca5e7 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -65,6 +65,7 @@ class _EngineTask:
     generation_start_time: float = 0.0
     in_thinking: bool = False
     reasoning_tokens: int = 0
+    prefill_tps: float = 0.0
 
 
 @dataclass(eq=False)
@@ -209,6 +210,7 @@ class ExoBatchGenerator:
             detokenizer=self.tokenizer.detokenizer,
             on_generation_token=on_generation_token,
             generation_start_time=time.perf_counter(),
+            prefill_tps=_prefill_tps,
         )
 
         return uid
@@ -285,21 +287,22 @@ class ExoBatchGenerator:
             stats: GenerationStats | None = None
             usage: Usage | None = None
             if is_done:
-                generation_elapsed = time.perf_counter() - state.generation_start_time
-                generation_tps = (
-                    state.completion_tokens / generation_elapsed
-                    if generation_elapsed > 0
-                    else 0.0
-                )
                 try:
                     mlx_stats = self._exo_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
                 except ZeroDivisionError:
-                    mlx_stats = None
+                    generation_elapsed = (
+                        time.perf_counter() - state.generation_start_time
+                    )
+                    generation_tps = (
+                        state.completion_tokens / generation_elapsed
+                        if generation_elapsed > 0
+                        else 0.0
+                    )
+
                 stats = GenerationStats(
-                    prompt_tps=float(mlx_stats.prompt_tps)
-                    if mlx_stats is not None and mlx_stats.prompt_time > 0
-                    else 0.0,
-                    generation_tps=float(generation_tps),
+                    prompt_tps=state.prefill_tps,
+                    generation_tps=generation_tps,
                     prompt_tokens=len(state.all_prompt_tokens),
                     generation_tokens=state.completion_tokens,
                     peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),

← ff4d20ee Fix image models through dashboard (#1746)  ·  back to Exo  ·  Fix exo bench again again (#1750) b713889f →