From fac6832e5f4bb7dddf8fc61fa4675b19f0059f00 Mon Sep 17 00:00:00 2001
From: Evan Quiney <evanev7@gmail.com>
Date: Tue, 24 Mar 2026 16:45:55 +0000
Subject: [PATCH] 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
---
 .../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