Add native Qwen MTP speculation

src/engine/metal.rs

@@ -14,7 +14,10 @@ use qwen::QwenExecutor;
 
 use super::gguf::{BF16, F16, F32, Gguf, IQ2_XXS, MXFP4, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor};
 use super::validation::{DsparkConfig, SupportKind, dspark_config};
-use super::{LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample};
+use super::{
+    LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample, exact_speculative_sample,
+    sample_from_logits,
+};
 use crate::model::ModelChoice;
 use crate::settings::{
     EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
@@ -2949,8 +2952,11 @@ pub(super) struct ExecutionStats {
     pub(super) speculative_cycles: u64,
     pub(super) drafted_tokens: u64,
     pub(super) accepted_draft_tokens: u64,
+    pub(super) accepted_depth_total: u64,
+    pub(super) accepted_depth_max: u64,
     pub(super) verifier_passes: u64,
     pub(super) verifier_ms: u64,
+    pub(super) repair_ms: u64,
     pub(super) ssd_enabled: bool,
     pub(super) ssd_resident_bytes: u64,
     pub(super) ssd_cache_bytes: u64,
@@ -4460,9 +4466,11 @@ impl Executor {
                 .map(Self::Glm)
             }
             LoadedModel::Gguf(_) => unreachable!("Qwen never uses a GGUF model"),
-            LoadedModel::Qwen(model) => QwenExecutor::open(*model, context)
-                .map(Box::new)
-                .map(Self::Qwen),
+            LoadedModel::Qwen(model) => {
+                { QwenExecutor::open_configured(*model, context, speculative) }
+                    .map(Box::new)
+                    .map(Self::Qwen)
+            }
         }
     }
 
@@ -4495,10 +4503,17 @@ impl Executor {
                 executor.eval(token)?;
                 Ok(vec![token])
             }
-            Self::Qwen(executor) => {
-                executor.eval(token)?;
-                Ok(vec![token])
-            }
+            Self::Qwen(executor) => executor.eval_speculative_sampled(
+                token,
+                max_tokens,
+                reasoning,
+                temperature,
+                top_p,
+                min_p,
+                top_k,
+                rng,
+                cancelled,
+            ),
         }
     }
 
@@ -4548,9 +4563,7 @@ impl Executor {
                 executor.eval_speculative_greedy(token, max_tokens, cancelled)
             }
             Self::Qwen(executor) => {
-                let _ = (max_tokens, reasoning, cancelled);
-                executor.eval(token)?;
-                Ok(vec![token])
+                executor.eval_speculative_greedy(token, max_tokens, reasoning, cancelled)
             }
         }
     }

src/engine/qwen.rs

@@ -30,6 +30,8 @@ const MTP_TOKEN_RESERVE: u64 = 3;
 const GDN_STATE_BYTES: u64 = 113_246_208;
 const GDN_CONV_BYTES: u64 = 2_211_840;
 const PLE_CONV_BYTES: u64 = 184_320;
+const MTP_CAPTURE_HIDDEN_BYTES: u64 = 10_240 * 4;
+const MTP_CAPTURE_LOGITS_BYTES: u64 = 248_320 * 4;
 
 #[derive(Deserialize)]
 struct Manifest {
@@ -125,10 +127,22 @@ pub(super) struct QwenModel {
 }
 
 impl QwenModel {
+    #[cfg(test)]
     pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
-        let loaded = load(root, context, false)?;
+        Self::open_configured(root, context, false)
+    }
+
+    pub(super) fn open_configured(
+        root: &Path,
+        context: u32,
+        enable_mtp: bool,
+    ) -> Result<Self, String> {
+        let loaded = load(root, context, enable_mtp)?;
         let mut bindings = loaded.bindings.core;
         bindings.extend(loaded.bindings.ple);
+        if enable_mtp {
+            bindings.extend(loaded.bindings.mtp);
+        }
         let mut paths = bindings
             .iter()
             .map(|binding| binding.file.clone())
@@ -176,7 +190,7 @@ impl QwenModel {
     pub(super) fn tensor(&self, name: &str) -> Result<&QwenTensor, String> {
         self.tensors
             .get(name)
-            .ok_or_else(|| format!("Qwen core tensor is missing: {name}"))
+            .ok_or_else(|| format!("Qwen tensor is missing: {name}"))
     }
 
     pub(super) fn map(&self, index: usize) -> (&[u8], &Path) {
@@ -424,10 +438,13 @@ fn validate_precision(name: &str, tensor: &ExpectedTensor) -> Result<(), String>
         (Some(bits @ (2 | 4 | 8)), Some(group @ (32 | 64)), Some("affine"))
             if tensor.dtype == "U32" || name.ends_with(".scales") || name.ends_with(".biases") =>
         {
-            if group == 32 && !name.starts_with("ngram.") {
+            if group == 32
+                && !name.starts_with("ngram.")
+                && !name.starts_with("mtp.layers.0.mlp.switch_mlp.")
+            {
                 return Err(format!("{name} unexpectedly uses 32-value groups"));
             }
-            if bits == 2 && !name.starts_with("mtp.") {
+            if bits == 2 {
                 return Err(format!("{name} unexpectedly uses 2-bit weights"));
             }
             Ok(())
@@ -666,7 +683,30 @@ pub(super) fn memory_plan(
         .and_then(|bytes| bytes.checked_add(QSA_FIXED_SCRATCH_BYTES))
         .and_then(|bytes| bytes.checked_add(topk_scratch))
         .ok_or_else(|| "Qwen prefill transient size overflows".to_owned())?;
-    let admitted_mtp = if enable_mtp { MTP_BYTES } else { 0 };
+    let admitted_mtp = if enable_mtp {
+        let capture_rows = MTP_TOKEN_RESERVE + 1;
+        let recurrent_capture = (GDN_STATE_BYTES
+            + GDN_CONV_BYTES
+            + PLE_CONV_BYTES
+            + MTP_CAPTURE_HIDDEN_BYTES
+            + MTP_CAPTURE_LOGITS_BYTES)
+            .checked_mul(capture_rows)
+            .ok_or_else(|| "Qwen MTP verifier capture size overflows".to_owned())?;
+        let mtp_attention = (KV_BYTES_PER_TOKEN + QSA_RAW_BYTES_PER_TOKEN)
+            .checked_mul(token_capacity)
+            .and_then(|bytes| {
+                QSA_POOLED_BYTES_PER_BLOCK
+                    .checked_mul(block_capacity)
+                    .and_then(|pooled| bytes.checked_add(pooled))
+            })
+            .ok_or_else(|| "Qwen MTP attention size overflows".to_owned())?;
+        MTP_BYTES
+            .checked_add(recurrent_capture)
+            .and_then(|bytes| bytes.checked_add(mtp_attention))
+            .ok_or_else(|| "Qwen MTP admission size overflows".to_owned())?
+    } else {
+        0
+    };
     let admission = CORE_BYTES
         .checked_add(admitted_mtp)
         .and_then(|bytes| bytes.checked_add(kv_and_recurrent))
@@ -706,7 +746,7 @@ mod tests {
         assert_eq!(plan.optional_mtp, MTP_BYTES);
         assert_eq!(plan.kv_and_recurrent, 7_564_812_288);
         assert_eq!(plan.prefill_transient, 28_056_068);
-        assert_eq!(plan.admission, 81_008_126_487);
+        assert_eq!(plan.admission, 88_924_002_839);
         let without_mtp = memory_plan(262_144, false, 512).unwrap();
         assert_eq!(without_mtp.optional_mtp, MTP_BYTES);
         assert_eq!(without_mtp.admission, 79_335_550_955);