feat: extend speculative_bench MTP path past Gemma 4 Unified #1613

Problem / Background

run_mtp hard-matches one LoadedModel variant at src/bin/speculative_bench.rs:436-443, so speculative_bench can produce MTP numbers for exactly one checkpoint:

let unified = match &model {
    LoadedModel::Gemma4Unified(u) => u,
    other => anyhow::bail!(
        "MTP bench currently supports a Gemma 4 Unified target; \
         load_model returned a different variant ({:?})",
        std::mem::discriminant(other)
    ),
};

Measured on the 2026-09-04 M5 Max sweep (mlxcel 0.6.0 at b2ff1eee, MLX pin 9a795735, speculative_bench --sweep --max-tokens 128, CSV benchmarks/metal_m5max_spec_2026-09-04.csv), that restriction costs two pairings:

1. The Gemma 4 31B MTP pairing (src/bin/speculative_bench.rs:260-266) fails at all three K with MTP run failed: MTP bench currently supports a Gemma 4 Unified target; load_model returned a different variant (Discriminant(26)). Not a missing checkpoint: models/gemma-4-31b-it-assistant-bf16 is on disk and the target measures 28.3 tok/s on its baseline row. models/gemma-4-31b-it-4bit/config.json is Gemma4ForConditionalGeneration (model_type: gemma4) and loads as LoadedModel::Gemma4VLM, while the 12B is Gemma4UnifiedForConditionalGeneration.
2. Qwen 3.8 MTP cannot be measured at all. models/qwen3.8-27b-4bit (Qwen3_5ForConditionalGeneration, qwen3_5), models/qwen3.8-27b-mtp-4bit and models/qwen3.8-27b-mtp-bf16 are all present, and the target measures 32.88 tok/s decode in benchmarks/metal_m5max_2026-09-03.csv, but REACHABLE_PAIRINGS (:236-284) has no Qwen 3.8 entry. A catalog entry alone would not help, since the same match rejects the target; lifting the restriction is the prerequisite.

The supported pairing is healthy on this host (Gemma 4 Unified 12B: 45.3 tok/s baseline, 62.9 / 70.9 / 70.0 tok/s at K=2/4/8, acceptance 55.6% / 35.0% / 34.6%, matching the GB10 matrix in docs/benchmark_results/model_tests.md exactly), so this is a harness limit, not missing model support.

Current Behavior

Every piece the bench needs already exists in the tree:

Proposed Solution

• Replace the Gemma4Unified unwrap with per-variant adapter selection mirroring speculative_burst.rs:1247-1360. Bind the drafter against &model as &dyn LanguageModel and release with model.release_sequence_state_by_id(seq_id) on the enum, so only the MtpTarget construction needs a match.
• MtpGenerator is generic over MtpTarget, so factor the timed run (warm-up burst, timed generate, acceptance capture, sequence release) into one function generic over T: MtpTarget and call it from each arm rather than copying the body per variant.
• Keep an explicit fallback arm that names the loaded variant and the supported families instead of printing std::mem::discriminant. Widen model_variant_label (src/server/batch/speculative_burst.rs:3034) to pub(crate) and reuse it.
• Add Qwen 3.8 to REACHABLE_PAIRINGS once the restriction is lifted: a qwen3.8-27b-4bit baseline with BenchKind::None, plus the same target with the qwen3.8-27b-mtp-4bit drafter, BenchKind::Mtp, block_size: Some(4).

Rejected: keeping the restriction and deleting the 31B pairing. The adapters exist and the server already dispatches them, so the bench would be the only place pretending the support is absent.

Scope

In scope: src/bin/speculative_bench.rs (run_mtp dispatch, REACHABLE_PAIRINGS); src/server/batch/speculative_burst.rs (visibility of model_variant_label only); docs/benchmark_results/model_tests.md (M5 Max table at :308-340, CSV index row at :25); docs/benchmarks.md:949.

Out of scope: the three --kind dflash rows, a separate known deferral (DFlash loader plus a public Qwen3NextCache API) that keeps its current status string; batched (B > 1) MTP adapters; Inkling pairings, which have no local checkpoint.

Implementation Notes

• Constraints: Qwen 3.5 MTP exactness is Metal-only in the server path, which declines when mlxcel_core::metal_is_available() is false (speculative_burst.rs:1123-1135). Mirror that guard so a CUDA sweep records the reason as a row status instead of emitting an unsound number.
• Edge cases: a drafter that loads but fails validate_target_compat (backbone hidden size or vocab mismatch) stays a per-row status string, never a sweep abort; a missing checkpoint directory keeps the existing resolve_model_dir fallback and skip; block_size < 2 keeps its early bail.
• Error handling: every failure stays a row status in the CSV and the Markdown table, and the sweep still completes all remaining rows.

Acceptance Criteria

• run_mtp accepts Gemma4, Gemma4VLM, Gemma4Unified, Qwen35, Qwen35Moe, Qwen35VLM and Qwen35MoeVLM targets, and an unsupported variant fails with a message naming the loaded variant and the supported families. (PR feat(bench): dispatch speculative_bench MTP past Gemma 4 Unified #1621; the check runs before the drafter load and also names the target directory, because the shared label table reports variants outside the speculative-capable set as other.)
• Gemma 4 31B + MTP assistant rows carry numeric tok/s, speedup, acceptance rate and mean accepted length at K=2/4/8, or the pairing is dropped from REACHABLE_PAIRINGS with the reason recorded in the docs. (Measured on M1 Ultra, not M5 Max: 18.4 / 14.9 / 14.9 tok/s, 0.93x / 0.75x / 0.75x against a 19.9 tok/s baseline.)
• A Qwen 3.8 27B baseline row and MTP pairing are in REACHABLE_PAIRINGS and both produce numbers. (24.9 tok/s baseline, 22.5 / 17.2 / 9.3 tok/s at K=2/4/8 on M1 Ultra.)
• Gemma 4 Unified 12B rows are unchanged within run-to-run noise (K=4 near 70.9 tok/s, acceptance 35.0%, mean accepted length 1.05). (Not assertable as written: 70.9 tok/s is an M5 Max figure and this run is M1 Ultra. The rows still produce numbers, 36.0 / 28.5 / 29.3 tok/s at K=2/4/8 against a 38.4 tok/s baseline, acceptance 54.5% / 39.6% / 39.6%. An M5 Max re-run against the fixed binary is still pending.)
• docs/benchmark_results/model_tests.md carries the new rows, drops the "unsupported target" note, and updates the CSV index row count at :25. (On bench/0.7.0-refresh at f4ff3752, not on main; those anchors exist only on that branch.)
• cargo clippy --all-targets -- -D warnings and cargo fmt --check are clean. (Run scoped to the touched target as cargo clippy --bin speculative_bench --features metal,accelerate -- -D warnings, plus cargo fmt --all -- --check; the unscoped workspace form exceeds this host's run budget.)

Verification

cargo build --release --features metal,accelerate --bin speculative_bench
cargo clippy --all-targets -- -D warnings && cargo fmt --check

# The two single pairings that fail today.
./target/release/speculative_bench --target models/gemma-4-31b-it-4bit \
    --draft models/gemma-4-31b-it-assistant-bf16 --kind mtp --block-size 4 --max-tokens 128
./target/release/speculative_bench --target models/qwen3.8-27b-4bit \
    --draft models/qwen3.8-27b-mtp-4bit --kind mtp --block-size 4 --max-tokens 128

# Full sweep, which regenerates the table pasted into the docs.
./target/release/speculative_bench --sweep --batch 1 --max-tokens 128 2>&1 | tee /tmp/bench-sweep.log

Pass: both single runs print a [bench/mtp] Done: line with non-zero tok/s and acceptance_rate, the sweep table shows numeric cells for Gemma 4 31B MTP and Qwen 3.8 MTP, and no row carries a status starting with MTP bench currently supports.

Related: #154 (Gemma 4 Unified MTP drafter), #1165 (qwen3_5_mtp drafter and its Metal exactness gate), #638 (the K sweep behind --k-values).