Synchronize before warmup (#1419)

8f3681cf7e81eab492d0771942db651b056e4a5c · 2026-02-07 00:14:19 +0000 · rltakashige

## Motivation

Maybe addresses #1303 

## Changes

Add an mx barrier before warmup

## Why It Works

It might, it might not. Shouldn't break anything that's not already
broken though.

## Test Plan

### Manual Testing
The two machines I tested on were fine on GLM 4.7 Flash 8bit (the one in
exo.log in the issue). Obviously not definitive for anything, however.

<img width="594" height="878" alt="image"
src="https://github.com/user-attachments/assets/534d3ad6-16ef-4cb5-b823-43c8d4e1d3c6"
/>

Files touched

M src/exo/worker/engines/mlx/generator/generate.py

Diff

commit 8f3681cf7e81eab492d0771942db651b056e4a5c
Author: rltakashige <rl.takashige@gmail.com>
Date:   Sat Feb 7 00:14:19 2026 +0000

    Synchronize before warmup (#1419)
    
    ## Motivation
    
    Maybe addresses #1303
    
    ## Changes
    
    Add an mx barrier before warmup
    
    ## Why It Works
    
    It might, it might not. Shouldn't break anything that's not already
    broken though.
    
    ## Test Plan
    
    ### Manual Testing
    The two machines I tested on were fine on GLM 4.7 Flash 8bit (the one in
    exo.log in the issue). Obviously not definitive for anything, however.
    
    <img width="594" height="878" alt="image"
    src="https://github.com/user-attachments/assets/534d3ad6-16ef-4cb5-b823-43c8d4e1d3c6"
    />
---
 src/exo/worker/engines/mlx/generator/generate.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index fa618d9c..7063f8de 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -145,6 +145,8 @@ def warmup_inference(
     # Use a default sampler for warmup
     sampler = make_sampler(temp=0.0)
 
+    mx_barrier(group)
+
     logger.info("Generating warmup tokens")
     for _r in stream_generate(
         model=model,

← 9dc4f786 Ciaran/image model listing (#1417) · back to Exo · Add image lightbox (#1414) 3f57416d →