Revert dumb merge mistake

612f58c78de71d39476e886f2beca177580c68f3 · 2025-11-06 18:39:08 -0800 · rltakashige

Files touched

M src/exo/worker/runner/generate.py

Diff

commit 612f58c78de71d39476e886f2beca177580c68f3
Author: rltakashige <rl.takashige@gmail.com>
Date:   Thu Nov 6 18:39:08 2025 -0800

    Revert dumb merge mistake
---
 src/exo/worker/runner/generate.py | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/src/exo/worker/runner/generate.py b/src/exo/worker/runner/generate.py
index 9fe58d40..e6821fd0 100644
--- a/src/exo/worker/runner/generate.py
+++ b/src/exo/worker/runner/generate.py
@@ -215,11 +215,9 @@ def generate_step(
 
         y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
 
-    prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
-
     mx.async_eval(y, logprobs)
-    next_y, next_logprobs = _step(y)
-    mx.async_eval(next_y, next_logprobs)
+    next_y: array | None = None
+    next_logprobs: array | None = None
     n = 0
 
     while True:
@@ -236,11 +234,8 @@ def generate_step(
         yield int(y.item()), logprobs
         if n % 256 == 0:
             mx.clear_cache()
-        if n == max_tokens:
-            break
-        y, logprobs = next_y, logprobs
-        next_y, next_logprobs = _step(y)
-        mx.async_eval(next_y, next_logprobs)
+        y, logprobs = next_y, next_logprobs
+        n += 1
 
 
 def stream_generate(

← 6bcac37d stop benching on all pushes · back to Exo · pipeline parallel fix 9058b117 →