Add MTP self speculative generation

Author: apstenku123

cppmega_mlx/inference/__init__.py

@@ -16,6 +16,7 @@
     GenerationChunk,
     build_prompt_cache,
     generate_tokens,
+    generate_tokens_mtp_self_speculative,
     generate_tokens_speculative,
     generate_tokens_with_prompt_cache,
     generate_tokens_with_kv_cache,
@@ -97,6 +98,7 @@
     "builtin_code_metadata_adapters",
     "clone_contiguous_kv_cache",
     "generate_tokens",
+    "generate_tokens_mtp_self_speculative",
     "generate_tokens_speculative",
     "generate_tokens_with_prompt_cache",
     "generate_tokens_with_kv_cache",

cppmega_mlx/inference/generation.py

@@ -247,6 +247,104 @@ def generate_tokens_speculative(
     return tokens
 
 
+def generate_tokens_mtp_self_speculative(
+    model: Any,
+    prompt_ids: mx.array,
+    *,
+    max_new_tokens: int,
+    draft_window: int | None = None,
+    model_kwargs: Mapping[str, mx.array] | None = None,
+    eos_token_id: int | None = None,
+    temperature: float = 1.0,
+    rng_key: Any | None = None,
+) -> mx.array:
+    """Generate with a model-owned MTP head as the speculative drafter.
+
+    This is the scoped FastMTP-aligned Stream I path: the same model computes
+    the last verified hidden state, its attached ``mtp_head`` drafts a bounded
+    token window, and the model verifies the candidate prefix with one normal
+    target forward. It is deliberately eager batch=1 and no-KV.
+    """
+
+    if len(prompt_ids.shape) != 2:
+        raise ValueError("prompt_ids must have shape (batch, sequence)")
+    if int(prompt_ids.shape[0]) != 1:
+        raise ValueError("MTP self-speculative generation currently supports batch=1")
+    if int(prompt_ids.shape[1]) <= 0:
+        raise ValueError("prompt_ids must contain at least one token")
+    if max_new_tokens < 0:
+        raise ValueError("max_new_tokens must be non-negative")
+    if max_new_tokens == 0:
+        return prompt_ids
+
+    mtp_head = _resolve_mtp_draft_head(model)
+    trained_depth = _mtp_head_trained_depth(mtp_head)
+    window_limit = trained_depth if draft_window is None else draft_window
+    if window_limit <= 0:
+        raise ValueError("draft_window must be positive")
+    if window_limit > trained_depth:
+        raise ValueError("draft_window must not exceed model.mtp_head.config.depth")
+
+    max_seq_length = _model_max_seq_length(model)
+    _validate_prefix_fits_max_length(prompt_ids, max_seq_length)
+
+    tokens = prompt_ids
+    generated = 0
+    key = rng_key
+    while generated < max_new_tokens:
+        remaining = max_new_tokens - generated
+        append_limit = remaining
+        window = min(window_limit, remaining)
+        target_slots = _available_generation_slots(tokens, max_seq_length)
+        if target_slots is not None:
+            window = min(window, target_slots)
+            append_limit = min(append_limit, target_slots)
+
+        draft_tokens, draft_logits, key = _propose_mtp_self_speculative_window(
+            model,
+            mtp_head,
+            tokens,
+            draft_window=window,
+            model_kwargs=model_kwargs,
+            eos_token_id=eos_token_id,
+            temperature=temperature,
+            rng_key=key,
+        )
+        candidate = mx.concatenate([tokens, draft_tokens[None, :].astype(tokens.dtype)], axis=1)
+        _validate_prefix_fits_max_length(candidate, max_seq_length)
+
+        target_logits = _standard_generation_logits(
+            model(
+                candidate,
+                **_model_kwargs_for_prefix(model_kwargs, candidate),
+            ),
+            candidate,
+        )
+        start = int(tokens.shape[1]) - 1
+        verifier_logits = target_logits[0, start : start + int(draft_tokens.shape[0]) + 1, :]
+
+        key, acceptance_key = _split_generation_key(key)
+        accepted, _n_accepted, next_token = speculative_acceptance(
+            draft_logits,
+            verifier_logits,
+            draft_tokens,
+            temperature=temperature,
+            rng_key=acceptance_key,
+        )
+        append_tokens, found_eos = _speculative_append_tokens(
+            accepted,
+            next_token,
+            remaining=append_limit,
+            eos_token_id=eos_token_id,
+        )
+        tokens = mx.concatenate([tokens, append_tokens[None, :].astype(tokens.dtype)], axis=1)
+        generated += int(append_tokens.shape[0])
+        if found_eos:
+            break
+
+    return tokens
+
+
 def generate_tokens_with_kv_cache(
     model: Any,
     prompt_ids: mx.array,
@@ -874,6 +972,147 @@ def _speculative_append_tokens(
     return proposed, False
 
 
+def _resolve_mtp_draft_head(model: Any) -> Any:
+    mtp_head = getattr(model, "mtp_head", None)
+    if mtp_head is None:
+        raise ValueError("model.mtp_head is required for MTP self-speculative generation")
+    if not callable(getattr(model, "decoder_hidden_states", None)):
+        raise TypeError("MTP self-speculative generation requires model.decoder_hidden_states")
+    return mtp_head
+
+
+def _mtp_head_trained_depth(mtp_head: Any) -> int:
+    config = getattr(mtp_head, "config", None)
+    depth = getattr(config, "depth", None)
+    if depth is None:
+        raise ValueError("model.mtp_head.config.depth is required")
+    depth_int = int(depth)
+    if depth_int <= 0:
+        raise ValueError("model.mtp_head.config.depth must be positive")
+    return depth_int
+
+
+def _propose_mtp_self_speculative_window(
+    model: Any,
+    mtp_head: Any,
+    tokens: mx.array,
+    *,
+    draft_window: int,
+    model_kwargs: Mapping[str, mx.array] | None,
+    eos_token_id: int | None,
+    temperature: float,
+    rng_key: Any | None,
+) -> tuple[mx.array, mx.array, Any | None]:
+    hidden_states = model.decoder_hidden_states(
+        tokens,
+        **_model_kwargs_for_prefix(model_kwargs, tokens),
+    )
+    if not isinstance(hidden_states, mx.array) or hidden_states.ndim != 3:
+        raise TypeError("model.decoder_hidden_states must return (batch, sequence, hidden)")
+    if hidden_states.shape[:2] != tokens.shape:
+        raise ValueError("decoder hidden states prefix shape must match tokens")
+
+    last_hidden = hidden_states[:, -1:, :]
+    last_token_ids = tokens[:, -1:]
+    draft_method = getattr(mtp_head, "draft", None)
+    if callable(draft_method):
+        draft_callable = cast(
+            Callable[..., tuple[mx.array, mx.array, Any | None]],
+            draft_method,
+        )
+        draft_rows, draft_logit_rows, key = draft_callable(
+            last_hidden,
+            last_token_ids,
+            num_draft_tokens=draft_window,
+            temperature=temperature,
+            rng_key=rng_key,
+        )
+    else:
+        draft_rows, draft_logit_rows, key = _draft_with_minimal_mtp_head(
+            mtp_head,
+            last_hidden,
+            last_token_ids,
+            num_draft_tokens=draft_window,
+            temperature=temperature,
+            rng_key=rng_key,
+        )
+    _validate_mtp_draft_rows(draft_rows, draft_logit_rows, draft_window=draft_window)
+
+    draft_tokens = draft_rows[0].astype(tokens.dtype)
+    draft_logits = draft_logit_rows[0]
+    if eos_token_id is not None:
+        for idx in range(int(draft_tokens.shape[0])):
+            if int(draft_tokens[idx].item()) == eos_token_id:
+                return draft_tokens[: idx + 1], draft_logits[: idx + 1], key
+    return draft_tokens, draft_logits, key
+
+
+def _draft_with_minimal_mtp_head(
+    mtp_head: Any,
+    last_hidden: mx.array,
+    last_token_ids: mx.array,
+    *,
+    num_draft_tokens: int,
+    temperature: float,
+    rng_key: Any | None,
+) -> tuple[mx.array, mx.array, Any | None]:
+    required = (
+        "token_embedding",
+        "hidden_norm",
+        "embedding_norm",
+        "proj",
+        "shared_block",
+        "output_norm",
+        "lm_head",
+    )
+    missing = [name for name in required if not callable(getattr(mtp_head, name, None))]
+    if missing:
+        raise TypeError(f"model.mtp_head is missing draft components: {', '.join(missing)}")
+
+    h = last_hidden
+    token_ids = last_token_ids
+    draft_tokens: list[mx.array] = []
+    draft_logits: list[mx.array] = []
+    key = rng_key
+    for _ in range(num_draft_tokens):
+        token_emb = mtp_head.token_embedding(token_ids)
+        h_mtp = mtp_head.proj(
+            mx.concatenate(
+                [mtp_head.hidden_norm(h), mtp_head.embedding_norm(token_emb)],
+                axis=-1,
+            )
+        )
+        h = mtp_head.output_norm(mtp_head.shared_block(h_mtp))
+        step_logits = mtp_head.lm_head(h)[:, -1, :]
+        key, step_key = _split_generation_key(key)
+        next_token = sample_next_token(
+            step_logits,
+            temperature=temperature,
+            rng_key=step_key,
+        ).astype(last_token_ids.dtype)
+        draft_logits.append(step_logits)
+        draft_tokens.append(next_token)
+        token_ids = next_token
+
+    return mx.concatenate(draft_tokens, axis=1), mx.stack(draft_logits, axis=1), key
+
+
+def _validate_mtp_draft_rows(
+    draft_tokens: mx.array,
+    draft_logits: mx.array,
+    *,
+    draft_window: int,
+) -> None:
+    if not isinstance(draft_tokens, mx.array):
+        raise TypeError("MTP draft tokens must be an mlx.core.array")
+    if not isinstance(draft_logits, mx.array):
+        raise TypeError("MTP draft logits must be an mlx.core.array")
+    if draft_tokens.shape != (1, draft_window):
+        raise ValueError("MTP draft tokens must have shape (1, draft_window)")
+    if len(draft_logits.shape) != 3 or draft_logits.shape[:2] != (1, draft_window):
+        raise ValueError("MTP draft logits must have shape (1, draft_window, vocab)")
+
+
 def _validate_prefix_fits_max_length(tokens: mx.array, max_seq_length: int | None) -> None:
     if max_seq_length is not None and tokens.shape[1] > max_seq_length:
         raise ValueError("tokens exceed model.config.max_seq_length")

docs/mlx_port_master_plan.md

@@ -43,7 +43,7 @@ This document is a planning synthesis from parallel research lanes (local-code a
 - **Sharded checkpoints** (model.safetensors.index.json).
 - **Sequence packing** (cumulative-doc-id attention mask).
 - **Tokenizer**: M0.1 is closed. The vendored GB10 BPE artifact and special-token contract exist for vocab=65536 with id7=<CODE_START>, id45=<FIM_INSTRUCTION>, id46=<SPACE>, id47=<NL>. Encode normalizes whitespace runs ([\r\n]+-><NL>, [ \t]+-><SPACE>); decode is plain token concat with sentinel substitution. Same wrapper deployed on Mac and gb10/CUDA side.
-- **Inference / generation**: temperature/top-k/top-p sampling, eager no-KV full-prefix generation, contiguous KV-cache generation/streaming, paged KV scheduler metadata with fail-closed model-integration guard, contiguous prompt-cache reuse, fail-closed standard MTP/draft-output guard, and prompt-only FIM/iFIM infilling construction exist locally. Still open: model-integrated paged attention/API serving, sliding-window/SSM prompt-cache safety, integrated FIM decode/postprocess loop, speculative/self-spec decode, q4 path, and KV-q4 path.
+- **Inference / generation**: temperature/top-k/top-p sampling, eager no-KV full-prefix generation, contiguous KV-cache generation/streaming, paged KV scheduler metadata with fail-closed model-integration guard, contiguous prompt-cache reuse, fail-closed standard MTP/draft-output guard, prompt-only FIM/iFIM infilling construction, eager vanilla speculative decoding, and eager MTP self-speculative decoding exist locally. Still open: model-integrated paged attention/API serving, sliding-window/SSM prompt-cache safety, integrated FIM decode/postprocess loop, EAGLE/token-recycling speculative paths, q4 path, and KV-q4 path.
 - **Quantization**: bf16 only; no mx.quantize integration, no q4 inference path.
 - **Structural parity anchors, not CUDA tensor parity**: existing tests/test_cppmega_parity_anchors.py checks NAM56R route/layer constants, DSA/MLA layer derivation, vocab/MoE anchors, fail-closed parity wording, and optional sibling cppmega source-anchor presence when that checkout exists. The source-file existence check is only one guard; the test is broader than file-presence coverage, but it still does not compare CUDA golden tensors or prove numerical agreement.
 
@@ -68,7 +68,7 @@ Full-stack research-grade LLM (40-45K LOC):
   - fim.py (Bavarian 2207.14255 PSM/SPM, plus AST-FIM 2506.00204).
   - ifim.py (Sun et al. arXiv 2509.24637, Sep 2025: instruction-aware FIM; paper/source proposal only, not the deployed GB10 token-id contract).
   - stp.py (Huang/LeCun arXiv 2602.22617, Feb 2026: JEPA geodesic regularization, ~100 LOC).
-- Speculative decoding (acceptance-rejection), EAGLE-2 draft head.
+- Speculative decoding references: vanilla acceptance-rejection and MTP self-spec now have scoped eager MLX loops; EAGLE-2 remains pattern-only/gated until simpler paths underperform.
 - Tokenizer/FIM intent is a source reference only: live local/GB10 artifacts use the deployed M0.1 contract documented below (id7=<CODE_START>, id45=<FIM_INSTRUCTION>, id46=<SPACE>, id47=<NL>), and M0.1 decode parity is covered by the explicit <SPACE>/<NL> token redesign rather than HF reversible round-trip decode.
 - Training: pretrain + midtrain + SFT + RL via Muon/AdamW, FA3 (CUDA) / Pallas (TPU) / SDPA fallback.
 - Inference: contiguous KV (engine.py) + paged KV scheduler (serving.py).
@@ -884,7 +884,7 @@ Excluded as Hopper-only dead-end on Metal: cppmega/megatron/cute_dsl_mimo/ (sm_9
 166. **SCOPED DONE**: add prompt-only FIM-aware infilling in cppmega_mlx/inference/infilling.py (PSM/SPM token routing plus iFIM id45 prefix). This constructs inference prompts ending at <FIM_MIDDLE> only; it does not append the target middle/EOT, run generation, post-process the generated span, or add a KV/paged decode loop.
 167. **SCOPED DONE**: add `generate_tokens_speculative(...)` in `cppmega_mlx/inference/generation.py`, exported from `cppmega_mlx.inference`. The eager batch=1 loop lets a draft model propose a bounded token window, verifies with one target-model full-prefix forward over K+1 logits, and applies the existing Leviathan-style acceptance-rejection helper before appending accepted tokens plus target fallback/tail. This is greenfield MLX work because cppmega CUDA does not implement speculative decoding. It is not KV-cache/paged speculative serving, EAGLE-2, MTP self-speculation, throughput/quality benchmarking, or GB10/CUDA parity.
 168. Add EAGLE-2-style draft head (separate small draft model, GQA-aware) referencing nanochat/experiments/sota_impl/eagle2/. Gated; defer if MTP self-speculation already meets the throughput target.
-169. Add self-speculative decoding via MTP (FastMTP-aligned: same model emits draft tokens via its MTP heads, target verifies in one extra forward) referencing nanochat/mtp_draft.py. Cheapest of the three since MTP head already lands in Stream H; benchmark first.
+169. **SCOPED DONE**: add `generate_tokens_mtp_self_speculative(...)` in `cppmega_mlx/inference/generation.py`, exported from `cppmega_mlx.inference`. The eager batch=1 loop uses the model-owned `mtp_head` to draft a bounded K-token window from the last verified hidden state, verifies the candidate prefix with one normal target forward over K+1 logits, and reuses the Leviathan-style acceptance-rejection helper. It supports the real `MinimalMTPHead` fallback plus explicit `mtp_head.draft(...)` adapters, and fails closed when no trained MTP head is attached or the requested draft window exceeds `mtp_head.config.depth`. This is not KV-cache/paged self-speculative serving, EAGLE-2, token recycling, throughput/quality benchmarking, or GB10/CUDA parity.
 170. **SCOPED DONE (local streaming slice)**: add `GenerationChunk` and `stream_generate_tokens(...)` in `cppmega_mlx/inference/generation.py`, exported from `cppmega_mlx.inference`. Covered eager full-prefix and contiguous-KV token streaming, batched rows, all-rows EOS stopping, optional per-token text decode, zero-new-token no-op, package export boundary, quantized-KV fail-closed behavior, and paged-attention non-integration guard. This is not full mlx-lm registry export, model-integrated paged attention, prompt cache, q4/KV-q4, speculative/self-spec decoding, OpenAI serving, throughput/quality/long-context benching, or GB10/CUDA parity.
 171. **SCOPED DONE (contiguous prompt-cache reuse)**: add PromptCacheEntry, clone_contiguous_kv_cache(...), build_prompt_cache(...), and generate_tokens_with_prompt_cache(...) for repeated-prefix reuse on the Mac-local contiguous KV path. Covered build/reuse, suffix decode, prefix mismatch fail-closed behavior, zero-new-token no-op, package exports, independent cache cloning, and real HybridTinyLM prompt-cache-vs-full-KV parity. This is not paged attention, OpenAI serving, sliding-window/SSM safety, q4/KV-q4, speculative decode, or CUDA/H200 parity.
 172. Validate prompt-cache safety with sliding-window/SSM against the installed mlx-lm prompt-cache APIs; do not assume a generic KV cache is safe for every route.