feat(speculative): Inkling multi-token-prediction drafter (chained MTP blocks reusing the decoder layer, snapshot-and-replay rollback) #1315

Summary

Inkling checkpoints ship a native MTP head (model.mtp.layers.{j}.*, 8 blocks on Inkling and Inkling-Small, 1 on the 0.6B checkpoint; Inkling-Small stores it in a separate 4.5 GB mtp.safetensors shard). Each block combines the previous hidden state with the next token's embedding through two RMSNorms and a 2H -> H projection, then runs one ordinary Inkling decoder layer with a dense MLP; the target's embedding, final norm and LM head are shared. This issue adds the drafter as an mlxcel Drafter, implements MtpTarget on the Inkling text model (verify with hidden capture, snapshot-and-replay rollback for the conv caches), and makes --draft-kind mtp --draft-model <dir with mtp.safetensors> work for B = 1.

Current behavior

• src/lib/mlxcel-core/src/speculative/mtp/target.rs:230 (MtpTarget) and src/lib/mlxcel-core/src/drafter/mod.rs:722 (Drafter) exist, with concrete pairs for Qwen 3.5 (drafter/qwen3_5_mtp/) and Gemma 4 (drafter/gemma4_assistant/). Neither drafter reuses a target decoder layer with a conv cache, and load_drafter (drafter/mod.rs:1097) has no Inkling arm.
• feat(models): add the Inkling text backbone (banded relative-position attention, short-conv residuals, sigmoid-routed MoE with shared experts) #1318 drops every model.mtp.* key at sanitize time and provides cache snapshot/restore but no verify/rollback hooks.
• The pre-converted mlx-community/Inkling-Small-mlx-4bit checkpoint contains no model.mtp.* tensors; the head is only in the original thinkingmachines/Inkling-Small (mtp.safetensors) and in the 0.6B checkpoint (model.safetensors).

Expected behavior

Drafter directory: any directory whose safetensors carry model.mtp.layers.* keys (the original checkpoint itself, or a directory holding mtp.safetensors plus the target's config.json). Depth n = mtp_config.num_nextn_predict_layers (also accept text_config.num_mtp_layers); mtp_config.local_layer_ids (also text_config.mtp_local_layer_ids) lists the sliding blocks ([0, 2, 4, 5, 6, 7] on Inkling and Inkling-Small; [0] on 0.6B). Default block_size = n + 2 (10 for Inkling-Small), overridable by the existing --draft-block-size.

Block j (0 <= j < n) weights, H = hidden_size:

model.mtp.layers.{j}.embed_norm.weight [H]        model.mtp.layers.{j}.hidden_norm.weight [H]        model.mtp.layers.{j}.input_proj.weight [H, 2H]
model.mtp.layers.{j}.transformer_block.<sub>   where <sub> is every per-layer key of the decoder (attn.wq_du/wk_dv/wv_dv/wr_du/wo_ud, attn.q_norm/k_norm,
    attn.k_sconv/v_sconv, attn.rel_logits_proj.proj [16, 512 if sliding else 1024], attn_norm, mlp_norm, attn_sconv, mlp_sconv, mlp.w13_dn [2*dense_width, H], mlp.w2_md, mlp.global_scale)

Block config = the target text config with num_hidden_layers = n, layer_types[j] = "hybrid_sliding" if j in local_layer_ids else "hybrid", every block dense (mlp_layer_types = ["dense"] * n), and the same sconv_kernel_size, d_rel, rel_extent, log_scaling_*, head counts. Sanitize: the same per-layer renames as the backbone applied under blocks.{j}.transformer_block.; embed_norm, hidden_norm, input_proj unchanged; model.llm.norm.weight, when present in the drafter dir, is ignored (the norm is bound from the target).

Forward of one block on (tokens: [B, L], hidden: [B, L, H]), with its own per-block cache (KVCache + 4 conv slots, as in the backbone):

h = input_proj( concat([ hidden_norm(hidden), embed_norm(embed_target(tokens)) ], -1) )       # [B, L, H]
h = transformer_block_j(h, cache_j)                                                             # pre-norm residual stream
logits(h) = lm_head_target( norm_target(h) / logits_mup_width_multiplier )[..., :unpadded_vocab_size]

Round loop (B = 1; K = block_size):

1. Prefill seed: after the target prefill of prompt[0..T] with pre-norm hidden states Ht [1, T, H] and first bonus token b0, run block 0 on tokens = prompt[1..T] ++ [b0] with hidden = Ht[:, :T]; seed_hidden = h[:, -1:], seed_token = argmax(logits(seed_hidden)) (greedy) or the sampler.
2. Draft: snapshot every block cache. tok = seed_token, h_prev = seed_hidden, drafts = [seed_token]; while len(drafts) < K - 1: h_prev = block_{min(r, n-1)}(tok, h_prev) where r counts blocks appended in this round starting at 0, tok = argmax(logits(h_prev)), push. Returns K - 1 draft tokens.
3. Verify on the target with [b0, d_0, ..., d_{K-2}] (one forward, capture pre-norm hidden Hv [1, K, H] and a cache snapshot taken before the forward); accept a drafts, new bonus b.
4. Accept: restore the block caches from the round snapshot; run block 0 over tokens = [d_0, ..., d_{a-1}, b], hidden = Hv[:, :a+1]; set the new seed from the last position. Target rollback: restore the pre-verify snapshot and replay the accepted prefix [b0, d_0, ..., d_{a-1}] (a + 1 tokens) through the target (conv states cannot be trimmed from the tail, so replay is the exact rollback; the replayed KV equals the kept prefix bit for bit because the conv and attention inputs are identical).

Greedy invariant: with temperature = 0 the speculative output equals the plain decode output token for token.

Implementation plan

1. src/lib/mlxcel-core/src/drafter/inkling_mtp/{config.rs, model.rs, drafter.rs, mod.rs}: InklingMtpConfig::from_dir (reads the target config.json, derives n, local_layer_ids, block config, block_size), InklingMtpBlock { embed_norm, hidden_norm: RMSNorm, input_proj: UnifiedLinear, layer: InklingDecoderLayer, cache } (move InklingDecoderLayer, InklingLayerCache and their from_weights into mlxcel-core or expose them through a small trait object so the drafter crate can hold them; follow whichever route qwen3_5_mtp took for its layer), InklingMtpDraftModel { blocks, bound embed / norm / lm_head, seed_token, seed_hidden, round_snapshot } implementing Drafter (bind, validate_target_compat = hidden size and vocab match, draft_block, accept hook, filter_batch). Register in load_drafter by detecting model.mtp.layers. keys (tensor presence, not a config flag) with DrafterKind::Mtp.
2. Loader: extend src/loading so a drafter directory may be the full original checkpoint; only shards containing model.mtp.* keys are read (use the index), so mtp.safetensors alone is loaded for Inkling-Small.
3. Target: implement MtpTarget on InklingModel (prefill_and_seed returning the last pre-norm hidden, verify_forward capturing Hv, verify_finalize = snapshot restore + replay of accepted + 1 tokens). Snapshot = per-layer KV copy plus the 4 conv states (ModelStateSnapshot from the backbone). Register Inkling in src/server/speculative_dispatch.rs and src/server/batch/mtp_policy.rs (B = 1 only; requires_uniform_batch_acceptance).
4. CLI: --draft-kind mtp --draft-model models/Inkling-Small-mtp where the directory holds config.json and mtp.safetensors; document mlxcel download thinkingmachines/Inkling-Small --include config.json mtp.safetensors as the way to fetch only the head (add --include glob support to src/downloader if it does not exist yet).
5. docs/supported-models.md: note the MTP drafter and the shard requirement.

Validation

(a) Unit tests, src/lib/mlxcel-core/src/drafter/inkling_mtp/tests.rs (tiny random configs, no checkpoint):

• block_config_marks_local_ids_sliding_and_dense: n = 2, local_layer_ids = [0] gives block 0 sliding (rel_proj [16, W]), block 1 global (rel_proj [16, rel_extent]), both dense.
• drafter_forward_shape: 5-token prompt -> logits [1, 1, vocab] finite.
• greedy_invariant_with_speculation: tiny random target + drafter, 40 greedy tokens with and without --draft-kind mtp are identical.
• rollback_replay_restores_conv_and_kv_exactly: run verify with K = 4, accept 1, rollback, then decode 3 tokens; compare every conv state and KV slab to a run that never speculated (bitwise).
• load_drafter_detects_mtp_tensors_not_config_flag: a directory with model.mtp.layers.0.* keys routes to the Inkling MTP drafter; one with only mtp_config in config.json does not.

(b) Real checkpoints: target models/Inkling-Small-mlx-4bit (about 153 GB) and drafter models/Inkling-Small-mtp (config.json + mtp.safetensors from thinkingmachines/Inkling-Small, 4.5 GB bf16):

./target/release/mlxcel generate -m models/Inkling-Small-mlx-4bit --draft-model models/Inkling-Small-mtp --draft-kind mtp -p "Write a haiku about tide pools, then explain its structure." -n 128
./target/release/mlxcel generate -m models/Inkling-Small-mlx-4bit -p "Write a haiku about tide pools, then explain its structure." -n 128

Acceptance: both outputs identical for 128 greedy tokens; mean accepted length above 2 on that prompt; decode tokens per second higher with the drafter (numbers in the PR). Also run the 0.6B checkpoint (models/Inkling-0.6B-A0.6B as both target and drafter directory) and require identical greedy output for 64 tokens as the cheap regression case.

Acceptance criteria

• The Inkling MTP head loads from mtp.safetensors (or the full checkpoint) and drafts through the chained-block procedure above
• Inkling implements MtpTarget with snapshot-and-replay rollback; greedy speculative output equals plain output on the 0.6B and Inkling-Small checkpoints
• docs/supported-models.md updated
• detection table in src/models/detection.rs updated with a test
• cargo test --workspace --profile test-fast --features metal,accelerate passes
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

Out of scope

• Batched (B > 1) acceptance and the tree-verify path.
• A standalone split tool writing a trimmed drafter directory (reading mtp.safetensors directly makes it unnecessary).

Dependencies

• Blocked by feat(models): add the Inkling text backbone (banded relative-position attention, short-conv residuals, sigmoid-routed MoE with shared experts) #1318.

• Part of epic feat(models): Inkling (Thinking Machines) omni model family, text + image + audio + video + MTP drafter #1313