Add Multi-Token Prediction (MTP) for DeepSeek V3 speculative decoding

Implement support for DeepSeek V3's MTP layer (layer 61) to enable
speculative decoding. Based on vLLM/SGLang research showing 81-82%
acceptance rate with k=1 and 1.5-2x speedup at low QPS.

Key changes:
- Add MTP module with MTPModule class and speculative decode logic
- Patch DeepSeek V3 sanitize() to preserve layer 61 weights
- Extract MTP weights and create MTPModule during model loading
- Integrate MTP generation path in mlx_generate()
- Add MTP_ENABLED and MTP_NUM_DRAFT_TOKENS configuration

Co-Authored-By: Claude Opus 4.5 <[email protected]>

src/exo/worker/engines/mlx/constants.py

@@ -13,3 +13,8 @@ KV_CACHE_BITS: int | None = None
 
 # TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True
 TRUST_REMOTE_CODE: bool = True
+
+# Multi-Token Prediction (MTP) configuration for DeepSeek V3
+# MTP enables speculative decoding using the model's built-in draft layer
+MTP_ENABLED: bool = True  # Feature flag to enable/disable MTP
+MTP_NUM_DRAFT_TOKENS: int = 1  # Number of tokens to draft (vLLM reports k=1 is optimal)

src/exo/worker/engines/mlx/generator/generate.py

@@ -19,7 +19,13 @@ from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS
+from exo.worker.engines.mlx.constants import (
+    KV_BITS,
+    KV_GROUP_SIZE,
+    MAX_TOKENS,
+    MTP_ENABLED,
+    MTP_NUM_DRAFT_TOKENS,
+)
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     make_kv_cache,
@@ -115,6 +121,11 @@ def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]:
     return eos
 
 
+def _has_mtp_module(model: Model) -> bool:
+    """Check if the model has an attached MTP module."""
+    return hasattr(model, "mtp_module") and model.mtp_module is not None  # type: ignore[attr-defined]
+
+
 def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
@@ -145,6 +156,43 @@ def mlx_generate(
     )
 
     max_tokens = task.max_tokens or MAX_TOKENS
+
+    # Check if we should use MTP speculative decoding
+    use_mtp = MTP_ENABLED and _has_mtp_module(model)
+
+    if use_mtp:
+        logger.info("Using MTP speculative decoding")
+        yield from _mlx_generate_with_mtp(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_tokens=max_tokens,
+            sampler=sampler,
+            logits_processors=logits_processors,
+            prompt_cache=caches,
+        )
+    else:
+        yield from _mlx_generate_standard(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_tokens=max_tokens,
+            sampler=sampler,
+            logits_processors=logits_processors,
+            prompt_cache=caches,
+        )
+
+
+def _mlx_generate_standard(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    prompt: str,
+    max_tokens: int,
+    sampler: Callable[[mx.array], mx.array],
+    logits_processors: list[Callable[[mx.array, mx.array], mx.array]],
+    prompt_cache: list[KVCache | Any],
+) -> Generator[GenerationResponse]:
+    """Standard generation path using mlx_lm stream_generate."""
     for out in stream_generate(
         model=model,
         tokenizer=tokenizer,
@@ -152,7 +200,7 @@ def mlx_generate(
         max_tokens=max_tokens,
         sampler=sampler,
         logits_processors=logits_processors,
-        prompt_cache=caches,
+        prompt_cache=prompt_cache,
         # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
         prefill_step_size=2048,
         kv_group_size=KV_GROUP_SIZE,
@@ -187,4 +235,64 @@ def mlx_generate(
         if out.finish_reason is not None:
             break
 
+
+def _mlx_generate_with_mtp(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    prompt: str,
+    max_tokens: int,
+    sampler: Callable[[mx.array], mx.array],
+    logits_processors: list[Callable[[mx.array, mx.array], mx.array]],
+    prompt_cache: list[KVCache | Any],
+) -> Generator[GenerationResponse]:
+    """MTP speculative decoding generation path.
+
+    Uses the model's attached MTP module for speculative decoding,
+    which can provide 1.5-2x speedup with ~81% acceptance rate.
+    """
+    from exo.worker.engines.mlx.mtp.speculative_decode import mtp_speculative_generate
+
+    mtp_module = model.mtp_module  # type: ignore[attr-defined]
+
+    for out in mtp_speculative_generate(
+        model=model,
+        mtp_module=mtp_module,
+        tokenizer=tokenizer,
+        prompt=prompt,
+        max_tokens=max_tokens,
+        sampler=sampler,
+        logits_processors=logits_processors,
+        prompt_cache=prompt_cache,
+        num_draft_tokens=MTP_NUM_DRAFT_TOKENS,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE if KV_GROUP_SIZE is not None else 64,
+        kv_bits=KV_BITS,
+    ):
+        logger.info(f"{out.text} (from_draft={out.from_draft})")
+
+        stats: GenerationStats | None = None
+        if out.finish_reason is not None:
+            stats = GenerationStats(
+                prompt_tps=float(out.prompt_tps),
+                generation_tps=float(out.generation_tps),
+                prompt_tokens=int(out.prompt_tokens),
+                generation_tokens=int(out.generation_tokens),
+                peak_memory_usage=Memory.from_gb(out.peak_memory),
+            )
+
+            if out.finish_reason not in get_args(FinishReason):
+                logger.warning(
+                    f"Model generated unexpected finish_reason: {out.finish_reason}"
+                )
+
+        yield GenerationResponse(
+            text=out.text,
+            token=out.token,
+            finish_reason=cast(FinishReason | None, out.finish_reason),
+            stats=stats,
+        )
+
+        if out.finish_reason is not None:
+            break
+
         # TODO: Do we want an mx_barrier?

src/exo/worker/engines/mlx/mtp/__init__.py

@@ -0,0 +1,6 @@
+"""Multi-Token Prediction (MTP) module for DeepSeek V3 speculative decoding."""
+
+from exo.worker.engines.mlx.mtp.module import MTPModule
+from exo.worker.engines.mlx.mtp.speculative_decode import mtp_speculative_generate
+
+__all__ = ["MTPModule", "mtp_speculative_generate"]

src/exo/worker/engines/mlx/mtp/module.py

@@ -0,0 +1,207 @@
+"""MTP Module for DeepSeek V3 Multi-Token Prediction.
+
+The MTP architecture predicts one additional token ahead using:
+1. hnorm - RMSNorm for hidden state normalization
+2. enorm - RMSNorm for embedding normalization
+3. eh_proj - Linear(2*hidden_size -> hidden_size) projection
+4. transformer_block - Single decoder layer (attention + MLP)
+5. Shared embedding/lm_head from main model
+
+Forward pass:
+    h_norm = hnorm(hidden_state)
+    e_norm = enorm(embed(token))
+    projected = eh_proj(concat([h_norm, e_norm]))
+    new_hidden = transformer_block(projected)
+    logits = lm_head(output_norm(new_hidden))
+"""
+
+from typing import Any
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models.cache import KVCache
+from mlx_lm.models.deepseek_v3 import (
+    DeepseekV3Attention,
+    DeepseekV3MLP,
+    ModelArgs,
+)
+
+MTP_LAYER_INDEX = 61
+
+
+class MTPModule(nn.Module):
+    """Multi-Token Prediction module for DeepSeek V3.
+
+    This module is initialized from the layer 61 weights that are normally
+    discarded during model loading. It enables speculative decoding by
+    predicting one token ahead using the hidden state from the main model.
+    """
+
+    def __init__(
+        self,
+        config: ModelArgs,
+        shared_embedding: nn.Embedding,
+        shared_lm_head: nn.Linear,
+        output_norm: nn.RMSNorm,
+    ) -> None:
+        super().__init__()
+        self.config = config
+
+        # MTP-specific normalization layers
+        self.hnorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.enorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Projection: concatenated [hidden, embedding] -> hidden_size
+        self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
+
+        # Single transformer block for MTP
+        # Use a dense MLP since this is just a single layer
+        self.transformer_block = MTPTransformerBlock(config)
+
+        # Share embedding and lm_head with main model
+        self._shared_embedding = shared_embedding
+        self._shared_lm_head = shared_lm_head
+        self._output_norm = output_norm
+
+    def __call__(
+        self,
+        hidden_state: mx.array,
+        draft_token: mx.array,
+        cache: KVCache | None = None,
+        mask: mx.array | None = None,
+    ) -> tuple[mx.array, mx.array]:
+        """Forward pass for MTP.
+
+        Args:
+            hidden_state: Hidden state from main model [batch, seq_len, hidden_size]
+            draft_token: Token to embed and combine with hidden state [batch, seq_len]
+            cache: Optional KV cache for the MTP transformer block
+            mask: Optional attention mask
+
+        Returns:
+            tuple of (logits, new_hidden_state)
+        """
+        # Get embedding of draft token
+        embedding = self._shared_embedding(draft_token)
+
+        # Normalize hidden state and embedding
+        h_norm = self.hnorm(hidden_state)
+        e_norm = self.enorm(embedding)
+
+        # Project concatenated representation
+        concatenated = mx.concatenate([h_norm, e_norm], axis=-1)
+        projected = self.eh_proj(concatenated)
+
+        # Pass through single transformer block
+        new_hidden = self.transformer_block(projected, mask=mask, cache=cache)
+
+        # Apply output norm and get logits
+        normed_hidden = self._output_norm(new_hidden)
+        logits = self._shared_lm_head(normed_hidden)
+
+        return logits, new_hidden
+
+
+class MTPTransformerBlock(nn.Module):
+    """Single transformer block for MTP.
+
+    This is similar to DeepseekV3DecoderLayer but uses a dense MLP
+    instead of MoE since this is just for the single MTP layer.
+    """
+
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.self_attn = DeepseekV3Attention(config)
+        # MTP uses dense MLP, not MoE
+        self.mlp = DeepseekV3MLP(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: mx.array | None = None,
+        cache: Any | None = None,
+    ) -> mx.array:
+        """Forward pass with residual connections."""
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+def extract_mtp_weights(weights: dict[str, mx.array]) -> dict[str, mx.array]:
+    """Extract MTP-specific weights from layer 61.
+
+    The MTP layer has these weight patterns:
+    - model.layers.61.enorm.weight -> MTP embedding normalization
+    - model.layers.61.hnorm.weight -> MTP hidden normalization
+    - model.layers.61.eh_proj.weight -> MTP projection layer
+    - model.layers.61.self_attn.* -> MTP attention
+    - model.layers.61.input_layernorm.* -> MTP layer norms
+    - model.layers.61.post_attention_layernorm.*
+    - model.layers.61.mlp.* -> MTP MLP (dense, not MoE)
+
+    Args:
+        weights: Full model weights dict
+
+    Returns:
+        Dict of MTP-specific weights with keys renamed for MTPModule
+    """
+    mtp_weights: dict[str, mx.array] = {}
+    mtp_prefix = f"model.layers.{MTP_LAYER_INDEX}."
+
+    for key, value in weights.items():
+        if key.startswith(mtp_prefix):
+            # Remove the layer prefix to get relative path
+            new_key = key[len(mtp_prefix) :]
+            mtp_weights[new_key] = value
+
+    return mtp_weights
+
+
+def load_mtp_weights_into_module(
+    mtp_module: MTPModule,
+    mtp_weights: dict[str, mx.array],
+) -> None:
+    """Load extracted MTP weights into the MTPModule.
+
+    Args:
+        mtp_module: The MTPModule instance to load weights into
+        mtp_weights: Extracted MTP weights from extract_mtp_weights()
+    """
+    # Map weight names to module attributes
+    weight_mapping: dict[str, str] = {
+        "enorm.weight": "enorm.weight",
+        "hnorm.weight": "hnorm.weight",
+        "eh_proj.weight": "eh_proj.weight",
+    }
+
+    # Load direct mappings
+    for src_name, dst_name in weight_mapping.items():
+        if src_name in mtp_weights:
+            parts = dst_name.split(".")
+            obj: Any = mtp_module
+            for part in parts[:-1]:
+                obj = getattr(obj, part)
+            setattr(obj, parts[-1], mtp_weights[src_name])
+
+    # Load transformer block weights (self_attn, mlp, layer norms)
+    transformer_prefixes = [
+        "self_attn",
+        "mlp",
+        "input_layernorm",
+        "post_attention_layernorm",
+    ]
+
+    for prefix in transformer_prefixes:
+        for key, value in mtp_weights.items():
+            if key.startswith(prefix):
+                # Navigate to the correct attribute
+                parts = key.split(".")
+                obj = mtp_module.transformer_block
+                for part in parts[:-1]:
+                    obj = getattr(obj, part)
+                setattr(obj, parts[-1], value)

src/exo/worker/engines/mlx/mtp/speculative_decode.py

@@ -0,0 +1,506 @@
+"""MTP Speculative Decoding for DeepSeek V3.
+
+This module implements speculative decoding using the Multi-Token Prediction (MTP)
+layer from DeepSeek V3. The key difference from standard speculative decoding is
+that MTP requires hidden states from the main model, not just token predictions.
+
+Based on vLLM/SGLang research:
+- 81-82% acceptance rate with k=1
+- 1.5-2x speedup at low QPS
+"""
+
+import functools
+import time
+from collections.abc import Callable, Generator
+from dataclasses import dataclass
+from typing import Any, cast
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models import cache
+from mlx_lm.models.cache import KVCache
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.worker.engines.mlx.mtp.module import MTPModule
+
+# Generation stream for async operations
+generation_stream = mx.new_stream(mx.default_device())
+
+
+@dataclass
+class MTPGenerationResponse:
+    """Response from MTP speculative generation.
+
+    Attributes:
+        text: The next segment of decoded text.
+        token: The next token.
+        logprobs: A vector of log probabilities.
+        from_draft: Whether the token was generated by the MTP draft module.
+        prompt_tokens: The number of tokens in the prompt.
+        prompt_tps: The prompt processing tokens-per-second.
+        generation_tokens: The number of generated tokens.
+        generation_tps: The tokens-per-second for generation.
+        peak_memory: The peak memory used so far in GB.
+        finish_reason: The reason the response is being sent: "length", "stop" or None.
+    """
+
+    text: str
+    token: int
+    logprobs: mx.array
+    from_draft: bool
+    prompt_tokens: int
+    prompt_tps: float
+    generation_tokens: int
+    generation_tps: float
+    peak_memory: float
+    finish_reason: str | None = None
+
+
+def maybe_quantize_kv_cache(
+    prompt_cache: list[Any],
+    quantized_kv_start: int,
+    kv_group_size: int,
+    kv_bits: int | None,
+) -> None:
+    """Quantize KV cache entries if needed."""
+    if kv_bits is None:
+        return
+    for e, c in enumerate(prompt_cache):
+        if (
+            hasattr(c, "to_quantized")
+            and hasattr(c, "offset")
+            and c.offset >= quantized_kv_start
+        ):
+            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
+
+
+class ModelWithHiddenStates(nn.Module):
+    """Wrapper to extract hidden states before lm_head.
+
+    This wrapper allows capturing the hidden states from the transformer
+    layers before the final lm_head projection, which is needed for MTP.
+    """
+
+    def __init__(self, base_model: nn.Module) -> None:
+        super().__init__()
+        self._base = base_model
+
+    def forward_with_hidden(
+        self,
+        inputs: mx.array,
+        model_cache: list[Any] | None = None,
+    ) -> tuple[mx.array, mx.array]:
+        """Forward pass that returns both logits and hidden states.
+
+        Args:
+            inputs: Input token ids
+            model_cache: KV cache
+
+        Returns:
+            Tuple of (logits, hidden_states)
+        """
+        # Call the inner model (transformer layers + norm)
+        hidden: mx.array = self._base.model(inputs, model_cache)
+        # Get logits from lm_head
+        logits: mx.array = self._base.lm_head(hidden)
+        return logits, hidden
+
+    def forward(
+        self,
+        inputs: mx.array,
+        model_cache: list[Any] | None = None,
+    ) -> mx.array:
+        """Standard forward pass returning only logits."""
+        return cast(mx.array, self._base(inputs, cache=model_cache))
+
+    @property
+    def layers(self) -> list[nn.Module]:
+        """Access layers for cache creation."""
+        return cast(list[nn.Module], self._base.layers)
+
+
+def mtp_speculative_generate_step(
+    prompt: mx.array,
+    model: nn.Module,
+    mtp_module: MTPModule,
+    *,
+    num_draft_tokens: int = 1,
+    max_tokens: int = 256,
+    sampler: Callable[[mx.array], mx.array] | None = None,
+    logits_processors: list[Callable[[mx.array, mx.array], mx.array]] | None = None,
+    prompt_cache: list[Any] | None = None,
+    mtp_cache: KVCache | None = None,
+    prefill_step_size: int = 512,
+    kv_bits: int | None = None,
+    kv_group_size: int = 64,
+    quantized_kv_start: int = 0,
+) -> Generator[tuple[int, mx.array, bool], None, None]:
+    """MTP speculative decoding generator.
+
+    Unlike standard speculative decoding where the draft model only needs tokens,
+    MTP requires the hidden states from the main model. This generator:
+
+    1. Runs the main model to get logits AND hidden states
+    2. Uses MTP module with hidden state + sampled token to predict next token
+    3. Verifies MTP predictions with the main model
+    4. Accepts/rejects based on matching
+
+    Args:
+        prompt: The input prompt as token ids
+        model: The main model (must support return_hidden=True)
+        mtp_module: The MTP module for draft prediction
+        num_draft_tokens: Number of tokens to draft (typically 1 for MTP)
+        max_tokens: Maximum number of tokens to generate
+        sampler: Optional sampler function for token selection
+        logits_processors: Optional list of logits processors
+        prompt_cache: KV cache for the main model
+        mtp_cache: KV cache for the MTP module
+        prefill_step_size: Step size for prompt processing
+        kv_bits: Bits for KV cache quantization
+        kv_group_size: Group size for KV cache quantization
+        quantized_kv_start: Step to begin cache quantization
+
+    Yields:
+        Tuple of (token, logprobs, from_draft)
+    """
+    y = prompt.astype(mx.uint32)
+    prev_tokens: mx.array | None = None
+
+    # Wrap model to get hidden states
+    wrapped_model = (
+        model
+        if isinstance(model, ModelWithHiddenStates)
+        else ModelWithHiddenStates(model)
+    )
+
+    # Create caches if needed
+    if prompt_cache is None:
+        prompt_cache = cache.make_prompt_cache(model)
+    if mtp_cache is None:
+        mtp_cache = KVCache()
+
+    final_sampler = (
+        sampler if sampler is not None else (lambda x: mx.argmax(x, axis=-1))
+    )
+
+    quantize_cache_fn = functools.partial(
+        maybe_quantize_kv_cache,
+        quantized_kv_start=quantized_kv_start,
+        kv_group_size=kv_group_size,
+        kv_bits=kv_bits,
+    )
+
+    def _process_and_sample(
+        tokens: mx.array | None,
+        logits: mx.array,
+    ) -> tuple[mx.array, mx.array]:
+        """Process logits and sample tokens."""
+        nonlocal logits_processors
+        processed_logits = logits
+        if logits_processors:
+            for processor in logits_processors:
+                processed_logits = processor(
+                    tokens if tokens is not None else mx.array([]), processed_logits
+                )
+
+        logprobs = processed_logits - mx.logsumexp(
+            processed_logits, axis=-1, keepdims=True
+        )
+        sampled = final_sampler(logprobs)
+        return sampled, logprobs
+
+    def _main_model_step_with_hidden(
+        input_y: mx.array,
+    ) -> tuple[mx.array, mx.array, mx.array]:
+        """Run main model step with hidden state return."""
+        nonlocal prev_tokens
+
+        with mx.stream(generation_stream):
+            logits, hidden = wrapped_model.forward_with_hidden(
+                input_y[None], prompt_cache
+            )
+            logits = logits[:, -1, :]
+            quantize_cache_fn(prompt_cache)
+
+            if logits_processors:
+                prev_tokens = (
+                    mx.concatenate([prev_tokens, input_y])
+                    if prev_tokens is not None
+                    else input_y
+                )
+
+            sampled, logprobs_result = _process_and_sample(prev_tokens, logits)
+            return sampled, logprobs_result.squeeze(0), hidden[:, -1:, :]
+
+    def _main_model_step(
+        input_y: mx.array,
+    ) -> tuple[mx.array, mx.array]:
+        """Run main model step without hidden state."""
+        nonlocal prev_tokens
+
+        with mx.stream(generation_stream):
+            logits = wrapped_model.forward(input_y[None], prompt_cache)
+            logits = logits[:, -1, :]
+            quantize_cache_fn(prompt_cache)
+
+            if logits_processors:
+                prev_tokens = (
+                    mx.concatenate([prev_tokens, input_y])
+                    if prev_tokens is not None
+                    else input_y
+                )
+
+            sampled, logprobs_result = _process_and_sample(prev_tokens, logits)
+            return sampled, logprobs_result.squeeze(0)
+
+    def _mtp_draft(
+        hidden_state: mx.array,
+        draft_token: mx.array,
+    ) -> tuple[mx.array, mx.array]:
+        """Generate draft token using MTP module."""
+        with mx.stream(generation_stream):
+            logits, new_hidden = mtp_module(
+                hidden_state,
+                draft_token,
+                cache=mtp_cache,
+            )
+            logits = logits[:, -1, :]
+            sampled, _ = _process_and_sample(None, logits)
+            return sampled, new_hidden
+
+    def _prefill(input_y: mx.array) -> mx.array:
+        """Prefill the prompt cache."""
+        result_y = input_y
+        while result_y.size > prefill_step_size:
+            _ = wrapped_model.forward(result_y[:prefill_step_size][None], prompt_cache)
+            quantize_cache_fn(prompt_cache)
+            mx.eval([c.state for c in prompt_cache])
+            result_y = result_y[prefill_step_size:]
+            mx.clear_cache()
+        return result_y
+
+    def _rewind_cache(num_draft: int, num_accept: int) -> None:
+        """Rewind caches after rejection."""
+        cache.trim_prompt_cache(prompt_cache, num_draft - num_accept)
+
+    # Prefill phase
+    with mx.stream(generation_stream):
+        y = _prefill(y)
+
+    ntoks = 0
+    num_draft = 0
+    n_accepted = 0
+    last_hidden: mx.array | None = None
+
+    try:
+        # Initial step to get first token and hidden state
+        sampled, logprobs, last_hidden = _main_model_step_with_hidden(y)
+        mx.eval(sampled, logprobs, last_hidden)
+
+        y = sampled
+        current_logprobs = logprobs
+
+        while ntoks < max_tokens:
+            # Draft phase: use MTP to predict next token
+            num_draft = min(max_tokens - ntoks - 1, num_draft_tokens)
+
+            if num_draft > 0 and last_hidden is not None:
+                # Use MTP to draft
+                draft_token, draft_hidden = _mtp_draft(last_hidden, y)
+                mx.eval(draft_token, draft_hidden)
+
+                # Verify with main model
+                # Feed the drafted token to main model
+                verify_input = mx.concatenate([y, draft_token.flatten()])
+                verify_sampled, verify_logprobs, new_hidden = (
+                    _main_model_step_with_hidden(verify_input)
+                )
+                mx.eval(verify_sampled, verify_logprobs, new_hidden)
+
+                # Check if draft matches verification
+                draft_token_val = int(draft_token.item())
+                verify_token_val = (
+                    int(verify_sampled[0].item())
+                    if verify_sampled.shape[0] > 1
+                    else int(verify_sampled.item())
+                )
+
+                # Yield the current token (not from draft)
+                ntoks += 1
+                yield int(y.item()), current_logprobs, False
+
+                if ntoks >= max_tokens:
+                    break
+
+                if draft_token_val == verify_token_val:
+                    # Draft accepted
+                    n_accepted += 1
+                    ntoks += 1
+                    draft_logprobs = (
+                        verify_logprobs[0]
+                        if verify_logprobs.ndim > 1
+                        else verify_logprobs
+                    )
+                    yield draft_token_val, draft_logprobs, True
+
+                    if ntoks >= max_tokens:
+                        break
+
+                    # Continue with the token after the draft
+                    y = (
+                        verify_sampled[-1:]
+                        if verify_sampled.ndim > 0 and verify_sampled.shape[0] > 1
+                        else verify_sampled
+                    )
+                    current_logprobs = (
+                        verify_logprobs[-1]
+                        if verify_logprobs.ndim > 1
+                        else verify_logprobs
+                    )
+                    last_hidden = new_hidden
+                else:
+                    # Draft rejected - rewind and use verified token
+                    _rewind_cache(1, 0)
+                    y = (
+                        verify_sampled[:1]
+                        if verify_sampled.ndim > 0 and verify_sampled.shape[0] > 1
+                        else verify_sampled
+                    )
+                    current_logprobs = (
+                        verify_logprobs[0]
+                        if verify_logprobs.ndim > 1
+                        else verify_logprobs
+                    )
+                    last_hidden = (
+                        new_hidden[:, :1, :] if new_hidden is not None else None
+                    )
+            else:
+                # No drafting, just do normal generation
+                ntoks += 1
+                yield int(y.item()), current_logprobs, False
+
+                if ntoks >= max_tokens:
+                    break
+
+                sampled, logprobs, last_hidden = _main_model_step_with_hidden(y)
+                mx.eval(sampled, logprobs, last_hidden)
+
+                y = sampled
+                current_logprobs = logprobs
+
+            if ntoks % 256 == 0:
+                mx.clear_cache()
+
+    finally:
+        _rewind_cache(num_draft, n_accepted)
+
+
+def mtp_speculative_generate(
+    model: nn.Module,
+    mtp_module: MTPModule,
+    tokenizer: TokenizerWrapper,
+    prompt: str | mx.array | list[int],
+    max_tokens: int = 256,
+    sampler: Callable[[mx.array], mx.array] | None = None,
+    logits_processors: list[Callable[[mx.array, mx.array], mx.array]] | None = None,
+    prompt_cache: list[Any] | None = None,
+    num_draft_tokens: int = 1,
+    prefill_step_size: int = 512,
+    kv_group_size: int = 64,
+    kv_bits: int | None = None,
+) -> Generator[MTPGenerationResponse, None, None]:
+    """High-level MTP speculative generation with text output.
+
+    Args:
+        model: The main model
+        mtp_module: The MTP module for draft prediction
+        tokenizer: Tokenizer for encoding/decoding
+        prompt: Input prompt (string, array, or token list)
+        max_tokens: Maximum tokens to generate
+        sampler: Optional sampler function
+        logits_processors: Optional logits processors
+        prompt_cache: Optional KV cache
+        num_draft_tokens: Number of draft tokens
+        prefill_step_size: Prefill step size
+        kv_group_size: KV group size
+        kv_bits: KV bits
+
+    Yields:
+        MTPGenerationResponse objects with text and metadata
+    """
+    if not isinstance(prompt, mx.array):
+        if isinstance(prompt, str):
+            bos_token = getattr(tokenizer, "bos_token", None)
+            add_special_tokens = bos_token is None or not prompt.startswith(
+                str(bos_token)
+            )
+            encoded: list[int] = tokenizer.encode(
+                prompt, add_special_tokens=add_special_tokens
+            )
+            prompt = mx.array(encoded)
+        else:
+            prompt = mx.array(prompt)
+
+    detokenizer = tokenizer.detokenizer
+    eos_token_ids: list[int] = getattr(tokenizer, "eos_token_ids", [])
+
+    token_generator = mtp_speculative_generate_step(
+        prompt,
+        model,
+        mtp_module,
+        max_tokens=max_tokens,
+        sampler=sampler,
+        logits_processors=logits_processors,
+        prompt_cache=prompt_cache,
+        num_draft_tokens=num_draft_tokens,
+        prefill_step_size=prefill_step_size,
+        kv_group_size=kv_group_size,
+        kv_bits=kv_bits,
+    )
+
+    tic = time.perf_counter()
+    prompt_tps = 0.0
+    token = 0
+    logprobs: mx.array = mx.array([0.0])
+    from_draft = False
+    n = 0
+
+    for n, (token, logprobs, from_draft) in enumerate(token_generator):
+        if n == 0:
+            prompt_time = time.perf_counter() - tic
+            prompt_tps = float(prompt.size) / prompt_time
+            tic = time.perf_counter()
+
+        if token in eos_token_ids:
+            break
+
+        detokenizer.add_token(token)
+        if (n + 1) == max_tokens:
+            break
+
+        yield MTPGenerationResponse(
+            text=str(detokenizer.last_segment),
+            token=token,
+            logprobs=logprobs,
+            from_draft=from_draft,
+            prompt_tokens=int(prompt.size),
+            prompt_tps=prompt_tps,
+            generation_tokens=n + 1,
+            generation_tps=(n + 1) / (time.perf_counter() - tic),
+            peak_memory=mx.get_peak_memory() / 1e9,
+            finish_reason=None,
+        )
+
+    detokenizer.finalize()
+    yield MTPGenerationResponse(
+        text=str(detokenizer.last_segment),
+        token=token,
+        logprobs=logprobs,
+        from_draft=from_draft,
+        prompt_tokens=int(prompt.size),
+        prompt_tps=prompt_tps,
+        generation_tokens=n + 1,
+        generation_tps=(n + 1) / (time.perf_counter() - tic),
+        peak_memory=mx.get_peak_memory() / 1e9,
+        finish_reason="stop" if token in eos_token_ids else "length",
+    )

src/exo/worker/engines/mlx/utils_mlx.py

@@ -27,6 +27,7 @@ from exo.shared.models.model_cards import ModelId
 from exo.worker.engines.mlx.constants import (
     CACHE_GROUP_SIZE,
     KV_CACHE_BITS,
+    MTP_ENABLED,
     TRUST_REMOTE_CODE,
 )
 
@@ -70,6 +71,67 @@ Group = mx.distributed.Group
 resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
 
 
+# MTP (Multi-Token Prediction) support for DeepSeek V3
+MTP_LAYER_INDEX = 61
+_original_deepseek_sanitize: Callable[..., dict[str, Any]] | None = None
+
+
+def _is_deepseek_v3_model(model: nn.Module) -> bool:
+    """Check if the model is DeepSeek V3."""
+    return hasattr(model, "model") and isinstance(model.model, DeepseekV3Model)
+
+
+def _patch_deepseek_sanitize_for_mtp() -> None:
+    """Patch DeepSeek V3 Model.sanitize to preserve MTP layer weights.
+
+    The original sanitize() method filters out layer 61 (MTP layer) weights.
+    This patch keeps them so we can extract and use the MTP module.
+    """
+    global _original_deepseek_sanitize
+    from mlx_lm.models.deepseek_v3 import Model as DeepSeekV3Model
+
+    if _original_deepseek_sanitize is not None:
+        # Already patched
+        return
+
+    _original_deepseek_sanitize = DeepSeekV3Model.sanitize
+
+    def sanitize_with_mtp(
+        self: DeepSeekV3Model, weights: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Modified sanitize that keeps MTP layer weights."""
+        # First, call the original sanitize to handle all the weight transformations
+        # (dequantization, expert stacking, etc.)
+        if _original_deepseek_sanitize is None:
+            raise RuntimeError(
+                "_original_deepseek_sanitize is None - patch not applied correctly"
+            )
+        original_result: dict[str, Any] = _original_deepseek_sanitize(self, weights)
+
+        # Re-add the MTP layer weights that were filtered out
+        mtp_weights = {
+            k: v
+            for k, v in weights.items()
+            if k.startswith(f"model.layers.{MTP_LAYER_INDEX}")
+        }
+
+        return {**original_result, **mtp_weights}
+
+    DeepSeekV3Model.sanitize = sanitize_with_mtp
+
+
+def _restore_deepseek_sanitize() -> None:
+    """Restore the original DeepSeek V3 sanitize method."""
+    global _original_deepseek_sanitize
+    if _original_deepseek_sanitize is None:
+        return
+
+    from mlx_lm.models.deepseek_v3 import Model as DeepSeekV3Model
+
+    DeepSeekV3Model.sanitize = _original_deepseek_sanitize
+    _original_deepseek_sanitize = None
+
+
 # TODO: Test this
 #  ALSO https://github.com/exo-explore/exo/pull/233#discussion_r2549683673
 def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
@@ -205,31 +267,164 @@ def load_mlx_items(
     group: Group | None,
     on_timeout: TimeoutCallback | None = None,
 ) -> tuple[Model, TokenizerWrapper]:
-    if group is None:
-        logger.info(f"Single device used for {bound_instance.instance}")
-        model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
-        start_time = time.perf_counter()
-        model, _ = load_model(model_path, strict=True)
-        end_time = time.perf_counter()
-        logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s")
-        tokenizer = get_tokenizer(model_path, bound_instance.bound_shard)
+    """Load MLX model and tokenizer.
 
-    else:
-        logger.info("Starting distributed init")
-        start_time = time.perf_counter()
-        model, tokenizer = shard_and_load(
-            bound_instance.bound_shard, group=group, on_timeout=on_timeout
-        )
-        end_time = time.perf_counter()
-        logger.info(
-            f"Time taken to shard and load model: {(end_time - start_time):.2f}s"
-        )
+    Returns:
+        Tuple of (model, tokenizer)
+    """
+    model_id = bound_instance.bound_shard.model_meta.model_id
+    mtp_module = None
+
+    # Patch sanitize for MTP if this might be DeepSeek V3
+    should_try_mtp = MTP_ENABLED and _might_be_deepseek_v3(model_id)
+    if should_try_mtp:
+        logger.info("Patching DeepSeek V3 sanitize for MTP weight preservation")
+        _patch_deepseek_sanitize_for_mtp()
+
+    try:
+        if group is None:
+            logger.info(f"Single device used for {bound_instance.instance}")
+            model_path = build_model_path(model_id)
+            start_time = time.perf_counter()
+            model, _ = load_model(model_path, strict=not should_try_mtp)
+            end_time = time.perf_counter()
+            logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s")
+            tokenizer = get_tokenizer(model_path, bound_instance.bound_shard)
+
+        else:
+            logger.info("Starting distributed init")
+            start_time = time.perf_counter()
+            model, tokenizer = shard_and_load(
+                bound_instance.bound_shard, group=group, on_timeout=on_timeout
+            )
+            end_time = time.perf_counter()
+            logger.info(
+                f"Time taken to shard and load model: {(end_time - start_time):.2f}s"
+            )
+
+        # Extract MTP module if available
+        if should_try_mtp and _is_deepseek_v3_model(model):
+            mtp_module = _extract_mtp_module(model)
+            if mtp_module is not None:
+                logger.info("Successfully extracted MTP module from DeepSeek V3")
+
+    finally:
+        # Restore original sanitize
+        if should_try_mtp:
+            _restore_deepseek_sanitize()
 
     set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
 
+    # Store MTP module on the model for later access
+    if mtp_module is not None:
+        model.mtp_module = mtp_module  # noqa: B010
+
     return cast(Model, model), tokenizer
 
 
+def _might_be_deepseek_v3(model_id: str) -> bool:
+    """Check if model ID suggests this might be DeepSeek V3."""
+    model_id_lower = model_id.lower()
+    return "deepseek" in model_id_lower and (
+        "v3" in model_id_lower or "r1" in model_id_lower
+    )
+
+
+def _flatten_params(
+    params: dict[str, Any],
+    prefix: str = "",
+) -> dict[str, mx.array]:
+    """Flatten nested parameter dict to flat dict with dot-separated keys."""
+    result: dict[str, mx.array] = {}
+    for key, value in params.items():
+        full_key = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, mx.array):
+            result[full_key] = value
+        elif isinstance(value, dict):
+            result.update(_flatten_params(value, full_key))
+    return result
+
+
+def _extract_mtp_module(model: nn.Module) -> Any | None:
+    """Extract MTP module from a loaded DeepSeek V3 model.
+
+    The MTP weights are stored in model.model.layers at index 61 (if preserved).
+    This function extracts them and creates an MTPModule.
+
+    Returns:
+        MTPModule if MTP weights were found and extracted, None otherwise.
+    """
+    from exo.worker.engines.mlx.mtp.module import (
+        MTPModule,
+        extract_mtp_weights,
+        load_mtp_weights_into_module,
+    )
+
+    try:
+        # Check if the model has the MTP layer
+        inner_model = getattr(model, "model", None)
+        if inner_model is None or not hasattr(inner_model, "layers"):
+            logger.debug("Model doesn't have expected structure for MTP extraction")
+            return None
+
+        layers: list[nn.Module] = inner_model.layers  # type: ignore[assignment]
+        if len(layers) <= MTP_LAYER_INDEX:
+            logger.debug(
+                f"Model has {len(layers)} layers, MTP layer {MTP_LAYER_INDEX} not found"
+            )
+            return None
+
+        # Get model config
+        config = getattr(model, "args", None)
+        if config is None:
+            logger.debug("Could not get model config for MTP module")
+            return None
+
+        # Create MTP module with shared weights
+        embed_tokens = getattr(inner_model, "embed_tokens", None)
+        lm_head = getattr(model, "lm_head", None)
+        norm = getattr(inner_model, "norm", None)
+
+        if embed_tokens is None or lm_head is None or norm is None:
+            logger.debug("Could not get required model components for MTP")
+            return None
+
+        mtp_module = MTPModule(
+            config=config,
+            shared_embedding=embed_tokens,
+            shared_lm_head=lm_head,
+            output_norm=norm,
+        )
+
+        # Extract MTP layer weights from the model's parameters
+        # The weights should be at model.model.layers.61.*
+        # model.parameters() returns a nested dict, we need to flatten it
+        raw_params: dict[str, Any] = dict(model.parameters())  # type: ignore[arg-type]
+        model_weights = _flatten_params(raw_params)
+        mtp_weights = extract_mtp_weights(model_weights)
+
+        if not mtp_weights:
+            logger.debug("No MTP weights found in model parameters")
+            return None
+
+        # Load weights into MTP module
+        load_mtp_weights_into_module(mtp_module, mtp_weights)
+
+        # Remove MTP layer from main model to avoid double computation
+        # Create new layers list without the MTP layer
+        new_layers = [layer for i, layer in enumerate(layers) if i != MTP_LAYER_INDEX]
+        inner_model.layers = new_layers  # noqa: B010
+
+        logger.info(
+            f"Extracted MTP module, main model now has {len(new_layers)} layers"
+        )
+        return mtp_module
+
+    except Exception as e:
+        logger.warning(f"Failed to extract MTP module: {e}")
+        return None
+
+
 def shard_and_load(
     shard_metadata: ShardMetadata,
     group: Group,