hugo/DS4Server
1
0
Fork 0
Code Issues 7 Pull Requests Actions 1 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

108
src/engine.rs
View File

@@ -15,6 +15,8 @@ use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Ten
#[cfg(target_os = "macos")]
use kvstore::{KvStore, StoreReason};
use sha2::{Digest, Sha256};
#[cfg(target_os = "macos")]
use std::collections::HashMap;
use std::path::{Path, PathBuf};
#[cfg(target_os = "macos")]
use std::sync::Arc;
@@ -363,6 +365,15 @@ pub(crate) struct Generator {
/// spaced like ds4's `continued_last_store_tokens`.
last_store_tokens: u32,
metrics: Arc<Metrics>,
resident_sessions: HashMap<PathBuf, ResidentSlot>,
resident_active: Option<PathBuf>,
resident_limit: usize,
}
#[cfg(target_os = "macos")]
struct ResidentSlot {
state: metal::ResidentState,
last_store_tokens: u32,
}
#[cfg(target_os = "macos")]
@@ -459,11 +470,6 @@ pub(crate) struct CompactionOutput {
#[cfg(target_os = "macos")]
impl Generator {
pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
if settings.speculative.glm_mtp {
return Err(
"GLM MTP requires the shared speculative verifier, which is not enabled".into(),
);
}
let simulated_memory =
SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?;
let model = Model::open(settings)?;
@@ -476,6 +482,7 @@ impl Generator {
settings.speculative,
settings.ssd,
settings.steering.clone(),
settings.diagnostics.expert_profile_path.as_deref(),
)?;
let stats = executor.execution_stats();
metrics.speculative_stats(
@@ -491,7 +498,18 @@ impl Generator {
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_cache_entries,
stats.ssd_preloaded_experts,
stats.ssd_cache_hits,
stats.ssd_cache_misses,
stats.ssd_cache_evictions,
stats.ssd_cache_wraps,
stats.ssd_buffer_allocs,
stats.ssd_buffer_reuses,
stats.ssd_pread_bytes,
stats.ssd_pread_ms,
stats.ssd_evict_advise_bytes,
stats.ssd_willneed_advise_bytes,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
@@ -502,6 +520,13 @@ impl Generator {
checkpoint: None,
last_store_tokens: 0,
metrics,
resident_sessions: HashMap::new(),
resident_active: None,
resident_limit: std::env::var("DS4_RESIDENT_SESSIONS")
.ok()
.and_then(|value| value.parse().ok())
.filter(|limit| *limit > 0)
.unwrap_or(4),
})
}
@@ -561,7 +586,18 @@ impl Generator {
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_cache_entries,
stats.ssd_preloaded_experts,
stats.ssd_cache_hits,
stats.ssd_cache_misses,
stats.ssd_cache_evictions,
stats.ssd_cache_wraps,
stats.ssd_buffer_allocs,
stats.ssd_buffer_reuses,
stats.ssd_pread_bytes,
stats.ssd_pread_ms,
stats.ssd_evict_advise_bytes,
stats.ssd_willneed_advise_bytes,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
@@ -597,7 +633,11 @@ impl Generator {
previous_checkpoint = None;
}
} else {
self.executor.reset()?;
let key = resident_key(directory, history_tag);
let restored = self.activate_resident(key)?;
if !restored || self.executor.checkpoint_tag() != history_tag {
self.executor.reset()?;
}
self.checkpoint = None;
self.last_store_tokens = 0;
previous_checkpoint = None;
@@ -629,6 +669,7 @@ impl Generator {
// live KV, so mark it and drop the stale file association.
self.executor.note_checkpoint_tag(completed_tag);
self.checkpoint = None;
self.resident_active = Some(resident_key(directory, completed_tag));
return Ok(output);
}
let completed_checkpoint = store.checkpoint_path(&completed_key);
@@ -653,6 +694,11 @@ impl Generator {
self.last_store_tokens = self.executor.position();
}
self.checkpoint = retained.then_some(completed_checkpoint);
if let Some(checkpoint) = &self.checkpoint {
self.resident_active = Some(checkpoint.clone());
} else {
self.resident_active = Some(resident_key(directory, completed_tag));
}
}
Ok(output)
}
@@ -669,6 +715,7 @@ impl Generator {
mut progress: impl FnMut(u32, u32, Option<f32>),
mut phase: impl FnMut(&'static str),
) -> Result<CompactionOutput, String> {
self.activate_resident(checkpoint.to_owned())?;
let _ = std::fs::remove_file(checkpoint);
self.executor.reset()?;
self.checkpoint = None;
@@ -806,6 +853,12 @@ impl Generator {
checkpoint: &Path,
expected_tag: [u8; 32],
) -> Result<bool, String> {
let resident_hit = self.activate_resident(checkpoint.to_owned())?;
if resident_hit && self.executor.checkpoint_tag() == expected_tag {
self.checkpoint = Some(checkpoint.to_owned());
self.metrics.kv_lookup(KvLookup::MemoryHit);
return Ok(true);
}
if self.checkpoint.as_deref() == Some(checkpoint) {
if !checkpoint.is_file() {
self.executor.reset()?;
@@ -844,6 +897,39 @@ impl Generator {
Ok(found)
}
fn activate_resident(&mut self, key: PathBuf) -> Result<bool, String> {
if self.resident_active.as_ref() == Some(&key) {
return Ok(true);
}
let restored = self.resident_sessions.contains_key(&key);
let (mut incoming_state, incoming_last_store) = self
.resident_sessions
.remove(&key)
.map_or((None, 0), |slot| (Some(slot.state), slot.last_store_tokens));
self.executor.swap_resident_state(&mut incoming_state)?;
if let (Some(previous), Some(outgoing)) = (self.resident_active.take(), incoming_state) {
// ponytail: four resident sessions bound memory; raise
// DS4_RESIDENT_SESSIONS when the machine can hold more KV state.
if self.resident_sessions.len() >= self.resident_limit {
let evicted = self.resident_sessions.keys().next().cloned();
if let Some(evicted) = evicted {
self.resident_sessions.remove(&evicted);
}
}
self.resident_sessions.insert(
previous,
ResidentSlot {
state: outgoing,
last_store_tokens: self.last_store_tokens,
},
);
}
self.resident_active = Some(key);
self.checkpoint = None;
self.last_store_tokens = incoming_last_store;
Ok(restored)
}
fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> {
self.metrics.kv_write_started();
let started = Instant::now();
@@ -1311,6 +1397,16 @@ fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn
Sha256::digest(conversation_key(system, reasoning, messages)).into()
}
#[cfg(target_os = "macos")]
fn resident_key(directory: &Path, tag: [u8; 32]) -> PathBuf {
let mut name = String::with_capacity(64);
for byte in tag {
use std::fmt::Write;
let _ = write!(name, "{byte:02x}");
}
directory.join("resident").join(name)
}
#[cfg(target_os = "macos")]
fn sample(
logits: &[f32],