Add native Qwen MTP speculation

This commit is contained in:
Georg Bauer committed 2026-09-03 22:38:52 +02:00
1 parent 714b39f7f3
commit bf82df77cb
14 files changed
+1302 -78

No files matched your search

+10 -4
View File
@@ -865,16 +865,22 @@ impl App {
self.preference_error = None;
}
Message::PreferenceGlmMtpChanged(value) => {
self.preference_draft.glm_mtp =
self.preference_draft.acceleration_model.supports_glm_mtp() && value;
self.preference_draft.glm_mtp = self
.preference_draft
.acceleration_model
.supports_integrated_mtp()
&& value;
if !self.preference_draft.glm_mtp {
self.preference_draft.glm_mtp_timing = false;
}
self.preference_error = None;
}
Message::PreferenceGlmMtpTimingChanged(value) => {
self.preference_draft.glm_mtp_timing =
self.preference_draft.acceleration_model.supports_glm_mtp() && value;
self.preference_draft.glm_mtp_timing = self
.preference_draft
.acceleration_model
.supports_integrated_mtp()
&& value;
if self.preference_draft.glm_mtp_timing {
self.preference_draft.glm_mtp = true;
}
+8 -8
View File
@@ -17,12 +17,12 @@ impl App {
let glm_mtp_toggle: Option<fn(bool) -> Message> = self
.preference_draft
.acceleration_model
.supports_glm_mtp()
.supports_integrated_mtp()
.then_some(Message::PreferenceGlmMtpChanged);
let glm_mtp_timing_toggle: Option<fn(bool) -> Message> = self
.preference_draft
.acceleration_model
.supports_glm_mtp()
.supports_integrated_mtp()
.then_some(Message::PreferenceGlmMtpTimingChanged);
let keep_vision_loaded_toggle: Option<fn(bool) -> Message> =
(self.preference_draft.acceleration_model == ModelChoice::Glm53Flash)
@@ -608,13 +608,13 @@ impl App {
text("SPECULATIVE DECODING").size(11).color(muted_text()),
hint(
toggle(self.preference_draft.glm_mtp)
.label("Enable integrated GLM MTP")
.label("Enable integrated MTP")
.on_toggle_maybe(glm_mtp_toggle),
"Uses the prediction head built into GLM for speculative decoding, so no separate draft model is loaded.",
"Uses the selected model's managed prediction head for speculative decoding. Qwen loads its pinned MTP sidecar; GLM uses its embedded head.",
),
hint(
toggle(self.preference_draft.glm_mtp_timing)
.label("Log GLM MTP timing counters")
.label("Log MTP timing counters")
.on_toggle_maybe(glm_mtp_timing_toggle),
"Records per-stage timings of the speculative path to the log, to show where the acceleration actually goes. A diagnostic aid that costs a little throughput.",
),
@@ -644,8 +644,8 @@ impl App {
),
text(if self.preference_draft.acceleration_model.supports_dspark() {
"DeepSeek V4 Flash 0731 uses its managed DSpark support artifact."
} else if self.preference_draft.acceleration_model.supports_glm_mtp() {
"GLM MTP is integrated; DSpark is unavailable for this model."
} else if self.preference_draft.acceleration_model.supports_integrated_mtp() {
"Integrated MTP is available; DSpark is unavailable for this model."
} else {
"No speculative-decoding support is available for this model."
})
@@ -656,7 +656,7 @@ impl App {
|engine| {
let settings = engine.speculative;
format!(
"Engine: GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {} • exact sampling {}",
"Engine: integrated MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {} • exact sampling {}",
if settings.glm_mtp { "on" } else { "off" },
if settings.glm_mtp_timing { "on" } else { "off" },
if settings.dspark { "on" } else { "off" },
+17
View File
@@ -369,6 +369,7 @@ impl App {
match stats.speculative_mode {
2 => "DSpark",
3 => "GLM MTP",
4 => "Qwen MTP",
_ => "Off",
},
),
@@ -376,11 +377,27 @@ impl App {
metric_row("Drafted", format_count(stats.drafted_tokens)),
metric_row("Accepted", format_count(stats.accepted_draft_tokens)),
metric_row("Acceptance", format!("{:.1}%", acceptance * 100.0)),
metric_row(
"Mean accepted depth",
if stats.speculative_cycles == 0 {
"0.00".into()
} else {
format!(
"{:.2}",
stats.accepted_depth_total as f64 / stats.speculative_cycles as f64
)
}
),
metric_row(
"Maximum accepted depth",
format_count(stats.accepted_depth_max)
),
metric_row(
"Target verifier passes",
format_count(stats.verifier_passes)
),
metric_row("Verifier wall time", format_milliseconds(stats.verifier_ms)),
metric_row("Repair wall time", format_milliseconds(stats.repair_ms)),
metric_row(
"Effective target-pass speedup",
format!("{effective_speedup:.2}×")
+103 -5
View File
@@ -329,9 +329,13 @@ impl LoadedModel {
)?;
let context = u32::try_from(settings.context_tokens)
.map_err(|_| "Qwen context must be a positive whole number")?;
qwen::QwenModel::open(&settings.artifacts.model, context)
.map(Box::new)
.map(Self::Qwen)
qwen::QwenModel::open_configured(
&settings.artifacts.model,
context,
settings.speculative.glm_mtp,
)
.map(Box::new)
.map(Self::Qwen)
} else {
Model::open(settings).map(Box::new).map(Self::Gguf)
}
@@ -815,8 +819,11 @@ impl Generator {
stats.speculative_cycles,
stats.drafted_tokens,
stats.accepted_draft_tokens,
stats.accepted_depth_total,
stats.accepted_depth_max,
stats.verifier_passes,
stats.verifier_ms,
stats.repair_ms,
);
self.metrics.ssd_stats(SsdStats {
enabled: stats.ssd_enabled,
@@ -1988,7 +1995,7 @@ fn sample_probabilities(
.sum();
let mut choice = rng.unit() * total;
for (token, probability) in probabilities {
if Some(*token) == excluded {
if Some(*token) == excluded || !probability.is_finite() || *probability <= 0.0 {
continue;
}
choice -= probability;
@@ -1999,7 +2006,9 @@ fn sample_probabilities(
probabilities
.iter()
.rev()
.find(|(token, _)| Some(*token) != excluded)
.find(|(token, probability)| {
Some(*token) != excluded && probability.is_finite() && *probability > 0.0
})
.map_or(0, |(token, _)| *token as i32)
}
@@ -2028,6 +2037,67 @@ fn exact_delta_sample(
)
}
#[cfg(any(target_os = "macos", test))]
fn sample_from_logits(
logits: &[f32],
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
) -> i32 {
sample_probabilities(
&sampling_probabilities(logits, temperature, top_p, min_p, top_k),
rng,
None,
)
}
#[cfg(any(target_os = "macos", test))]
#[allow(clippy::too_many_arguments)]
fn exact_speculative_sample(
target_logits: &[f32],
draft_logits: &[f32],
draft: i32,
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
) -> (i32, bool) {
let mut target = sampling_probabilities(target_logits, temperature, top_p, min_p, top_k);
let mut proposal = sampling_probabilities(draft_logits, temperature, top_p, min_p, top_k);
target.sort_unstable_by_key(|(token, _)| *token);
proposal.sort_unstable_by_key(|(token, _)| *token);
let probability = |values: &[(usize, f32)], token: usize| {
values
.binary_search_by_key(&token, |(candidate, _)| *candidate)
.ok()
.map_or(0.0, |index| values[index].1)
};
let target_probability = probability(&target, draft as usize);
let draft_probability = probability(&proposal, draft as usize);
if draft_probability > 0.0 && rng.unit() <= (target_probability / draft_probability).min(1.0) {
return (draft, true);
}
let target_best = target
.iter()
.max_by(|left, right| left.1.total_cmp(&right.1))
.map_or(draft, |(token, _)| *token as i32);
let mut residual = Vec::with_capacity(target.len());
for (token, target_probability) in target {
let remaining = target_probability - probability(&proposal, token);
if remaining > 0.0 {
residual.push((token, remaining));
}
}
if residual.is_empty() {
return (target_best, false);
}
(sample_probabilities(&residual, rng, None), false)
}
#[cfg(any(target_os = "macos", test))]
struct Rng(u64);
@@ -2113,6 +2183,34 @@ mod sampling_tests {
assert_eq!((replacement, was_draft), (1, false));
}
#[test]
fn probability_ratio_speculation_preserves_the_target_distribution() {
let target = [0.75_f32.ln(), 0.25_f32.ln()];
let proposal = [0.25_f32.ln(), 0.75_f32.ln()];
let mut rng = Rng::new(0x51a7_1c5e);
let mut counts = [0_u32; 2];
for _ in 0..40_000 {
let draft = sample_from_logits(&proposal, 1.0, 1.0, 0.0, 0, &mut rng);
let (token, _) =
exact_speculative_sample(&target, &proposal, draft, 1.0, 1.0, 0.0, 0, &mut rng);
counts[token as usize] += 1;
}
let observed = counts[0] as f32 / counts.iter().sum::<u32>() as f32;
assert!((observed - 0.75).abs() < 0.015, "observed {observed}");
let run = |seed| {
let mut rng = Rng::new(seed);
(0..32)
.map(|_| {
let draft = sample_from_logits(&proposal, 1.0, 1.0, 0.0, 0, &mut rng);
exact_speculative_sample(&target, &proposal, draft, 1.0, 1.0, 0.0, 0, &mut rng)
})
.collect::<Vec<_>>()
};
assert_eq!(run(77), run(77));
assert!(run(77).iter().any(|(_, accepted)| !accepted));
}
#[test]
fn sampling_probabilities_match_ds4_filter_order() {
let probabilities = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 0.8, 0.2, 0);
+24 -11
View File
@@ -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)
}
}
}
+1048 -23
View File
File diff suppressed because it is too large. Load diff
+46 -6
View File
@@ -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);
+22 -1
View File
@@ -113,8 +113,11 @@ pub(crate) struct MetricsSnapshot {
pub(crate) speculative_cycles: u64,
pub(crate) drafted_tokens: u64,
pub(crate) accepted_draft_tokens: u64,
pub(crate) accepted_depth_total: u64,
pub(crate) accepted_depth_max: u64,
pub(crate) verifier_passes: u64,
pub(crate) verifier_ms: u64,
pub(crate) repair_ms: u64,
pub(crate) ssd_enabled: bool,
pub(crate) ssd_resident_bytes: u64,
pub(crate) ssd_cache_bytes: u64,
@@ -206,8 +209,11 @@ pub(crate) struct Metrics {
speculative_cycles: AtomicU64,
drafted_tokens: AtomicU64,
accepted_draft_tokens: AtomicU64,
accepted_depth_total: AtomicU64,
accepted_depth_max: AtomicU64,
verifier_passes: AtomicU64,
verifier_ms: AtomicU64,
repair_ms: AtomicU64,
ssd_enabled: AtomicBool,
ssd_resident_bytes: AtomicU64,
ssd_cache_bytes: AtomicU64,
@@ -304,8 +310,11 @@ impl Metrics {
speculative_cycles: AtomicU64::new(0),
drafted_tokens: AtomicU64::new(0),
accepted_draft_tokens: AtomicU64::new(0),
accepted_depth_total: AtomicU64::new(0),
accepted_depth_max: AtomicU64::new(0),
verifier_passes: AtomicU64::new(0),
verifier_ms: AtomicU64::new(0),
repair_ms: AtomicU64::new(0),
ssd_enabled: AtomicBool::new(false),
ssd_resident_bytes: AtomicU64::new(0),
ssd_cache_bytes: AtomicU64::new(0),
@@ -472,23 +481,32 @@ impl Metrics {
self.source.store(WorkSource::None as u8, Ordering::Relaxed);
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn speculative_stats(
&self,
mode: u8,
cycles: u64,
drafted: u64,
accepted: u64,
accepted_depth_total: u64,
accepted_depth_max: u64,
verifier_passes: u64,
verifier_ms: u64,
repair_ms: u64,
) {
self.speculative_mode.store(mode, Ordering::Relaxed);
self.speculative_cycles.store(cycles, Ordering::Relaxed);
self.drafted_tokens.store(drafted, Ordering::Relaxed);
self.accepted_draft_tokens
.store(accepted, Ordering::Relaxed);
self.accepted_depth_total
.store(accepted_depth_total, Ordering::Relaxed);
self.accepted_depth_max
.store(accepted_depth_max, Ordering::Relaxed);
self.verifier_passes
.store(verifier_passes, Ordering::Relaxed);
self.verifier_ms.store(verifier_ms, Ordering::Relaxed);
self.repair_ms.store(repair_ms, Ordering::Relaxed);
}
pub(crate) fn ssd_stats(&self, stats: SsdStats) {
@@ -540,7 +558,7 @@ impl Metrics {
self.decode_tps.store(0, Ordering::Relaxed);
self.prefill_tps.store(0, Ordering::Relaxed);
self.prefill_sample.store(0, Ordering::Relaxed);
self.speculative_stats(0, 0, 0, 0, 0, 0);
self.speculative_stats(0, 0, 0, 0, 0, 0, 0, 0, 0);
self.ssd_stats(SsdStats::default());
self.model_unloads.fetch_add(1, Ordering::Relaxed);
}
@@ -695,8 +713,11 @@ impl Metrics {
speculative_cycles: self.speculative_cycles.load(Ordering::Relaxed),
drafted_tokens: self.drafted_tokens.load(Ordering::Relaxed),
accepted_draft_tokens: self.accepted_draft_tokens.load(Ordering::Relaxed),
accepted_depth_total: self.accepted_depth_total.load(Ordering::Relaxed),
accepted_depth_max: self.accepted_depth_max.load(Ordering::Relaxed),
verifier_passes: self.verifier_passes.load(Ordering::Relaxed),
verifier_ms: self.verifier_ms.load(Ordering::Relaxed),
repair_ms: self.repair_ms.load(Ordering::Relaxed),
ssd_enabled: self.ssd_enabled.load(Ordering::Relaxed),
ssd_resident_bytes: self.ssd_resident_bytes.load(Ordering::Relaxed),
ssd_cache_bytes: self.ssd_cache_bytes.load(Ordering::Relaxed),
+2 -2
View File
@@ -128,8 +128,8 @@ impl ModelChoice {
self == Self::Qwen38FlashNext
}
pub(crate) fn supports_glm_mtp(self) -> bool {
self.is_glm()
pub(crate) fn supports_integrated_mtp(self) -> bool {
self.is_glm() || self.is_qwen38()
}
pub(crate) fn main_artifact_size(self) -> u64 {
+1 -1
View File
@@ -13,7 +13,7 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
pub(super) const LABEL: &str = "Qwen3.8 Flash Next Bare Speed artifact set";
pub(super) const REPOSITORY: &str = "Youssofal/Qwen3.8-Flash-Next-MTPLX-Bare-Speed";
pub(super) const REVISION: &str = "74559cdf34fbfc0b593de72d17e93f37fd4f9ea7";
const MANIFEST_SHA256: &str = "eeec490fd3d0c0b1c9093be0a7fe7b9be4b49389ed5276de7530c2d3c11290ca";
const MANIFEST_SHA256: &str = "6f1172de47fa30b9602e13fc7ad14e578a813bb5b1fab8f0ef320041aac6c19a";
const MANIFEST_BYTES: &[u8] =
include_bytes!("../../assets/models/qwen38-flash-next-bare-speed.json");
+4 -4
View File
@@ -45,10 +45,10 @@ pub(crate) struct SpeculativePreferences {
impl SpeculativePreferences {
pub(crate) fn validate(&self, model: ModelChoice) -> Result<(), String> {
if self.glm_mtp_timing && !self.glm_mtp {
return Err("GLM MTP timing requires GLM MTP.".into());
return Err("MTP timing requires integrated MTP.".into());
}
if !model.supports_glm_mtp() && (self.glm_mtp || self.glm_mtp_timing) {
return Err("GLM MTP is available only for GLM models.".into());
if !model.supports_integrated_mtp() && (self.glm_mtp || self.glm_mtp_timing) {
return Err("Integrated MTP is unavailable for the selected model.".into());
}
if self.keep_vision_loaded && model != ModelChoice::Glm53Flash {
return Err("Persistent vision weights are available only for GLM 5.3 Flash.".into());
@@ -83,7 +83,7 @@ impl SpeculativePreferences {
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub(crate) struct EngineSpeculativeSettings {
pub(crate) glm_mtp: bool,
pub(crate) glm_mtp_timing: bool,