# Copyright © 2023 Apple Inc. import argparse import gc import json import logging import pickle import platform import socket import time import uuid import warnings from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from queue import Empty as QueueEmpty from queue import Queue from threading import Thread from typing import ( Any, Callable, Dict, List, Literal, Optional, Tuple, Union, ) import mlx.core as mx from huggingface_hub import scan_cache_dir from ._version import __version__ from .generate import ( DEFAULT_QUANTIZED_KV_START, BatchGenerator, TextStateMachine, make_stop_sequences, make_text_state_machine, stream_generate, ) from .models.cache import LRUPromptCache, make_prompt_cache from .sample_utils import make_logits_processors, make_sampler from .utils import ( _parse_size, load, maybe_set_recommended_wired_limit, sharded_load, ) def get_system_fingerprint(): gpu_arch = mx.device_info()["architecture"] return f"{__version__}-{mx.__version__}-{platform.platform()}-{gpu_arch}" class ToolCallFormatter: def __init__(self, tool_parser, tools, streaming=False): self._idx = 0 self._tool_parser = tool_parser self._tools = tools self._streaming = streaming def _format(self, tc): tc_id = tc.pop("id", None) or str(uuid.uuid4()) tc["arguments"] = json.dumps(tc["arguments"], ensure_ascii=False) out = { "function": tc, "type": "function", "id": tc_id, } if self._streaming: out["index"] = self._idx self._idx += 1 return out def __call__(self, tool_calls): if not tool_calls: return [] result = [] for tool_text in tool_calls: try: parsed = self._tool_parser(tool_text, self._tools) except (ValueError, json.JSONDecodeError) as e: logging.warning( f"Failed to parse tool call ({type(e).__name__}: {e}) — " f"tool text was likely truncated mid-generation." ) continue if not isinstance(parsed, list): parsed = [parsed] result.extend(self._format(tc) for tc in parsed) return result def convert_chat(messages: List[dict], role_mapping: Optional[dict] = None): default_role_mapping = { "system_prompt": ( "A chat between a curious user and an artificial intelligence " "assistant. The assistant follows the given rules no matter what." ), "system": "ASSISTANT's RULE: ", "user": "USER: ", "assistant": "ASSISTANT: ", "stop": "\n", } role_mapping = role_mapping or default_role_mapping prompt = "" for line in messages: role_prefix = role_mapping.get(line["role"], "") stop = role_mapping.get("stop", "") content = line.get("content", "") prompt += f"{role_prefix}{content}{stop}" prompt += role_mapping.get("assistant", "") return prompt.rstrip() def process_message_content(messages): """ Convert message content to a format suitable for `apply_chat_template`. The function operates on messages in place. It converts the 'content' field to a string instead of a list of text fragments. Args: message_list (list): A list of dictionaries, where each dictionary may have a 'content' key containing a list of dictionaries with 'type' and 'text' keys. Raises: ValueError: If the 'content' type is not supported or if 'text' is missing. """ for message in messages: content = message.get("content") if isinstance(content, list): text_fragments = [ fragment["text"] for fragment in content if fragment["type"] == "text" ] if len(text_fragments) != len(content): raise ValueError("Only 'text' content type is supported.") message["content"] = "".join(text_fragments) elif content is None: message["content"] = "" if tool_calls := message.get("tool_calls"): for tool_call in tool_calls: if func := tool_call.get("function"): if args := func.get("arguments"): func["arguments"] = json.loads(args) @dataclass class ModelDescription: model: str draft: str adapter: str @dataclass class SamplingArguments: temperature: float top_p: float top_k: int min_p: float xtc_probability: float xtc_threshold: float @dataclass class LogitsProcessorArguments: logit_bias: Optional[Dict[int, float]] repetition_penalty: float repetition_context_size: int presence_penalty: float presence_context_size: int frequency_penalty: float frequency_context_size: int @dataclass class GenerationArguments: model: ModelDescription sampling: SamplingArguments logits: LogitsProcessorArguments stop_words: List[str] max_tokens: int num_draft_tokens: int logprobs: bool top_logprobs: int seed: Optional[int] chat_template_kwargs: Optional[Dict[str, Any]] @dataclass class CompletionRequest: request_type: Literal["chat", "text"] prompt: str messages: List[Any] tools: Optional[List[Any]] role_mapping: Optional[Dict[str, Any]] @dataclass class GenerationContext: has_tool_calling: bool has_thinking: bool tool_parser: Callable[[str, Any], Dict] text_sm: TextStateMachine initial_state: str prompt: List[int] prompt_cache_count: int = -1 _should_stop: bool = False def stop(self): self._should_stop = True @dataclass class Response: text: str token: int logprob: float finish_reason: Optional[str] top_tokens: Tuple[Dict[str, Any]] class TimeBudget: def __init__(self, budget=0.5, iterations=25, sync_frequency=10): self._is_distributed = mx.distributed.init().size() > 1 self._budget = budget self._iterations = iterations self._sync_frequency = sync_frequency self._start = None self._current_iterations = None self._loops = 0 self._time_spent = 0 def __iter__(self): self._start = time.time() self._current_iterations = 0 return self def __next__(self): if not self._is_distributed: if time.time() - self._start > self._budget: raise StopIteration() return None self._current_iterations += 1 if self._current_iterations <= self._iterations: return None self._loops += 1 self._time_spent += time.time() - self._start if self._loops % self._sync_frequency == 0: loop_time = mx.distributed.all_sum(self._time_spent).item() avg_loop_time = loop_time / ( mx.distributed.init().size() * self._sync_frequency ) factor = self._budget / avg_loop_time self._iterations = max(round(self._iterations * factor), 1) self._loops = 0 self._time_spent = 0 raise StopIteration() class ModelProvider: def __init__(self, cli_args: argparse.Namespace): """Load models on demand and persist them across the whole process.""" self.cli_args = cli_args self.model_key = None self.model = None self.tokenizer = None self.draft_model = None self.is_batchable = False group = mx.distributed.init() self.pipeline_group = group if group.size() > 1 and cli_args.pipeline else None self.tensor_group = ( group if group.size() > 1 and not cli_args.pipeline else None ) self.is_distributed = group.size() > 1 # 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 # Build the tokenizer config for later use in load self._tokenizer_config = {"trust_remote_code": cli_args.trust_remote_code} if cli_args.chat_template: self._tokenizer_config["chat_template"] = cli_args.chat_template def reset(self) -> None: self.model_key = None self.model = None self.tokenizer = None self.draft_model = None self.is_batchable = False 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. Dropping the refs returns the # weights to MLX's buffer pool, not the OS; clear that pool so a # later load of a different model does not keep every previous one. self.reset() gc.collect() mx.clear_cache() # 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, trust_remote_code=self.cli_args.trust_remote_code, ) else: model, tokenizer = load( model_path, adapter_path=adapter_path, tokenizer_config=self._tokenizer_config, trust_remote_code=self.cli_args.trust_remote_code, ) # 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 # 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." ) # Compute batchability is_batchable = draft_model is None is_batchable = is_batchable and all( hasattr(c, "merge") for c in make_prompt_cache(model) ) # 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 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): adapter_path = self._adapter_map.get(model_path, adapter_path) model_path = self._model_map.get(model_path, model_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 def _make_sampler(args, tokenizer): return make_sampler( args.sampling.temperature, top_p=args.sampling.top_p, top_k=args.sampling.top_k, min_p=args.sampling.min_p, xtc_probability=args.sampling.xtc_probability, xtc_threshold=args.sampling.xtc_threshold, xtc_special_tokens=tokenizer.encode("\n") + list(tokenizer.eos_token_ids), ) def _make_logits_processors(args): return make_logits_processors( args.logits.logit_bias, args.logits.repetition_penalty, args.logits.repetition_context_size, args.logits.presence_penalty, args.logits.presence_context_size, args.logits.frequency_penalty, args.logits.frequency_context_size, ) def _format_top_logprobs(logprobs, top_n, tokenizer) -> Tuple[Dict[str, Any]]: """Returns info dicts for the top `top_n` tokens from `logprobs`""" if top_n <= 0: return () sorted_indices = mx.argpartition(-logprobs, kth=top_n - 1) top_indices = sorted_indices[:top_n].tolist() top_probs = logprobs[top_indices].tolist() txts = tokenizer.convert_ids_to_tokens(top_indices) return tuple( {"id": i, "token": s, "logprob": g} for i, s, g in zip(top_indices, txts, top_probs) ) class ResponseGenerator: def __init__(self, model_provider: ModelProvider, prompt_cache: LRUPromptCache): self.model_provider = model_provider self.prompt_cache = prompt_cache self.requests = Queue() self._state_machine_cache = {} self._time_budget = TimeBudget() self._is_distributed = mx.distributed.init().size() > 1 self._rank = mx.distributed.init().rank() self._stop = False self._generation_failed = False self._generation_thread = Thread(target=self._run_generate) self._generation_thread.start() def _run_generate(self): try: self._generate() except Exception as e: logging.exception(f"mlx_lm.server generation thread died: {e}") self._generation_failed = True def generation_available(self): return self._generation_thread.is_alive() and not self._generation_failed def stop_and_join(self): self._stop = True self._generation_thread.join() def join(self): self._generation_thread.join() @property def is_healthy(self): return self.generation_available() def _log_cache_stats(self): n_sequences = len(self.prompt_cache) n_bytes = self.prompt_cache.nbytes logging.info(f"Prompt Cache: {n_sequences} sequences, {n_bytes / 1e9:.2f} GB") for cache_type, stats in self.prompt_cache.stats_by_type().items(): n_sequences = stats["n_sequences"] n_bytes = stats["n_bytes"] logging.info( f"- {cache_type}: {n_sequences} sequences, {n_bytes / 1e9:.2f} GB" ) def _next_request(self, timeout=None): request = None if not self._is_distributed or self._rank == 0: try: if timeout is not None: request = self.requests.get(timeout=timeout) else: request = self.requests.get_nowait() except QueueEmpty: pass return self._share_request(request) def _share_object(self, obj): if not self._is_distributed: return obj 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: return request shareable = request[1:] if request is not None else None shareable = self._share_object(shareable) if shareable is None: return None rq = request[0] if request is not None else Queue() return rq, *shareable def _tokenize(self, tokenizer, request, args): """Tokenize a request and split the prompt into segments. Returns a tuple * prompt - Full list of tokens * segments - A list of lists of tokens. Up to 3 segments that correspond to system prompt, context, thinking tail. * segment_types - A string per segment indicating if the segment is a system prompt or a user prompt or nothing special. * initial state - A string that contains the initial state of the state machine (normal or thinking depending on whether we have tail or not) """ if request.request_type == "chat": messages = request.messages tools = request.tools role_mapping = request.role_mapping if tokenizer.has_chat_template: process_message_content(messages) if tools and not tokenizer.has_tool_calling: logging.warning( "Received tools but model does not support tool calling. " "If you think this is an error, file an issue here: " "https://github.com/ml-explore/mlx-lm/issues" ) chat_template_args = self.model_provider.cli_args.chat_template_args if args.chat_template_kwargs: chat_template_args = chat_template_args.copy() chat_template_args.update(args.chat_template_kwargs) template_kwargs = dict( tools=tools, tokenize=True, **chat_template_args, ) prompt = tokenizer.apply_chat_template( messages, add_generation_prompt=True, **template_kwargs, ) else: prompt = tokenizer.encode(convert_chat(messages, role_mapping)) return prompt, [prompt], ["assistant"], "normal" else: prompt = tokenizer.encode(request.prompt) return prompt, [prompt], ["assistant"], "normal" # If we are here it means we have a chat request so we need to search # for segments for better cache management. # Choose the initial state among only reasoning or normal initial_state = "normal" if tokenizer.has_thinking: think_start = tokenizer.rfind_think_start(prompt) think_end = tokenizer.rfind_think_end(prompt) if think_start > think_end: initial_state = "reasoning" # It is not a user message so no segmentation needed. if messages[-1]["role"] != "user": return prompt, [prompt], ["assistant"], initial_state segments = [] segment_types = [] # Find where the system prompt ends and add it as a segment. num_system = 0 sys_end = 0 for m in messages: if m["role"] == "system": num_system += 1 else: break if num_system > 0: sys_tokens = tokenizer.apply_chat_template( messages[:num_system] + [{"role": "user", "content": ""}], add_generation_prompt=False, **template_kwargs, ) for i, (a, b) in enumerate(zip(sys_tokens, prompt)): if a != b: sys_end = i break if sys_end > 0 and sys_end < len(prompt): segments.append(prompt[:sys_end]) segment_types.append("system") # Find a tail segment that contains thinking tokens (small up to 11 # tokens) tail_start = len(prompt) if tokenizer.has_thinking: think_start = tokenizer.rfind_think_start(prompt, start=tail_start - 11) if think_start >= 0: tail_start = think_start # Finalize the segments and return if sys_end < tail_start: segments.append(prompt[sys_end:tail_start]) segment_types.append("user") if tail_start < len(prompt): segments.append(prompt[tail_start:]) segment_types.append("assistant") if not segments: segments = [prompt] segment_types = ["assistant"] return prompt, segments, segment_types, initial_state def _make_state_machine(self, model_key, tokenizer, stop_words): """Make (and cache) a StopSequences and TextStateMachine.""" cache_key = (model_key, tuple(stop_words)) rs = self._state_machine_cache.get(cache_key) if rs is not None: return rs stop_sequences = make_stop_sequences(tokenizer, stop_words) text_sm = make_text_state_machine(tokenizer, stop_words) if len(self._state_machine_cache) > 100: self._state_machine_cache.clear() self._state_machine_cache[cache_key] = (stop_sequences, text_sm) return stop_sequences, text_sm def _is_batchable(self, args): return ( self.model_provider.is_batchable and args.seed is None and self.cli_args.kv_bits 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_tokenizer = None current_model_key = None batch_generator = None drain_batch = False batch_results = {} unprocessed_requests = [] def get_next_request(timeout=None): if unprocessed_requests: return unprocessed_requests.pop() else: return self._next_request(timeout) if self._is_distributed: seed = mx.distributed.all_sum(mx.random.state[0]).view(mx.uint64).item() mx.random.seed(seed) while not self._stop: request = None if not drain_batch: timeout = ( None if (batch_generator is not None and len(batch_results) > 0) else 0.1 ) request = get_next_request(timeout=timeout) # We got a request if request is not None: rqueue, request, args = request # Can it be added to the current batch? if ( batch_generator is not None and current_model == args.model and self._is_batchable(args) ): try: prompt, segments, segment_types, initial_state = self._tokenize( current_tokenizer, request, args ) except Exception as e: rqueue.put(e) continue stop_sequences, text_sm = self._make_state_machine( self.model_provider.model_key, current_tokenizer, args.stop_words, ) self._log_cache_stats() cache, rest = self.prompt_cache.fetch_nearest_cache( current_model_key, prompt ) prompt_cache_count = len(prompt) - len(rest) N = prompt_cache_count while N > 0: if N >= len(segments[0]): N -= len(segments.pop(0)) segment_types.pop(0) else: segments[0] = segments[0][N:] break ctx = GenerationContext( has_tool_calling=current_tokenizer.has_tool_calling, has_thinking=current_tokenizer.has_thinking, tool_parser=current_tokenizer.tool_parser, text_sm=text_sm, initial_state=initial_state, prompt=prompt, prompt_cache_count=prompt_cache_count, ) rqueue.put(ctx) (uid,) = batch_generator.insert_segments( segments=[segments], max_tokens=[args.max_tokens], caches=[cache], all_tokens=[prompt[:prompt_cache_count]], samplers=[_make_sampler(args, current_tokenizer)], logits_processors=[_make_logits_processors(args)], stop_sequences=[stop_sequences], ) batch_results[uid] = { "ctx": ctx, "rqueue": rqueue, "detokenizer": current_tokenizer.detokenizer, "segment_types": segment_types[::-1], "top_logprobs": args.top_logprobs, } # just making sure we don't leave a reference around del cache if self.model_provider.cli_args.prompt_cache_bytes is not None: total = self.model_provider.cli_args.prompt_cache_bytes active = batch_generator.prompt_cache_nbytes self.prompt_cache.trim_to(n_bytes=total - active) continue # No batch generator. Load the model and if it's not # batchable serve sequential, o/w make a batch generaotr and # serve batched elif batch_generator is None: try: model, tokenizer = self.model_provider.load( args.model.model, args.model.adapter, args.model.draft ) except Exception as e: rqueue.put(e) continue if not self._is_batchable(args): self._serve_single((rqueue, request, args), generation_stream) continue current_model = args.model current_tokenizer = tokenizer current_model_key = self.model_provider.model_key batch_results = {} batch_generator = BatchGenerator( model, 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 # We have a batch but this request cannot be added to the # batch so drain it to process the request. else: drain_batch = True unprocessed_requests.append((rqueue, request, args)) continue # No request so serve from the current batch elif batch_generator is not None: if len(batch_results) == 0: if drain_batch: current_model = None current_tokenizer = None current_model_key = None batch_generator.close() batch_generator = None drain_batch = False continue uids_to_remove = [] for _ in self._time_budget: prompt_responses, gen_responses = batch_generator.next() if not prompt_responses and not gen_responses: break # Progress report for prompt processing for r in prompt_responses: result = batch_results[r.uid] result["rqueue"].put(r.progress) if result["ctx"]._should_stop: uids_to_remove.append(r.uid) # Save the caches at end of segments eos_ids = [ r.uid for r in prompt_responses if r.end_of_segment and not r.end_of_prompt and batch_results[r.uid]["segment_types"] ] caches = batch_generator.extract_cache(eos_ids) for uid, (cache, cache_key) in caches.items(): self.prompt_cache.insert_cache( self.model_provider.model_key, cache_key[:], cache, cache_type=batch_results[uid]["segment_types"].pop(), ) del caches for r in gen_responses: result = batch_results[r.uid] # Don't decode the final stop token if r.finish_reason == "stop": result["detokenizer"].finalize() text = result["detokenizer"].last_segment elif r.finish_reason == "length": result["detokenizer"].add_token(r.token) result["detokenizer"].finalize() text = result["detokenizer"].last_segment else: result["detokenizer"].add_token(r.token) text = result["detokenizer"].last_segment result["rqueue"].put( Response( text, r.token, r.logprobs[r.token].item(), r.finish_reason, _format_top_logprobs( r.logprobs, result["top_logprobs"], current_tokenizer, ), ) ) if r.finish_reason is not None: result["rqueue"].put(None) self.prompt_cache.insert_cache( current_model_key, r.all_tokens[:], r.prompt_cache, cache_type="assistant", ) del batch_results[r.uid] if result["ctx"]._should_stop: uids_to_remove.append(r.uid) uids_to_remove = self._share_object(uids_to_remove) if uids_to_remove: 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) # Make sure the model and prompt cache are destroyed in the generation # thread under same stream. self.model_provider.reset() del self.prompt_cache gc.collect() def _serve_single(self, request, stream): rqueue, request, args = request # Define the progress callback def progress(tokens_processed, tokens_total): rqueue.put((tokens_processed, tokens_total)) try: # Load the model and tokenizer model = self.model_provider.model tokenizer = self.model_provider.tokenizer draft_model = self.model_provider.draft_model # Prepare the prompt and state machine prompt, _, _, initial_state = self._tokenize(tokenizer, request, args) stop_sequences, text_sm = self._make_state_machine( self.model_provider.model_key, tokenizer, args.stop_words, ) # Start the generation context ctx = GenerationContext( has_thinking=tokenizer.has_thinking, has_tool_calling=tokenizer.has_tool_calling, tool_parser=tokenizer.tool_parser, text_sm=text_sm, initial_state=initial_state, prompt=prompt, ) rqueue.put(ctx) # Seed if requested if args.seed is not None: mx.random.seed(args.seed) # Make the sampler and logit processor sampler = _make_sampler(args, tokenizer) logits_processors = _make_logits_processors(args) # Load the KV cache self._log_cache_stats() cache, rest = self.prompt_cache.fetch_nearest_cache( self.model_provider.model_key, prompt ) ctx.prompt_cache_count = len(prompt) - len(rest) cache_key = prompt[:] if cache is None: cache = make_prompt_cache(self.model_provider.model) if self.model_provider.draft_model is not None: cache += make_prompt_cache(self.model_provider.draft_model) # Process the prompt and generate tokens # Own matcher: the automaton is cached across requests. stop_matcher = stop_sequences.matcher() for gen in stream_generate( model=model, stream=stream, tokenizer=tokenizer, prompt=rest, max_tokens=args.max_tokens, sampler=sampler, logits_processors=logits_processors, prompt_cache=cache, draft_model=draft_model, num_draft_tokens=args.num_draft_tokens, prompt_progress_callback=progress, prefill_step_size=self.cli_args.prefill_step_size, kv_bits=self.cli_args.kv_bits, kv_group_size=self.cli_args.kv_group_size, quantized_kv_start=self.cli_args.quantized_kv_start, ): finish_reason = gen.finish_reason # Token-level stop word detection if stop_matcher.advance(gen.token): finish_reason = "stop" rqueue.put( Response( gen.text, gen.token, gen.logprobs[gen.token].item(), finish_reason, _format_top_logprobs( gen.logprobs, args.top_logprobs, tokenizer ), ) ) cache_key.append(gen.token) if ctx._should_stop: if self._is_distributed: raise NotImplementedError() break if finish_reason is not None: break rqueue.put(None) # Save the KV cache again self.prompt_cache.insert_cache( self.model_provider.model_key, cache_key, cache ) except Exception as e: rqueue.put(e) def _await_response(self, response_queue): # Wherever the request was when the thread died, nothing more will be # put on its queue, so give up rather than block forever. while True: try: return response_queue.get(timeout=1.0) except QueueEmpty: if not self.generation_available(): raise RuntimeError("generation thread died") from None def generate( self, request: CompletionRequest, generation_args: GenerationArguments, progress_callback: Optional[Callable[[int, int], None]] = None, ): response_queue = Queue() if not self.generation_available(): raise RuntimeError("generation thread died") self.requests.put((response_queue, request, generation_args)) def _inner(): while True: response = self._await_response(response_queue) if response is None: break if isinstance(response, Exception): raise response if isinstance(response, tuple): if progress_callback is not None: progress_callback(*response) continue yield response ctx = self._await_response(response_queue) if isinstance(ctx, Exception): raise ctx return ctx, _inner() @property def cli_args(self): return self.model_provider.cli_args class APIHandler(BaseHTTPRequestHandler): def __init__( self, response_generator: ResponseGenerator, *args, system_fingerprint: Optional[str] = None, **kwargs, ): """ Create static request specific metadata """ self.created = int(time.time()) self.response_generator = response_generator self.system_fingerprint = system_fingerprint or get_system_fingerprint() super().__init__(*args, **kwargs) def _set_cors_headers(self): allowed_origins = self.response_generator.cli_args.allowed_origins origin = self.headers.get("Origin") if "*" in allowed_origins: self.send_header("Access-Control-Allow-Origin", "*") elif origin in allowed_origins: self.send_header("Access-Control-Allow-Origin", origin) self.send_header("Vary", "Origin") self.send_header("Access-Control-Allow-Methods", "*") self.send_header("Access-Control-Allow-Headers", "*") def _set_completion_headers(self, status_code: int = 200): self.send_response(status_code) self.send_header("Content-type", "application/json") self._set_cors_headers() def _set_stream_headers(self, status_code: int = 200): self.send_response(status_code) self.send_header("Content-type", "text/event-stream") self.send_header("Cache-Control", "no-cache") self._set_cors_headers() def do_OPTIONS(self): self._set_completion_headers(204) self.end_headers() def do_POST(self): """ Respond to a POST request from a client. """ request_factories = { "/v1/completions": self.handle_text_completions, "/v1/chat/completions": self.handle_chat_completions, "/chat/completions": self.handle_chat_completions, } if self.path not in request_factories: self._set_completion_headers(404) self.end_headers() self.wfile.write(b"Not Found") return # Fetch and parse request body content_length = self.headers.get("Content-Length") if content_length is None: self._set_completion_headers(411) self.end_headers() self.wfile.write( json.dumps({"error": "Content-Length header is required"}).encode() ) return try: content_length = int(content_length) except ValueError: self._set_completion_headers(400) self.end_headers() self.wfile.write( json.dumps({"error": "Invalid Content-Length header"}).encode() ) return raw_body = self.rfile.read(content_length) try: self.body = json.loads(raw_body.decode()) except json.JSONDecodeError as e: logging.error(f"JSONDecodeError: {e} - Raw body: {raw_body.decode()}") self._set_completion_headers(400) self.end_headers() self.wfile.write( json.dumps({"error": f"Invalid JSON in request body: {e}"}).encode() ) return if logging.getLogger().isEnabledFor(logging.DEBUG): debug_body = json.dumps(self.body, indent="\t") logging.debug(f"Incoming Request Body: {debug_body}") if not isinstance(self.body, dict): debug_body = json.dumps(self.body, indent="\t") logging.error(f"Invalid Request Body: {debug_body}") self._set_completion_headers(400) self.end_headers() self.wfile.write( json.dumps({"error": "Request should be a JSON dictionary"}).encode() ) return # Extract request parameters from the body self.stream = self.body.get("stream", False) self.stream_options = self.body.get("stream_options", None) self.requested_model = self.body.get("model", "default_model") self.requested_draft_model = self.body.get("draft_model", "default_model") self.num_draft_tokens = self.body.get( "num_draft_tokens", self.response_generator.cli_args.num_draft_tokens ) self.adapter = self.body.get("adapters", None) self.max_tokens = self.body.get("max_completion_tokens", None) if self.max_tokens is None: self.max_tokens = self.body.get( "max_tokens", self.response_generator.cli_args.max_tokens ) self.temperature = self.body.get( "temperature", self.response_generator.cli_args.temp ) self.top_p = self.body.get("top_p", self.response_generator.cli_args.top_p) self.top_k = self.body.get("top_k", self.response_generator.cli_args.top_k) self.min_p = self.body.get("min_p", self.response_generator.cli_args.min_p) self.repetition_penalty = self.body.get("repetition_penalty", 0.0) self.repetition_context_size = self.body.get("repetition_context_size", 20) self.presence_penalty = self.body.get("presence_penalty", 0.0) self.presence_context_size = self.body.get("presence_context_size", 20) self.frequency_penalty = self.body.get("frequency_penalty", 0.0) self.frequency_context_size = self.body.get("frequency_context_size", 20) self.xtc_probability = self.body.get("xtc_probability", 0.0) self.xtc_threshold = self.body.get("xtc_threshold", 0.1) self.logit_bias = self.body.get("logit_bias", None) self.logprobs = self.body.get("logprobs", False) self.top_logprobs = self.body.get("top_logprobs", -1) self.seed = self.body.get("seed", None) self.chat_template_kwargs = self.body.get("chat_template_kwargs") self.validate_model_parameters() # Get stop sequences stop_words = self.body.get("stop") stop_words = stop_words or [] stop_words = [stop_words] if isinstance(stop_words, str) else stop_words # Create the completion request request = request_factories[self.path]() self.handle_completion(request, stop_words) def _validate( self, name, expected_type, min_val=None, max_val=None, optional=False, whitelist=None, ): value = getattr(self, name) if optional and value is None: return if not isinstance(value, expected_type): try: allowed = tuple(et.__name__ for et in expected_type) except TypeError: allowed = expected_type.__name__ raise ValueError(f"{name} must be of type {allowed}") if whitelist is not None and value in whitelist: return if min_val is not None and value < min_val: raise ValueError(f"{name} must be at least {min_val}") if max_val is not None and value > max_val: raise ValueError(f"{name} must be at most {max_val}") def validate_model_parameters(self): """Validate that the passed model parameters have correct types and values.""" self._validate("stream", bool) self._validate("max_tokens", int, min_val=0) self._validate("temperature", (float, int), min_val=0) self._validate("top_p", (float, int), min_val=0, max_val=1) self._validate("top_k", int, min_val=0) self._validate("min_p", (float, int), min_val=0, max_val=1) self._validate("num_draft_tokens", int, min_val=0) self._validate("repetition_penalty", (float, int), min_val=0) self._validate("repetition_context_size", int, min_val=0) self._validate("presence_penalty", (float, int)) self._validate("presence_context_size", int, min_val=0) self._validate("frequency_penalty", (float, int)) self._validate("frequency_context_size", int, min_val=0) self._validate("logprobs", bool) self._validate("top_logprobs", int, min_val=0, max_val=11, whitelist=[-1]) self._validate("xtc_probability", float, min_val=0, max_val=1) self._validate("xtc_threshold", float, min_val=0, max_val=1) self._validate("requested_model", str) self._validate("adapter", str, optional=True) self._validate("seed", int, optional=True) self._validate("logit_bias", dict, optional=True) if self.logit_bias is not None: try: self.logit_bias = {int(k): float(v) for k, v in self.logit_bias.items()} except ValueError as e: raise ValueError("logit_bias must be a dict of int to float") from e def generate_response( self, text: str, finish_reason: Union[Literal["length", "stop"], None], prompt_token_count: Optional[int] = None, completion_token_count: Optional[int] = None, prompt_cache_count: Optional[int] = None, token_logprobs: Optional[List[float]] = None, top_tokens: Optional[List[Tuple[Dict[str, Any]]]] = None, tokens: Optional[List[int]] = None, tool_calls: Optional[List[str]] = None, reasoning_text: Optional[str] = None, ) -> dict: """ Generate a single response packet based on response type (stream or not), completion type and parameters. Args: text (str): Text generated by model finish_reason (Union[Literal["length", "stop"], None]): The reason the response is being sent: "length", "stop" or `None`. prompt_token_count (Optional[int]): The number of tokens in the prompt, used to populate the "usage" field (not used when stream). completion_token_count (Optional[int]): The number of tokens in the response, used to populate the "usage" field (not used when stream). prompt_cache_count (Optional[int]): The portion of prompt_token_count that was found in the cache when servicing the request. token_logprobs (Optional[List[float]]): The log probabilities per token, in token order. top_tokens (Optional[List[Tuple[Dict[str, Any]]]]): List of outputs from _format_top_logprobs, giving info on the top N tokens at each token position. tokens (Optional[List[int]]): List of tokens to return with logprobs structure tool_calls (Optional[List[str]]): List of tool calls. reasoning_text (Optional[str]): The reasoning text generated by the model. Returns: dict: A dictionary containing the response, in the same format as OpenAI's API. """ token_logprobs = token_logprobs or [] top_logprobs = top_tokens or [] tool_calls = tool_calls or [] # Static response response = { "id": self.request_id, "system_fingerprint": self.system_fingerprint, "object": self.object_type, "model": self.requested_model, "created": self.created, "choices": [ { "index": 0, "finish_reason": finish_reason, }, ], } if top_logprobs: response["choices"][0]["logprobs"] = { "content": [ dict(i[0], top_logprobs=i) if i else {} for i in top_logprobs ] } elif token_logprobs: response["choices"][0]["logprobs"] = { "content": [ dict(id=i, logprob=g) for i, g in zip(tokens, token_logprobs) ] } if not self.stream: if not ( isinstance(prompt_token_count, int) and isinstance(completion_token_count, int) ): raise ValueError( "Response type is complete, but token counts not provided" ) response["usage"] = { "prompt_tokens": prompt_token_count, "completion_tokens": completion_token_count, "total_tokens": prompt_token_count + completion_token_count, } if prompt_cache_count is not None and prompt_cache_count >= 0: response["usage"]["prompt_tokens_details"] = { "cached_tokens": prompt_cache_count, } choice = response["choices"][0] # Add dynamic response if self.object_type.startswith("chat.completion"): key_name = "delta" if self.stream else "message" choice[key_name] = {"role": "assistant"} if not self.stream: # The schema requires "content" field to be present choice[key_name]["content"] = text if text else None elif text: choice[key_name]["content"] = text if reasoning_text: choice[key_name]["reasoning"] = reasoning_text if tool_calls: choice[key_name]["tool_calls"] = tool_calls elif self.object_type == "text_completion": choice.update(text=text) else: raise ValueError(f"Unsupported response type: {self.object_type}") return response def handle_completion(self, request: CompletionRequest, stop_words: List[str]): """ Generate a response to a prompt and send it to the client in a single batch. Args: prompt (List[int]): The tokenized prompt. stop_words (List[str]): A list of stop words """ args = GenerationArguments( model=ModelDescription( model=self.requested_model, draft=self.requested_draft_model, adapter=self.adapter, ), sampling=SamplingArguments( temperature=self.temperature, top_p=self.top_p, top_k=self.top_k, min_p=self.min_p, xtc_probability=self.xtc_probability, xtc_threshold=self.xtc_threshold, ), logits=LogitsProcessorArguments( logit_bias=self.logit_bias, repetition_penalty=self.repetition_penalty, repetition_context_size=self.repetition_context_size, presence_penalty=self.presence_penalty, presence_context_size=self.presence_context_size, frequency_penalty=self.frequency_penalty, frequency_context_size=self.frequency_context_size, ), stop_words=stop_words, max_tokens=self.max_tokens, num_draft_tokens=self.num_draft_tokens, logprobs=self.logprobs, top_logprobs=self.top_logprobs, seed=self.seed, chat_template_kwargs=self.chat_template_kwargs, ) # Keep connection allive during long prompt processing (and also log # the progress) def keepalive_callback(processed, total): logging.info(f"Prompt processing progress: {processed}/{total}") if self.stream: msg = f": keepalive {processed}/{total}\n\n".encode() self.wfile.write(msg) self.wfile.flush() # Create the token generator try: ctx, response = self.response_generator.generate( request, args, progress_callback=keepalive_callback, ) except Exception as e: self._set_completion_headers(404) self.end_headers() self.wfile.write(json.dumps({"error": str(e)}).encode()) return # Prepare the headers if self.stream: self._set_stream_headers(200) self.end_headers() logging.debug("Starting stream:") else: self._set_completion_headers(200) logging.debug("Starting completion:") # Tool call formatter tool_formatter = ToolCallFormatter(ctx.tool_parser, request.tools, self.stream) # Initialize the text state machine sm_state = ctx.text_sm.make_state(ctx.initial_state) # Variables to save the generated text, tokens, logprobs, tools etc prev_state = ctx.initial_state finish_reason = "stop" reasoning_text = "" made_tool_call = False tool_text = "" tool_calls = [] text = "" tokens = [] token_logprobs = [] top_tokens = [] try: for gen in response: logging.debug(gen.text) # Advance the text state machine to strip control sequences if gen.finish_reason == "stop": sm_state, current_state = TextStateMachine.discard(sm_state) clean_text = "" elif gen.finish_reason == "length": sm_state, clean_text, current_state = TextStateMachine.step( sm_state, gen.text ) sm_state, flushed, current_state = TextStateMachine.flush(sm_state) clean_text += flushed else: sm_state, clean_text, current_state = TextStateMachine.step( sm_state, gen.text ) # Collect the clean text by state: reasoning, tool, or normal if current_state == "reasoning": reasoning_text += clean_text elif current_state == "tool": tool_text += clean_text elif current_state == "normal": if prev_state == "tool": tool_calls.append(tool_text) tool_text = "" made_tool_call = True text += clean_text # Add the tokens and logprobs to the vars. tokens.append(gen.token) if args.logprobs: token_logprobs.append(gen.logprob) if args.top_logprobs > 0: top_tokens.append(gen.top_tokens) if ( self.stream and current_state != "tool" and (text or tool_calls or reasoning_text) ): resp = self.generate_response( text, None, tool_calls=tool_formatter(tool_calls), reasoning_text=reasoning_text, ) self.wfile.write(f"data: {json.dumps(resp)}\n\n".encode()) self.wfile.flush() reasoning_text = "" text = "" tool_calls = [] if gen.finish_reason is not None: finish_reason = gen.finish_reason prev_state = current_state if prev_state == "tool" and tool_text: tool_calls.append(tool_text) made_tool_call = True if finish_reason == "stop" and made_tool_call: finish_reason = "tool_calls" if self.stream: resp = self.generate_response( text, finish_reason, tool_calls=tool_formatter(tool_calls), reasoning_text=reasoning_text, ) self.wfile.write(f"data: {json.dumps(resp)}\n\n".encode()) self.wfile.flush() if ( self.stream_options is not None and self.stream_options["include_usage"] ): resp = self.completion_usage_response( len(ctx.prompt), len(tokens), ctx.prompt_cache_count, ) self.wfile.write(f"data: {json.dumps(resp)}\n\n".encode()) self.wfile.flush() self.wfile.write("data: [DONE]\n\n".encode()) self.wfile.flush() else: resp = self.generate_response( text, finish_reason, len(ctx.prompt), len(tokens), ctx.prompt_cache_count, token_logprobs=token_logprobs, top_tokens=top_tokens, tokens=tokens, reasoning_text=reasoning_text, tool_calls=tool_formatter(tool_calls), ) if logging.getLogger().isEnabledFor(logging.DEBUG): response_debug = json.dumps(resp, indent="\t") logging.debug(f"Outgoing Response: {response_debug}") response_json = json.dumps(resp).encode() self.send_header("Content-Length", str(len(response_json))) self.end_headers() self.wfile.write(response_json) self.wfile.flush() finally: ctx.stop() def completion_usage_response( self, prompt_token_count: Optional[int] = None, completion_token_count: Optional[int] = None, prompt_cache_count: Optional[int] = None, ): response = { "id": self.request_id, "system_fingerprint": self.system_fingerprint, "object": "chat.completion", "model": self.requested_model, "created": self.created, "choices": [], "usage": { "prompt_tokens": prompt_token_count, "completion_tokens": completion_token_count, "total_tokens": prompt_token_count + completion_token_count, }, } if prompt_cache_count is not None and prompt_cache_count >= 0: response["usage"]["prompt_tokens_details"] = { "cached_tokens": prompt_cache_count, } return response def handle_chat_completions(self) -> CompletionRequest: """ Handle a chat completion request. Returns: mx.array: A mx.array of the tokenized prompt from the request body """ body = self.body assert "messages" in body, "Request did not contain messages" # Determine response type self.request_id = f"chatcmpl-{uuid.uuid4()}" self.object_type = "chat.completion.chunk" if self.stream else "chat.completion" return CompletionRequest( "chat", "", body["messages"], body.get("tools") or None, body.get("role_mapping"), ) def handle_text_completions(self) -> CompletionRequest: """ Handle a text completion request. Returns: mx.array: A mx.array of the tokenized prompt from the request body """ # Determine response type self.request_id = f"cmpl-{uuid.uuid4()}" self.object_type = "text_completion" assert "prompt" in self.body, "Request did not contain a prompt" return CompletionRequest( "text", self.body["prompt"], [], None, None, ) def do_GET(self): """ Respond to a GET request from a client. """ if self.path.startswith("/v1/models"): self.handle_models_request() elif self.path == "/health": self.handle_health_check() else: self._set_completion_headers(404) self.end_headers() self.wfile.write(b"Not Found") def handle_health_check(self): """ Handle a GET request for the /health endpoint. """ is_healthy = self.response_generator.is_healthy status_code = 200 if is_healthy else 503 status = "ok" if is_healthy else "unavailable" self._set_completion_headers(status_code) self.end_headers() self.wfile.write(json.dumps({"status": status}).encode()) self.wfile.flush() def handle_models_request(self): """ Handle a GET request for the /v1/models endpoint. """ self._set_completion_headers(200) self.end_headers() files = ["config.json", "model.safetensors.index.json", "tokenizer_config.json"] parts = self.path.split("/") filter_repo_id = None if len(parts) > 3: filter_repo_id = "/".join(parts[3:]) def probably_mlx_lm(repo): if repo.repo_type != "model": return False if "main" not in repo.refs: return False if filter_repo_id is not None and repo.repo_id != filter_repo_id: return False file_names = {f.file_path.name for f in repo.refs["main"].files} return all(f in file_names for f in files) # Scan the cache directory for downloaded mlx models hf_cache_info = scan_cache_dir() downloaded_models = [ repo for repo in hf_cache_info.repos if probably_mlx_lm(repo) ] # Create a list of available models models = [ { "id": repo.repo_id, "object": "model", "created": self.created, } for repo in downloaded_models ] if self.response_generator.cli_args.model: model_path = Path(self.response_generator.cli_args.model) if model_path.exists(): model_id = str(model_path.resolve()) models.append( { "id": model_id, "object": "model", "created": self.created, } ) response = {"object": "list", "data": models} response_json = json.dumps(response).encode() self.wfile.write(response_json) self.wfile.flush() def _run_http_server( host: str, port: int, response_generator, server_class=ThreadingHTTPServer, handler_class=APIHandler, ): server_address = (host, port) infos = socket.getaddrinfo( *server_address, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE ) server_class.address_family, _, _, _, server_address = next(iter(infos)) httpd = server_class( server_address, lambda *args, **kwargs: handler_class( response_generator, system_fingerprint=get_system_fingerprint(), *args, **kwargs, ), ) warnings.warn( "mlx_lm.server is not recommended for production as " "it only implements basic security checks." ) logging.info(f"Starting httpd at {host} on port {port}...") try: httpd.serve_forever() except KeyboardInterrupt: httpd.shutdown() response_generator.stop_and_join() def run( host: str, port: int, model_provider: ModelProvider, server_class=ThreadingHTTPServer, handler_class=APIHandler, ): group = mx.distributed.init() prompt_cache = LRUPromptCache(model_provider.cli_args.prompt_cache_size) response_generator = ResponseGenerator(model_provider, prompt_cache) if group.rank() == 0: _run_http_server(host, port, response_generator) else: response_generator.join() def main(): parser = argparse.ArgumentParser(description="MLX Http Server.") parser.add_argument( "--model", type=str, help="The path to the MLX model weights, tokenizer, and config", ) parser.add_argument( "--adapter-path", type=str, help="Optional path for the trained adapter weights and config.", ) parser.add_argument( "--host", type=str, default="127.0.0.1", help="Host for the HTTP server (default: 127.0.0.1)", ) parser.add_argument( "--port", type=int, default=8080, help="Port for the HTTP server (default: 8080)", ) parser.add_argument( "--allowed-origins", type=lambda x: x.split(","), default="*", help="Allowed origins (default: *)", ) parser.add_argument( "--draft-model", type=str, help="A model to be used for speculative decoding.", default=None, ) parser.add_argument( "--num-draft-tokens", type=int, help="Number of tokens to draft when using speculative decoding.", default=3, ) parser.add_argument( "--trust-remote-code", action="store_true", help="Enable trusting remote code for tokenizer", ) parser.add_argument( "--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Set the logging level (default: INFO)", ) parser.add_argument( "--chat-template", type=str, default="", help="Specify a chat template for the tokenizer", required=False, ) parser.add_argument( "--use-default-chat-template", action="store_true", help="Use the default chat template", ) parser.add_argument( "--temp", type=float, default=0.0, help="Default sampling temperature (default: 0.0)", ) parser.add_argument( "--top-p", type=float, default=1.0, help="Default nucleus sampling top-p (default: 1.0)", ) parser.add_argument( "--top-k", type=int, default=0, help="Default top-k sampling (default: 0, disables top-k)", ) parser.add_argument( "--min-p", type=float, default=0.0, help="Default min-p sampling (default: 0.0, disables min-p)", ) parser.add_argument( "--max-tokens", type=int, default=512, help="Default maximum number of tokens to generate (default: 512)", ) parser.add_argument( "--chat-template-args", type=json.loads, help="""A JSON formatted string of arguments for the tokenizer's apply_chat_template, e.g. '{"enable_thinking":false}'""", default="{}", ) parser.add_argument( "--decode-concurrency", type=int, default=32, help="When a request is batchable then decode that many requests in parallel", ) parser.add_argument( "--prompt-concurrency", type=int, default=8, help="When a request is batchable then process that many prompts in parallel", ) parser.add_argument( "--prefill-step-size", type=int, default=2048, help="Step size for prefill processing (default: 2048)", ) parser.add_argument( "--prompt-cache-size", type=int, default=10, help="Maximum number of distinct KV caches to hold in the prompt cache", ) parser.add_argument( "--prompt-cache-bytes", type=_parse_size, help="Maximum size in bytes of the KV caches", ) parser.add_argument( "--kv-bits", type=int, default=None, help="Number of bits for KV cache quantization (e.g., 4 or 8). " "Reduces memory usage for long contexts. Disables batching, so " "requests are served one at a time. Default: None (full precision)", ) parser.add_argument( "--kv-group-size", type=int, default=64, help="Group size for KV cache quantization (default: 64)", ) parser.add_argument( "--quantized-kv-start", type=int, default=DEFAULT_QUANTIZED_KV_START, help="Token position to start KV cache quantization " f"(default: {DEFAULT_QUANTIZED_KV_START})", ) parser.add_argument( "--pipeline", action="store_true", help="Use pipelining instead of tensor parallelism", ) args = parser.parse_args() _ = maybe_set_recommended_wired_limit() logging.basicConfig( level=getattr(logging, args.log_level.upper(), None), format="%(asctime)s - %(levelname)s - %(message)s", ) run(args.host, args.port, ModelProvider(args)) if __name__ == "__main__": print( "Calling `python -m mlx_lm.server...` directly is deprecated." " Use `mlx_lm.server...` or `python -m mlx_lm server ...` instead." ) main()