From ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Wed, 22 Apr 2026 00:34:09 -0700 Subject: [PATCH] Thread local generation stream (#1090) --- mlx_lm/generate.py | 16 +++- mlx_lm/server.py | 187 ++++++++++++++++++++++--------------------- mlx_lm/utils.py | 4 +- setup.py | 2 +- tests/test_gguf.py | 3 + tests/test_server.py | 3 + 6 files changed, 117 insertions(+), 98 deletions(-) diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index d0913b5..3573b26 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -223,7 +223,7 @@ def setup_arg_parser(): # A stream on the default device just for generation -generation_stream = mx.new_stream(mx.default_device()) +generation_stream = mx.new_thread_local_stream(mx.default_device()) @contextlib.contextmanager @@ -1497,6 +1497,7 @@ class BatchGenerator: def __init__( self, model: nn.Module, + *, max_tokens: int = 128, stop_tokens: Optional[Sequence[Sequence[int]]] = None, sampler: Optional[Callable[[mx.array], mx.array]] = None, @@ -1507,6 +1508,7 @@ class BatchGenerator: prefill_batch_size: int = 8, prefill_step_size: int = 2048, max_kv_size: Optional[int] = None, + stream=None, ): self.model = model self.max_tokens = max_tokens @@ -1518,6 +1520,8 @@ class BatchGenerator: self.completion_batch_size = max(completion_batch_size, prefill_batch_size) self.max_kv_size = max_kv_size + self._stream = stream or generation_stream + self._default_state_machine = SequenceStateMachine( {"normal": [(seq, None) for seq in stop_tokens]} if stop_tokens else {}, initial="normal", @@ -1544,9 +1548,13 @@ class BatchGenerator: else: self._old_wired_limit = None + @property + def stream(self): + return self._stream + def close(self): if self._old_wired_limit is not None: - mx.synchronize(generation_stream) + mx.synchronize(self._stream) mx.set_wired_limit(self._old_wired_limit) self._old_wired_limit = None @@ -1843,7 +1851,7 @@ class BatchGenerator: Returns: Tuple of prompt processing responses and generation responses. """ - with mx.stream(generation_stream): + with mx.stream(self._stream): return self._next() def next_generated(self): @@ -1853,7 +1861,7 @@ class BatchGenerator: Returns: List of GenerationBatch.Response objects """ - with mx.stream(generation_stream): + with mx.stream(self._stream): while True: prompt_responses, generation_responses = self._next() if not generation_responses and prompt_responses: diff --git a/mlx_lm/server.py b/mlx_lm/server.py index 29b7dd7..ce8d958 100644 --- a/mlx_lm/server.py +++ b/mlx_lm/server.py @@ -36,7 +36,6 @@ from ._version import __version__ from .generate import ( BatchGenerator, SequenceStateMachine, - generation_stream, stream_generate, ) from .models.cache import ( @@ -279,8 +278,7 @@ class TimeBudget: self._loops += 1 self._time_spent += time.time() - self._start if self._loops % self._sync_frequency == 0: - with mx.stream(generation_stream): - loop_time = mx.distributed.all_sum(self._time_spent).item() + loop_time = mx.distributed.all_sum(self._time_spent).item() avg_loop_time = loop_time / ( mx.distributed.init().size() * self._sync_frequency ) @@ -308,94 +306,92 @@ class ModelProvider: ) self.is_distributed = group.size() > 1 - # Preload the default model if it is provided - self.default_model_map = {} - if self.cli_args.model is not None: - self.default_model_map[self.cli_args.model] = "default_model" - self.load(self.cli_args.model, draft_model_path="default_model") + # Maps model and adapter paths the actual paths to be used. Used to + # map 'default_model' to the provided model by cli argument but could + # be used for more in the future. + self._model_map = {} + self._adapter_map = {} + self._draft_model_map = {} + self._model_map["default_model"] = self.cli_args.model + self._adapter_map["default_model"] = self.cli_args.adapter_path + self._draft_model_map["default_model"] = self.cli_args.draft_model - # Added in adapter_path to load dynamically - def load(self, model_path, adapter_path=None, draft_model_path=None): - model_path = self.default_model_map.get(model_path, model_path) - if self.model_key == (model_path, adapter_path, draft_model_path): - return self.model, self.tokenizer + # Build the tokenizer config for later use in load + self._tokenizer_config = { + "trust_remote_code": True if cli_args.trust_remote_code else None + } + if cli_args.chat_template: + self._tokenizer_config["chat_template"] = cli_args.chat_template + + def _load(self, model_path, adapter_path=None, draft_model_path=None): + if self.is_distributed and ( + adapter_path is not None or draft_model_path is not None + ): + raise ValueError( + "Loading with adapters or draft models not supported in distributed mode" + ) # Remove the old model if it exists. + self.model_key = None self.model = None self.tokenizer = None - self.model_key = None self.draft_model = None - # Building tokenizer_config - tokenizer_config = { - "trust_remote_code": True if self.cli_args.trust_remote_code else None - } - if self.cli_args.chat_template: - tokenizer_config["chat_template"] = self.cli_args.chat_template - - if model_path == "default_model": - if self.cli_args.model is None: - raise ValueError( - "A model path has to be given as a CLI " - "argument or in the HTTP request" - ) - adapter_path = adapter_path or self.cli_args.adapter_path - # TODO: Generalize distributed load - if self.is_distributed: - model, tokenizer = sharded_load( - self.cli_args.model, self.pipeline_group, self.tensor_group - ) - else: - model, tokenizer = load( - self.cli_args.model, - adapter_path=adapter_path, - tokenizer_config=tokenizer_config, - ) + # Load the model and tokenizer + if self.is_distributed: + model, tokenizer = sharded_load( + model_path, + pipeline_group=self.pipeline_group, + tensor_group=self.tensor_group, + tokenizer_config=self._tokenizer_config, + ) else: - # TODO: Generalize distributed load - if self.is_distributed: - model, tokenizer = sharded_load( - model_path, self.pipeline_group, self.tensor_group - ) - else: - model, tokenizer = load( - model_path, - adapter_path=adapter_path, - tokenizer_config=tokenizer_config, - ) + model, tokenizer = load( + model_path, + adapter_path=adapter_path, + tokenizer_config=self._tokenizer_config, + ) + # Use the default chat template if needed if self.cli_args.use_default_chat_template: if tokenizer.chat_template is None: tokenizer.chat_template = tokenizer.default_chat_template - self.model_key = (model_path, adapter_path, draft_model_path) - self.model = model - self.tokenizer = tokenizer - - def validate_draft_tokenizer(draft_tokenizer): - # Check if tokenizers are compatible + # Load the draft model for speculative decoding + draft_model = None + if draft_model_path is not None: + draft_model, draft_tokenizer = load(draft_model_path) if draft_tokenizer.vocab_size != tokenizer.vocab_size: logging.warning( "Draft model tokenizer does not match model tokenizer. " "Speculative decoding may not work as expected." ) - # Load draft model if specified - if ( - draft_model_path == "default_model" - and self.cli_args.draft_model is not None - ): - self.draft_model, draft_tokenizer = load(self.cli_args.draft_model) - validate_draft_tokenizer(draft_tokenizer) + # Compute batchability + is_batchable = draft_model is None + is_batchable = is_batchable and all( + hasattr(c, "merge") for c in make_prompt_cache(model) + ) - elif draft_model_path is not None and draft_model_path != "default_model": - self.draft_model, draft_tokenizer = load(draft_model_path) - validate_draft_tokenizer(draft_tokenizer) + # Update the member variables + self.model_key = (model_path, adapter_path, draft_model_path) + self.model = model + self.tokenizer = tokenizer + self.draft_model = draft_model + self.is_batchable = is_batchable - if self.draft_model is None: - self.is_batchable = all( - hasattr(c, "merge") for c in make_prompt_cache(self.model) - ) + def load_default(self): + if self._model_map["default_model"] is not None: + self.load("default_model", None, "default_model") + + def load(self, model_path, adapter_path=None, draft_model_path=None): + model_path = self._model_map.get(model_path, model_path) + adapter_path = self._adapter_map.get(model_path, adapter_path) + draft_model_path = self._draft_model_map.get(draft_model_path, draft_model_path) + + model_key = (model_path, adapter_path, draft_model_path) + if self.model_key != model_key: + self._load(*model_key) return self.model, self.tokenizer @@ -489,22 +485,21 @@ class ResponseGenerator: if not self._is_distributed: return obj - with mx.stream(generation_stream): - if self._rank == 0: - if obj is None: - mx.eval(mx.distributed.all_sum(0)) - return None - data = mx.array(pickle.dumps(obj)) - mx.eval(mx.distributed.all_sum(data.size)) - mx.eval(mx.distributed.all_sum(data)) - return obj - else: - size = mx.distributed.all_sum(0).item() - if size == 0: - return None - data = mx.zeros(size, dtype=mx.uint8) - data = mx.distributed.all_sum(data) - return pickle.loads(data) + if self._rank == 0: + if obj is None: + mx.eval(mx.distributed.all_sum(0)) + return None + data = mx.array(pickle.dumps(obj)) + mx.eval(mx.distributed.all_sum(data.size)) + mx.eval(mx.distributed.all_sum(data)) + return obj + else: + size = mx.distributed.all_sum(0).item() + if size == 0: + return None + data = mx.zeros(size, dtype=mx.uint8) + data = mx.distributed.all_sum(data) + return pickle.loads(data) def _share_request(self, request): if not self._is_distributed: @@ -691,6 +686,14 @@ class ResponseGenerator: return self.model_provider.is_batchable and args.seed is None def _generate(self): + # Local thread stream that we 'll pass to the BatchGenerator to make + # sure that all generation runs in the same stream as the + # synchronization messages. + generation_stream = mx.default_stream(mx.default_device()) + + # Load the default model if it is given + self.model_provider.load_default() + current_model = None current_sampling = None current_tokenizer = None @@ -820,6 +823,7 @@ class ResponseGenerator: completion_batch_size=self.cli_args.decode_concurrency, prefill_batch_size=self.cli_args.prompt_concurrency, prefill_step_size=self.cli_args.prefill_step_size, + stream=generation_stream, ) unprocessed_requests.append((rqueue, request, args)) continue @@ -909,12 +913,11 @@ class ResponseGenerator: uids_to_remove = self._share_object(uids_to_remove) if uids_to_remove: - with mx.stream(generation_stream): - batch_generator.remove(uids_to_remove) - for uid in uids_to_remove: - # It may have already been removed during - # generation - batch_results.pop(uid, None) + batch_generator.remove(uids_to_remove) + for uid in uids_to_remove: + # It may have already been removed during + # generation + batch_results.pop(uid, None) def _serve_single(self, request): rqueue, request, args = request diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 70bf8c8..ef3d266 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -507,6 +507,8 @@ def sharded_load( pipeline_group: Optional[mx.distributed.Group] = None, tensor_group: Optional[mx.distributed.Group] = None, return_config: bool = False, + *, + tokenizer_config: Optional[Dict[str, Any]] = None, ): # Get model path with everything but weight safetensors model_path = _download( @@ -571,7 +573,7 @@ def sharded_load( # Load and shard the model, and load the weights tokenizer = load_tokenizer( model_path, - {"trust_remote_code": True}, + tokenizer_config or {"trust_remote_code": True}, eos_token_ids=config.get("eos_token_id", None), ) model, _ = load_model(model_path, lazy=True, strict=False) diff --git a/setup.py b/setup.py index b4d45da..dd769a7 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ sys.path.append(str(package_dir)) from _version import __version__ -MIN_MLX_VERSION = "0.30.4" +MIN_MLX_VERSION = "0.31.2" setup( name="mlx-lm", diff --git a/tests/test_gguf.py b/tests/test_gguf.py index c7774c6..f7e789a 100644 --- a/tests/test_gguf.py +++ b/tests/test_gguf.py @@ -35,6 +35,9 @@ class TestConvertToGGUFWithoutMocks(unittest.TestCase): mock_tokenizer.get_vocab.return_value = {"": 0, "hello": 1, "world": 2} mock_tokenizer.all_special_tokens = [""] mock_tokenizer.all_special_ids = [0] + mock_tokenizer.bos_token_id = None + mock_tokenizer.eos_token_id = None + mock_tokenizer.unk_token_id = None mock_from_pretrained.return_value = mock_tokenizer model_path = Path(self.test_dir) diff --git a/tests/test_server.py b/tests/test_server.py index 0ed8e84..9a8a2ad 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -68,6 +68,9 @@ class DummyModelProvider: assert model in ["default_model", "chat_model"] return self.model, self.tokenizer + def load_default(self): + return self.load("default_model", None, "default_model") + class MockCache: def __init__(self, value, is_trimmable: bool = True):