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

Finish DS4 execution parity

Browse Source
This commit is contained in:
Georg Bauer
2026-07-26 20:44:23 +02:00
parent fd3f8e45dc
commit 0d80c217c4
15 changed files with 16905 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

446
src/engine/metal.rs
View File

@@ -2,11 +2,13 @@ mod checkpoint;
mod glm;
mod gpu;
mod hotlist;
mod profile;
use glm::GlmExecutor;
use gpu::*;
use profile::ExpertProfile;
use super::gguf::{F16, F32, Gguf, Q4_K, Q8_0, Tensor as GgufTensor};
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 crate::model::ModelChoice;
@@ -1991,7 +1993,12 @@ impl SsdPlan {
}
let mut by_layer = vec![Vec::<(i32, u32)>::new(); model.shape.layers as usize];
let mut loaded = 0_u32;
for &(layer, expert) in hotlist::FLASH {
let hotlist = match model.shape.model {
ModelChoice::DeepSeekV4Flash => hotlist::FLASH,
ModelChoice::DeepSeekV4Pro => hotlist::PRO,
ModelChoice::Glm52 => unreachable!("GLM uses its dedicated executor"),
};
for &(layer, expert) in hotlist {
if loaded == self.preload_experts {
break;
}
@@ -2163,6 +2170,31 @@ fn estimated_deepseek_runtime_bytes(shape: super::Shape, context: u32, prefill:
raw.saturating_add(compressed).saturating_add(scratch)
}
fn resident_deepseek_admission_bytes(
model: &Model,
context: u32,
prefill: u32,
) -> Result<u64, String> {
resident_deepseek_admission_for_weights(
model.main.len() - model.main.data_offset(),
model.shape,
context,
prefill,
)
}
fn resident_deepseek_admission_for_weights(
weight_bytes: u64,
shape: super::Shape,
context: u32,
prefill: u32,
) -> Result<u64, String> {
weight_bytes
.checked_add(estimated_deepseek_runtime_bytes(shape, context, prefill))
.and_then(|bytes| bytes.checked_add(512 * 1024 * 1024))
.ok_or_else(|| "DeepSeek runtime memory size overflow".into())
}
fn effective_prefill_cap(context: u32, requested: u32) -> u32 {
if requested == 0 {
context.clamp(1, DEFAULT_PREFILL_CHUNK)
@@ -2512,7 +2544,18 @@ pub(super) struct ExecutionStats {
pub(super) ssd_resident_bytes: u64,
pub(super) ssd_cache_bytes: u64,
pub(super) ssd_cache_experts: u64,
pub(super) ssd_cache_entries: u64,
pub(super) ssd_preloaded_experts: u64,
pub(super) ssd_cache_hits: u64,
pub(super) ssd_cache_misses: u64,
pub(super) ssd_cache_evictions: u64,
pub(super) ssd_cache_wraps: u64,
pub(super) ssd_buffer_allocs: u64,
pub(super) ssd_buffer_reuses: u64,
pub(super) ssd_pread_bytes: u64,
pub(super) ssd_pread_ms: u64,
pub(super) ssd_evict_advise_bytes: u64,
pub(super) ssd_willneed_advise_bytes: u64,
pub(super) ssd_selected_requests: u64,
pub(super) ssd_requested_bytes: u64,
pub(super) ssd_wait_ms: u64,
@@ -2525,6 +2568,7 @@ pub(super) struct DeepSeekExecutor {
dspark: Option<Dspark>,
steering: Option<Steering>,
ssd: Option<SsdPlan>,
profile: Option<ExpertProfile>,
logits: Vec<f32>,
tokens: Vec<i32>,
quality: bool,
@@ -2539,9 +2583,20 @@ pub(super) struct DeepSeekExecutor {
model_identity: [u8; 32],
_context: Context,
model: Model,
speculative: EngineSpeculativeSettings,
}
pub(super) struct DeepSeekResidentState {
session: Session,
legacy_mtp: Option<LegacyMtp>,
dspark: Option<Dspark>,
logits: Vec<f32>,
tokens: Vec<i32>,
checkpoint_tag: [u8; 32],
}
impl DeepSeekExecutor {
#[allow(dead_code)]
#[allow(clippy::too_many_arguments)]
pub(super) fn open(
model: Model,
@@ -2552,6 +2607,31 @@ impl DeepSeekExecutor {
speculative: EngineSpeculativeSettings,
ssd: EngineSsdSettings,
steering: EngineSteeringSettings,
) -> Result<Self, String> {
Self::open_profile(
model,
context,
quality,
prefill_chunk,
power_percent,
speculative,
ssd,
steering,
None,
)
}
#[allow(clippy::too_many_arguments)]
pub(super) fn open_profile(
model: Model,
context: u32,
quality: bool,
prefill_chunk: u32,
power_percent: u8,
speculative: EngineSpeculativeSettings,
ssd: EngineSsdSettings,
steering: EngineSteeringSettings,
expert_profile_path: Option<&str>,
) -> Result<Self, String> {
let weights = Weights::bind(&model)?;
let mut ssd_plan = ssd
@@ -2559,7 +2639,11 @@ impl DeepSeekExecutor {
.then(|| SsdPlan::new(&model, &weights, ssd, context, prefill_chunk))
.transpose()?;
let spans = ssd_plan.as_ref().map(|plan| plan.model_spans.as_slice());
let admission = ssd_plan.as_ref().map_or(0, |plan| plan.admission_bytes);
let admission = if let Some(plan) = &ssd_plan {
plan.admission_bytes
} else {
resident_deepseek_admission_bytes(&model, context, prefill_chunk)?
};
let context_handle = Context::open(&model, quality, ssd.enabled, admission, spans)?;
let steering = Steering::load(&model, steering)?;
let session = Session::new(
@@ -2620,6 +2704,13 @@ impl DeepSeekExecutor {
} else {
model.checkpoint_identity()
};
let profile = ExpertProfile::new(
expert_profile_path,
model.shape.model,
model.shape.layers,
model.shape.experts,
model.shape.experts_used,
)?;
Ok(Self {
weights,
session,
@@ -2627,6 +2718,7 @@ impl DeepSeekExecutor {
dspark,
steering,
ssd: ssd_plan,
profile,
logits: vec![0.0; model.shape.vocab as usize],
tokens: Vec::new(),
quality,
@@ -2645,6 +2737,7 @@ impl DeepSeekExecutor {
model_identity,
_context: context_handle,
model,
speculative,
})
}
@@ -2681,6 +2774,9 @@ impl DeepSeekExecutor {
self.session.scratch.logits.read_f32(&mut self.logits)?;
self.session.position += 1;
self.tokens.push(token);
if let Some(profile) = &self.profile {
profile.write()?;
}
throttle(
&mut self.decode_average,
started.elapsed(),
@@ -3335,6 +3431,16 @@ impl DeepSeekExecutor {
self.steering.as_ref(),
self.ssd.as_ref(),
)?;
if let Some(profile) = &mut self.profile {
profile.record(
index,
pos,
&batch.router_selected,
&batch.router_weights,
rows,
index < shape.hash_layers as usize,
)?;
}
if let Some(dspark) = &mut self.dspark {
dspark.capture_batch(
index as u32,
@@ -3388,6 +3494,9 @@ impl DeepSeekExecutor {
}
self.session.position += rows;
self.tokens.extend_from_slice(tokens);
if let Some(profile) = &self.profile {
profile.write()?;
}
Ok(tops)
}
@@ -3414,16 +3523,29 @@ impl DeepSeekExecutor {
..ExecutionStats::default()
};
if let Some(ssd) = &self.ssd {
let mut native = StreamExpertCacheStats::default();
unsafe { ds4_gpu_stream_expert_cache_get_stats(&mut native) };
let selected = ssd
.selected_experts
.load(std::sync::atomic::Ordering::Relaxed);
stats.ssd_enabled = true;
stats.ssd_resident_bytes = ssd.resident_bytes;
stats.ssd_cache_experts = u64::from(ssd.cache_experts);
stats.ssd_cache_entries = u64::from(native.current_count);
stats.ssd_cache_bytes = ssd
.per_expert_bytes
.saturating_mul(u64::from(ssd.cache_experts));
stats.ssd_preloaded_experts = u64::from(ssd.preload_experts);
stats.ssd_cache_hits = native.hits;
stats.ssd_cache_misses = native.misses;
stats.ssd_cache_evictions = native.evictions;
stats.ssd_cache_wraps = native.wraps;
stats.ssd_buffer_allocs = native.buffer_allocs;
stats.ssd_buffer_reuses = native.buffer_reuses;
stats.ssd_pread_bytes = native.pread_bytes;
stats.ssd_pread_ms = native.pread_ms.max(0.0) as u64;
stats.ssd_evict_advise_bytes = native.evict_advise_bytes;
stats.ssd_willneed_advise_bytes = native.willneed_advise_bytes;
stats.ssd_selected_requests = ssd
.selected_requests
.load(std::sync::atomic::Ordering::Relaxed);
@@ -3466,6 +3588,67 @@ impl DeepSeekExecutor {
Ok(())
}
fn blank_resident_state(&self) -> Result<DeepSeekResidentState, String> {
let session = Session::new(&self.model, self.session.context, self.session.prefill_cap)?;
let legacy_mtp = match (self.model.support_kind, self.model.support.as_ref()) {
(Some(SupportKind::LegacyMtp), Some(support)) => {
let hc_dim = self.model.shape.hc * self.model.shape.embd;
Some(LegacyMtp {
weights: LegacyMtpWeights::bind(support, self.model.shape)?,
layer: LayerState::allocate(&self.model, 1, session.context, session.raw_cap)?,
state_hc: Buffer::floats(hc_dim)?,
next_hc: Buffer::floats(hc_dim)?,
draft_token: None,
raw_rows: 0,
draft_limit: self.speculative.mtp_draft_tokens.max(1) as u32,
margin: self.speculative.mtp_margin,
drafted: 0,
accepted: 0,
})
}
_ => None,
};
let dspark = match (
self.model.support_kind,
self.model.support.as_ref(),
self.speculative.dspark,
) {
(Some(SupportKind::DSpark), Some(support), true) => Some(Dspark::new(
&self.model,
support,
&session,
self.speculative,
self.quality,
)?),
_ => None,
};
Ok(DeepSeekResidentState {
session,
legacy_mtp,
dspark,
logits: vec![0.0; self.model.shape.vocab as usize],
tokens: Vec::new(),
checkpoint_tag: [0; 32],
})
}
pub(super) fn swap_resident_state(
&mut self,
state: &mut Option<DeepSeekResidentState>,
) -> Result<(), String> {
let mut incoming = state
.take()
.map_or_else(|| self.blank_resident_state(), Ok)?;
std::mem::swap(&mut self.session, &mut incoming.session);
std::mem::swap(&mut self.legacy_mtp, &mut incoming.legacy_mtp);
std::mem::swap(&mut self.dspark, &mut incoming.dspark);
std::mem::swap(&mut self.logits, &mut incoming.logits);
std::mem::swap(&mut self.tokens, &mut incoming.tokens);
std::mem::swap(&mut self.checkpoint_tag, &mut incoming.checkpoint_tag);
*state = Some(incoming);
Ok(())
}
pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
if !tokens.starts_with(&self.tokens) {
self.reset()?;
@@ -3523,6 +3706,16 @@ impl DeepSeekExecutor {
self.steering.as_ref(),
self.ssd.as_ref(),
)?;
if let Some(profile) = &mut self.profile {
profile.record(
index,
self.session.position,
&scratch.router_selected,
&scratch.router_weights,
1,
index < shape.hash_layers as usize,
)?;
}
if let Some(dspark) = &mut self.dspark {
dspark.capture_decode(index as u32, &scratch.current_hc, shape)?;
}
@@ -3537,6 +3730,11 @@ pub(super) enum Executor {
Glm(Box<GlmExecutor>),
}
pub(super) enum ResidentState {
DeepSeek(Box<DeepSeekResidentState>),
Glm(Box<glm::GlmResidentState>),
}
impl Executor {
#[allow(dead_code)]
pub(super) fn open(
@@ -3575,6 +3773,7 @@ impl Executor {
ffn_scale: 0.0,
attention_scale: 0.0,
},
None,
)
}
@@ -3588,9 +3787,10 @@ impl Executor {
speculative: EngineSpeculativeSettings,
ssd: EngineSsdSettings,
steering: EngineSteeringSettings,
expert_profile_path: Option<&str>,
) -> Result<Self, String> {
match model.shape.family {
ModelFamily::DeepSeek => DeepSeekExecutor::open(
ModelFamily::DeepSeek => DeepSeekExecutor::open_profile(
model,
context,
quality,
@@ -3599,12 +3799,20 @@ impl Executor {
speculative,
ssd,
steering,
expert_profile_path,
)
.map(Box::new)
.map(Self::DeepSeek),
ModelFamily::Glm => GlmExecutor::open(model, context, quality, ssd)
.map(Box::new)
.map(Self::Glm),
ModelFamily::Glm => GlmExecutor::open_profile(
model,
context,
quality,
ssd,
speculative,
expert_profile_path,
)
.map(Box::new)
.map(Self::Glm),
}
}
@@ -3627,8 +3835,8 @@ impl Executor {
executor.eval_speculative_greedy(token, max_tokens, reasoning, cancelled)
}
Self::Glm(executor) => {
executor.eval(token)?;
Ok(vec![token])
let _ = reasoning;
executor.eval_speculative_greedy(token, max_tokens, cancelled)
}
}
}
@@ -3654,7 +3862,7 @@ impl Executor {
pub(super) fn execution_stats(&self) -> ExecutionStats {
match self {
Self::DeepSeek(executor) => executor.execution_stats(),
Self::Glm(_) => ExecutionStats::default(),
Self::Glm(executor) => executor.execution_stats(),
}
}
@@ -3686,6 +3894,37 @@ impl Executor {
}
}
pub(super) fn swap_resident_state(
&mut self,
state: &mut Option<ResidentState>,
) -> Result<(), String> {
match self {
Self::DeepSeek(executor) => {
let mut inner = match state.take() {
Some(ResidentState::DeepSeek(state)) => Some(*state),
Some(ResidentState::Glm(_)) => {
return Err("resident session belongs to a different model family".into());
}
None => None,
};
executor.swap_resident_state(&mut inner)?;
*state = inner.map(|state| ResidentState::DeepSeek(Box::new(state)));
}
Self::Glm(executor) => {
let mut inner = match state.take() {
Some(ResidentState::Glm(state)) => Some(*state),
Some(ResidentState::DeepSeek(_)) => {
return Err("resident session belongs to a different model family".into());
}
None => None,
};
executor.swap_resident_state(&mut inner)?;
*state = inner.map(|state| ResidentState::Glm(Box::new(state)));
}
}
Ok(())
}
pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
match self {
Self::DeepSeek(executor) => executor.align_prompt(tokens),
@@ -6628,6 +6867,30 @@ mod tests {
);
}
#[test]
fn pro_resident_admission_includes_weights_context_and_scratch() {
let weights = 464_627_334_560_u64;
let runtime = estimated_deepseek_runtime_bytes(PRO, 32_768, 4_096);
assert_eq!(
super::resident_deepseek_admission_for_weights(weights, PRO, 32_768, 4_096).unwrap(),
weights + runtime + 512 * 1024 * 1024
);
}
#[test]
fn generated_hotlists_match_each_model_shape() {
for (hotlist, shape) in [
(super::hotlist::FLASH, FLASH),
(super::hotlist::PRO, PRO),
(super::hotlist::GLM52, crate::engine::GLM),
] {
assert!(hotlist.len() >= 4_096);
assert!(hotlist.iter().all(|&(layer, expert)| {
u32::from(layer) < shape.layers && u64::from(expert) < shape.experts
}));
}
}
#[test]
#[ignore = "requires the installed 81 GiB Flash and legacy MTP GGUF fixtures"]
fn legacy_mtp_runs_a_target_owned_greedy_cycle() {
@@ -7086,6 +7349,91 @@ mod tests {
assert_eq!(resident, run(path, true));
}
#[test]
#[ignore = "requires the installed 81 GiB Flash GGUF fixture and a Metal device"]
fn resident_multi_session_switching_preserves_each_kv_frontier() {
use super::{DeepSeekExecutor, argmax, configure_sources};
use crate::engine::Model;
use crate::model::ModelChoice;
use crate::settings::{
EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings, ReasoningMode,
};
use std::path::Path;
configure_sources().unwrap();
let model = Model::open_main(
Path::new(
"../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf",
),
ModelChoice::DeepSeekV4Flash,
)
.unwrap();
let prompts = ["Reply with A.", "Reply with B."].map(|content| {
model.render_conversation(
"",
&[crate::engine::ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: content.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: false,
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
},
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(&prompts[0], |_| true).unwrap();
let first = (executor.tokens().to_vec(), argmax(executor.logits()));
let mut inactive = None;
executor.swap_resident_state(&mut inactive).unwrap();
executor.prefill(&prompts[1], |_| true).unwrap();
let second = (executor.tokens().to_vec(), argmax(executor.logits()));
executor.swap_resident_state(&mut inactive).unwrap();
assert_eq!(
(executor.tokens(), argmax(executor.logits())),
(&*first.0, first.1)
);
executor.swap_resident_state(&mut inactive).unwrap();
assert_eq!(
(executor.tokens(), argmax(executor.logits())),
(&*second.0, second.1)
);
}
#[test]
#[ignore = "requires the installed Flash GGUF, DS4 steering fixture, and a Metal device"]
fn directional_steering_matches_the_ds4_token_oracle() {
@@ -7163,4 +7511,82 @@ mod tests {
assert_eq!(tokens, expected);
}
}
#[test]
#[ignore = "requires the installed DeepSeek V4 Pro GGUF and Apple Metal"]
fn pro_resident_and_ssd_streaming_choose_the_same_tokens() {
use super::{DeepSeekExecutor, argmax, configure_sources};
use crate::engine::Model;
use crate::model::ModelChoice;
use crate::settings::{
EngineSpeculativeSettings, EngineSsdSettings, EngineSteeringSettings,
};
use std::path::Path;
configure_sources().unwrap();
let path = std::env::var("DS4_PRO_MODEL").unwrap_or_else(|_| {
"../ds4/models/DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf".into()
});
if !Path::new(&path).is_file() {
eprintln!("skipping unavailable Pro fixture: {path}");
return;
}
let run = |streaming| {
let model = Model::open_main(Path::new(&path), ModelChoice::DeepSeekV4Pro).unwrap();
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(),
}],
crate::settings::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: false,
dspark_confidence_threshold: 0.9,
dspark_confidence_threshold_set: false,
dspark_strict: false,
},
EngineSsdSettings {
enabled: streaming,
cold: true,
cache_experts: if streaming { 32 } else { 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();
(0..4)
.map(|_| {
let token = argmax(executor.logits());
executor.eval(token).unwrap();
token
})
.collect::<Vec<_>>()
};
assert_eq!(run(false), run(true));
}
}