diff --git a/src/exo/worker/engines/mlx/constants.py b/src/exo/worker/engines/mlx/constants.py index 2544e9be..acb01ccc 100644 --- a/src/exo/worker/engines/mlx/constants.py +++ b/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) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index b4ea2a6c..f44e38af 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/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? diff --git a/src/exo/worker/engines/mlx/mtp/__init__.py b/src/exo/worker/engines/mlx/mtp/__init__.py new file mode 100644 index 00000000..15ab7877 --- /dev/null +++ b/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"] diff --git a/src/exo/worker/engines/mlx/mtp/module.py b/src/exo/worker/engines/mlx/mtp/module.py new file mode 100644 index 00000000..1e62d81f --- /dev/null +++ b/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) diff --git a/src/exo/worker/engines/mlx/mtp/speculative_decode.py b/src/exo/worker/engines/mlx/mtp/speculative_decode.py new file mode 100644 index 00000000..2772003c --- /dev/null +++ b/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", + ) diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 5bab9926..53b374e3 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/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,