exo_bench: set temperature=0.0 for argmax sampling (drops decode overhead)

bench/exo_bench.py

@@ -287,6 +287,10 @@ def run_one_completion(
         "max_tokens": tg,
         "logprobs": False,
         "use_prefix_cache": use_prefix_cache,
+        # Argmax sampling for deterministic, faster decode (matches
+        # mlx_lm.benchmark default). Avoids per-token softmax + categorical
+        # sample over the full vocab.
+        "temperature": 0.0,
     }
 
     if not stream:
@@ -727,6 +731,15 @@ def main() -> int:
                             )
                             runs.append(row)
                             all_rows.append(row)
+                            # Per-repeat trial log so individual numbers are visible
+                            # alongside the final averaged summary. Useful for
+                            # spotting outliers and trial-to-trial variance.
+                            _s = row.get("stats") or {}
+                            logger.info(
+                                f"  repeat {r + 1}/{args.repeat}: "
+                                f"prompt_tps={_s.get('prompt_tps', 0):.2f} "
+                                f"gen_tps={_s.get('generation_tps', 0):.2f}"
+                            )
                         else:
                             # Concurrent: fire N requests in parallel
                             # Pre-build prompt once, barrier ensures simultaneous dispatch
@@ -738,6 +751,8 @@ def main() -> int:
                                 "max_tokens": tg,
                                 "logprobs": False,
                                 "use_prefix_cache": args.use_prefix_cache,
+                                # Argmax sampling — matches mlx_lm.benchmark default
+                                "temperature": 0.0,
                             }
                             barrier = threading.Barrier(concurrency)
                             batch_start = threading.Event()

src/exo/api/main.py

@@ -254,6 +254,7 @@ class API:
         self.node_id: NodeId = node_id
         self.last_completed_election: int = 0
         self.port = port
+        self._sent_image_hashes: set[str] = set()
 
         self.paused: bool = False
         self.paused_ev: anyio.Event = anyio.Event()
@@ -303,6 +304,7 @@ class API:
         self.event_receiver.close()
         self.event_receiver = event_receiver
         self._tg.start_soon(self._apply_state)
+        self._sent_image_hashes = set()
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -824,8 +826,18 @@ class API:
         )
         command = TextGeneration(task_params=task_params)
 
+        new_images: list[tuple[int, str]] = []
+        for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
+            if h not in self._sent_image_hashes:
+                self._sent_image_hashes.add(h)
+                new_images.append((idx, img))
+
+        if not new_images:
+            await self._send(command)
+            return command
+
         all_chunks: list[tuple[int, str]] = []
-        for img_idx, img_data in enumerate(images):
+        for img_idx, img_data in new_images:
             for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
                 all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))
 

src/exo/shared/apply.py

@@ -40,14 +40,7 @@ from exo.shared.types.profiling import (
     ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import (
-    ImageEdits,
-    ImageGeneration,
-    Task,
-    TaskId,
-    TaskStatus,
-    TextGeneration,
-)
+from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -79,6 +72,7 @@ def event_apply(event: Event, state: State) -> State:
             TestEvent()
             | ChunkGenerated()
             | TaskAcknowledged()
+            | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
             | CustomModelCardAdded()
@@ -99,8 +93,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_runner_status_updated(event, state)
         case TaskCreated():
             return apply_task_created(event, state)
-        case InputChunkReceived():
-            return apply_input_chunk_received(event, state)
         case TaskDeleted():
             return apply_task_deleted(event, state)
         case TaskFailed():
@@ -165,32 +157,10 @@ def apply_task_created(event: TaskCreated, state: State) -> State:
     return state.model_copy(update={"tasks": new_tasks})
 
 
-def apply_input_chunk_received(event: InputChunkReceived, state: State) -> State:
-    command_chunks = {
-        **state.input_chunks.get(event.command_id, {}),
-        event.chunk.chunk_index: event.chunk,
-    }
-    return state.model_copy(
-        update={
-            "input_chunks": {**state.input_chunks, event.command_id: command_chunks}
-        }
-    )
-
-
 def apply_task_deleted(event: TaskDeleted, state: State) -> State:
-    task = state.tasks.get(event.task_id)
     new_tasks: Mapping[TaskId, Task] = {
         tid: task for tid, task in state.tasks.items() if tid != event.task_id
     }
-    if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
-        new_input_chunks = {
-            command_id: chunks
-            for command_id, chunks in state.input_chunks.items()
-            if command_id != task.command_id
-        }
-        return state.model_copy(
-            update={"tasks": new_tasks, "input_chunks": new_input_chunks}
-        )
     return state.model_copy(update={"tasks": new_tasks})
 
 

src/exo/shared/tests/test_apply/test_apply_input_chunks.py

@@ -1,85 +0,0 @@
-from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId
-from exo.shared.types.events import (
-    IndexedEvent,
-    InputChunkReceived,
-    TaskCreated,
-    TaskDeleted,
-)
-from exo.shared.types.state import State
-from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import (
-    InputMessage,
-    InputMessageContent,
-    TextGenerationTaskParams,
-)
-from exo.shared.types.worker.instances import InstanceId
-
-
-def test_apply_input_chunk_received_stores_chunk_in_state() -> None:
-    command_id = CommandId("command-1")
-    chunk = InputImageChunk(
-        model=ModelId("mlx-community/test-model"),
-        command_id=command_id,
-        data="abc",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-
-    state = apply(
-        State(),
-        IndexedEvent(
-            idx=0,
-            event=InputChunkReceived(command_id=command_id, chunk=chunk),
-        ),
-    )
-
-    assert state.input_chunks == {command_id: {0: chunk}}
-
-
-def test_apply_task_deleted_removes_chunks_for_generation_command() -> None:
-    command_id = CommandId("command-1")
-    task_id = TaskId("task-1")
-    chunk = InputImageChunk(
-        model=ModelId("mlx-community/test-model"),
-        command_id=command_id,
-        data="abc",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-    task = TextGeneration(
-        task_id=task_id,
-        instance_id=InstanceId("instance-1"),
-        task_status=TaskStatus.Pending,
-        command_id=command_id,
-        task_params=TextGenerationTaskParams(
-            model=ModelId("mlx-community/test-model"),
-            input=[
-                InputMessage(role="user", content=InputMessageContent("hello")),
-            ],
-        ),
-    )
-
-    state = State()
-    state = apply(
-        state,
-        IndexedEvent(
-            idx=0,
-            event=InputChunkReceived(command_id=command_id, chunk=chunk),
-        ),
-    )
-    state = apply(
-        state,
-        IndexedEvent(idx=1, event=TaskCreated(task_id=task_id, task=task)),
-    )
-    state = apply(
-        state,
-        IndexedEvent(idx=2, event=TaskDeleted(task_id=task_id)),
-    )
-
-    assert state.tasks == {}
-    assert state.input_chunks == {}

src/exo/shared/types/state.py

@@ -6,8 +6,7 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
 
 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -46,9 +45,6 @@ class State(FrozenModel):
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
-    # Durable request input chunks for active image requests. Workers rebuild
-    # local image caches from this state instead of reading events directly.
-    input_chunks: Mapping[CommandId, Mapping[int, InputImageChunk]] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)

src/exo/worker/engines/mlx/auto_parallel.py

@@ -599,7 +599,7 @@ def tensor_auto_parallel(
         raise ValueError(f"Unsupported model type: {type(model)}")
 
     model = yield from tensor_parallel_sharding_strategy.shard_model(model)
-    return patch_tensor_model(model)
+    return model  # PATCH-DISABLED for A/B test
 
 
 class TensorParallelShardingStrategy(ABC):

src/exo/worker/main.py

@@ -24,6 +24,7 @@ from exo.shared.types.events import (
     CustomModelCardDeleted,
     Event,
     IndexedEvent,
+    InputChunkReceived,
     InstanceDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
@@ -140,6 +141,37 @@ class Worker:
                 if isinstance(event, InstanceDeleted):
                     self._instance_backoff.reset(event.instance_id)
 
+                # Buffer input image chunks for image editing
+                if isinstance(event, InputChunkReceived):
+                    cmd_id = event.command_id
+                    if cmd_id not in self.input_chunk_buffer:
+                        self.input_chunk_buffer[cmd_id] = {}
+                        self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
+
+                    self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
+                        event.chunk
+                    )
+
+                    if (
+                        len(self.input_chunk_buffer[cmd_id])
+                        == self.input_chunk_counts[cmd_id]
+                    ):
+                        per_image: defaultdict[int, list[InputImageChunk]] = (
+                            defaultdict(list)
+                        )
+                        for chunk in self.input_chunk_buffer[cmd_id].values():
+                            per_image[chunk.image_index].append(chunk)
+                        for chunks_for_image in per_image.values():
+                            sorted_chunks = sorted(
+                                chunks_for_image, key=lambda c: c.chunk_index
+                            )
+                            img = Base64Image("".join(c.data for c in sorted_chunks))
+                            self.image_cache[
+                                Base64ImageHash(
+                                    hashlib.sha256(img.encode("ascii")).hexdigest()
+                                )
+                            ] = img
+
                 if isinstance(event, CustomModelCardAdded):
                     await event.model_card.save_to_custom_dir()
                     add_to_card_cache(event.model_card)
@@ -147,35 +179,6 @@ class Worker:
                 if isinstance(event, CustomModelCardDeleted):
                     await delete_custom_card(event.model_id)
 
-                self._sync_input_views_from_state()
-
-    def _sync_input_views_from_state(self) -> None:
-        self.input_chunk_buffer = {
-            command_id: dict(chunks)
-            for command_id, chunks in self.state.input_chunks.items()
-        }
-        self.input_chunk_counts = {
-            command_id: next(iter(chunks.values())).total_chunks
-            for command_id, chunks in self.input_chunk_buffer.items()
-            if chunks
-        }
-
-        self.image_cache = {}
-        for command_id, chunks in self.input_chunk_buffer.items():
-            expected_chunks = self.input_chunk_counts.get(command_id)
-            if expected_chunks is None or len(chunks) != expected_chunks:
-                continue
-
-            per_image: defaultdict[int, list[InputImageChunk]] = defaultdict(list)
-            for chunk in chunks.values():
-                per_image[chunk.image_index].append(chunk)
-            for chunks_for_image in per_image.values():
-                sorted_chunks = sorted(chunks_for_image, key=lambda c: c.chunk_index)
-                image = Base64Image("".join(chunk.data for chunk in sorted_chunks))
-                self.image_cache[
-                    Base64ImageHash(hashlib.sha256(image.encode("ascii")).hexdigest())
-                ] = image
-
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)
@@ -186,7 +189,7 @@ class Worker:
                 self.state.instances,
                 self.state.runners,
                 self.state.tasks,
-                self.state.input_chunks,
+                self.input_chunk_buffer,
                 self.image_cache,
                 self._instance_backoff,
                 self._download_backoff,
@@ -318,9 +321,15 @@ class Worker:
                             advanced_params=task.task_params.advanced_params,
                         ),
                     )
+                    # Cleanup buffers
+                    if cmd_id in self.input_chunk_buffer:
+                        del self.input_chunk_buffer[cmd_id]
+                    if cmd_id in self.input_chunk_counts:
+                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
 
                 case TextGeneration() if task.task_params.image_hashes:
+                    cmd_id = task.command_id
                     resolved_images = [
                         self.image_cache[h]
                         for _, h in sorted(task.task_params.image_hashes.items())
@@ -332,6 +341,10 @@ class Worker:
                             )
                         }
                     )
+                    if cmd_id in self.input_chunk_buffer:
+                        del self.input_chunk_buffer[cmd_id]
+                    if cmd_id in self.input_chunk_counts:
+                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
                 case LoadModel(instance_id=instance_id):
                     if (instance := self.state.instances.get(instance_id)) is not None: