hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Add native Qwen MTP speculation

Browse Source
This commit is contained in:
Georg Bauer
2026-09-03 22:38:52 +02:00
parent 714b39f7f3
commit bf82df77cb
14 changed files with 1302 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
assets/models/qwen38-flash-next-bare-speed-tensors.tsv
View File

@@ -2492,15 +2492,15 @@ mtp.safetensors mtp.layers.0.mlp.shared_expert.up_proj.weight U32 640x320 4 64 a
mtp.safetensors mtp.layers.0.mlp.shared_expert_gate.biases BF16 1x40 4 64 affine 140045196 140045276
mtp.safetensors mtp.layers.0.mlp.shared_expert_gate.scales BF16 1x40 4 64 affine 140045276 140045356
mtp.safetensors mtp.layers.0.mlp.shared_expert_gate.weight U32 1x320 4 64 affine 620915756 620917036
mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.biases BF16 512x2560x20 2 64 affine 140966956 193395756
mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.scales BF16 512x2560x20 2 64 affine 641468716 693897516
mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.weight U32 512x2560x80 2 64 affine 1253145132 1672575532
mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.biases BF16 512x640x80 2 64 affine 693897772 746326572
mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.scales BF16 512x640x80 2 64 affine 20277900 72706700
mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.weight U32 512x640x320 2 64 affine 760314412 1179744812
mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.biases BF16 512x640x80 2 64 affine 87616396 140045196
mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.scales BF16 512x640x80 2 64 affine 1192852012 1245280812
mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.weight U32 512x640x320 2 64 affine 194808876 614239276
mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.biases BF16 512x2560x20 4 32 affine 140966956 193395756
mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.scales BF16 512x2560x20 4 32 affine 641468716 693897516
mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.weight U32 512x2560x80 4 32 affine 1253145132 1672575532
mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.biases BF16 512x640x80 4 32 affine 693897772 746326572
mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.scales BF16 512x640x80 4 32 affine 20277900 72706700
mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.weight U32 512x640x320 4 32 affine 760314412 1179744812
mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.biases BF16 512x640x80 4 32 affine 87616396 140045196
mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.scales BF16 512x640x80 4 32 affine 1192852012 1245280812
mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.weight U32 512x640x320 4 32 affine 194808876 614239276
mtp.safetensors mtp.layers.0.mlp_hyper_connection.block_inject_weight.weight BF16 4x10240 74345100 74427020
mtp.safetensors mtp.layers.0.mlp_hyper_connection.hc_norm.weight BF16 10240 752921132 752941612
mtp.safetensors mtp.layers.0.mlp_hyper_connection.input_mix_weight_down.weight BF16 320x10240 81062540 87616140
1 file name dtype shape quant_bits group_size quant_mode data_start data_end
2492 mtp.safetensors mtp.layers.0.mlp.shared_expert_gate.biases BF16 1x40 4 64 affine 140045196 140045276
2493 mtp.safetensors mtp.layers.0.mlp.shared_expert_gate.scales BF16 1x40 4 64 affine 140045276 140045356
2494 mtp.safetensors mtp.layers.0.mlp.shared_expert_gate.weight U32 1x320 4 64 affine 620915756 620917036
2495 mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.biases BF16 512x2560x20 2 4 64 32 affine 140966956 193395756
2496 mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.scales BF16 512x2560x20 2 4 64 32 affine 641468716 693897516
2497 mtp.safetensors mtp.layers.0.mlp.switch_mlp.down_proj.weight U32 512x2560x80 2 4 64 32 affine 1253145132 1672575532
2498 mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.biases BF16 512x640x80 2 4 64 32 affine 693897772 746326572
2499 mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.scales BF16 512x640x80 2 4 64 32 affine 20277900 72706700
2500 mtp.safetensors mtp.layers.0.mlp.switch_mlp.gate_proj.weight U32 512x640x320 2 4 64 32 affine 760314412 1179744812
2501 mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.biases BF16 512x640x80 2 4 64 32 affine 87616396 140045196
2502 mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.scales BF16 512x640x80 2 4 64 32 affine 1192852012 1245280812
2503 mtp.safetensors mtp.layers.0.mlp.switch_mlp.up_proj.weight U32 512x640x320 2 4 64 32 affine 194808876 614239276
2504 mtp.safetensors mtp.layers.0.mlp_hyper_connection.block_inject_weight.weight BF16 4x10240 74345100 74427020
2505 mtp.safetensors mtp.layers.0.mlp_hyper_connection.hc_norm.weight BF16 10240 752921132 752941612
2506 mtp.safetensors mtp.layers.0.mlp_hyper_connection.input_mix_weight_down.weight BF16 320x10240 81062540 87616140

4
assets/models/qwen38-flash-next-bare-speed.json
View File

@@ -11,7 +11,7 @@
"converter": "qwen38-artifacts-v1-identity"
},
"tensor_inventory": "qwen38-flash-next-bare-speed-tensors.tsv",
"tensor_inventory_sha256": "369cdc7e53c09eaa5eff5cdcc49f0d0eab0ddc175f34965ba2ed1be2673d5858",
"tensor_inventory_sha256": "b5731e6febcf865d276a0e7b144da02375f2d7e3129594f271ee9c6c351f4c8f",
"config": {
"/architectures/0": "Qwen4ExpForConditionalGeneration",
"/model_type": "qwen4_exp",
@@ -73,7 +73,7 @@
{"class":"qsa","file":"model-00016-of-00017.safetensors","tensor":"language_model.model.layers.11.self_attn.indexer.index_qk_proj.weight","row":0,"values":64,"bits":8,"group_size":64,"sha256":"c803ea77148621a5d6dfa4060a52c75f57185234a0bec4771671f69336bb8346"},
{"class":"gdn","file":"model-00016-of-00017.safetensors","tensor":"language_model.model.layers.0.linear_attn.A_log","row":0,"values":48,"sha256":"88c53a2a04bda1d96ee1ade6fa7dfa9c49d3245b1cc1e83673f02b7a36e07a85"},
{"class":"ple","file":"ngram-table.safetensors","tensor":"ngram.weight","row":0,"values":64,"bits":4,"group_size":32,"sha256":"2243dd9766046bb80d98e3baf5e59ba246958cfe4f340ebce8c1fff57a2810d9"},
{"class":"mtp","file":"mtp.safetensors","tensor":"mtp.layers.0.mlp.switch_mlp.gate_proj.weight","row":0,"values":64,"bits":2,"group_size":64,"sha256":"6293cce5e1b25a4f35cb8af46e6e13b8378ec4ca84ec7f68859066dda48049cc"}
{"class":"mtp","file":"mtp.safetensors","tensor":"mtp.layers.0.mlp.switch_mlp.gate_proj.weight","row":0,"values":64,"bits":4,"group_size":32,"sha256":"393d82675ae7e25275243fd4907d3894146ac1c5662b601aae998c489d10af5c"}
],
"files": [
{"path":"model-00001-of-00017.safetensors","role":"core","size":4666167150,"sha256":"a27232c9434b9d8961f198cf36f44e228b3f23a5f8d874ba4ca8af32a1b23ffc"},

14
src/app/preferences.rs
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;
}

16
src/app/view/preferences.rs
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
src/app/view/stats.rs
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}×")

104
src/engine.rs
View File

@@ -329,7 +329,11 @@ 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)
qwen::QwenModel::open_configured(
&settings.artifacts.model,
context,
settings.speculative.glm_mtp,
)
.map(Box::new)
.map(Self::Qwen)
} else {
@@ -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);

33
src/engine/metal.rs
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)
LoadedModel::Qwen(model) => {
{ QwenExecutor::open_configured(*model, context, speculative) }
.map(Box::new)
.map(Self::Qwen),
.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)
}
}
}

1071
src/engine/metal/qwen.rs
View File

File diff suppressed because it is too large Load Diff

52
src/engine/qwen.rs
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);

23
src/metrics.rs
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),

4
src/model.rs
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 {

2
src/model/qwen.rs
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");

8
src/settings.rs
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,

8
tools/qwen38-artifacts.rs
View File

@@ -657,7 +657,7 @@ fn precision_assignments(
reject_unassigned_quantized(&core, &assignments)?;
let mtp = tensor_headers(manifest, source, Some("mtp"), complete)?;
add_inferred_sidecar(&mtp, 64, &mut assignments)?;
add_inferred_sidecar(&mtp, &mut assignments)?;
let ple = tensor_headers(manifest, source, Some("ple"), complete)?;
let ple_group = config
.pointer("/mtplx_recipe/ngram/group_size")
@@ -711,7 +711,6 @@ fn validate_unquantized(base: &str, tensors: &BTreeMap<String, Tensor>) -> Resul
fn add_inferred_sidecar(
tensors: &BTreeMap<String, Tensor>,
group_size: u64,
assignments: &mut BTreeMap<String, Quantization>,
) -> Result<(), String> {
for (name, tensor) in tensors {
@@ -732,6 +731,11 @@ fn add_inferred_sidecar(
.shape
.last()
.ok_or_else(|| format!("{base}.scales has no dimensions"))?;
let group_size = if base.contains(".switch_mlp.") {
32
} else {
64
};
let pack = groups
.checked_mul(group_size)
.and_then(|columns| columns.checked_div(packed))