feat(v4): vendor mlx-lm PR #990 MTPModule + SequentialMTPHead adapter

PR #990 adds native MTP speculative decoding to mlx-lm. The Qwen3.5-
specific driver (generate.py / sample_utils.py / server.py / cache.py
changes) is out of scope, but the MTPModule itself is the reusable
"DeepSeek V3 SequentialMTPHead" pattern that any backbone can plug in:

    e     = pre_fc_norm_embedding(embed_tokens(next_token_ids))
    h     = pre_fc_norm_hidden(hidden_states)
    fused = fc(concat([e, h], -1))
    for layer in layers: fused = layer(fused, mask, cache)
    return norm(fused)

Two landings:

1. cppmega_v4/nn/_external/_mlx_lm_mtp_module_vendored.py — vendored
   MTPModule + create_attention_mask, refactored to accept a
   decoder_layer_factory injection (instead of qwen3_5-specific
   Attention/MLP/MoE imports). This is the *minimum* edit needed to
   make the module usable inside the v4 plugin without dragging the
   whole Qwen3.5 model into scope. MIT, © 2025-2026 Apple Inc.

2. cppmega_v4/nn/mtp_speculative_adapter.py — wraps an existing
   SequentialMTPHead's depth block (or builds a fresh one) and exposes
   the (hidden, next_token_ids, embed_tokens, cache=None) → fused_h
   signature that PR #990's driver expects. Lets our training head
   serve speculative inference without a separate inference-only
   module. Key bridging:
     - PR #990 pre_fc_norm_hidden/_embedding ↔ depth_block.hidden_norm/embedding_norm
     - PR #990 fc(concat)                    ↔ depth_block.proj
     - PR #990 layers + final norm           ↔ depth_block.transformer + .output_norm

Tests (6 new): vendored MTPModule shape contract; adapter built fresh;
adapter from existing head reuses depth-0 weights; rejection of
depth-0 head; rejection of out-of-range depth index; determinism.

v4 suite: 190 passed / 2 skipped.

Author: apstenku123

cppmega_v4/nn/_external/_mlx_lm_mtp_module_vendored.py

@@ -0,0 +1,95 @@
+# Vendored from ml-explore/mlx-lm PR #990 (mlx_lm/models/qwen3_5.py)
+# Upstream license: MIT, © 2025-2026 Apple Inc.
+#
+# Excerpt only — the MTPModule class + its decoder layer.  The full PR also
+# touches generate.py / sample_utils.py / server.py / cache.py to wire
+# speculative decoding end-to-end through mlx_lm.generate; those drivers are
+# out of scope for the v4 plugin (we expose the module via our own dispatch).
+#
+# Public surface used by cppmega_v4:
+#   - MTPDecoderLayer   — full-attention decoder block for the MTP head
+#   - MTPModule         — fuses hidden + next-token embedding via (pre_norm,
+#                         pre_norm, fc, layers, post_norm) — the "DeepSeek
+#                         V3 SequentialMTPHead" pattern that Qwen3.5 inherited.
+
+"""Multi-Token Prediction module (Qwen3.5 native speculative decoding).
+
+Vendored verbatim except for:
+  - removed unused imports (Attention/MLP/SparseMoeBlock/SwitchLinear) that
+    pull the whole qwen3_5 model into scope. The MTPModule itself only
+    constructs sub-layers via *injected* nn.Module classes — callers pass
+    in the decoder-layer factory rather than relying on the qwen3_5 module
+    globals.
+
+That refactor (factory pattern) is the *minimum* edit to make the vendored
+class usable without dragging Qwen3.5-specific Attention/MLP/MoE code into
+the v4 plugin.
+"""
+
+from typing import Any, Callable, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def create_attention_mask(x: mx.array, cache: Any) -> Optional[mx.array]:
+    """Causal mask helper (lifted from mlx_lm.models.base for self-containment)."""
+    L = x.shape[1]
+    if L <= 1:
+        return None
+    offset = 0 if cache is None else cache.offset
+    rinds = mx.arange(offset + L)
+    linds = mx.arange(offset, offset + L) if offset else rinds
+    mask = linds[:, None] < rinds[None]
+    return mask * -1e9
+
+
+class MTPModule(nn.Module):
+    """Multi-Token Prediction head (Qwen3.5 native speculative decoding).
+
+    Predicts token t+2 from the backbone hidden state h_t and the sampled
+    token t+1, using a shared lm_head with the backbone.
+
+    Architecture:
+        e     = pre_fc_norm_embedding(embed_tokens(next_token_ids))   # (B, N, H)
+        h     = pre_fc_norm_hidden(hidden_states)                     # (B, N, H)
+        fused = fc(concat([e, h], -1))                                # (B, N, H)
+        for layer in layers: fused = layer(fused, mask, cache)
+        out   = norm(fused)
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_layers: int,
+        rms_norm_eps: float,
+        decoder_layer_factory: Callable[[], nn.Module],
+    ):
+        super().__init__()
+        self.pre_fc_norm_hidden = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.pre_fc_norm_embedding = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.fc = nn.Linear(hidden_size * 2, hidden_size, bias=False)
+        self.layers = [decoder_layer_factory() for _ in range(num_layers)]
+        self.norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        next_token_ids: mx.array,
+        embed_tokens: nn.Embedding,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        embeds = embed_tokens(next_token_ids)  # (B, N, H)
+        e = self.pre_fc_norm_embedding(embeds)
+        h = self.pre_fc_norm_hidden(hidden_states)
+        fused = self.fc(mx.concatenate([e, h], axis=-1))  # (B, N, H)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        mask = create_attention_mask(fused, cache[0])
+        for layer, c in zip(self.layers, cache):
+            fused = layer(fused, mask, c)
+
+        return self.norm(fused)  # (B, N, H)

cppmega_v4/nn/mtp_speculative_adapter.py

@@ -0,0 +1,93 @@
+"""Speculative-decoding adapter: SequentialMTPHead → mlx-lm PR #990 MTPModule API.
+
+PR #990 introduces native MTP speculative decoding into mlx-lm with a
+specific module contract:
+
+    MTPModule(hidden_states, next_token_ids, embed_tokens, cache=None) -> fused_h
+    # fused_h then fed into the shared lm_head to produce token-t+2 logits.
+
+Our cppmega_v4 SequentialMTPHead exposes a *training-time* loss interface
+(returns per-depth logits + loss). This adapter wraps a SequentialMTPHead
+(or just its depth-0 block) and presents the (hidden, next_token_ids,
+embed_tokens) signature that PR #990's driver expects, so the same head
+can serve both training and speculative inference without a separate
+inference-only module.
+
+Key behavioural mapping:
+  - PR #990's pre_fc_norm_hidden / pre_fc_norm_embedding ↔ our
+    SequentialMTPDepthBlock.hidden_norm / embedding_norm
+  - PR #990's fc(concat([e, h]))                       ↔ our
+    SequentialMTPDepthBlock.proj
+  - PR #990's stack of MTPDecoderLayer                 ↔ our
+    SequentialMTPDepthBlock.transformer + .output_norm
+
+The adapter uses depth-0's block by default (depth-1+ are for deeper
+look-ahead in training; at inference time we predict one token ahead).
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from cppmega_v4.nn.mtp_v4 import SequentialMTPDepthBlock, SequentialMTPHead
+
+
+class SequentialMTPHeadAsMTPModule(nn.Module):
+    """Wraps a SequentialMTPDepthBlock to expose the PR #990 MTPModule API.
+
+    Construct directly from a SequentialMTPHead via ``from_head``, or build
+    fresh from hidden_size if you only want the inference contract without
+    sharing weights with a training head.
+    """
+
+    def __init__(self, depth_block: SequentialMTPDepthBlock):
+        super().__init__()
+        self.depth_block = depth_block
+
+    @classmethod
+    def from_head(
+        cls,
+        head: SequentialMTPHead,
+        depth_index: int = 0,
+    ) -> "SequentialMTPHeadAsMTPModule":
+        if head.config.depth <= 0:
+            raise ValueError(
+                "Cannot build MTPModule adapter from a depth=0 SequentialMTPHead"
+            )
+        if not 0 <= depth_index < head.config.depth:
+            raise ValueError(
+                f"depth_index {depth_index} out of range [0, {head.config.depth})"
+            )
+        return cls(head.depth_blocks[depth_index])
+
+    @classmethod
+    def fresh(cls, hidden_size: int) -> "SequentialMTPHeadAsMTPModule":
+        return cls(SequentialMTPDepthBlock(hidden_size))
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        next_token_ids: mx.array,
+        embed_tokens: nn.Embedding,
+        cache: Optional[Any] = None,  # accepted for PR #990 API compat; unused
+    ) -> mx.array:
+        """Run one MTP fusion + transformer step.
+
+        Args:
+            hidden_states: (B, N, H) backbone pre-norm hidden state.
+            next_token_ids: (B, N) sampled next-token ids (the speculative draft).
+            embed_tokens: the backbone's token embedding (shared with lm_head
+                or independent — caller decides).
+            cache: KV cache slot list, ignored here (SequentialMTPDepthBlock's
+                transformer is a single self-contained block).
+
+        Returns:
+            (B, N, H) fused hidden state; caller applies the shared lm_head.
+        """
+        del cache  # PR #990 API compat; our block has no per-layer KV cache
+        embeds = embed_tokens(next_token_ids)
+        return self.depth_block(hidden_states, embeds)
+
+
+__all__ = ["SequentialMTPHeadAsMTPModule"]

tests/v4/test_mtp_speculative_adapter.py

@@ -0,0 +1,106 @@
+"""Tests for SequentialMTPHead → mlx-lm PR #990 MTPModule adapter."""
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+import pytest
+
+from cppmega_v4.nn._external._mlx_lm_mtp_module_vendored import MTPModule
+from cppmega_v4.nn.mtp_speculative_adapter import SequentialMTPHeadAsMTPModule
+from cppmega_v4.nn.mtp_v4 import SequentialMTPDepthBlock, SequentialMTPHead
+
+pytest.importorskip("mlx_lm", reason="mlx_lm not available")
+from cppmega_mlx.training.mtp import MTPLossConfig  # noqa: E402
+
+
+def _make_emb(vocab=32, hidden=16) -> nn.Embedding:
+    return nn.Embedding(vocab, hidden)
+
+
+def test_vendored_mtp_module_factory_shape_contract():
+    """Vendored MTPModule must produce (B, N, H) from (B, N, H) + (B, N)."""
+    B, N, H = 1, 4, 16
+    vocab = 32
+    embed = _make_emb(vocab=vocab, hidden=H)
+
+    def decoder_layer():
+        # Minimal pass-through layer — just RMSNorm. The MTPModule API only
+        # requires layer(x, mask, cache) → x with same shape.
+        norm = nn.RMSNorm(H)
+
+        class PassThrough(nn.Module):
+            def __call__(self, x, mask=None, cache=None):
+                return norm(x)
+        return PassThrough()
+
+    mod = MTPModule(
+        hidden_size=H, num_layers=1, rms_norm_eps=1e-5,
+        decoder_layer_factory=decoder_layer,
+    )
+    hidden = mx.random.normal((B, N, H))
+    next_ids = mx.array(np.random.randint(0, vocab, (B, N)).astype(np.int32))
+    out = mod(hidden, next_ids, embed)
+    assert out.shape == (B, N, H)
+
+
+def test_adapter_from_fresh_matches_shape_contract():
+    """Adapter built fresh must accept (hidden, next_ids, embed) and return (B,N,H)."""
+    B, N, H = 1, 4, 16
+    vocab = 32
+    embed = _make_emb(vocab=vocab, hidden=H)
+    mod = SequentialMTPHeadAsMTPModule.fresh(H)
+    hidden = mx.random.normal((B, N, H))
+    next_ids = mx.array(np.random.randint(0, vocab, (B, N)).astype(np.int32))
+    out = mod(hidden, next_ids, embed)
+    assert out.shape == (B, N, H)
+    assert not bool(mx.any(mx.isnan(out)).item())
+
+
+def test_adapter_from_training_head_reuses_weights():
+    """Building from an existing SequentialMTPHead must share parameters with depth-0."""
+    B, N, H = 1, 4, 16
+    vocab = 32
+    embed = _make_emb(vocab=vocab, hidden=H)
+    head_lm = nn.Linear(H, vocab, bias=False)
+    head = SequentialMTPHead(
+        embed, head_lm, config=MTPLossConfig(depth=2),
+    )
+    adapter = SequentialMTPHeadAsMTPModule.from_head(head, depth_index=0)
+    # adapter.depth_block should be the same object as head.depth_blocks[0]
+    assert adapter.depth_block is head.depth_blocks[0]
+    # And calling it should produce a valid fused hidden state.
+    hidden = mx.random.normal((B, N, H))
+    next_ids = mx.array(np.random.randint(0, vocab, (B, N)).astype(np.int32))
+    out = adapter(hidden, next_ids, embed)
+    assert out.shape == (B, N, H)
+
+
+def test_adapter_rejects_zero_depth_head():
+    H = 16
+    embed = _make_emb(hidden=H)
+    head_lm = nn.Linear(H, 32, bias=False)
+    head = SequentialMTPHead(embed, head_lm, config=MTPLossConfig(depth=0))
+    with pytest.raises(ValueError, match="depth=0"):
+        SequentialMTPHeadAsMTPModule.from_head(head)
+
+
+def test_adapter_rejects_out_of_range_depth_index():
+    H = 16
+    embed = _make_emb(hidden=H)
+    head_lm = nn.Linear(H, 32, bias=False)
+    head = SequentialMTPHead(embed, head_lm, config=MTPLossConfig(depth=2))
+    with pytest.raises(ValueError, match="out of range"):
+        SequentialMTPHeadAsMTPModule.from_head(head, depth_index=2)
+
+
+def test_adapter_call_is_deterministic():
+    """Same inputs → same outputs (no hidden RNG)."""
+    B, N, H = 1, 3, 8
+    vocab = 16
+    embed = _make_emb(vocab=vocab, hidden=H)
+    mod = SequentialMTPHeadAsMTPModule.fresh(H)
+    hidden = mx.random.normal((B, N, H))
+    next_ids = mx.array(np.random.randint(0, vocab, (B, N)).astype(np.int32))
+    out_a = mod(hidden, next_ids, embed)
+    out_b = mod(hidden, next_ids, embed)
+    np.testing.assert_array_equal(np.array(out_a), np.array(out_b))