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

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

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