feat(mlx_lm): basic speculative decoding support in mlx_lm.generate / mlx_lm.server

llms/mlx_lm/generate.py

@@ -23,6 +23,12 @@ def setup_arg_parser():
         default="mlx_model",
         help="The path to the local model directory or Hugging Face repo.",
     )
+    parser.add_argument(
+        "--draft-model",
+        type=str,
+        required=False,
+        help="The path to the local model directory or Hugging Face repo for speculative decoding.",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,
@@ -81,7 +87,7 @@ def setup_arg_parser():
         "--max-kv-size",
         type=int,
         default=1024,
-        help="Set the maximum key-value cache size",
+        help="Set the maximum key-value cache size (0 for unlimited)",
     )
     return parser
 
@@ -132,6 +138,7 @@ def main():
         adapter_path=args.adapter_path,
         tokenizer_config=tokenizer_config,
     )
+    draft_model = load(args.draft_model)[0] if args.draft_model is not None else None
 
     if args.use_default_chat_template:
         if tokenizer.chat_template is None:
@@ -159,7 +166,8 @@ def main():
         formatter=formatter,
         temp=args.temp,
         top_p=args.top_p,
-        max_kv_size=args.max_kv_size,
+        max_kv_size=args.max_kv_size if args.max_kv_size > 0 else None,
+        draft_model=draft_model,
     )
 
 

llms/mlx_lm/models/base.py

@@ -9,7 +9,6 @@
 
 
 class KVCache:
-
     def __init__(self, head_dim, n_kv_heads):
         self.n_kv_heads = n_kv_heads
         if isinstance(head_dim, int):
@@ -23,6 +22,13 @@ def __init__(self, head_dim, n_kv_heads):
         self.offset = 0
         self.step = 256
 
+    def drop(self, n):
+        if n >= self.offset:
+            self.keys = self.values = None
+            self.offset = 0
+        elif n > 0:
+            self.offset -= n
+
     def update_and_fetch(self, keys, values):
         prev = self.offset
         if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
@@ -33,11 +39,10 @@ def update_and_fetch(self, keys, values):
             new_k = mx.zeros(k_shape, keys.dtype)
             new_v = mx.zeros(v_shape, values.dtype)
             if self.keys is not None:
-                if prev % self.step != 0:
-                    self.keys = self.keys[..., :prev, :]
-                    self.values = self.values[..., :prev, :]
-                self.keys = mx.concatenate([self.keys, new_k], axis=2)
-                self.values = mx.concatenate([self.values, new_v], axis=2)
+                self.keys = mx.concatenate([self.keys[..., :prev, :], new_k], axis=2)
+                self.values = mx.concatenate(
+                    [self.values[..., :prev, :], new_v], axis=2
+                )
             else:
                 self.keys, self.values = new_k, new_v
 
@@ -51,7 +56,6 @@ def state(self):
 
 
 class RotatingKVCache:
-
     def __init__(self, head_dim, n_kv_heads, max_size, keep=0, step=256):
         self.n_kv_heads = n_kv_heads
         if isinstance(head_dim, int):

llms/mlx_lm/sample_utils.py

@@ -26,7 +26,10 @@ def min_p_sampling(
             0.99-0.8 range.
         min_tokens_to_keep (int, optional): Minimum number of tokens that cannot
             be filtered. Default: ``1``.
-
+        temperature: Temperature parameter for softmax distribution reshaping.
+    Returns:
+        token(s) selected based on the min-p criterion.
+        Shape: same as logits, but with the last dimension having size 1.
     """
     if not (0 <= min_p <= 1.0):
         raise ValueError(
@@ -39,14 +42,14 @@ def min_p_sampling(
     # reference implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L531-L605
 
     # Softmax probabilities
-    probs = mx.softmax(logits * (1 / temperature), axis=-1)
+    probs = mx.softmax(logits / temperature, axis=-1)
 
     # Indices sorted in decreasing order
-    sorted_indices = mx.argsort(-logits).squeeze(0)
-    sorted_probs = probs[..., sorted_indices]
+    sorted_indices = mx.argsort(-logits)
+    sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)
 
     # Top probability
-    top_probs = probs[..., sorted_indices[0]]
+    top_probs = mx.expand_dims(sorted_probs[..., 0], axis=-1)
 
     # Calculate the min_p threshold
     scaled_min_p = min_p * top_probs
@@ -58,43 +61,54 @@ def min_p_sampling(
     # Create pool of tokens with probability less than scaled min_p
     selected_probs = mx.where(tokens_to_remove, 0, sorted_probs)
 
-    # Return sampled token
-    sorted_token = mx.random.categorical(mx.log(selected_probs))
-    return sorted_indices[sorted_token]
+    # Return sampled token(s)
+    sampled_indices = mx.random.categorical(mx.log(selected_probs))
+    tokens = mx.take_along_axis(
+        sorted_indices, mx.expand_dims(sampled_indices, axis=-1), axis=-1
+    )
+    return tokens.squeeze(-1)
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.array:
+def top_p_sampling(
+    logits: mx.array, top_p: float, temperature: float, axis: int = -1
+) -> mx.array:
     """
     Apply top-p (nucleus) sampling to logits.
 
     Args:
         logits: The logits from the model's output.
         top_p: The cumulative probability threshold for top-p filtering.
         temperature: Temperature parameter for softmax distribution reshaping.
+        axis: The axis along which to apply top-p sampling.
     Returns:
-        token selected based on the top-p criterion.
+        token(s) selected based on the top-p criterion.
     """
-    # referenced implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L449-L460
-    probs = mx.softmax(logits * (1 / temperature), axis=-1)
+    # Apply temperature and compute softmax
+    probs = mx.softmax(logits / temperature, axis=axis)
 
-    # sort probs in ascending order
-    sorted_indices = mx.argsort(probs, axis=-1)
-    sorted_probs = probs[..., sorted_indices.squeeze(0)]
+    # Sort probs in descending order
+    sorted_indices = mx.argsort(-probs, axis=axis)
+    sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=axis)
 
-    cumulative_probs = mx.cumsum(sorted_probs, axis=-1)
+    # Compute cumulative probabilities
+    cumulative_probs = mx.cumsum(sorted_probs, axis=axis)
 
-    # select tokens with cumulative probs below threshold
-    top_probs = mx.where(
-        cumulative_probs > 1 - top_p,
-        sorted_probs,
-        0,
-    )
+    # Create a mask for probs above the threshold
+    mask = cumulative_probs <= top_p
+
+    # Apply the mask to the sorted probabilities
+    masked_probs = sorted_probs * mask
 
-    sorted_token = mx.random.categorical(mx.log(top_probs))
-    token = sorted_indices.squeeze(0)[sorted_token]
+    # Sample from the normalized probabilities
+    sampled_indices = mx.random.categorical(mx.log(masked_probs), axis=axis)
+
+    # Gather the original token indices
+    tokens = mx.take_along_axis(
+        sorted_indices, mx.expand_dims(sampled_indices, axis=axis), axis=axis
+    )
 
-    return token
+    return tokens.squeeze(axis)
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)

llms/mlx_lm/server.py

@@ -104,6 +104,8 @@ def __init__(self, cli_args: argparse.Namespace):
         # Preload the default model if it is provided
         if self.cli_args.model is not None:
             self.load("default_model")
+        if self.cli_args.draft_model is not None:
+            self.draft_model, _ = load(self.cli_args.model)
 
     def _validate_model_path(self, model_path: str):
         model_path = Path(model_path)
@@ -161,6 +163,7 @@ def __init__(self, model_provider: ModelProvider, *args, **kwargs):
         """
         self.created = int(time.time())
         self.model_provider = model_provider
+        self.draft_model = model_provider.draft_model
         super().__init__(*args, **kwargs)
 
     def _set_cors_headers(self):
@@ -410,8 +413,10 @@ def handle_completion(
         top_tokens = []
         for (token, logprobs), _ in zip(
             generate_step(
-                prompt=prompt,
+                prompts=prompt[None],
                 model=self.model,
+                draft_model=self.draft_model,
+                tokenizer=self.tokenizer,
                 temp=self.temperature,
                 top_p=self.top_p,
                 repetition_penalty=self.repetition_penalty,
@@ -420,6 +425,8 @@ def handle_completion(
             ),
             range(self.max_tokens),
         ):
+            token = token.item()
+            logprobs = logprobs.squeeze(0)
             detokenizer.add_token(token)
             logging.debug(detokenizer.text)
             tokens.append(token)
@@ -497,15 +504,18 @@ def handle_stream(
 
         for (token, _), _ in zip(
             generate_step(
-                prompt=prompt,
+                prompts=prompt[None],
                 model=self.model,
+                draft_model=self.draft_model,
+                tokenizer=self.tokenizer,
                 temp=self.temperature,
                 top_p=self.top_p,
                 repetition_penalty=self.repetition_penalty,
                 repetition_context_size=self.repetition_context_size,
             ),
             range(self.max_tokens),
         ):
+            token = token.item()
             detokenizer.add_token(token)
             logging.debug(detokenizer.text)
             tokens.append(token)
@@ -646,6 +656,11 @@ def main():
         type=str,
         help="The path to the MLX model weights, tokenizer, and config",
     )
+    parser.add_argument(
+        "--draft-model",
+        type=str,
+        help="The path to the MLX model weights and config for speculative decoding",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,