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

Integrate DS4 execution parity in Rust

Browse Source
This commit is contained in:
Georg Bauer
2026-07-26 17:58:05 +02:00
parent c9f0c3661c
commit 4420b81117
20 changed files with 11643 additions and 358 deletions
Download Patch File Download Diff File Expand all files Collapse all files

335
src/engine.rs
View File

@@ -22,7 +22,22 @@ use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(target_os = "macos")]
use std::time::Instant;
use tokenizer::Tokenizer;
use validation::{validate_dspark, validate_main};
use validation::{SupportKind, validate_main, validate_support};
#[cfg(target_os = "macos")]
unsafe extern "C" {
fn mmap(
address: *mut std::ffi::c_void,
length: usize,
protection: i32,
flags: i32,
fd: i32,
offset: i64,
) -> *mut std::ffi::c_void;
fn mlock(address: *const std::ffi::c_void, length: usize) -> i32;
fn munlock(address: *const std::ffi::c_void, length: usize) -> i32;
fn munmap(address: *mut std::ffi::c_void, length: usize) -> i32;
}
pub(crate) use validation::validate_model_artifact;
@@ -189,6 +204,7 @@ const GLM: Shape = Shape {
pub(crate) struct Model {
main: Gguf,
support: Option<Gguf>,
support_kind: Option<SupportKind>,
shape: Shape,
tokenizer: Tokenizer,
}
@@ -206,10 +222,17 @@ impl Model {
#[allow(dead_code)]
pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
let mut model = Self::open_main(&settings.artifacts.model, settings.model)?;
if settings.execution.warm_weights {
model.main.warm()?;
}
if let Some(path) = &settings.artifacts.mtp {
let support = Gguf::open(path)?;
validate_dspark(&support, &model.shape)?;
let kind = validate_support(&support, &model.shape)?;
if settings.execution.warm_weights {
support.warm()?;
}
model.support = Some(support);
model.support_kind = Some(kind);
}
Ok(model)
}
@@ -228,6 +251,7 @@ impl Model {
Ok(Self {
main,
support: None,
support_kind: None,
shape,
tokenizer,
})
@@ -333,6 +357,7 @@ impl Model {
#[cfg(target_os = "macos")]
pub(crate) struct Generator {
executor: metal::Executor,
_simulated_memory: Option<SimulatedMemory>,
checkpoint: Option<PathBuf>,
/// Token frontier of the last transient store, so continued checkpoints are
/// spaced like ds4's `continued_last_store_tokens`.
@@ -340,6 +365,68 @@ pub(crate) struct Generator {
metrics: Arc<Metrics>,
}
#[cfg(target_os = "macos")]
struct SimulatedMemory {
address: std::ptr::NonNull<std::ffi::c_void>,
bytes: usize,
}
#[cfg(target_os = "macos")]
impl SimulatedMemory {
fn acquire(bytes: u64) -> Result<Option<Self>, String> {
if bytes == 0 {
return Ok(None);
}
let bytes = usize::try_from(bytes).map_err(|_| "simulated memory size is too large")?;
let address = unsafe { mmap(std::ptr::null_mut(), bytes, 1 | 2, 2 | 0x1000, -1, 0) };
let Some(address) = std::ptr::NonNull::new(address) else {
return Err(format!(
"cannot reserve simulated used memory: {}",
std::io::Error::last_os_error()
));
};
if address.as_ptr() as isize == -1 {
return Err(format!(
"cannot reserve simulated used memory: {}",
std::io::Error::last_os_error()
));
}
let chunk = 256 * 1024 * 1024;
let mut locked = 0;
while locked < bytes {
let length = (bytes - locked).min(chunk);
let start = unsafe { address.as_ptr().cast::<u8>().add(locked) };
for page in (0..length).step_by(16 * 1024) {
unsafe { start.add(page).write((page / (16 * 1024)) as u8) };
}
unsafe { start.add(length - 1).write(1) };
if unsafe { mlock(start.cast(), length) } != 0 {
if locked != 0 {
unsafe { munlock(address.as_ptr(), locked) };
}
unsafe { munmap(address.as_ptr(), bytes) };
return Err(format!(
"cannot lock simulated used memory after {:.2} GiB: {}",
locked as f64 / 1_073_741_824.0,
std::io::Error::last_os_error()
));
}
locked += length;
}
Ok(Some(Self { address, bytes }))
}
}
#[cfg(target_os = "macos")]
impl Drop for SimulatedMemory {
fn drop(&mut self) {
unsafe {
munlock(self.address.as_ptr(), self.bytes);
munmap(self.address.as_ptr(), self.bytes);
}
}
}
#[derive(Clone)]
pub(crate) struct ChatTurn {
pub(crate) user: bool,
@@ -377,25 +464,41 @@ impl Generator {
"GLM MTP requires the shared speculative verifier, which is not enabled".into(),
);
}
if settings.speculative.dspark
|| (settings.ssd.enabled && settings.model != ModelChoice::Glm52)
|| settings.steering.file.is_some()
{
return Err(
"DSpark, SSD streaming, and steering are not yet available in the Rust executor"
.into(),
);
}
let simulated_memory =
SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?;
let model = Model::open(settings)?;
let executor = metal::Executor::open_configured(
model,
settings.context_tokens.max(1) as u32,
settings.execution.quality,
settings.execution.prefill_chunk,
settings.execution.power_percent,
settings.speculative,
settings.ssd,
settings.steering.clone(),
)?;
let stats = executor.execution_stats();
metrics.speculative_stats(
stats.speculative_mode,
stats.speculative_cycles,
stats.drafted_tokens,
stats.accepted_draft_tokens,
stats.verifier_passes,
stats.verifier_ms,
);
metrics.ssd_stats(
stats.ssd_enabled,
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_preloaded_experts,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
);
Ok(Self {
executor,
_simulated_memory: simulated_memory,
checkpoint: None,
last_store_tokens: 0,
metrics,
@@ -422,8 +525,9 @@ impl Generator {
checkpoint,
conversation_tag(&settings.system_prompt, settings.reasoning_mode, history),
)?;
let (mut output, prompt_complete) =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes =
@@ -442,6 +546,28 @@ impl Generator {
Ok(output)
}
fn publish_execution_stats(&self) {
let stats = self.executor.execution_stats();
self.metrics.speculative_stats(
stats.speculative_mode,
stats.speculative_cycles,
stats.drafted_tokens,
stats.accepted_draft_tokens,
stats.verifier_passes,
stats.verifier_ms,
);
self.metrics.ssd_stats(
stats.ssd_enabled,
stats.ssd_resident_bytes,
stats.ssd_cache_bytes,
stats.ssd_cache_experts,
stats.ssd_preloaded_experts,
stats.ssd_selected_requests,
stats.ssd_requested_bytes,
stats.ssd_wait_ms,
);
}
pub(crate) fn generate_transient(
&mut self,
directory: &Path,
@@ -481,8 +607,9 @@ impl Generator {
self.metrics.kv_lookup(KvLookup::MemoryHit);
}
let (mut output, prompt_complete) =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes = previous_checkpoint
@@ -809,6 +936,7 @@ impl Generator {
}
completed
};
self.publish_execution_stats();
if completed != suffix.len() {
return Ok((
GenerationOutput {
@@ -826,11 +954,12 @@ impl Generator {
progress(self.executor.position(), self.executor.context(), Some(0.0));
let generation_started = Instant::now();
let mut generated_tokens = 0_u32;
for _ in 0..settings
let generation_limit = settings
.max_generated_tokens
.max(0)
.min((max_context - self.executor.position() as usize) as i32)
{
as u32;
while generated_tokens < generation_limit {
if cancelled.load(Ordering::Relaxed) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
flush_generated(
@@ -887,86 +1016,132 @@ impl Generator {
true,
));
}
if self.executor.model().is_think_start_token(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
reasoning = true;
generated.reasoning.get_or_insert_default();
} else if self.executor.model().is_think_end_token(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
if let Some(reasoning_text) = &mut generated.reasoning
&& emit_safe_text(
reasoning_text,
&mut emitted_reasoning,
&settings.stops,
true,
true,
emit,
)
let cycle = if settings.temperature <= 0.0 {
self.executor.eval_speculative_greedy(
token,
generation_limit - generated_tokens,
settings.reasoning_mode,
cancelled,
)?
} else {
self.executor.eval(token)?;
vec![token]
};
self.publish_execution_stats();
for token in cycle {
if generated_tokens >= generation_limit
|| self
.executor
.model()
.is_stop_token_for_reasoning(token, settings.reasoning_mode)
{
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
reasoning = false;
generated.reasoning_complete = true;
emit(false, String::new());
} else if let Some(bytes) = self.executor.model().token_bytes(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &bytes, false);
let stopped = if reasoning {
let text = generated.reasoning.get_or_insert_default();
emit_safe_text(
text,
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
flush_generated(
&mut generated,
&mut emitted_reasoning,
&settings.stops,
false,
true,
emit,
)
} else {
generated.reasoning_complete = true;
emit_safe_text(
&mut generated.content,
&mut emitted_content,
&settings.stops,
false,
false,
emit,
)
};
if stopped {
);
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
completion_tokens: generated_tokens,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
true,
));
}
if self.executor.model().is_think_start_token(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
reasoning = true;
generated.reasoning.get_or_insert_default();
} else if self.executor.model().is_think_end_token(token) {
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
if let Some(reasoning_text) = &mut generated.reasoning
&& emit_safe_text(
reasoning_text,
&mut emitted_reasoning,
&settings.stops,
true,
true,
emit,
)
{
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
reasoning = false;
generated.reasoning_complete = true;
emit(false, String::new());
} else if let Some(bytes) = self.executor.model().token_bytes(token) {
append_generated_bytes(
&mut generated,
reasoning,
&mut pending_utf8,
&bytes,
false,
);
let stopped = if reasoning {
let text = generated.reasoning.get_or_insert_default();
emit_safe_text(
text,
&mut emitted_reasoning,
&settings.stops,
false,
true,
emit,
)
} else {
generated.reasoning_complete = true;
emit_safe_text(
&mut generated.content,
&mut emitted_content,
&settings.stops,
false,
false,
emit,
)
};
if stopped {
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
}
generated_tokens += 1;
progress(
self.executor.position(),
self.executor.context(),
Some(
generated_tokens as f32
/ generation_started.elapsed().as_secs_f32().max(1.0e-6),
),
);
}
self.executor.eval(token)?;
generated_tokens += 1;
progress(
self.executor.position(),
self.executor.context(),
Some(
generated_tokens as f32
/ generation_started.elapsed().as_secs_f32().max(1.0e-6),
),
);
}
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
flush_generated(