electron-rare/mlx-lm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Thread local generation stream (#1090)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-04-22 00:34:09 -07:00
committed by GitHub
parent 4f5cbd2a4f
commit ed1fca4cef
6 changed files with 117 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx_lm/generate.py
+12 -4
View File
@@ -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:
mlx_lm/server.py
+95 -92
View File
@@ -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
mlx_lm/utils.py
+3 -1
View File
@@ -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)
setup.py
+1 -1
View File
@@ -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",
tests/test_gguf.py
+3
View File
@@ -35,6 +35,9 @@ class TestConvertToGGUFWithoutMocks(unittest.TestCase):
mock_tokenizer.get_vocab.return_value = {"<pad>": 0, "hello": 1, "world": 2}
mock_tokenizer.all_special_tokens = ["<pad>"]
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)
tests/test_server.py
+3
View File
@@ -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):