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

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

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.

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],