← back to Exo

fix warmup consistency for slow machines (#1748)

fac6832e5f4bb7dddf8fc61fa4675b19f0059f00 · 2026-03-24 16:45:55 +0000 · Evan Quiney

a fix from pr #1643 which is now stale - should make prefill more
consistent on very slow machines

## testing
qwen-3.5-35b-a3b loads normally
gpt-oss-120b-mxfp4-q8 loads normally

Files touched

• M src/exo/worker/engines/mlx/generator/generate.py
• M tests/auto_bench.sh

Diff

commit fac6832e5f4bb7dddf8fc61fa4675b19f0059f00
Author: Evan Quiney <evanev7@gmail.com>
Date:   Tue Mar 24 16:45:55 2026 +0000

    fix warmup consistency for slow machines (#1748)
    
    a fix from pr #1643 which is now stale - should make prefill more
    consistent on very slow machines
    
    ## testing
    qwen-3.5-35b-a3b loads normally
    gpt-oss-120b-mxfp4-q8 loads normally
---
 src/exo/worker/engines/mlx/generator/generate.py | 33 +++++++++++++-----------
 tests/auto_bench.sh                              |  2 +-
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 16de44d8..82b06ef5 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -337,21 +337,24 @@ def warmup_inference(
     mx_barrier(group)
 
     logger.info("Generating warmup tokens")
-    for _r in stream_generate(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=warmup_prompt,
-        max_tokens=50,
-        sampler=sampler,
-        prompt_cache=cache,
-        prefill_step_size=2048,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
-    ):
-        logger.info("Generated warmup token: " + str(_r.text))
-        tokens_generated += 1
-
-    logger.info("Generated ALL warmup tokens")
+    try:
+        # for slow warmups, pipeline prefill=True tends to be more likely to succeed within the 5s gpu timeout window
+        # as we don't block on the last all gather.
+        set_pipeline_prefill(model, is_prefill=True)
+        for _r in stream_generate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=warmup_prompt,
+            max_tokens=50,
+            sampler=sampler,
+            prompt_cache=cache,
+            prefill_step_size=2048,
+            kv_group_size=KV_GROUP_SIZE,
+            kv_bits=KV_BITS,
+        ):
+            tokens_generated += 1
+    finally:
+        set_pipeline_prefill(model, is_prefill=False)
 
     mx_barrier(group)
 
diff --git a/tests/auto_bench.sh b/tests/auto_bench.sh
index 3f3d6b5d..c0debf2d 100755
--- a/tests/auto_bench.sh
+++ b/tests/auto_bench.sh
@@ -50,6 +50,6 @@ bench_runner="${hosts[0]}"
 mkdir -p "./bench/$commit"
 nix run .#exo-get-all-models-on-cluster -- "$bench_runner" | while IFS= read -r model; do
   echo "running bench for $model" 1>&2
-  ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$bench_runner@$bench_runner" "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-bench -- --model $model --pp 128 4096 --tg 128 --stdout --skip-tensor-ring" >>"./bench/$commit/${model//\//--}.json"
+  ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$bench_runner@$bench_runner" "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-bench -- --model $model --pp 128 4096 --tg 128 --concurrency 1 3 8 --stdout --skip-tensor-ring" >>"./bench/$commit/${model//\//--}.json"
   echo
 done

← 7df3774c Improve batch performance and stats reporting (#1777)  ·  back to Exo  ·  Add HF_ENDPOINT in the app settings (#1783) 6cdfbb7e →