diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 2e3d051251..c45e3751dd 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -6,6 +6,9 @@ from typing import Callable, Generator, cast, get_args import mlx.core as mx +from mlx_lm.generate import ( + generation_stream as mlx_lm_generation_stream, +) from mlx_lm.generate import ( maybe_quantize_kv_cache, stream_generate, @@ -326,6 +329,9 @@ def combined_progress_callback(processed: int, total: int) -> None: set_pipeline_prefill(model, is_prefill=True) + # The batch's next decode step may still be running: let it finish before this prefill's + # collectives, as for the batch's own prompts (see opt_batch_gen._prompt_after_decode_step) + mx.synchronize(mlx_lm_generation_stream) mx_barrier(group) logger.info("Starting prefill") diff --git a/src/exo/worker/engines/mlx/patches/opt_batch_gen.py b/src/exo/worker/engines/mlx/patches/opt_batch_gen.py index e7bff41966..0fad4c6d02 100644 --- a/src/exo/worker/engines/mlx/patches/opt_batch_gen.py +++ b/src/exo/worker/engines/mlx/patches/opt_batch_gen.py @@ -1,8 +1,9 @@ +from collections.abc import Callable from dataclasses import dataclass, field from typing import cast import mlx.core as mx -from mlx_lm.generate import GenerationBatch +from mlx_lm.generate import GenerationBatch, PromptProcessingBatch, generation_stream _PRECOMPUTE_TOP_K = 20 @@ -132,5 +133,24 @@ def _patched_step(self: GenerationBatch) -> tuple[list[int], list[mx.array]]: return token_list, current_lp +_original_prompt = cast( + Callable[[PromptProcessingBatch, list[list[int]]], None], + PromptProcessingBatch.prompt, +) + + +def _prompt_after_decode_step( + self: PromptProcessingBatch, tokens: list[list[int]] +) -> None: + """BatchGenerator._next async-evaluates the batch's next decode step and then processes new + prompts. Wait for that step first. With MLX_METAL_FAST_SYNCH, a tensor-parallel prompt eval + started while the step's collectives are still running can deadlock: its GPU work spins + waiting for the communication stream, which is still waiting for the step's GPU work.""" + if tokens: + mx.synchronize(generation_stream) + _original_prompt(self, tokens) + + def apply_batch_gen_patch() -> None: GenerationBatch._step = _patched_step + PromptProcessingBatch.prompt = _prompt_after_decode_step diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prompt_waits_for_decode_step.py b/src/exo/worker/tests/unittests/test_mlx/test_prompt_waits_for_decode_step.py new file mode 100644 index 0000000000..35dd8f171a --- /dev/null +++ b/src/exo/worker/tests/unittests/test_mlx/test_prompt_waits_for_decode_step.py @@ -0,0 +1,35 @@ +from typing import cast + +import mlx.core as mx +import pytest +from mlx_lm.generate import PromptProcessingBatch + +from exo.worker.engines.mlx.patches import opt_batch_gen + + +def test_prompt_waits_for_the_decode_step_before_new_tokens( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + def synchronize(stream: object = None) -> None: + calls.append("wait") + + def original_prompt(batch: PromptProcessingBatch, tokens: list[list[int]]) -> None: + calls.append(f"prompt {len(tokens)}") + + monkeypatch.setattr(mx, "synchronize", synchronize) + monkeypatch.setattr(opt_batch_gen, "_original_prompt", original_prompt) + batch = PromptProcessingBatch.__new__(PromptProcessingBatch) + + opt_batch_gen._prompt_after_decode_step(batch, [[1, 2, 3]]) # pyright: ignore[reportPrivateUsage] + # mlx-lm calls it every step, mostly with nothing to process: no need to wait then + opt_batch_gen._prompt_after_decode_step(batch, []) # pyright: ignore[reportPrivateUsage] + + assert calls == ["wait", "prompt 1", "prompt 0"] + + +def test_batch_gen_patch_installs_it() -> None: + opt_batch_gen.apply_batch_gen_patch() + installed = cast(object, PromptProcessingBatch.prompt) + assert installed is opt_batch_gen._prompt_after_decode_step # pyright: ignore[reportPrivateUsage]