← back to Exo
Fix pipeline mismatched send after 1587 (#1629)
152a27ea5dd68674519f65491e6e314983be126e · 2026-02-26 16:48:34 +0000 · rltakashige
## Motivation
Tests caught a bug. It was a real bug.
Files touched
M src/exo/worker/engines/mlx/auto_parallel.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
Diff
commit 152a27ea5dd68674519f65491e6e314983be126e
Author: rltakashige <rl.takashige@gmail.com>
Date: Thu Feb 26 16:48:34 2026 +0000
Fix pipeline mismatched send after 1587 (#1629)
## Motivation
Tests caught a bug. It was a real bug.
---
src/exo/worker/engines/mlx/auto_parallel.py | 9 ++++++-
src/exo/worker/engines/mlx/generator/generate.py | 4 ++++
.../test_mlx/test_pipeline_prefill_callbacks.py | 28 +++++++++++++++-------
3 files changed, 31 insertions(+), 10 deletions(-)
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index ab65e552..9ada458f 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -165,6 +165,7 @@ class PipelineLastLayer(CustomMlxLayer):
self.group = group
self.original_layer_signature = signature(self.original_layer.__call__)
self.is_prefill: bool = False
+ self.queue_sends: bool = False
def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
cache = self.original_layer_signature.bind_partial(
@@ -178,7 +179,7 @@ class PipelineLastLayer(CustomMlxLayer):
mx.eval(output)
if self.r != self.s - 1:
- if self.is_prefill:
+ if self.queue_sends:
_pending_prefill_sends.append(
(output, (self.r + 1) % self.s, self.group)
)
@@ -210,6 +211,12 @@ def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None:
layer.is_prefill = is_prefill
+def set_pipeline_queue_sends(model: nn.Module, queue_sends: bool) -> None:
+ for layer in model.layers: # type: ignore
+ if isinstance(layer, PipelineLastLayer):
+ layer.queue_sends = queue_sends
+
+
def get_inner_model(model: nn.Module) -> nn.Module:
inner = getattr(model, "model", None)
if isinstance(inner, nn.Module):
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index cd84f644..0ecaf8f0 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -34,6 +34,7 @@ from exo.worker.engines.mlx.auto_parallel import (
clear_prefill_sends,
flush_prefill_sends,
set_pipeline_prefill,
+ set_pipeline_queue_sends,
)
from exo.worker.engines.mlx.cache import (
CacheSnapshot,
@@ -245,6 +246,7 @@ def prefill(
try:
if is_pipeline and num_tokens >= prefill_step_size:
+ set_pipeline_queue_sends(model, queue_sends=True)
assert group is not None, "Pipeline prefill requires a distributed group"
pipeline_parallel_prefill(
model=model,
@@ -274,9 +276,11 @@ def prefill(
):
break # Stop after first iteration - cache is now filled
except PrefillCancelled:
+ set_pipeline_queue_sends(model, queue_sends=False)
set_pipeline_prefill(model, is_prefill=False)
raise
+ set_pipeline_queue_sends(model, queue_sends=False)
set_pipeline_prefill(model, is_prefill=False)
# stream_generate added 1 extra generated token to the cache, so we should trim it.
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py b/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
index 1c156655..89f1b66a 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
@@ -377,10 +377,7 @@ class TestPipelinePrefillCallbacks:
ids=["short_50", "medium_500", "long_5000"],
)
def test_callbacks_match(self, prompt_tokens: int) -> None:
- """Pipeline and stream_generate must produce identical callback sequences."""
- # Run single-device (stream_generate path)
- single = _run_single_device_test(prompt_tokens, timeout=180)
-
+ """All pipeline ranks must produce identical callback sequences."""
# Run 4-rank pipeline
pipeline_results = _run_pipeline_test(
layer_splits=LAYER_SPLITS_4WAY,
@@ -389,10 +386,11 @@ class TestPipelinePrefillCallbacks:
timeout=180,
)
- single_callbacks = single["callbacks"]
- prefill_count = single["prefill_token_count"]
+ # All ranks must agree on prefill token count and callback sequence
+ rank0_data = pipeline_results[0]
+ rank0_callbacks = rank0_data["callbacks"]
+ prefill_count = rank0_data["prefill_token_count"]
- # Every rank must produce the same callback sequence as stream_generate
for rank, pipe_data in sorted(pipeline_results.items()):
pipe_callbacks = pipe_data["callbacks"]
@@ -401,13 +399,25 @@ class TestPipelinePrefillCallbacks:
f"{pipe_data['prefill_token_count']} vs {prefill_count}"
)
- assert pipe_callbacks == single_callbacks, (
+ assert pipe_callbacks == rank0_callbacks, (
f"Rank {rank} callback mismatch for {prompt_tokens} prompt tokens "
f"(prefill M={prefill_count}):\n"
- f" stream_generate ({len(single_callbacks)} callbacks): {single_callbacks}\n"
+ f" pipeline R0 ({len(rank0_callbacks)} callbacks): {rank0_callbacks}\n"
f" pipeline R{rank} ({len(pipe_callbacks)} callbacks): {pipe_callbacks}"
)
+ # Structural checks: starts with (0, M), ends with (M, M), monotonically increasing
+ assert rank0_callbacks[0] == (0, prefill_count), (
+ f"First callback should be (0, {prefill_count}), got {rank0_callbacks[0]}"
+ )
+ assert rank0_callbacks[-1] == (prefill_count, prefill_count), (
+ f"Last callback should be ({prefill_count}, {prefill_count}), got {rank0_callbacks[-1]}"
+ )
+ for i in range(1, len(rank0_callbacks)):
+ assert rank0_callbacks[i][0] >= rank0_callbacks[i - 1][0], (
+ f"Callbacks not monotonically increasing at index {i}: {rank0_callbacks}"
+ )
+
@pytest.mark.parametrize(
"prompt_tokens",
[50, 500],
← db36bd5a Add custom prefill for pipeline (#1587)
·
back to Exo
·
fmt: add taplo TOML formatter to treefmt configuration dc0bb5e1 →