fix(mlx_lm.server): fail fast when --draft-model set with non-trimmable cache

[tejkas]

Summary

Fixes #1446.

mlx_lm.server with --draft-model accepts startup and /v1/models, then crashes during the first completion with

"ValueError: Speculative decoding requires a trimmable prompt cache".

This issue would affect hybrid-attention models such as qwen3_5_moe which use ArraysCache on their linear-attention layers, which cannot be trimmed. Speculative decoding needs trimming to roll back rejected draft tokens.

Since the user provides --draft-model explicitly, opting to fail fast here makes more sense (as opposed to say, silently disabling speculation). There are two other options named in the issue which don't strike me as particularly feasible at this time:

• use/select a trimmable cache for this path: The model's linear attention layer requires recurrent state https://sebastianraschka.com/llm-architecture-gallery/hybrid-attention/ which cannot be simply replaced by a trimmable cache structure.
• support trimming the cache type used by these Qwen hybrid-attention models: Similar to above, a cache for linear attention models cannot be easily trimmed (it would probably require checkpointing of the recurrent state, resulting in memory blowup).

Implementation

Built the check into a small helper _ensure_speculation_supported(prompt_cache, model_path) with related unit tests using a mock cache (optionally, could inline the check but this seemed more testable).

Tests

Added 3 unit tests (ensure check rejects non-trimmable caches, allows trimmable, and that the error message names the cache type in question). No model download.
Happy path cases remain covered by existing TestServer and TestServerWithDraftModel test classes.

Verified Locally

With this fix, the server now fails at startup, before serving any request, instead of crashing mid-generation.

❯ python -m mlx_lm.server \
    --model mlx-community/mamba-130m-hf-bf16 \
    --draft-model mlx-community/Qwen2.5-0.5B-Instruct-4bit \
    --port 18087 --log-level ERROR

Exception in thread Thread-1 (_generate):
Traceback (most recent call last):
  File "/opt/miniconda3/lib/python3.13/threading.py", line 1043, in _bootstrap_inner
    self.run()
    ~~~~~~~~^^
  File "/opt/miniconda3/lib/python3.13/threading.py", line 994, in run
    self._target(*self._args, **self._kwargs)
    ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tejaskasturi/Documents/mlx-lm/mlx_lm/server.py", line 708, in _generate
    self.model_provider.load_default()
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/Users/tejaskasturi/Documents/mlx-lm/mlx_lm/server.py", line 398, in load_default
    self.load("default_model", None, "default_model")
    ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tejaskasturi/Documents/mlx-lm/mlx_lm/server.py", line 407, in load
    self._load(*model_key)
    ~~~~~~~~~~^^^^^^^^^^^^
  File "/Users/tejaskasturi/Documents/mlx-lm/mlx_lm/server.py", line 373, in _load
    _ensure_speculation_supported(prompt_cache, model_path)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tejaskasturi/Documents/mlx-lm/mlx_lm/server.py", line 53, in _ensure_speculation_supported
    raise ValueError(
    ...<2 lines>...
    )
ValueError: --draft-model was set, but speculative decoding requires a trimmable prompt cache; 'mlx-community/mamba-130m-hf-bf16' uses non-trimmable cache(s): {'ArraysCache'}