# How to Enable and Use Speculative Decoding with Draft Models in MLX Omni Server

> Enable and use speculative decoding with draft models in MLX Omni Server to boost LLM generation speed. Learn how to implement this efficient technique.

- Repository: [madroid/mlx-omni-server](https://github.com/madroidmaq/mlx-omni-server)
- Tags: how-to-guide
- Published: 2026-03-06

---

**Speculative decoding accelerates MLX LLM generation by using a small draft model to predict tokens that a larger target model verifies, and you enable it in mlx-omni-server simply by providing a `draft_model_id` when initializing the chat generator.**

Speculative decoding reduces inference latency by drafting candidate tokens with a cheap, small model and accepting or rejecting them with the full target model in a single forward pass. In the **mlx-omni-server** repository, this optimization is exposed through the `ChatGenerator` class and the OpenAI-compatible HTTP endpoint without requiring additional configuration flags.

## How Speculative Decoding Works in MLX Omni Server

The implementation delegates the actual speculative algorithm to the underlying **mlx-lm** library while managing model lifecycle and caching within the server. When you supply a draft model identifier, the server loads both models and automatically triggers accelerated generation.

### Model Loading and Validation

The `load_mlx_model` factory in [`src/mlx_omni_server/chat/mlx/model_types.py`](https://github.com/madroidmaq/mlx-omni-server/blob/main/src/mlx_omni_server/chat/mlx/model_types.py) (lines 38-52) handles dual-model initialization. It loads the primary model with `mlx_lm.load` and, if `draft_model_id` is provided, loads the draft model into the same `MLXModel` instance. The `MLXModel.has_draft_model` property (lines 94-96) validates that both the draft model and its tokenizer are present before allowing speculative decoding to proceed.

The `ChatGenerator` class exposes this state through `has_draft_model()` (lines 55-58 in [`chat_generator.py`](https://github.com/madroidmaq/mlx-omni-server/blob/main/chat_generator.py)), letting downstream code verify availability before generation begins.

### Automatic Triggering

When a draft model is present, the `stream_generate` routine from **mlx-lm** automatically executes speculative decoding. The server does not require explicit flags or parameters beyond the initial model ID. The `ChatGenerator.create` and `ChatGenerator.get_or_create` methods (lines 45-58) accept the `draft_model_id` argument and pass it directly to the model loader.

## Enabling Draft Models in Your Code

You can enable speculative decoding through three interfaces: direct Python instantiation, the OpenAI-compatible adapter, or the HTTP API.

### Direct Python Usage

Instantiate `ChatGenerator` with both a target and draft model ID to activate speculative decoding:

```python
from mlx_omni_server.chat.mlx.chat_generator import ChatGenerator

# Main model (large, high-quality)

model_id = "mlx-community/Llama-3.1-8B-Instruct-4bit"

# Draft model (tiny, cheap)

draft_id = "mlx-community/Qwen3-0.6B-4bit"

generator = ChatGenerator.create(
    model_id=model_id,
    draft_model_id=draft_id,
)

# Verify speculative decoding is available

assert generator.has_draft_model()

```

The `create` method chains into `MLXModel.load` (referenced in [`model_types.py`](https://github.com/madroidmaq/mlx-omni-server/blob/main/model_types.py)), which fetches both models from the Hugging Face Hub or local cache.

### OpenAI-Compatible Adapter

When using the adapter pattern, pass the draft model during wrapper construction. The `OpenAIAdapter` delegates generation to the underlying `ChatGenerator`, inheriting its draft capabilities:

```python
from mlx_omni_server.chat.openai.openai_adapter import OpenAIAdapter
from mlx_omni_server.chat.mlx.chat_generator import ChatGenerator

# Initialize with draft model

wrapper = ChatGenerator.get_or_create(
    model_id="mlx-community/Llama-3.1-8B-Instruct-4bit",
    draft_model_id="mlx-community/Qwen3-0.6B-4bit",
)

adapter = OpenAIAdapter(wrapper=wrapper)

# Generate with automatic speculative decoding

response = adapter.create_chat_completion(
    messages=[{"role": "user", "content": "Explain speculative decoding"}],
    stream=False,
)
print(response.choices[0].message.content)

```

The adapter itself remains agnostic to draft models; the optimization happens transparently within the wrapped `ChatGenerator`.

### HTTP API Requests

Send requests to the `/v1/chat/completions` endpoint with the `extra_params.draft_model` field:

```bash
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
        "model": "mlx-community/Llama-3.1-8B-Instruct-4bit",
        "messages": [{"role": "user", "content": "What is speculative decoding?"}],
        "stream": false,
        "extra_params": {"draft_model": "mlx-community/Qwen3-0.6B-4bit"}
      }'

```

In [`src/mlx_omni_server/chat/openai/router.py`](https://github.com/madroidmaq/mlx-omni-server/blob/main/src/mlx_omni_server/chat/openai/router.py) (lines 22-26), the `_create_text_model` function extracts `draft_model` from `extra_params` and forwards it to `ChatGenerator.get_or_create`, ensuring the draft model is loaded before generation begins.

## Caching and Performance

### Generator Instance Caching

The [`wrapper_cache.py`](https://github.com/madroidmaq/mlx-omni-server/blob/main/wrapper_cache.py) module (lines 28-42) maintains a shared cache of `ChatGenerator` instances keyed by the tuple `(model_id, adapter_path, draft_model_id)`. This prevents redundant model loading when the same target-draft pair is requested repeatedly, reducing startup latency for subsequent requests.

### Draft Model Selection

Choose a draft model that is significantly smaller than your target model—typically 10-20x fewer parameters. The draft model must share the same tokenizer vocabulary as the target model for speculative decoding to function correctly. The server validates tokenizer compatibility during the `load_mlx_model` call.

## Summary

- **Provide a `draft_model_id`** when calling `ChatGenerator.create()` or `ChatGenerator.get_or_create()` to enable speculative decoding.
- **Use the `extra_params.draft_model` JSON field** in OpenAI-compatible HTTP requests to specify the draft model dynamically.
- **Check `generator.has_draft_model()`** to verify that both models loaded successfully before relying on accelerated generation.
- **Generator instances are cached** by the combination of target model, adapter path, and draft model ID to avoid redundant loading.
- **mlx-lm handles the actual speculation**; mlx-omni-server only manages model lifecycle and routing.

## Frequently Asked Questions

### What performance improvement can I expect from speculative decoding?

Performance gains depend on the draft model's acceptance rate. When the small model predicts tokens that the large target model accepts, you save the full forward-pass cost of the large model for those positions. Typical speedups range from 1.5x to 2.5x for well-matched model pairs, though results vary based on task complexity and model compatibility.

### Can I use any model as a draft model?

No, the draft model must use the same tokenizer as the target model. In [`model_types.py`](https://github.com/madroidmaq/mlx-omni-server/blob/main/model_types.py), the loading sequence assumes vocabulary compatibility between the primary and draft models. Attempting to pair incompatible tokenizers will result in generation errors or corrupted output. Verify that both models use identical tokenizer configurations before deployment.

### Does the draft model stay loaded between requests?

Yes. The [`wrapper_cache.py`](https://github.com/madroidmaq/mlx-omni-server/blob/main/wrapper_cache.py) implementation caches `ChatGenerator` instances using a key that includes the `draft_model_id`. As long as subsequent requests specify the same target and draft model combination, the server reuses the existing instance without reloading weights from disk.

### Is speculative decoding supported in streaming mode?

Yes. The underlying `mlx_lm.stream_generate` function supports speculative decoding for both streaming and non-streaming outputs. When you set `stream=True` in the HTTP API or adapter calls, the server continues to use the draft model for token prediction while delivering chunks incrementally.