MLX Swift MTP

Native Multi-Token Prediction (MTP) speculative decoding for Apple Silicon, built on MLX Swift.

Up to 2-3x faster inference with zero quality degradation. The drafter uses the model's own built-in MTP heads — no second model loaded into RAM.

How It Works

Standard LLM inference generates one token at a time. MTP speculative decoding uses lightweight draft heads (built into models like Gemma 4) to predict multiple tokens simultaneously, then verifies them in a single backbone forward pass.

Draft Phase          Verify Phase         Emit
MTP heads predict    Backbone verifies    Accepted tokens
K tokens (fast)  →   all K in parallel →  + bonus token

The verification uses Leviathan-Chen rejection sampling, guaranteeing the output distribution is identical to standard autoregressive generation — same quality, just faster.

Supported Models

Backbone	MTP Drafter	Drafter Params	Best For
gemma-4-e2b-it (4-bit)	gemma-4-E2B-it-assistant-bf16	78M	8GB devices
gemma-4-e4b-it (4-bit)	gemma-4-E4B-it-assistant-bf16	78.8M	16GB+ devices
gemma-4-26B-A4B-it	gemma-4-26B-A4B-it-assistant-bf16	0.4B	32GB+ devices
gemma-4-31B-it	gemma-4-31B-it-assistant-bf16	0.5B	64GB+ devices

Quick Start

Add the package to your Package.swift:

.package(url: "https://github.com/joelnishanth/mlx-swift-mtp", branch: "main"),

Simple Generation

import MLXSwiftMTP

let modelPair = try await MTPModelPair.load(
    backboneId: "mlx-community/gemma-4-e4b-it-4bit",
    assistantId: "mlx-community/gemma-4-E4B-it-assistant-bf16"
)

let output = try await MTPGenerator.generate(
    prompt: "Explain quantum computing in simple terms.",
    modelPair: modelPair,
    parameters: .default
)
print(output)

Streaming Generation

for try await chunk in MTPGenerator.generateStream(
    prompt: "Write a haiku about Swift programming.",
    modelPair: modelPair
) {
    print(chunk, terminator: "")
}

Architecture

The MTP drafter is a tiny transformer (~78M params for E4B) that:

1. Shares input embeddings with the backbone model
2. Consumes backbone's last-layer activations (concatenated with token embeddings, down-projected)
3. Shares KV cache with the backbone (no duplicate memory)
4. Uses an efficient centroid embedder to avoid full-vocab logit computation

Speculative Decoding Cycle

Each generation cycle:

1. Draft: MTP heads run K forward passes using backbone activations (tiny, fast)
2. Verify: Backbone runs ONE batched forward pass over all K draft tokens
3. Accept/Reject: Probability-ratio acceptance per position: alpha = min(1, p_target(x) / p_draft(x))
4. Residual correction on reject: sample from max(0, p_target - p_draft) (normalized)
5. Bonus token: if all K accepted, target's logits give a free extra token

CLI Benchmark

swift run mtp-cli --backbone mlx-community/gemma-4-e4b-it-4bit \
                  --assistant mlx-community/gemma-4-E4B-it-assistant-bf16 \
                  --prompt "Explain the theory of relativity"

Requirements

• macOS 14.0+ (Sonoma) or iOS 16.0+
• Apple Silicon (M1/M2/M3/M4/M5)
• Xcode 15.0+
• Swift 5.9+

References

• Multi-token prediction in Gemma 4 — Google blog post
• Fast Inference from Transformers via Speculative Decoding — Leviathan et al., 2023
• MLX Swift — Apple's ML framework for Apple Silicon
• MLX Swift LM — LLM/VLM support for MLX Swift

License

Apache License 2.0. See LICENSE.