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

Support DeepSeek V4 Flash 0731

Browse Source
This commit is contained in:
Georg Bauer
2026-08-29 20:28:50 +02:00
parent ad855b321e
commit f1c177b754
23 changed files with 10510 additions and 2324 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/app.rs
View File

@@ -398,6 +398,7 @@ pub(crate) enum Message {
PreferenceGlmMtpTimingChanged(bool),
PreferenceDsparkConfidenceChanged(String),
PreferenceDsparkStrictChanged(bool),
PreferenceDsparkExactSamplingChanged(bool),
PreferenceSsdChanged(bool),
PreferenceSsdColdChanged(bool),
PreferenceSsdCacheChanged(String),
@@ -2719,7 +2720,7 @@ mod tests {
Some(Message::FocusPrevious)
));
assert!(ModelChoice::DeepSeekV4Flash.supports_dspark());
assert!(!ModelChoice::DeepSeekV4Flash0731.supports_dspark());
assert!(ModelChoice::DeepSeekV4Flash0731.supports_dspark());
assert!(!ModelChoice::DeepSeekV4Pro.supports_dspark());
assert!(!ModelChoice::Glm52.supports_dspark());
}

26
src/app/preferences.rs
View File

@@ -45,6 +45,7 @@ pub(super) struct PreferenceDraft {
pub(super) glm_mtp_timing: bool,
pub(super) dspark_confidence_threshold: String,
pub(super) dspark_strict: bool,
pub(super) dspark_exact_sampling: bool,
pub(super) ssd_streaming: bool,
pub(super) ssd_streaming_cold: bool,
pub(super) ssd_cache: String,
@@ -111,6 +112,7 @@ impl PreferenceDraft {
glm_mtp_timing: speculative.glm_mtp_timing,
dspark_confidence_threshold: optional_string(speculative.dspark_confidence_threshold),
dspark_strict: speculative.dspark_strict,
dspark_exact_sampling: speculative.dspark_exact_sampling,
ssd_streaming: runtime.ssd.enabled,
ssd_streaming_cold: runtime.ssd.cold,
ssd_cache: optional_string(runtime.ssd.cache),
@@ -226,6 +228,7 @@ impl PreferenceDraft {
&self.dspark_confidence_threshold,
)?,
dspark_strict: self.dspark_strict,
dspark_exact_sampling: self.dspark_exact_sampling,
})
}
@@ -267,6 +270,7 @@ impl PreferenceDraft {
self.glm_mtp_timing = speculative.glm_mtp_timing;
self.dspark_confidence_threshold = optional_string(speculative.dspark_confidence_threshold);
self.dspark_strict = speculative.dspark_strict;
self.dspark_exact_sampling = speculative.dspark_exact_sampling;
self.ssd_streaming = ssd.enabled;
self.ssd_streaming_cold = ssd.cold;
self.ssd_cache = optional_string(ssd.cache);
@@ -693,12 +697,16 @@ impl App {
self.preference_error = None;
}
Message::PreferenceLegacyMtpChanged(enabled) => {
self.preference_draft.legacy_mtp_enabled =
self.preference_draft.acceleration_model.supports_dspark() && enabled;
self.preference_draft.legacy_mtp_enabled = self
.preference_draft
.acceleration_model
.supports_legacy_mtp()
&& enabled;
if self.preference_draft.legacy_mtp_enabled {
self.preference_draft.dspark_enabled = false;
self.preference_draft.dspark_confidence_threshold.clear();
self.preference_draft.dspark_strict = false;
self.preference_draft.dspark_exact_sampling = false;
}
self.preference_error = None;
}
@@ -708,6 +716,7 @@ impl App {
if !self.preference_draft.dspark_enabled {
self.preference_draft.dspark_confidence_threshold.clear();
self.preference_draft.dspark_strict = false;
self.preference_draft.dspark_exact_sampling = false;
} else {
self.preference_draft.legacy_mtp_enabled = false;
}
@@ -903,6 +912,15 @@ impl App {
}
self.preference_error = None;
}
Message::PreferenceDsparkExactSamplingChanged(value) => {
self.preference_draft.dspark_exact_sampling =
self.preference_draft.acceleration_model.supports_dspark() && value;
if self.preference_draft.dspark_exact_sampling {
self.preference_draft.dspark_enabled = true;
self.preference_draft.legacy_mtp_enabled = false;
}
self.preference_error = None;
}
Message::PreferenceSsdChanged(value) => {
self.preference_draft.ssd_streaming = value;
self.preference_error = None;
@@ -1011,14 +1029,14 @@ mod tests {
assert_eq!(draft.context_tokens, "32768");
draft.context_tokens = "456".into();
draft
.select_generation(ModelChoice::DeepSeekV4Flash, ReasoningMode::High)
.select_generation(ModelChoice::DeepSeekV4Flash0731, ReasoningMode::High)
.unwrap();
assert_eq!(draft.context_tokens, "123");
draft.select_acceleration(ModelChoice::Glm52).unwrap();
assert!(!draft.ssd_streaming);
draft
.select_acceleration(ModelChoice::DeepSeekV4Flash)
.select_acceleration(ModelChoice::DeepSeekV4Flash0731)
.unwrap();
assert!(draft.ssd_streaming);
}

20
src/app/view/preferences.rs
View File

@@ -6,7 +6,7 @@ impl App {
let legacy_mtp_toggle: Option<fn(bool) -> Message> = self
.preference_draft
.acceleration_model
.supports_dspark()
.supports_legacy_mtp()
.then_some(Message::PreferenceLegacyMtpChanged);
let legacy_mtp = hint(
toggle(self.preference_draft.legacy_mtp_enabled)
@@ -36,6 +36,11 @@ impl App {
.acceleration_model
.supports_dspark()
.then_some(Message::PreferenceDsparkStrictChanged);
let dspark_exact_toggle: Option<fn(bool) -> Message> = self
.preference_draft
.acceleration_model
.supports_dspark()
.then_some(Message::PreferenceDsparkExactSamplingChanged);
let effective = self
.preference_draft
.effective_for(
@@ -74,7 +79,7 @@ impl App {
text_input("Automatic", &self.preference_draft.directional_steering_ffn);
let mut steering_attn = text_input("0", &self.preference_draft.directional_steering_attn);
let mut dspark_confidence = text_input(
"0.9 (DS4 default)",
"0.6 (DS4 default)",
&self.preference_draft.dspark_confidence_threshold,
);
if self.preference_draft.model != ModelChoice::Glm52 {
@@ -509,7 +514,7 @@ impl App {
dspark,
preference_input_row(
"DSpark confidence threshold",
"How sure the draft model must be, from 0 to 1, before its token is handed to the verifier. Lower forwards more guesses for more speed and more rejected work; blank uses DS4's 0.9.",
"How sure the draft model must be, from 0 to 1, before its token is handed to the verifier. Lower forwards more guesses for more speed and more rejected work; blank uses DS4's 0.6, or 0.8 for exact sampling.",
dspark_confidence,
),
hint(
@@ -518,6 +523,12 @@ impl App {
.on_toggle_maybe(dspark_strict_toggle),
"Lets the draft model only propose, never decide: every token is sampled by the full model. Gives up some of the speedup in exchange for output identical to non-speculative decoding.",
),
hint(
toggle(self.preference_draft.dspark_exact_sampling)
.label("Use exact DSpark sampling")
.on_toggle_maybe(dspark_exact_toggle),
"For non-zero temperatures, applies DS4's exact acceptance and corrected rejection sampling. Off uses the faster opportunistic mode: sample a boundary token, then accept DSpark tokens only while they match the target's greedy path.",
),
text(if self.preference_draft.acceleration_model.supports_dspark() {
"Legacy MTP and DSpark use separate managed support artifacts; entering a DSpark threshold or enabling strict mode selects DSpark."
} else if self.preference_draft.acceleration_model == ModelChoice::Glm52 {
@@ -532,7 +543,7 @@ impl App {
|engine| {
let settings = engine.speculative;
format!(
"Engine: MTP draft {} • margin {} • legacy MTP {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}",
"Engine: MTP draft {} • margin {} • legacy MTP {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {} • exact sampling {}",
settings.mtp_draft_tokens,
settings.mtp_margin,
if self.preference_draft.legacy_mtp_enabled { "on" } else { "off" },
@@ -542,6 +553,7 @@ impl App {
settings.dspark_confidence_threshold,
if settings.dspark_confidence_threshold_set { " explicit" } else { " default" },
if settings.dspark_strict { "on" } else { "off" },
if settings.dspark_exact_sampling { "on" } else { "off" },
)
},
))

185
src/engine.rs
View File

@@ -8,7 +8,7 @@ mod validation;
#[cfg(target_os = "macos")]
use crate::metrics::{KvLookup, Metrics, SsdStats};
use crate::model::ModelChoice;
use crate::model::{ModelChoice, validate_engine_artifacts};
#[cfg(target_os = "macos")]
use crate::settings::TurnSettings;
use crate::settings::{EngineSettings, ReasoningMode};
@@ -232,6 +232,12 @@ pub(crate) struct ModelSummary {
impl Model {
#[allow(dead_code)]
pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
validate_engine_artifacts(
settings.model,
settings.artifacts.mtp.is_some() && !settings.speculative.dspark,
settings.speculative.dspark,
&settings.artifacts,
)?;
let mut model = Self::open_main(&settings.artifacts.model, settings.model)?;
if settings.execution.warm_weights {
model.main.warm()?;
@@ -1225,8 +1231,17 @@ impl Generator {
cancelled,
)?
} else {
self.executor.eval(token)?;
vec![token]
self.executor.eval_speculative_sampled(
token,
generation_limit - generated_tokens,
settings.reasoning_mode,
settings.temperature,
settings.top_p,
settings.min_p,
settings.top_k,
&mut rng,
cancelled,
)?
};
self.publish_execution_stats();
for token in cycle {
@@ -1545,12 +1560,30 @@ fn sample(
top_k: i32,
rng: &mut Rng,
) -> i32 {
let probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k);
sample_probabilities(&probabilities, rng, None)
}
#[cfg(any(target_os = "macos", test))]
fn sampling_probabilities(
logits: &[f32],
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
) -> Vec<(usize, f32)> {
let greedy = || {
vec![(
logits
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.map_or(0, |(index, _)| index),
1.0,
)]
};
if temperature <= 0.0 {
return logits
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.map_or(0, |(index, _)| index as i32);
return greedy();
}
let maximum = logits
.iter()
@@ -1558,7 +1591,7 @@ fn sample(
.filter(|value| value.is_finite())
.fold(f32::NEG_INFINITY, f32::max);
if !maximum.is_finite() {
return 0;
return greedy();
}
let top_p = if top_p <= 0.0 || top_p > 1.0 {
1.0
@@ -1571,45 +1604,96 @@ fn sample(
.enumerate()
.filter(|(_, logit)| logit.is_finite())
.map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp()))
.filter(|(_, probability)| *probability >= min_p)
.collect();
if probabilities.is_empty() {
return logits
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.map_or(0, |(index, _)| index as i32);
return greedy();
}
if top_p < 1.0 || top_k > 0 {
if top_p < 1.0 || top_k > 0 || min_p > 0.0 {
probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
if top_k > 0 {
probabilities.truncate(probabilities.len().min(top_k as usize));
probabilities.truncate(probabilities.len().min((top_k as usize).min(1024)));
}
}
if top_p < 1.0 {
let total: f32 = probabilities
.iter()
.map(|(_, probability)| probability)
.sum();
let mut kept = 0.0;
let count = probabilities
.iter()
.position(|(_, probability)| {
kept += *probability;
kept / total >= top_p
})
.map_or(probabilities.len(), |index| index + 1);
probabilities.truncate(count);
let total: f32 = probabilities
.iter()
.map(|(_, probability)| probability)
.sum();
let mut kept = 0.0;
let mut count = 0;
for (_, probability) in &probabilities {
if count > 0 && *probability < min_p {
break;
}
kept += *probability;
count += 1;
if kept / total >= top_p {
break;
}
}
let kept_total: f32 = probabilities.iter().map(|(_, p)| p).sum();
let mut choice = rng.unit() * kept_total;
for (token, probability) in &probabilities {
probabilities.truncate(count);
if probabilities.is_empty() || !kept.is_finite() || kept <= 0.0 {
return greedy();
}
for (_, probability) in &mut probabilities {
*probability /= kept;
}
if top_p >= 1.0 && top_k <= 0 {
probabilities.sort_unstable_by_key(|(token, _)| *token);
}
probabilities
}
#[cfg(any(target_os = "macos", test))]
fn sample_probabilities(
probabilities: &[(usize, f32)],
rng: &mut Rng,
excluded: Option<usize>,
) -> i32 {
let total: f32 = probabilities
.iter()
.filter(|(token, _)| Some(*token) != excluded)
.map(|(_, probability)| probability)
.sum();
let mut choice = rng.unit() * total;
for (token, probability) in probabilities {
if Some(*token) == excluded {
continue;
}
choice -= probability;
if choice <= 0.0 {
return *token as i32;
}
}
probabilities.last().map_or(0, |(token, _)| *token as i32)
probabilities
.iter()
.rev()
.find(|(token, _)| Some(*token) != excluded)
.map_or(0, |(token, _)| *token as i32)
}
#[cfg(any(target_os = "macos", test))]
fn exact_delta_sample(
logits: &[f32],
draft: i32,
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
) -> (i32, bool) {
let mut probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k);
let draft_probability = probabilities
.iter()
.find(|(token, _)| *token == draft as usize)
.map_or(0.0, |(_, probability)| *probability);
if rng.unit() <= draft_probability {
return (draft, true);
}
probabilities.sort_unstable_by_key(|(token, _)| *token);
(
sample_probabilities(&probabilities, rng, Some(draft as usize)),
false,
)
}
#[cfg(any(target_os = "macos", test))]
@@ -1665,6 +1749,37 @@ mod sampling_tests {
assert_eq!(chunks, ["hello "]);
}
#[test]
fn exact_delta_sampling_accepts_or_corrects_the_draft() {
let mut accept_rng = Rng::new(2);
let (accepted, was_draft) =
exact_delta_sample(&[10.0, 0.0], 0, 1.0, 1.0, 0.0, 0, &mut accept_rng);
assert_eq!((accepted, was_draft), (0, true));
let mut reject_rng = Rng::new(1);
let (replacement, was_draft) =
exact_delta_sample(&[0.0, 10.0], 0, 1.0, 1.0, 0.0, 0, &mut reject_rng);
assert_eq!((replacement, was_draft), (1, false));
}
#[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);
assert_eq!(probabilities.len(), 2);
assert_eq!(probabilities[0].0, 1);
assert_eq!(probabilities[1].0, 2);
assert!((probabilities.iter().map(|(_, value)| value).sum::<f32>() - 1.0).abs() < 1e-6);
let min_p_only = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 1.0, 0.2, 0);
assert_eq!(
min_p_only
.iter()
.map(|(token, _)| *token)
.collect::<Vec<_>>(),
[1, 2]
);
}
#[test]
fn split_utf8_token_bytes_are_joined_before_decoding() {
let mut generated = ChatTurn {

680
src/engine/metal.rs
View File

@@ -10,7 +10,7 @@ use profile::ExpertProfile;
use super::gguf::{F16, F32, Gguf, IQ2_XXS, Q4_K, Q8_0, Tensor as GgufTensor};
use super::validation::{DsparkConfig, SupportKind, dspark_config};
use super::{Model, ModelFamily};
use super::{Model, ModelFamily, Rng, exact_delta_sample};
use crate::model::ModelChoice;
use crate::settings::{
EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
@@ -407,7 +407,13 @@ impl Dspark {
capture_mask: 0,
cache_start: 0,
cache_len: 0,
confidence_threshold: settings.dspark_confidence_threshold,
confidence_threshold: if settings.dspark_exact_sampling
&& !settings.dspark_confidence_threshold_set
{
settings.dspark_confidence_threshold.max(0.8)
} else {
settings.dspark_confidence_threshold
},
strict: settings.dspark_strict || quality,
drafted: 0,
accepted: 0,
@@ -2600,6 +2606,48 @@ struct SpecFrontier {
dspark_cache_len: u32,
}
struct SpecPrefixFrontier {
layers: Vec<LayerFrontier>,
}
struct BatchVerification {
tops: Vec<i32>,
logits: Vec<Vec<f32>>,
prefixes: Vec<SpecPrefixFrontier>,
}
fn capture_compression_frontier(
state: &CompressionState,
bytes: u64,
purpose: &str,
) -> Result<CompressionFrontier, String> {
let state_kv = Buffer::bytes(bytes)?;
let state_score = Buffer::bytes(bytes)?;
state_kv.copy_from(0, &state.state_kv, 0, bytes, purpose)?;
state_score.copy_from(0, &state.state_score, 0, bytes, purpose)?;
Ok(CompressionFrontier {
state_kv,
state_score,
bytes,
rows: state.rows,
})
}
fn restore_compression_frontier(
state: &mut CompressionState,
saved: &CompressionFrontier,
purpose: &str,
) -> Result<(), String> {
state
.state_kv
.copy_from(0, &saved.state_kv, 0, saved.bytes, purpose)?;
state
.state_score
.copy_from(0, &saved.state_score, 0, saved.bytes, purpose)?;
state.rows = saved.rows;
Ok(())
}
impl LayerState {
fn allocate(model: &Model, index: u32, context: u32, raw_cap: u32) -> Result<Self, String> {
let shape = model.shape;
@@ -2671,7 +2719,7 @@ impl Session {
// SAFETY: declaration order is required because Rust drops fields in order.
// `session` must release every Buffer before `_context` calls ds4_gpu_cleanup(),
// and `_context` must drop before `model` unmaps memory wrapped without copying
// by native/metal/ds4_metal.m:10329. This intentionally differs from
// by `ds4_gpu_cleanup` in native/metal/ds4_metal.m. This intentionally differs from
// DS4's `ds4.c` consumes this exact field order; do not reorder it.
#[derive(Clone, Copy, Default)]
pub(super) struct ExecutionStats {
@@ -2934,31 +2982,6 @@ impl DeepSeekExecutor {
}
fn snapshot_spec_frontier(&self) -> Result<SpecFrontier, String> {
fn snapshot(state: &CompressionState, bytes: u64) -> Result<CompressionFrontier, String> {
let state_kv = Buffer::bytes(bytes)?;
let state_score = Buffer::bytes(bytes)?;
state_kv.copy_from(
0,
&state.state_kv,
0,
bytes,
"saving speculative compressor KV state",
)?;
state_score.copy_from(
0,
&state.state_score,
0,
bytes,
"saving speculative compressor score state",
)?;
Ok(CompressionFrontier {
state_kv,
state_score,
bytes,
rows: state.rows,
})
}
let shape = self.model.shape;
let commands = Commands::begin()?;
let layers = self
@@ -2971,9 +2994,10 @@ impl DeepSeekExecutor {
.as_ref()
.map(|state| {
let coefficient = if state.ratio == 4 { 2 } else { 1 };
snapshot(
capture_compression_frontier(
state,
coefficient * coefficient * state.ratio as u64 * shape.head_dim * 4,
"saving speculative compressor state",
)
})
.transpose()?;
@@ -2981,7 +3005,11 @@ impl DeepSeekExecutor {
.indexer
.as_ref()
.map(|state| {
snapshot(state, 4 * state.ratio as u64 * shape.indexer_head_dim * 4)
capture_compression_frontier(
state,
4 * state.ratio as u64 * shape.indexer_head_dim * 4,
"saving speculative indexer state",
)
})
.transpose()?;
Ok::<_, String>(LayerFrontier {
@@ -3020,40 +3048,26 @@ impl DeepSeekExecutor {
}
fn restore_spec_frontier(&mut self, frontier: &SpecFrontier) -> Result<(), String> {
fn restore(
state: &mut CompressionState,
saved: &CompressionFrontier,
) -> Result<(), String> {
state.state_kv.copy_from(
0,
&saved.state_kv,
0,
saved.bytes,
"restoring speculative compressor KV state",
)?;
state.state_score.copy_from(
0,
&saved.state_score,
0,
saved.bytes,
"restoring speculative compressor score state",
)?;
state.rows = saved.rows;
Ok(())
}
if frontier.layers.len() != self.session.layers.len() {
return Err("speculative frontier layer count changed".into());
}
let commands = Commands::begin()?;
for (layer, saved) in self.session.layers.iter_mut().zip(&frontier.layers) {
match (&mut layer.compression, &saved.compression) {
(Some(state), Some(saved)) => restore(state, saved)?,
(Some(state), Some(saved)) => restore_compression_frontier(
state,
saved,
"restoring speculative compressor state",
)?,
(None, None) => {}
_ => return Err("speculative compressor layout changed".into()),
}
match (&mut layer.indexer, &saved.indexer) {
(Some(state), Some(saved)) => restore(state, saved)?,
(Some(state), Some(saved)) => restore_compression_frontier(
state,
saved,
"restoring speculative indexer state",
)?,
(None, None) => {}
_ => return Err("speculative indexer layout changed".into()),
}
@@ -3078,6 +3092,63 @@ impl DeepSeekExecutor {
Ok(())
}
fn commit_spec_prefix(
&mut self,
baseline: &SpecFrontier,
prefix: &SpecPrefixFrontier,
proposals: &[i32],
logits: &[f32],
) -> Result<(), String> {
let count =
u32::try_from(proposals.len()).map_err(|_| "speculative prefix is too large")?;
if count == 0 || prefix.layers.len() != self.session.layers.len() {
return Err("invalid speculative prefix frontier".into());
}
let commands = Commands::begin()?;
for (layer, saved) in self.session.layers.iter_mut().zip(&prefix.layers) {
match (&mut layer.compression, &saved.compression) {
(Some(state), Some(saved)) => restore_compression_frontier(
state,
saved,
"committing speculative compressor prefix",
)?,
(None, None) => {}
_ => return Err("speculative compressor prefix layout changed".into()),
}
match (&mut layer.indexer, &saved.indexer) {
(Some(state), Some(saved)) => restore_compression_frontier(
state,
saved,
"committing speculative indexer prefix",
)?,
(None, None) => {}
_ => return Err("speculative indexer prefix layout changed".into()),
}
}
if let Some(dspark) = &mut self.dspark {
let row = u64::from(count - 1);
for slot in 0..dspark.config.target_layers.len() as u64 {
dspark.target_hidden.copy_from(
slot * self.model.shape.embd * 4,
&dspark.target_hidden_batch,
(slot * u64::from(self.session.prefill_cap) + row) * self.model.shape.embd * 4,
self.model.shape.embd * 4,
"committing speculative DSpark target prefix",
)?;
}
dspark.capture_mask = (1_u32 << dspark.config.target_layers.len()) - 1;
dspark.cache_start = baseline.dspark_cache_start;
dspark.cache_len = baseline.dspark_cache_len;
dspark.commit_proposed_prefix(count, self.session.raw_cap);
}
commands.finish()?;
self.session.position = baseline.position + count;
self.tokens.truncate(baseline.token_len);
self.tokens.extend_from_slice(proposals);
self.logits.clone_from_slice(logits);
Ok(())
}
fn verify_target_suffix(
&mut self,
proposals: &[i32],
@@ -3115,10 +3186,10 @@ impl DeepSeekExecutor {
}
let frontier = self.snapshot_spec_frontier()?;
let row_tops = match self.eval_batch_tops(proposals) {
Ok(tops) => {
let verification = match self.eval_batch_tops(proposals) {
Ok(verification) => {
self.verifier_passes += 1;
tops
verification
}
Err(error) => {
self.restore_spec_frontier(&frontier)?;
@@ -3129,7 +3200,7 @@ impl DeepSeekExecutor {
}
};
let mut commit = 1_usize;
while commit < proposals.len() && row_tops[commit - 1] == proposals[commit] {
while commit < proposals.len() && verification.tops[commit - 1] == proposals[commit] {
commit += 1;
}
if commit == proposals.len() {
@@ -3139,6 +3210,17 @@ impl DeepSeekExecutor {
return Ok(proposals.to_vec());
}
if let (Some(prefix), Some(logits)) = (
verification.prefixes.get(commit - 1),
verification.logits.get(commit - 1),
) {
self.commit_spec_prefix(&frontier, prefix, &proposals[..commit], logits)?;
self.verifier_ns = self
.verifier_ns
.saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
return Ok(proposals[..commit].to_vec());
}
self.restore_spec_frontier(&frontier)?;
if let Some(dspark) = &mut self.dspark {
dspark.commit_proposed_prefix(1, self.session.raw_cap);
@@ -3157,6 +3239,133 @@ impl DeepSeekExecutor {
Ok(proposals[..commit].to_vec())
}
#[allow(clippy::too_many_arguments)]
fn verify_target_suffix_stochastic(
&mut self,
proposals: &[i32],
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
cancelled: &std::sync::atomic::AtomicBool,
) -> Result<(Vec<i32>, usize), String> {
if proposals.is_empty() || cancelled.load(std::sync::atomic::Ordering::Relaxed) {
return Ok((Vec::new(), 0));
}
let started = Instant::now();
if self.quality
|| proposals.len() == 1
|| self
.ssd
.as_ref()
.is_some_and(|ssd| u64::from(ssd.cache_experts) < self.model.shape.experts)
{
let mut emitted = Vec::new();
let mut accepted = 0;
for &proposal in proposals {
let (token, was_draft) = exact_delta_sample(
&self.logits,
proposal,
temperature,
top_p,
min_p,
top_k,
rng,
);
self.eval_target(token)?;
self.verifier_passes += 1;
emitted.push(token);
if !was_draft {
break;
}
accepted += 1;
if cancelled.load(std::sync::atomic::Ordering::Relaxed) {
break;
}
}
self.verifier_ns = self
.verifier_ns
.saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
return Ok((emitted, accepted));
}
let (first, accepted_first) = exact_delta_sample(
&self.logits,
proposals[0],
temperature,
top_p,
min_p,
top_k,
rng,
);
if !accepted_first {
if let Some(dspark) = &mut self.dspark {
dspark.commit_proposed_prefix(1, self.session.raw_cap);
}
self.eval_target(first)?;
self.verifier_passes += 1;
self.verifier_ns = self
.verifier_ns
.saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
return Ok((vec![first], 0));
}
let frontier = self.snapshot_spec_frontier()?;
let verification = match self.eval_batch_tops(proposals) {
Ok(verification) => {
self.verifier_passes += 1;
verification
}
Err(error) => {
self.restore_spec_frontier(&frontier)?;
return Err(error);
}
};
let mut accepted = 1;
let mut replacement = None;
for (index, &proposal) in proposals.iter().enumerate().skip(1) {
let (token, was_draft) = exact_delta_sample(
&verification.logits[index - 1],
proposal,
temperature,
top_p,
min_p,
top_k,
rng,
);
if !was_draft {
replacement = Some(token);
break;
}
accepted += 1;
}
if replacement.is_none() {
self.verifier_ns = self
.verifier_ns
.saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
return Ok((proposals.to_vec(), accepted));
}
let prefix = verification
.prefixes
.get(accepted - 1)
.ok_or("missing stochastic verifier prefix")?;
let logits = verification
.logits
.get(accepted - 1)
.ok_or("missing stochastic verifier logits")?;
self.commit_spec_prefix(&frontier, prefix, &proposals[..accepted], logits)?;
let replacement = replacement.expect("replacement disappeared");
self.eval_target(replacement)?;
self.verifier_passes += 1;
let mut emitted = proposals[..accepted].to_vec();
emitted.push(replacement);
self.verifier_ns = self
.verifier_ns
.saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
Ok((emitted, accepted))
}
pub(super) fn eval_speculative_greedy(
&mut self,
first_token: i32,
@@ -3269,6 +3478,81 @@ impl DeepSeekExecutor {
Ok(accepted)
}
#[allow(clippy::too_many_arguments)]
fn eval_speculative_sampled(
&mut self,
first_token: i32,
max_tokens: u32,
reasoning: ReasoningMode,
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
cancelled: &std::sync::atomic::AtomicBool,
) -> Result<Vec<i32>, String> {
if self.dspark.is_none() {
self.eval_target(first_token)?;
return Ok(vec![first_token]);
}
if !self.speculative.dspark_exact_sampling {
return self.eval_speculative_greedy(first_token, max_tokens, reasoning, cancelled);
}
self.speculative_cycles += 1;
self.eval_target(first_token)?;
let mut emitted = vec![first_token];
if self.dspark.as_ref().is_some_and(|dspark| dspark.strict)
|| max_tokens <= 1
|| cancelled.load(std::sync::atomic::Ordering::Relaxed)
{
return Ok(emitted);
}
if self.ssd.is_some() {
install_speculative_model_maps(&self.model, "DSpark support mapping")?;
}
let mut dspark = self.dspark.take().expect("DSpark disappeared");
let proposals = dspark.propose(
&self.model,
&self.weights,
first_token,
self.session.position.saturating_sub(1),
self.session.raw_cap,
);
self.dspark = Some(dspark);
let mut proposals = proposals?;
proposals.truncate(
max_tokens
.saturating_sub(1)
.min(self.session.context.saturating_sub(self.session.position))
as usize,
);
if let Some(stop) = proposals
.iter()
.position(|token| self.model.is_stop_token_for_reasoning(*token, reasoning))
{
proposals.truncate(stop + 1);
}
if proposals.len() < 2 {
self.dspark
.as_mut()
.expect("DSpark disappeared")
.commit_proposed_prefix(1, self.session.raw_cap);
return Ok(emitted);
}
let (verified, accepted) = self.verify_target_suffix_stochastic(
&proposals,
temperature,
top_p,
min_p,
top_k,
rng,
cancelled,
)?;
emitted.extend_from_slice(&verified);
self.dspark.as_mut().expect("DSpark disappeared").accepted += accepted as u64;
Ok(emitted)
}
fn legacy_mtp_draft(&mut self, token: i32, target_hc: bool) -> Result<(i32, f32), String> {
let support = self
.model
@@ -3511,11 +3795,15 @@ impl DeepSeekExecutor {
self.eval_batch_inner(tokens, false).map(|_| ())
}
fn eval_batch_tops(&mut self, tokens: &[i32]) -> Result<Vec<i32>, String> {
fn eval_batch_tops(&mut self, tokens: &[i32]) -> Result<BatchVerification, String> {
self.eval_batch_inner(tokens, true)
}
fn eval_batch_inner(&mut self, tokens: &[i32], collect_tops: bool) -> Result<Vec<i32>, String> {
fn eval_batch_inner(
&mut self,
tokens: &[i32],
collect_tops: bool,
) -> Result<BatchVerification, String> {
let rows = u32::try_from(tokens.len()).map_err(|_| "prefill batch is too large")?;
if rows == 0 || rows > self.session.prefill_cap {
return Err("prefill batch exceeds the configured prefill workspace".into());
@@ -3535,17 +3823,16 @@ impl DeepSeekExecutor {
let size = self.model.main.len();
let shape = self.model.shape;
let pos = self.session.position;
let batch_selected_addr = self.ssd.is_some()
&& self.weights.layers.first().is_some_and(|layer| unsafe {
ds4_gpu_stream_prefill_batch_selected_addr_enabled(
rows,
shape.experts as u32,
shape.experts_used as u32,
layer.expert_gate.kind,
layer.expert_down.kind,
) != 0
});
let mut prefixes = (0..if collect_tops { rows } else { 0 })
.map(|_| SpecPrefixFrontier {
layers: (0..shape.layers)
.map(|_| LayerFrontier {
compression: None,
indexer: None,
})
.collect(),
})
.collect::<Vec<_>>();
if self.ssd.is_some() {
install_deepseek_model_spans(
&self.model,
@@ -3580,6 +3867,16 @@ impl DeepSeekExecutor {
.enumerate()
{
let started = Instant::now();
let layer_selected_addr = self.ssd.is_some()
&& unsafe {
ds4_gpu_stream_prefill_batch_selected_addr_enabled(
rows,
shape.experts as u32,
shape.experts_used as u32,
weights.expert_gate.kind,
weights.expert_down.kind,
) != 0
};
if let Some(ssd) = &self.ssd {
install_deepseek_model_spans(
&self.model,
@@ -3587,7 +3884,7 @@ impl DeepSeekExecutor {
&self.model,
weights,
index as u32,
batch_selected_addr,
layer_selected_addr,
ssd.per_expert_bytes,
)?,
"DeepSeek prefill layer mapping",
@@ -3606,6 +3903,7 @@ impl DeepSeekExecutor {
rows,
self.session.raw_cap,
self.steering.as_ref(),
collect_tops.then_some(prefixes.as_mut_slice()),
)?;
if let Some(profile) = &mut self.profile {
profile.record(
@@ -3659,6 +3957,7 @@ impl DeepSeekExecutor {
let output_rows = if collect_tops { rows } else { 1 };
let first_output = rows - output_rows;
let mut tops = Vec::with_capacity(output_rows as usize);
let mut output_logits = Vec::with_capacity(output_rows as usize);
for row in first_output..rows {
let commands = Commands::begin()?;
self.session.scratch.current_hc.copy_from(
@@ -3672,13 +3971,20 @@ impl DeepSeekExecutor {
commands.finish()?;
self.session.scratch.logits.read_f32(&mut self.logits)?;
tops.push(argmax(&self.logits));
if collect_tops {
output_logits.push(self.logits.clone());
}
}
self.session.position += rows;
self.tokens.extend_from_slice(tokens);
if let Some(profile) = &self.profile {
profile.write()?;
}
Ok(tops)
Ok(BatchVerification {
tops,
logits: output_logits,
prefixes,
})
}
pub(super) fn logits(&self) -> &[f32] {
@@ -4003,6 +4309,7 @@ impl Executor {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
crate::settings::EngineSsdSettings {
enabled: false,
@@ -4061,6 +4368,38 @@ impl Executor {
}
}
#[allow(clippy::too_many_arguments)]
pub(super) fn eval_speculative_sampled(
&mut self,
token: i32,
max_tokens: u32,
reasoning: ReasoningMode,
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
cancelled: &std::sync::atomic::AtomicBool,
) -> Result<Vec<i32>, String> {
match self {
Self::DeepSeek(executor) => executor.eval_speculative_sampled(
token,
max_tokens,
reasoning,
temperature,
top_p,
min_p,
top_k,
rng,
cancelled,
),
Self::Glm(executor) => {
executor.eval(token)?;
Ok(vec![token])
}
}
}
pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
match self {
Self::DeepSeek(executor) => executor.eval(token),
@@ -4272,10 +4611,13 @@ fn compress_attention_batch(
freq_scale: f32,
ext: f32,
attn_factor: f32,
mut prefixes: Option<&mut [SpecPrefixFrontier]>,
layer: usize,
) -> Result<u32, String> {
let ratio = state.ratio;
let chunk = rows / ratio;
if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) {
if prefixes.is_none() && (pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)))
{
let before = if pos == 0 { 0 } else { state.rows };
let target = s
.compressed_stage
@@ -4439,6 +4781,15 @@ fn compress_attention_batch(
)?;
state.rows += 1;
}
if let Some(prefixes) = prefixes.as_deref_mut() {
let coefficient = if ratio == 4 { 2 } else { 1 };
prefixes[row as usize].layers[layer].compression =
Some(capture_compression_frontier(
state,
coefficient * coefficient * u64::from(ratio) * shape.head_dim * 4,
"capturing speculative compressor prefix",
)?);
}
}
}
Ok(state.rows)
@@ -4459,9 +4810,12 @@ fn compress_index_batch(
freq_scale: f32,
ext: f32,
attn_factor: f32,
mut prefixes: Option<&mut [SpecPrefixFrontier]>,
layer: usize,
) -> Result<(), String> {
let ratio = state.ratio;
if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) {
if prefixes.is_none() && (pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)))
{
let before = if pos == 0 { 0 } else { state.rows };
let chunk = rows / ratio;
let target = state.cache.view(
@@ -4613,6 +4967,13 @@ fn compress_index_batch(
)?;
state.rows += 1;
}
if let Some(prefixes) = prefixes.as_deref_mut() {
prefixes[row as usize].layers[layer].indexer = Some(capture_compression_frontier(
state,
4 * u64::from(ratio) * shape.indexer_head_dim * 4,
"capturing speculative indexer prefix",
)?);
}
}
}
Ok(())
@@ -4631,6 +4992,7 @@ fn encode_batch_layer(
rows: u32,
raw_cap: u32,
steering: Option<&Steering>,
mut prefixes: Option<&mut [SpecPrefixFrontier]>,
) -> Result<(), String> {
let hc_dim = shape.hc * shape.embd;
let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
@@ -4897,6 +5259,8 @@ fn encode_batch_layer(
freq_scale,
ext,
attn_factor,
prefixes.as_deref_mut(),
layer as usize,
)?;
}
@@ -4943,6 +5307,8 @@ fn encode_batch_layer(
freq_scale,
ext,
attn_factor,
prefixes,
layer as usize,
)?;
matmul_rows(
&s.indexer_q,
@@ -6558,48 +6924,32 @@ fn encode_output(
},
"output HC weights",
)?;
let fused_sum_norm = unsafe {
ds4_gpu_hc_weighted_sum_norm_tensor(
s.output_embedding.raw(),
s.output_norm.raw(),
s.current_hc.raw(),
s.output_weights.raw(),
map,
size,
w.output_norm.offset,
shape.embd as u32,
shape.hc as u32,
shape.rms_epsilon,
)
} != 0;
if !fused_sum_norm {
call(
unsafe {
ds4_gpu_hc_weighted_sum_tensor(
s.output_embedding.raw(),
s.current_hc.raw(),
s.output_weights.raw(),
shape.embd as u32,
shape.hc as u32,
)
},
"output HC collapse",
)?;
call(
unsafe {
ds4_gpu_rms_norm_weight_tensor(
s.output_norm.raw(),
s.output_embedding.raw(),
map,
size,
w.output_norm.offset,
shape.embd as u32,
shape.rms_epsilon,
)
},
"output norm",
)?;
}
call(
unsafe {
ds4_gpu_hc_weighted_sum_tensor(
s.output_embedding.raw(),
s.current_hc.raw(),
s.output_weights.raw(),
shape.embd as u32,
shape.hc as u32,
)
},
"output HC collapse",
)?;
call(
unsafe {
ds4_gpu_rms_norm_weight_tensor(
s.output_norm.raw(),
s.output_embedding.raw(),
map,
size,
w.output_norm.offset,
shape.embd as u32,
shape.rms_epsilon,
)
},
"output norm",
)?;
q8(
&s.logits,
w.output,
@@ -7186,6 +7536,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
EngineSsdSettings {
enabled: false,
@@ -7323,6 +7674,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
EngineSsdSettings {
enabled: false,
@@ -7414,6 +7766,94 @@ mod tests {
);
}
#[test]
#[ignore = "requires the installed 0731 target and checkpoint-specific DSpark GGUF fixtures"]
fn flash_0731_runs_exact_sampled_dspark() {
use super::{DeepSeekExecutor, argmax, configure_sources};
use crate::engine::gguf::Gguf;
use crate::engine::validation::validate_support;
use crate::engine::{Model, Rng};
use crate::model::{ModelChoice, validate_engine_artifacts};
use crate::settings::{
EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
};
use std::sync::atomic::AtomicBool;
configure_sources().unwrap();
let artifacts = installed_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true);
validate_engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true, &artifacts)
.unwrap();
let mut model =
Model::open_main(&artifacts.model, ModelChoice::DeepSeekV4Flash0731).unwrap();
let support = Gguf::open(artifacts.mtp.as_ref().unwrap()).unwrap();
model.support_kind = Some(validate_support(&support, &model.shape).unwrap());
model.support = Some(support);
let prompt = model.render_conversation(
"",
&[crate::engine::ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "hi".into(),
}],
ReasoningMode::Direct,
);
let mut executor = DeepSeekExecutor::open(
model,
64,
false,
64,
100,
EngineSpeculativeSettings {
mtp_draft_tokens: 1,
mtp_margin: 3.0,
glm_mtp: false,
glm_mtp_timing: false,
dspark: true,
dspark_confidence_threshold: 0.6,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: true,
},
EngineSsdSettings {
enabled: false,
cold: false,
cache_experts: 0,
cache_bytes: 0,
full_layers: 0,
full_layers_set: false,
preload_experts: 0,
},
EngineSteeringSettings {
file: None,
ffn_scale: 0.0,
attention_scale: 0.0,
},
)
.unwrap();
executor.prefill(&prompt, |_| true).unwrap();
let first = argmax(executor.logits());
let cycle = executor
.eval_speculative_sampled(
first,
4,
ReasoningMode::Direct,
0.8,
0.95,
0.0,
0,
&mut Rng::new(7),
&AtomicBool::new(false),
)
.unwrap();
assert!(!cycle.is_empty());
assert!(executor.logits().iter().all(|logit| logit.is_finite()));
assert!(executor.session.position >= prompt.len() as u32 + cycle.len() as u32);
}
#[test]
#[ignore = "requires the installed Flash, legacy MTP, and DSpark GGUF fixtures"]
fn ssd_streaming_supports_legacy_mtp_and_dspark() {
@@ -7468,6 +7908,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
EngineSsdSettings {
enabled: true,
@@ -7544,6 +7985,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
};
let mut executor = DeepSeekExecutor::open(
model,
@@ -7627,6 +8069,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
EngineSsdSettings {
enabled: true,
@@ -7691,6 +8134,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
EngineSsdSettings {
enabled: false,
@@ -7775,6 +8219,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
EngineSsdSettings {
enabled: false,
@@ -7850,6 +8295,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
EngineSsdSettings {
enabled: true,

2
src/engine/metal/glm.rs
View File

@@ -401,6 +401,7 @@ impl GlmExecutor {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
None,
)
@@ -3412,6 +3413,7 @@ mod tests {
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
dspark_exact_sampling: false,
},
None,
)

18
src/engine/metal/gpu.rs
View File

@@ -1283,18 +1283,6 @@ unsafe extern "C" {
hc: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_hc_weighted_sum_norm_tensor(
out: *mut GpuTensor,
norm: *mut GpuTensor,
residual: *const GpuTensor,
weights: *const GpuTensor,
map: *const c_void,
size: u64,
norm_weight: u64,
embd: u32,
hc: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_hc_weighted_sum_tensor(
out: *mut GpuTensor,
residual: *const GpuTensor,
@@ -1443,9 +1431,9 @@ impl Buffer {
}
pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
// SAFETY: native/metal/ds4_metal.m:7916-7940 bounds-checks the view,
// ARC-retains base_obj.buffer at :7926, and marks the view non-owning at
// :7929, so it may outlive and be freed independently of this wrapper.
// SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks
// the view, ARC-retains base_obj.buffer, and marks the view non-owning, so
// it may outlive and be freed independently of this wrapper.
// Recheck those guarantees whenever the vendored Metal file is re-synced.
NonNull::new(unsafe { ds4_gpu_tensor_view(self.raw(), offset, bytes) })
.map(Self)

14
src/engine/validation.rs
View File

@@ -7,7 +7,14 @@ pub(crate) fn validate_model_artifact(
) -> Result<(), String> {
if support {
let model = Gguf::open(path)?;
validate_support(&model, &FLASH).map(|_| ())
let shape = match expected {
ModelChoice::DeepSeekV4Flash => FLASH,
ModelChoice::DeepSeekV4Flash0731 => FLASH_0731,
ModelChoice::DeepSeekV4Pro | ModelChoice::Glm52 => {
return Err(format!("{expected} does not use an external support GGUF"));
}
};
validate_support(&model, &shape).map(|_| ())
} else {
let model = Model::open_main(path, expected)?;
let summary = model.summary();
@@ -759,7 +766,10 @@ fn validate_glm_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
}
pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String> {
if shape.model != ModelChoice::DeepSeekV4Flash {
if !matches!(
shape.model,
ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Flash0731
) {
return Err("DSpark support is available only for DeepSeek V4 Flash".into());
}
let DsparkConfig {

134
src/model.rs
View File

@@ -10,22 +10,22 @@ use std::fs;
use std::path::{Path, PathBuf};
pub(crate) const MODEL_CHOICES: [ModelChoice; 4] = [
ModelChoice::DeepSeekV4Flash,
ModelChoice::DeepSeekV4Flash0731,
ModelChoice::DeepSeekV4Flash,
ModelChoice::DeepSeekV4Pro,
ModelChoice::Glm52,
];
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 6] = [
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [
ManagedArtifactId::DeepSeekV4Flash0731,
ManagedArtifactId::DeepSeekV4Flash0731Dspark,
ManagedArtifactId::DeepSeekV4Flash,
ManagedArtifactId::DeepSeekV4FlashMtp,
ManagedArtifactId::DeepSeekV4FlashDspark,
ManagedArtifactId::DeepSeekV4Flash0731,
ManagedArtifactId::DeepSeekV4Pro,
ManagedArtifactId::Glm52,
];
const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf";
const DEEPSEEK_FLASH_0731_REPOSITORY: &str = "Rednalreden/DeepSeek-V4-Flash-0731-dwarfstar-q2-gguf";
const GLM_REPOSITORY: &str = "antirez/glm-5.2-gguf";
const FLASH: Artifact = Artifact {
@@ -54,18 +54,26 @@ const FLASH_MTP: Artifact = Artifact {
};
const FLASH_0731: Artifact = Artifact {
label: "DeepSeek V4 Flash 0731 model",
file_name: "DeepSeek-V4-Flash-0731-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-imatrix.gguf",
repository: DEEPSEEK_FLASH_0731_REPOSITORY,
size: 86_720_111_520,
sha256: "0b39f9c337d6b49c77db2190556b8563abf3c5fbb98be3b58cf8d3a1db191e5f",
file_name: "DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf",
repository: DEEPSEEK_REPOSITORY,
size: 86_720_111_488,
sha256: "ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0",
support: Some(false),
};
const FLASH_0731_DSPARK: Artifact = Artifact {
label: "DeepSeek V4 Flash 0731 DSpark support",
file_name: "DeepSeek-V4-Flash-DSpark-support-0731.gguf",
repository: DEEPSEEK_REPOSITORY,
size: 5_989_114_272,
sha256: "7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360",
support: Some(true),
};
const PRO: Artifact = Artifact {
label: "DeepSeek V4 Pro model",
file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf",
label: "DeepSeek V4 Pro 0813 model",
file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf",
repository: DEEPSEEK_REPOSITORY,
size: 464_627_334_560,
sha256: "a0314d9c0e16122cd60071079124a2d17185d317c55a8f95ecb3ed3506278a96",
sha256: "c4d997ab9894b6c78b759f7869fe1726b6314b6515f6ff82607df3797c5eb193",
support: Some(false),
};
const GLM: Artifact = Artifact {
@@ -79,9 +87,9 @@ const GLM: Artifact = Artifact {
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
pub(crate) enum ModelChoice {
#[default]
#[serde(rename = "deepseek-v4-flash")]
DeepSeekV4Flash,
#[default]
#[serde(rename = "deepseek-v4-flash-0731")]
DeepSeekV4Flash0731,
#[serde(rename = "deepseek-v4-pro")]
@@ -104,10 +112,14 @@ impl ModelChoice {
MODEL_CHOICES.into_iter().find(|model| model.id() == id)
}
pub(crate) fn supports_dspark(self) -> bool {
pub(crate) fn supports_legacy_mtp(self) -> bool {
self == Self::DeepSeekV4Flash
}
pub(crate) fn supports_dspark(self) -> bool {
matches!(self, Self::DeepSeekV4Flash | Self::DeepSeekV4Flash0731)
}
fn main_artifact(self) -> &'static Artifact {
match self {
Self::DeepSeekV4Flash => &FLASH,
@@ -117,6 +129,14 @@ impl ModelChoice {
}
}
fn dspark_artifact(self) -> Option<&'static Artifact> {
match self {
Self::DeepSeekV4Flash => Some(&FLASH_DSPARK),
Self::DeepSeekV4Flash0731 => Some(&FLASH_0731_DSPARK),
Self::DeepSeekV4Pro | Self::Glm52 => None,
}
}
#[cfg(test)]
fn artifacts(
self,
@@ -125,8 +145,8 @@ impl ModelChoice {
) -> impl Iterator<Item = &'static Artifact> {
[
Some(self.main_artifact()),
(self.supports_dspark() && legacy_mtp_enabled).then_some(&FLASH_MTP),
(self.supports_dspark() && dspark_enabled).then_some(&FLASH_DSPARK),
(self.supports_legacy_mtp() && legacy_mtp_enabled).then_some(&FLASH_MTP),
dspark_enabled.then(|| self.dspark_artifact()).flatten(),
]
.into_iter()
.flatten()
@@ -147,22 +167,61 @@ pub(crate) fn engine_artifacts(
) -> EngineArtifacts {
EngineArtifacts {
model: model.main_artifact().path(model, models_path),
mtp: if model.supports_dspark() && legacy_mtp_enabled {
mtp: if model.supports_legacy_mtp() && legacy_mtp_enabled {
Some(FLASH_MTP.path(model, models_path))
} else if model.supports_dspark() && dspark_enabled {
Some(FLASH_DSPARK.path(model, models_path))
} else if dspark_enabled {
model
.dspark_artifact()
.map(|artifact| artifact.path(model, models_path))
} else {
None
},
}
}
pub(crate) fn validate_engine_artifacts(
model: ModelChoice,
legacy_mtp_enabled: bool,
dspark_enabled: bool,
artifacts: &EngineArtifacts,
) -> Result<(), String> {
if legacy_mtp_enabled && dspark_enabled {
return Err("Legacy MTP and DSpark cannot be enabled together".into());
}
if legacy_mtp_enabled && !model.supports_legacy_mtp() {
return Err(format!("Legacy MTP is not compatible with {model}"));
}
if dspark_enabled && !model.supports_dspark() {
return Err(format!("DSpark is not compatible with {model}"));
}
model
.main_artifact()
.validate_installed_path(&artifacts.model)?;
let expected_support = if legacy_mtp_enabled {
model.supports_legacy_mtp().then_some(&FLASH_MTP)
} else if dspark_enabled {
model.dspark_artifact()
} else {
None
};
match (expected_support, artifacts.mtp.as_deref()) {
(Some(expected), Some(path)) => expected.validate_installed_path(path),
(None, None) => Ok(()),
(Some(_), None) => Err(format!("{model} is missing its required support GGUF")),
(None, Some(path)) => Err(format!(
"{} is not compatible with the selected {model} checkpoint",
path.display()
)),
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ManagedArtifactId {
DeepSeekV4Flash,
DeepSeekV4FlashMtp,
DeepSeekV4FlashDspark,
DeepSeekV4Flash0731,
DeepSeekV4Flash0731Dspark,
DeepSeekV4Pro,
Glm52,
}
@@ -173,7 +232,9 @@ impl ManagedArtifactId {
Self::DeepSeekV4Flash | Self::DeepSeekV4FlashMtp | Self::DeepSeekV4FlashDspark => {
ModelChoice::DeepSeekV4Flash
}
Self::DeepSeekV4Flash0731 => ModelChoice::DeepSeekV4Flash0731,
Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => {
ModelChoice::DeepSeekV4Flash0731
}
Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro,
Self::Glm52 => ModelChoice::Glm52,
}
@@ -185,6 +246,7 @@ impl ManagedArtifactId {
Self::DeepSeekV4FlashMtp => &FLASH_MTP,
Self::DeepSeekV4FlashDspark => &FLASH_DSPARK,
Self::DeepSeekV4Flash0731 => &FLASH_0731,
Self::DeepSeekV4Flash0731Dspark => &FLASH_0731_DSPARK,
Self::DeepSeekV4Pro => &PRO,
Self::Glm52 => &GLM,
}
@@ -278,9 +340,9 @@ pub(crate) enum DownloadOutcome {
impl fmt::Display for ModelChoice {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::DeepSeekV4Flash => "DeepSeek V4 Flash",
Self::DeepSeekV4Flash => "DeepSeek V4 Flash (deprecated preview)",
Self::DeepSeekV4Flash0731 => "DeepSeek V4 Flash 0731",
Self::DeepSeekV4Pro => "DeepSeek V4 Pro",
Self::DeepSeekV4Pro => "DeepSeek V4 Pro 0813",
Self::Glm52 => "GLM 5.2",
})
}
@@ -296,6 +358,36 @@ struct Artifact {
}
impl Artifact {
fn validate_installed_path(&self, path: &Path) -> Result<(), String> {
if path.file_name().and_then(|name| name.to_str()) != Some(self.file_name) {
return Err(format!(
"{} is not the expected {} artifact",
path.display(),
self.label
));
}
let size = path
.metadata()
.map_err(|error| format!("Could not inspect {}: {error}", path.display()))?
.len();
if size != self.size {
return Err(format!(
"{} has {size} bytes, expected {}",
path.display(),
self.size
));
}
let verification = fs::read_to_string(path.with_extension("gguf.sha256"))
.map_err(|_| format!("{} has not passed checksum verification", path.display()))?;
if verification.trim() != self.sha256 {
return Err(format!(
"{} has the wrong checkpoint identity",
path.display()
));
}
Ok(())
}
fn path(&self, model: ModelChoice, models_path: &Path) -> PathBuf {
models_path.join(model.id()).join(self.file_name)
}

42
src/model/transfer.rs
View File

@@ -337,7 +337,7 @@ mod tests {
assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448);
assert_eq!(
ModelChoice::DeepSeekV4Flash0731.main_artifact().size,
86_720_111_520
86_720_111_488
);
assert_eq!(
ModelChoice::DeepSeekV4Flash.artifacts(true, true).count(),
@@ -348,7 +348,7 @@ mod tests {
ModelChoice::DeepSeekV4Flash0731
.artifacts(true, true)
.count(),
1
2
);
let id = SystemTime::now()
@@ -365,8 +365,14 @@ mod tests {
engine.mtp.as_deref().and_then(Path::file_name),
Some(std::ffi::OsStr::new(FLASH_DSPARK.file_name))
);
let flash_0731 =
engine_artifacts(ModelChoice::DeepSeekV4Flash0731, false, true, &models_path);
assert_eq!(
flash_0731.mtp.as_deref().and_then(Path::file_name),
Some(std::ffi::OsStr::new(FLASH_0731_DSPARK.file_name))
);
assert!(
engine_artifacts(ModelChoice::DeepSeekV4Flash0731, true, true, &models_path)
engine_artifacts(ModelChoice::DeepSeekV4Flash0731, true, false, &models_path)
.mtp
.is_none()
);
@@ -389,6 +395,36 @@ mod tests {
.unwrap(),
empty.sha256
);
let installed = empty.path(ModelChoice::DeepSeekV4Flash, &models_path);
assert!(empty.validate_installed_path(&installed).is_ok());
let wrong_name = installed.with_file_name("wrong-checkpoint.gguf");
fs::write(&wrong_name, []).unwrap();
fs::write(wrong_name.with_extension("gguf.sha256"), empty.sha256).unwrap();
assert!(empty.validate_installed_path(&wrong_name).is_err());
assert!(
validate_engine_artifacts(
ModelChoice::DeepSeekV4Flash0731,
true,
false,
&EngineArtifacts {
model: installed.clone(),
mtp: None,
},
)
.is_err()
);
assert!(
validate_engine_artifacts(
ModelChoice::DeepSeekV4Pro,
false,
true,
&EngineArtifacts {
model: installed,
mtp: None,
},
)
.is_err()
);
fs::remove_dir_all(models_path).unwrap();
}

24
src/settings.rs
View File

@@ -30,6 +30,7 @@ pub(crate) struct SpeculativePreferences {
pub(crate) dspark_enabled: bool,
pub(crate) dspark_confidence_threshold: Option<f32>,
pub(crate) dspark_strict: bool,
pub(crate) dspark_exact_sampling: bool,
}
impl Default for SpeculativePreferences {
@@ -43,6 +44,7 @@ impl Default for SpeculativePreferences {
dspark_enabled: false,
dspark_confidence_threshold: None,
dspark_strict: false,
dspark_exact_sampling: false,
}
}
}
@@ -62,13 +64,15 @@ impl SpeculativePreferences {
if self.dspark_enabled && !model.supports_dspark() {
return Err("DSpark is not available for the selected model.".into());
}
if self.legacy_mtp_enabled && !model.supports_dspark() {
if self.legacy_mtp_enabled && !model.supports_legacy_mtp() {
return Err("Legacy MTP is not available for the selected model.".into());
}
if self.legacy_mtp_enabled && self.dspark_enabled {
return Err("Legacy MTP and DSpark use different support artifacts.".into());
}
if (self.dspark_confidence_threshold.is_some() || self.dspark_strict)
if (self.dspark_confidence_threshold.is_some()
|| self.dspark_strict
|| self.dspark_exact_sampling)
&& !self.dspark_enabled
{
return Err("DSpark tuning requires DSpark to be enabled.".into());
@@ -86,9 +90,10 @@ impl SpeculativePreferences {
glm_mtp: self.glm_mtp,
glm_mtp_timing: self.glm_mtp_timing,
dspark: self.dspark_enabled,
dspark_confidence_threshold: self.dspark_confidence_threshold.unwrap_or(0.9),
dspark_confidence_threshold: self.dspark_confidence_threshold.unwrap_or(0.6),
dspark_confidence_threshold_set: self.dspark_confidence_threshold.is_some(),
dspark_strict: self.dspark_strict,
dspark_exact_sampling: self.dspark_exact_sampling,
}
}
}
@@ -103,6 +108,7 @@ pub(crate) struct EngineSpeculativeSettings {
pub(crate) dspark_confidence_threshold: f32,
pub(crate) dspark_confidence_threshold_set: bool,
pub(crate) dspark_strict: bool,
pub(crate) dspark_exact_sampling: bool,
}
/// An expert count, or a whole GiB budget. Written as `4` or `64GB`, the same
@@ -785,7 +791,7 @@ mod tests {
let defaults = SpeculativePreferences::default();
let engine = defaults.engine_settings();
assert_eq!((engine.mtp_draft_tokens, engine.mtp_margin), (1, 3.0));
assert_eq!(engine.dspark_confidence_threshold, 0.9);
assert_eq!(engine.dspark_confidence_threshold, 0.6);
assert!(!engine.dspark_confidence_threshold_set);
let tuned = SpeculativePreferences {
@@ -793,6 +799,7 @@ mod tests {
dspark_enabled: true,
dspark_confidence_threshold: Some(0.7),
dspark_strict: true,
dspark_exact_sampling: true,
..defaults
};
assert!(tuned.validate(ModelChoice::DeepSeekV4Flash).is_ok());
@@ -812,6 +819,7 @@ mod tests {
..SpeculativePreferences::default()
};
assert!(legacy.validate(ModelChoice::DeepSeekV4Flash).is_ok());
assert!(legacy.validate(ModelChoice::DeepSeekV4Flash0731).is_err());
assert!(legacy.validate(ModelChoice::DeepSeekV4Pro).is_err());
assert!(
SpeculativePreferences {
@@ -821,6 +829,14 @@ mod tests {
.validate(ModelChoice::DeepSeekV4Flash)
.is_err()
);
assert!(
SpeculativePreferences {
dspark_exact_sampling: true,
..SpeculativePreferences::default()
}
.validate(ModelChoice::DeepSeekV4Flash0731)
.is_err()
);
}
#[test]