feat: make --mtp work for model families that ship MTP heads separately

[solderzzc]

Implements the wiring discussed on #109, in the shape that avoids a second generation branch.

The gap

--mtp is currently silently a no-op for Gemma 4. The gate at Server.swift:1629/:1640 is context.model is any MTPLanguageModel, and only Qwen35Model, Qwen35TextModel and DeepseekV4Model conform — those carry MTP heads inside the checkpoint (mtp.* weights).

Gemma 4 doesn't. Google ships the heads as a separate assistant checkpoint, which is why DualModelMTP exists — MTPLanguageModel plus a mainModelRef back-pointer — and Gemma4AssistantModel conforms to it. But nothing in Sources/ ever set that reference except Sources/Gemma4MTPBench/main.swift:173, and Gemma4MTPBench is not a target in Package.swift, so it cannot build. The whole path was unreachable.

The change

--mtp-assistant-model <path-or-id> loads the assistant, injects mainModelRef, and routes through the existing generateMTP call. Instead of a second branch, mtpContext() decides which context generateMTP runs against:

• in-checkpoint MTP → the main context, exactly as today
• separate assistant → a derived context whose model is the assistant, while tokenizer, processor, configuration — and the KV cache passed alongside — remain the trunk's

That mirrors the reference usage in Gemma4MTPBench (assistant as model, main model's cache) and keeps one code path, so the prompt cache is untouched.

An explicit flag rather than #109's id table: that table maps gemma-4-e4b-it → the E2B assistant and gemma-4-31b-it → the 26B one, which look like transcription slips, and a wrong guess silently drafts from the wrong model.

Measured — and it is not a speedup yet

Correctness is fine: output prefixes are identical to baseline. Throughput is worse on both pairs available here:

trunk + assistant	baseline	with MTP
gemma-4-e2b-it-4bit + E2B assistant	136.8 tok/s	117.2
gemma-4-26b-a4b-it-4bit + 26B assistant	74.1 tok/s	63.6

And flat across depth on the 26B pair: --num-mtp-tokens 1 / 2 / 3 → 63.3 / 64.1 / 63.6 tok/s.

Invariance to draft depth is the interesting part — it points at a fixed per-round cost rather than per-drafted-token cost, which is exactly what the unlanded maxSharedKV=16 cap in #109 targets (O(T) → O(16) cross-attention in runMTPHead). Both assistants also ship bf16 against 4-bit trunks, so every drafted token costs more than the trunk token it replaces.

What this is and isn't

It makes the flag mean something instead of silently doing nothing, and gives the perf work in #109 something to be measured against. It is not a speedup on its own, and I would not enable MTP by default on these numbers.

With no --mtp-assistant-model, behaviour is byte-identical to before — the fallback is main.model is any MTPLanguageModel ? main : nil, the original condition.

259 tests pass.

Refs #109.

🤖 Generated with Claude Code

[solderzzc]

Long-context testing of this branch turned up a crash in the feature it adds, plus a correction to the performance picture I posted earlier.

Crash (blocks merge). On a cold server, any prompt over prefillStepSize (512 tokens) aborts:

Fatal error: Layer 0 is a KV-shared layer but received no sharedKV

MTPTokenIterator.prepare() prefills through context.model, and this PR makes that the assistant. A Gemma 4 assistant is entirely KV-shared layers and cannot run standalone — only callMTP supplies it sharedKV. Below 512 tokens prepare() returns .tokens without a forward pass, which is exactly why every test I ran on this branch passed.

Fixed in SharpAI/mlx-swift-lm#46 (delegate the assistant's callAsFunction to the trunk). This PR needs its submodule pointer bumped to that before merging.

Performance, corrected. My earlier ~14% loss was measured only at 150 tokens. At ~9k tokens the sign flips:

context	no-MTP	dual-model MTP
150 tok	74.1 tok/s	63.6 tok/s (−14%)
~9k tok	17.4 tok/s	18.5 tok/s (+6%)

Which is the expected shape: the fixed per-round cost of speculation is a big fraction of a cheap short-context step and a small fraction of an expensive long-context one. So this feature helps where it matters and hurts on trivial prompts — worth saying in the docs rather than treating the short-prompt number as the headline.

On the maxSharedKV window from #109: I ported it and it makes no measurable difference — 62.6 vs 63.6 tok/s at 150 tokens, 18.5 vs 18.6 at ~9k. The claimed 2–4× was at 40K–100K context, which I can't reach on the 26B pair within a reasonable run. Not landing it on unreproduced numbers; the port is parked on perf/mtp-shared-kv-window if anyone wants to take it to that scale.