Add progress and fixes

- Fix access to per sample logprobs
- Compute more than 1 token at a time in server batch generate

Author: angeloskath

mlx_lm/generate.py

@@ -308,7 +308,7 @@ def generate_step(
     kv_bits: Optional[int] = None,
     kv_group_size: int = 64,
     quantized_kv_start: int = 0,
-    prompt_progress_callback: Optional[Callable[[int], int]] = None,
+    prompt_progress_callback: Optional[Callable[[int, int], None]] = None,
     input_embeddings: Optional[mx.array] = None,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
@@ -334,7 +334,7 @@ def generate_step(
         kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
         quantized_kv_start (int): Step to begin using a quantized KV cache.
            when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[[int], int]): A call-back which takes the
+        prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
         input_embeddings (mx.array, optional): Input embeddings to use instead of or in
           conjunction with prompt tokens. Default: ``None``.
@@ -846,18 +846,18 @@ def __len__(self):
 
     def filter(self, keep_idx: List[int]):
         self.uids = [self.uids[k] for k in keep_idx]
+        self.logprobs = [self.logprobs[k] for k in keep_idx]
         self.max_tokens = [self.max_tokens[k] for k in keep_idx]
         self.num_tokens = [self.num_tokens[k] for k in keep_idx]
         keep_idx = mx.array(keep_idx, mx.int32)
         self.y = self.y[keep_idx]
-        self.logprobs = self.logprobs[keep_idx]
         for c in self.cache:
             c.filter(keep_idx)
 
     def extend(self, other):
         self.uids.extend(other.uids)
         self.y = mx.concatenate([self.y, other.y])
-        self.logprobs = mx.concatenate([self.logprobs, other.logprobs])
+        self.logprobs.extend(other.logprobs)
         self.num_tokens.extend(other.num_tokens)
         self.max_tokens.extend(other.max_tokens)
         for c, o in zip(self.cache, other.cache):
@@ -930,6 +930,9 @@ def __init__(
         completion_batch_size: int = 32,
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
+        prompt_progress_callback: Optional[
+            Callable[[List[Tuple[int, int, int]]], None]
+        ] = None,
     ):
         self.model = model
         self.unprocessed_prompts = []
@@ -940,6 +943,7 @@ def __init__(
         self.prefill_step_size = prefill_step_size
         self.prefill_batch_size = prefill_batch_size
         self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
+        self.prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
         self._stats = BatchStats()
 
         self.active_batch = None
@@ -968,6 +972,20 @@ def insert(
         )
         return uids
 
+    def remove(self, uids: List[int]):
+        uids = set(uids)
+        if self.active_batch is not None:
+            batch = self.active_batch
+            keep_idx = [e for e, uid in enumerate(batch.uids) if uid not in uids]
+            if len(keep_idx) > 0:
+                batch.filter(keep_idx)
+            else:
+                self.active_batch = None
+
+        for i in reversed(range(len(self.unprocessed_prompts))):
+            if self.unprocessed_prompts[i][0] in uids:
+                self.unprocessed_prompts.pop(i)
+
     def _process_prompts(self, prompts):
         uids, inputs, max_tokens, caches = zip(*prompts)
 
@@ -979,6 +997,8 @@ def _process_prompts(self, prompts):
 
         self._stats.prompt_tokens += sum(lengths)
 
+        processed_tokens = 0
+
         # New prompts so
         #   1. Left-pad the inputs
         #   2. Process
@@ -991,6 +1011,13 @@ def _process_prompts(self, prompts):
                 self.model(inputs[:, :n_to_process], cache=prompt_cache)
                 mx.eval([c.state for c in prompt_cache])
                 inputs = inputs[:, n_to_process:]
+                processed_tokens += n_to_process
+                self.prompt_progress_callback(
+                    [
+                        (uid, processed_tokens, length)
+                        for uid, length in zip(uids, lengths)
+                    ]
+                )
                 mx.clear_cache()
 
         # Further prompt processing so we need to
@@ -1011,6 +1038,13 @@ def _process_prompts(self, prompts):
                 self.model(inputs[:, :n_to_process], cache=prompt_cache)
                 mx.eval([c.state for c in prompt_cache])
                 inputs = inputs[:, n_to_process:]
+                processed_tokens += n_to_process
+                self.prompt_progress_callback(
+                    [
+                        (uid, processed_tokens, length)
+                        for uid, length in zip(uids, lengths)
+                    ]
+                )
                 mx.clear_cache()
 
             for c in prompt_cache:
@@ -1030,7 +1064,7 @@ def _step(self, input_tokens: mx.array, prompt_cache: List[Any]):
         logits = logits[:, -1, :]
         logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
         sampled = self.sampler(logprobs)
-        return sampled, logprobs
+        return sampled, list(logprobs)
 
     def stats(self):
         self._stats.prompt_tps = self._stats.prompt_tokens / self._stats.prompt_time

mlx_lm/server.py

@@ -18,6 +18,7 @@
 from threading import Condition, Lock, Thread
 from typing import (
     Any,
+    Callable,
     Dict,
     List,
     Literal,
@@ -357,6 +358,11 @@ class GenerationContext:
     stop_token_sequences: List[List[int]]
     prompt: List[int]
 
+    _should_stop: bool = False
+
+    def stop(self):
+        self._should_stop = True
+
 
 @dataclass
 class Response:
@@ -512,9 +518,14 @@ def get_next_request():
                 except QueueEmpty:
                     return None
 
+        def progress_callback(info):
+            for uid, processed, total in info:
+                if uid in batch_results:
+                    batch_results[uid]["rqueue"].put((min(processed, total), total))
+
         while True:
             request = None
-            if not drain_batch and len(batch_results) < 100:
+            if not drain_batch:
                 request = get_next_request()
 
             # We got a request
@@ -554,6 +565,7 @@ def get_next_request():
                         [rest], args.max_tokens, caches=[cache]
                     )
                     batch_results[uid] = {
+                        "ctx": ctx,
                         "cache_key": prompt[:],
                         "rqueue": rqueue,
                         "detokenizer": tokenizer.detokenizer,
@@ -596,6 +608,7 @@ def get_next_request():
                                 tokenizer.encode("\n"),
                             ],
                         ),
+                        prompt_progress_callback=progress_callback,
                     )
                     unprocessed_requests.append((rqueue, request, args))
                     continue
@@ -619,43 +632,63 @@ def get_next_request():
                         drain_batch = False
                     continue
 
-                responses = batch_generator.next()
-                for r in responses:
-                    result = batch_results[r.uid]
-                    result["cache_key"].append(r.token)
-                    result["detokenizer"].add_token(r.token)
-
-                    top_tokens = None
-                    if args.logprobs > 0:
-                        sorted_indices = mx.argpartition(
-                            -gen.logprobs, kth=args.logprobs - 1
-                        )
-                        top_indices = sorted_indices[: args.logprobs]
-                        top_logprobs = gen.logprobs[top_indices]
-                        top_token_info = zip(
-                            top_indices.tolist(), top_logprobs.tolist()
-                        )
-                        top_tokens = tuple(top_token_info)
-                    result["rqueue"].put(
-                        Response(
-                            result["detokenizer"].last_segment,
-                            r.token,
-                            r.logprobs[r.token].item(),
-                            r.finish_reason,
-                            top_tokens,
+                uids_to_remove = []
+                time_budget = 0.5
+                start = time.time()
+                while True:
+                    if time.time() - start > time_budget:
+                        break
+
+                    responses = batch_generator.next()
+                    if not responses:
+                        break
+
+                    for r in responses:
+                        result = batch_results[r.uid]
+                        result["cache_key"].append(r.token)
+                        result["detokenizer"].add_token(r.token)
+
+                        top_tokens = None
+                        if args.logprobs > 0:
+                            sorted_indices = mx.argpartition(
+                                -gen.logprobs, kth=args.logprobs - 1
+                            )
+                            top_indices = sorted_indices[: args.logprobs]
+                            top_logprobs = gen.logprobs[top_indices]
+                            top_token_info = zip(
+                                top_indices.tolist(), top_logprobs.tolist()
+                            )
+                            top_tokens = tuple(top_token_info)
+                        result["rqueue"].put(
+                            Response(
+                                result["detokenizer"].last_segment,
+                                r.token,
+                                r.logprobs[r.token].item(),
+                                r.finish_reason,
+                                top_tokens,
+                            )
                         )
-                    )
 
-                    if r.finish_reason is not None:
-                        result["rqueue"].put(None)
-                        self.prompt_cache.insert_cache(
-                            current_model_key, result["cache_key"], r.prompt_cache
-                        )
-                        del batch_results[r.uid]
+                        if r.finish_reason is not None:
+                            result["rqueue"].put(None)
+                            self.prompt_cache.insert_cache(
+                                current_model_key, result["cache_key"], r.prompt_cache
+                            )
+                            del batch_results[r.uid]
+
+                        if result["ctx"]._should_stop:
+                            uids_to_remove.append(r.uid)
+
+                    if uids_to_remove:
+                        batch_generator.remove(uids_to_remove)
 
     def _serve_single(self, request):
         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, tokenizer = self.model_provider.load(
@@ -724,7 +757,7 @@ def _serve_single(self, request):
                 prompt_cache=cache,
                 draft_model=draft_model,
                 num_draft_tokens=args.num_draft_tokens,
-                # TODO: prompt progress callback
+                prompt_progress_callback=progress,
             ):
                 top_tokens = None
                 if args.logprobs > 0:
@@ -746,6 +779,10 @@ def _serve_single(self, request):
                     )
                 )
                 cache_key.append(gen.token)
+
+                if ctx._should_stop:
+                    break
+
             rqueue.put(None)
 
             # Save the KV cache again
@@ -760,6 +797,7 @@ def generate(
         self,
         request: CompletionRequest,
         generation_args: GenerationArguments,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
     ):
         response_queue = Queue()
         self.requests.put((response_queue, request, generation_args))
@@ -771,6 +809,10 @@ def _inner():
                     break
                 if isinstance(response, Exception):
                     raise response
+                if isinstance(response, tuple):
+                    if progress_callback is not None:
+                        progress_callback(*response)
+                    continue
                 yield response
 
         ctx = response_queue.get()
@@ -1097,9 +1139,29 @@ def handle_completion(self, request: CompletionRequest, stop_words: List[str]):
             seed=self.seed,
         )
 
+        # Create keepalive callback to send SSE comments during long prompt processing
+        def keepalive_callback(processed_tokens, total_tokens):
+            logging.info(
+                f"Prompt processing progress: {processed_tokens}/{total_tokens}"
+            )
+            if self.stream:
+                try:
+                    # Send SSE comment for keepalive - invisible to clients but keeps connection alive
+                    self.wfile.write(
+                        f": keepalive {processed_tokens}/{total_tokens}\n\n".encode()
+                    )
+                    self.wfile.flush()
+                except (BrokenPipeError, ConnectionResetError, OSError):
+                    # Client disconnected, ignore
+                    pass
+
         # Create the token generator
         try:
-            ctx, response = self.response_generator.generate(request, args)
+            ctx, response = self.response_generator.generate(
+                request,
+                args,
+                progress_callback=keepalive_callback,
+            )
         except Exception as e:
             self._set_completion_headers(404)
             self.end_headers()
@@ -1160,14 +1222,12 @@ def handle_completion(self, request: CompletionRequest, stop_words: List[str]):
                 top_tokens.append(gen.top_tokens)
 
             # Check if we should stop early
-            # TODO: This doesn't actually stop generation in the generation
-            #       thread we should probably have a way to do that via the ctx
-            #       object
             stop_condition = stopping_criteria(
                 tokens, ctx.stop_token_sequences, stop_words, ctx.eos_token_id
             )
             if stop_condition.stop_met:
                 finish_reason = "stop"
+                ctx.stop()
                 tokens = tokens[: len(tokens) - stop_condition.trim_length]
                 text = text[: len(text) - stop_condition.trim_text_length]
                 segment = ""