Only patch for certain models

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

@@ -229,10 +229,12 @@ def pipeline_auto_parallel(
         "Expected a list of layers after auto-parallel initialisation"
     )
 
-    return patch_pipeline_model(model, group)
+    if isinstance(model, (GptOssModel, Qwen3MoeModel, Qwen3NextModel)):
+        model = patch_distributed_model(model)
+    return model
 
 
-def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
+def patch_distributed_model[T](model: T) -> T:
     # Patch __call__ on the model's class
     cls = model.__class__
     original_call = cls.__call__  # type :ignore
@@ -258,32 +260,6 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
     return model
 
 
-def patch_tensor_model[T](model: T) -> T:
-    """Patch model's __call__ to ensure distributed ops sync during inference."""
-    cls = model.__class__
-    original_call = cls.__call__
-    call_signature = signature(original_call)
-
-    def patched_call(
-        self: T,
-        *args: object,
-        **kwargs: object,
-    ) -> mx.array:
-        logits: mx.array = original_call(self, *args, **kwargs)  # pyright: ignore[reportAny]
-        cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
-            "cache", None
-        )
-
-        # Add dependency to last cache entry to ensure distributed ops are evaluated
-        if cache is not None and len(cache) > 0:  # pyright: ignore[reportAny]
-            cache[-1].state = mx.depends(cache[-1].state, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
-
-        return logits
-
-    cls.__call__ = patched_call
-    return model
-
-
 def tensor_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
@@ -333,7 +309,7 @@ def tensor_auto_parallel(
     if hasattr(model, "shard"):
         try:
             model.shard(group)  # type: ignore
-            return patch_tensor_model(model)
+            return patch_distributed_model(model)
         except (AttributeError, TypeError, NameError):
             pass
 
@@ -386,7 +362,7 @@ def tensor_auto_parallel(
     model = tensor_parallel_sharding_strategy.shard_model(
         model, timeout_seconds, on_timeout
     )
-    return patch_tensor_model(model)
+    return patch_distributed_model(model)
 
 
 class TensorParallelShardingStrategy(ABC):

src/exo/worker/tests/unittests/test_mlx/conftest.py

@@ -18,7 +18,7 @@ from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.generator.generate import mlx_generate
-from exo.worker.engines.mlx.utils_mlx import shard_and_load, apply_chat_template
+from exo.worker.engines.mlx.utils_mlx import apply_chat_template, shard_and_load
 
 
 class MockLayer(nn.Module):
@@ -121,10 +121,7 @@ def run_gpt_oss_pipeline_device(
 
         generated_text = ""
         for response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt
+            model=model, tokenizer=tokenizer, task=task, prompt=prompt
         ):
             generated_text += response.text
             if response.finish_reason is not None:
@@ -190,10 +187,7 @@ def run_gpt_oss_tensor_parallel_device(
 
         generated_text = ""
         for response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt
+            model=model, tokenizer=tokenizer, task=task, prompt=prompt
         ):
             generated_text += response.text
             if response.finish_reason is not None:

src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py

@@ -12,7 +12,7 @@ from exo.worker.engines.mlx.auto_parallel import (
     CustomMlxLayer,
     PipelineFirstLayer,
     PipelineLastLayer,
-    patch_pipeline_model,
+    patch_distributed_model,
 )
 from exo.worker.tests.unittests.test_mlx.conftest import MockLayer
 
@@ -55,7 +55,7 @@ def run_pipeline_device(
 
         # Wrap in a mock model, then wrap in PipelineParallelModel for all_gather
         inner_model = MockModel([composed])
-        model = patch_pipeline_model(inner_model, group)
+        model = patch_distributed_model(inner_model)
 
         x = mx.ones((1, 4))
         result = model(x)
@@ -143,7 +143,7 @@ def test_composed_call_works() -> None:
             # With world_size=2 and each layer doing x*2:
             #   - Rank 0: 1.0 * 2 = 2.0 (sends to rank 1)
             #   - Rank 1: 2.0 * 2 = 4.0 (last rank, final result)
-            expected = 2.0 * (2 ** rank)  # 2.0 for rank 0, 4.0 for rank 1
+            expected = 2.0 * (2**rank)  # 2.0 for rank 0, 4.0 for rank 1
             assert (result_array == expected).all(), (
                 f"Device {rank}: expected {expected}, got {result_array}"
             )