Repository: ml-explore/mlx-lm

Issue:
BUG in 0.24.1
nohup mlx_lm.server --model /Users/wws/models/Qwen3-30B-A3B-8bit --host 0.0.0.0 --port  8000 --max-tokens 60000  &

first curl any text "test": work ✅
curl "test" again: notwork ❌

try many times any same string ，will raise error ：
```
10.0.0.101 - - [17/May/2025 00:43:14] "POST /v1/chat/completions HTTP/1.1" 200 -
----------------------------------------
Exception occurred during processing of request from ('10.0.0.101', 52088)
Traceback (most recent call last):
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 318, in _handle_request_noblock
    self.process_request(request, client_address)
    ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 349, in process_request
    self.finish_request(request, client_address)
    ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 362, in finish_request
    self.RequestHandlerClass(request, client_address, self)
    ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 867, in <lambda>
    lambda *args, **kwargs: handler_class(
                            ~~~~~~~~~~~~~^
        model_provider,
        ^^^^^^^^^^^^^^^
    ...<3 lines>...
        **kwargs,
        ^^^^^^^^^
    ),
    ^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 257, in __init__
    super().__init__(*args, **kwargs)
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 766, in __init__
    self.handle()
    ~~~~~~~~~~~^^
  File "/Users/wws/miniconda3/lib/python3.13/http/server.py", line 436, in handle
    self.handle_one_request()
    ~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/Users/wws/miniconda3/lib/python3.13/http/server.py", line 424, in handle_one_request
    method()
    ~~~~~~^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 363, in do_POST
    self.handle_completion(prompt, stop_id_sequences)
    ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 638, in handle_completion
    for gen_response in stream_generate(
                        ~~~~~~~~~~~~~~~^
        model=self.model,
        ^^^^^^^^^^^^^^^^^
    ...<7 lines>...
        num_draft_tokens=self.num_draft_tokens,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ):
    ^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 636, in stream_generate
    for n, (token, logprobs, from_draft) in enumerate(token_generator):
                                            ~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 626, in <genexpr>
    (token, logprobs, False) for token, logprobs in token_generator
                                                    ^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 385, in generate_step
    y, logprobs = _step(y)
                  ~~~~~^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 357, in _step
    logits = model(y[None], cache=prompt_cache)
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/models/qwen3_moe.py", line 224, in __call__
    out = self.model(inputs, mask, cache)
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/models/qwen3_moe.py", line 196, in __call__
    h = self.embed_tokens(inputs)
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx/nn/layers/quantized.py", line 102, in __call__
    self["weight"][x],
    ~~~~~~~~~~~~~~^^^
ValueError: [gather] Got indices with invalid dtype. Indices must be integral.
----------------------------------------
```

Discussion:
@trillionmonster: conda install conda-forge::mlx-lm=0.24.0

no error  ✅

but without --max-tokens @trillionmonster: and bf16

nohup mlx_lm.server --model /Users/wws/models/Qwen3-30B-A3B-bf16 --host 0.0.0.0 --port  8000 --max-tokens 60000  & 

```
10.0.0.101 - - [17/May/2025 00:54:43] "POST /v1/chat/completions HTTP/1.1" 200 -
----------------------------------------
Exception occurred during processing of request from ('10.0.0.101', 52343)
Traceback (most recent call last):
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 318, in _handle_request_noblock
    self.process_request(request, client_address)
    ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 349, in process_request
    self.finish_request(request, client_address)
    ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 362, in finish_request
    self.RequestHandlerClass(request, client_address, self)
    ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 867, in <lambda>
    lambda *args, **kwargs: handler_class(
                            ~~~~~~~~~~~~~^
        model_provider,
        ^^^^^^^^^^^^^^^
    ...<3 lines>...
        **kwargs,
        ^^^^^^^^^
    ),
    ^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 257, in __init__
    super().__init__(*args, **kwargs)
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/socketserver.py", line 766, in __init__
    self.handle()
    ~~~~~~~~~~~^^
  File "/Users/wws/miniconda3/lib/python3.13/http/server.py", line 436, in handle
    self.handle_one_request()
    ~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/Users/wws/miniconda3/lib/python3.13/http/server.py", line 424, in handle_one_request
    method()
    ~~~~~~^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 363, in do_POST
    self.handle_completion(prompt, stop_id_sequences)
    ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/server.py", line 638, in handle_completion
    for gen_response in stream_generate(
                        ~~~~~~~~~~~~~~~^
        model=self.model,
        ^^^^^^^^^^^^^^^^^
    ...<7 lines>...
        num_draft_tokens=self.num_draft_tokens,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ):
    ^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 636, in stream_generate
    for n, (token, logprobs, from_draft) in enumerate(token_generator):
                                            ~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 626, in <genexpr>
    (token, logprobs, False) for token, logprobs in token_generator
                                                    ^^^^^^^^^^^^^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 385, in generate_step
    y, logprobs = _step(y)
                  ~~~~~^^^
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/generate.py", line 357, in _step
    logits = model(y[None], cache=prompt_cache)
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/models/qwen3_moe.py", line 224, in __call__
    out = self.model(inputs, mask, cache)
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx_lm/models/qwen3_moe.py", line 196, in __call__
    h = self.embed_tokens(inputs)
  File "/Users/wws/miniconda3/lib/python3.13/site-packages/mlx/nn/layers/embedding.py", line 31, in __call__
    return self.weight[x]
           ~~~~~~~~~~~^^^
ValueError: [gather] Got indices with invalid dtype. Indices must be integral.
----------------------------------------

```@trillionmonster: same error after:

pip install git+https://github.com/ml-explore/mlx-lm.git  @awni: Reproduces for me as well. Will investigate.@jpangburn: @awni worked for me, resolved this issue, thanks!

Solution patch:
diff --git a/mlx_lm/models/llama.py b/mlx_lm/models/llama.py
index 117adf0..39f550c 100644
--- a/mlx_lm/models/llama.py
+++ b/mlx_lm/models/llama.py
@@ -157,8 +157,12 @@ class LlamaModel(nn.Module):
         inputs: mx.array,
         mask: mx.array = None,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        h = self.embed_tokens(inputs)
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
 
         if mask is None:
             mask = create_attention_mask(h, cache)
@@ -186,8 +190,9 @@ class Model(nn.Module):
         inputs: mx.array,
         mask: mx.array = None,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, mask, cache, input_embeddings)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:
diff --git a/mlx_lm/models/pixtral.py b/mlx_lm/models/pixtral.py
new file mode 100644
index 0000000..85e8ed3
--- /dev/null
+++ b/mlx_lm/models/pixtral.py
@@ -0,0 +1,52 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import llama
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+        self.text_config["num_attention_heads"] = self.text_config.get(
+            "num_attention_heads", 32
+        )
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("vision_tower", None)
+        weights.pop("multi_modal_projector", None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
diff --git a/mlx_lm/server.py b/mlx_lm/server.py
index b0fd1e0..f7a3f0a 100644
--- a/mlx_lm/server.py
+++ b/mlx_lm/server.py
@@ -547,6 +547,9 @@ class APIHandler(BaseHTTPRequestHandler):
         prompt_len = len(prompt)
         com_prefix_len = common_prefix_len(self.prompt_cache.tokens, prompt)
 
+        # Leave at least one token in the prompt
+        com_prefix_len = min(com_prefix_len, len(prompt) - 1)
+
         # Condition 1: Model changed or no common prefix at all. Reset cache.
         if (
             self.prompt_cache.model_key != self.model_provider.model_key


Test patch:
diff --git a/tests/test_server.py b/tests/test_server.py
index 53826a3..73b43f1 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -357,7 +357,9 @@ class TestGetPromptCache(unittest.TestCase):
         self.assertEqual(self.handler.prompt_cache.model_key, ("model_v1", None, None))
         mock_make_cache.assert_called_once()
 
-    def test_identical_request_full_hit(self):
+    @patch("mlx_lm.server.trim_prompt_cache")
+    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=True)
+    def test_identical_request_full_hit(self, mock_can_trim, mock_trim_cache):
         """Test when the new prompt is identical to the cached one."""
         self.handler.prompt_cache.tokens = [1, 2, 3]
         self.handler.prompt_cache.model_key = ("model_v1", None, None)
@@ -368,10 +370,9 @@ class TestGetPromptCache(unittest.TestCase):
         with patch("mlx_lm.server.common_prefix_len", return_value=3):
             processed_prompt = self.handler.get_prompt_cache(prompt)
 
-        # Should process nothing, cache remains unchanged
-        self.assertEqual(processed_prompt, [])
+        mock_trim_cache.assert_called_once_with("existing_cache_obj", 1)
+        self.assertEqual(processed_prompt, [3])
         self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3])
-        self.assertEqual(self.handler.prompt_cache.cache, "existing_cache_obj")
 
     def test_cache_is_prefix(self):
         """Test when the cached prompt is a prefix of the new prompt."""