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

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 {