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

Implement batched server parity and observability

Browse Source
This commit is contained in:
Georg Bauer
2026-07-25 08:32:35 +02:00
parent d554b77b9d
commit 2fbe605b5d
15 changed files with 6478 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files

195
src/engine.rs
View File

@@ -3,12 +3,16 @@ mod gguf;
mod metal;
mod tokenizer;
#[cfg(target_os = "macos")]
use crate::metrics::{KvLookup, Metrics};
use crate::model::ModelChoice;
use crate::settings::TurnSettings;
use crate::settings::{EngineSettings, ReasoningMode};
use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value};
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
#[cfg(target_os = "macos")]
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(target_os = "macos")]
use std::time::Instant;
@@ -256,8 +260,14 @@ impl Model {
.encode_conversation(system, messages, reasoning)
}
fn render_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec<i32> {
self.tokenizer.encode_continuation(prompt, reasoning)
fn render_continuation(
&self,
prompt: &str,
reasoning: ReasoningMode,
skip_previous_eos: bool,
) -> Vec<i32> {
self.tokenizer
.encode_continuation(prompt, reasoning, skip_previous_eos)
}
pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
@@ -295,11 +305,13 @@ impl Model {
pub(crate) struct Generator {
executor: metal::Executor,
checkpoint: Option<PathBuf>,
metrics: Arc<Metrics>,
}
#[derive(Clone)]
pub(crate) struct ChatTurn {
pub(crate) user: bool,
pub(crate) skip_previous_eos: bool,
pub(crate) reasoning: Option<String>,
pub(crate) reasoning_complete: bool,
pub(crate) content: String,
@@ -311,11 +323,13 @@ pub(crate) struct GenerationOutput {
pub(crate) cached_tokens: u32,
pub(crate) completion_tokens: u32,
pub(crate) finish_reason: &'static str,
pub(crate) previous_checkpoint_bytes: Option<u64>,
pub(crate) checkpoint_bytes: u64,
}
#[cfg(target_os = "macos")]
impl Generator {
pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
if settings.model != ModelChoice::DeepSeekV4Flash {
return Err("local generation currently supports DeepSeek V4 Flash only".into());
}
@@ -330,13 +344,19 @@ impl Generator {
model,
settings.context_tokens.max(1) as u32,
settings.execution.quality,
settings.execution.prefill_chunk,
)?;
Ok(Self {
executor,
checkpoint: None,
metrics,
})
}
pub(crate) fn summary(&self) -> ModelSummary {
self.executor.model().summary()
}
pub(crate) fn generate(
&mut self,
checkpoint: &Path,
@@ -346,12 +366,20 @@ impl Generator {
mut emit: impl FnMut(bool, String),
mut progress: impl FnMut(u32, u32, Option<f32>),
) -> Result<GenerationOutput, String> {
self.select_checkpoint(checkpoint)?;
let (output, prompt_complete) =
let history = messages
.split_last()
.map_or(messages, |(_, history)| history);
self.select_checkpoint(
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 mut completed = messages.to_vec();
completed.push(output.message.clone());
self.executor.save_checkpoint(
output.previous_checkpoint_bytes =
std::fs::metadata(checkpoint).ok().map(|item| item.len());
self.save_checkpoint(
checkpoint,
if prompt_complete {
conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
@@ -359,6 +387,9 @@ impl Generator {
[0; 32]
},
)?;
output.checkpoint_bytes = std::fs::metadata(checkpoint)
.map(|item| item.len())
.unwrap_or(0);
Ok(output)
}
@@ -378,15 +409,17 @@ impl Generator {
conversation_tag(&settings.system_prompt, settings.reasoning_mode, history);
let checkpoint = directory.join(format!("{}.bin", hex_tag(history_tag)));
if self.executor.checkpoint_tag() != history_tag {
self.select_checkpoint(&checkpoint)?;
self.select_checkpoint(&checkpoint, history_tag)?;
if self.executor.checkpoint_tag() != history_tag {
self.executor.reset()?;
self.checkpoint = None;
let _ = std::fs::remove_file(&checkpoint);
}
} else {
self.metrics.kv_lookup(KvLookup::MemoryHit);
}
let (output, prompt_complete) =
let (mut output, prompt_complete) =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
let mut completed = messages.to_vec();
completed.push(output.message.clone());
@@ -399,7 +432,10 @@ impl Generator {
)
})?;
let completed_checkpoint = directory.join(format!("{}.bin", hex_tag(completed_tag)));
self.executor.save_checkpoint(
output.previous_checkpoint_bytes = std::fs::metadata(&completed_checkpoint)
.ok()
.map(|item| item.len());
self.save_checkpoint(
&completed_checkpoint,
if prompt_complete {
completed_tag
@@ -407,23 +443,62 @@ impl Generator {
[0; 32]
},
)?;
output.checkpoint_bytes = std::fs::metadata(&completed_checkpoint)
.map(|item| item.len())
.unwrap_or(0);
self.checkpoint = Some(completed_checkpoint);
Ok(output)
}
fn select_checkpoint(&mut self, checkpoint: &Path) -> Result<(), String> {
fn select_checkpoint(
&mut self,
checkpoint: &Path,
expected_tag: [u8; 32],
) -> Result<(), String> {
if self.checkpoint.as_deref() == Some(checkpoint) {
self.metrics
.kv_lookup(if self.executor.checkpoint_tag() == expected_tag {
KvLookup::MemoryHit
} else {
KvLookup::Miss
});
return Ok(());
}
self.executor.reset()?;
if self.executor.load_checkpoint(checkpoint).is_err() {
self.executor.reset()?;
let _ = std::fs::remove_file(checkpoint);
}
self.metrics.kv_read_started();
let started = Instant::now();
let loaded = self
.executor
.load_checkpoint(checkpoint, &mut |bytes| self.metrics.kv_read_bytes(bytes));
self.metrics
.kv_read_finished(started.elapsed(), loaded.is_err());
let lookup = match loaded {
Ok(true) if self.executor.checkpoint_tag() == expected_tag => KvLookup::DiskHit,
Ok(true) | Ok(false) => KvLookup::Miss,
Err(_) => {
self.executor.reset()?;
let _ = std::fs::remove_file(checkpoint);
KvLookup::Invalid
}
};
self.metrics.kv_lookup(lookup);
self.checkpoint = Some(checkpoint.to_owned());
Ok(())
}
fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> {
self.metrics.kv_write_started();
let started = Instant::now();
let result = self
.executor
.save_checkpoint(checkpoint, tag, &mut |bytes| {
self.metrics.kv_write_bytes(bytes);
});
self.metrics
.kv_write_finished(started.elapsed(), result.is_err());
result
}
fn generate_inner(
&mut self,
messages: &[ChatTurn],
@@ -443,11 +518,11 @@ impl Generator {
) =>
{
let mut tokens = self.executor.tokens().to_vec();
tokens.extend(
self.executor
.model()
.render_continuation(&latest.content, settings.reasoning_mode),
);
tokens.extend(self.executor.model().render_continuation(
&latest.content,
settings.reasoning_mode,
latest.skip_previous_eos,
));
tokens
}
_ => self.executor.model().render_conversation(
@@ -467,11 +542,13 @@ impl Generator {
));
}
let reused = self.executor.align_prompt(&tokens)?;
self.metrics.kv_prefix_reused(reused);
progress(self.executor.position(), self.executor.context(), None);
let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
let mut generated = ChatTurn {
user: false,
skip_previous_eos: false,
reasoning: reasoning.then(String::new),
reasoning_complete: !reasoning,
content: String::new(),
@@ -479,24 +556,40 @@ impl Generator {
let mut emitted_reasoning = 0;
let mut emitted_content = 0;
let prompt_tokens = tokens.len();
for (index, token) in tokens.into_iter().enumerate().skip(reused) {
if cancelled.load(Ordering::Relaxed) {
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: 0,
finish_reason: "stop",
},
false,
));
}
self.executor.eval(token)?;
if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens {
let suffix = &tokens[reused..];
let completed = if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
let context = self.executor.context();
self.executor.prefill(suffix, |used| {
progress(used, context, None);
!cancelled.load(Ordering::Relaxed)
})?
} else {
let mut completed = 0;
for &token in suffix {
if cancelled.load(Ordering::Relaxed) {
break;
}
self.executor.eval(token)?;
completed += 1;
progress(self.executor.position(), self.executor.context(), None);
}
completed
};
if completed != suffix.len() {
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: 0,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
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
@@ -519,6 +612,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
true,
));
@@ -550,6 +645,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
true,
));
@@ -575,6 +672,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
@@ -615,6 +714,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
@@ -645,6 +746,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "length",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
true,
))
@@ -891,6 +994,7 @@ mod sampling_tests {
fn checkpoint_tag_covers_the_canonical_chat_state() {
let messages = [ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
@@ -924,7 +1028,7 @@ mod sampling_tests {
ReasoningMode::Direct,
);
assert_eq!(tokens.len(), 10);
let mut executor = metal::Executor::open(model, 32_768, false).unwrap();
let mut executor = metal::Executor::open(model, 32_768, false, 0).unwrap();
assert_eq!(executor.context(), 32_768);
for &token in &tokens {
executor.eval(token).unwrap();
@@ -947,7 +1051,9 @@ mod sampling_tests {
let next = argmax.0 as i32;
let checkpoint =
std::env::temp_dir().join(format!("ds4-rust-kv-{}.bin", std::process::id()));
executor.save_checkpoint(&checkpoint, [7; 32]).unwrap();
executor
.save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
.unwrap();
executor.eval(next).unwrap();
let continued_logit = executor.logits()[0];
let continued_argmax = executor
@@ -960,8 +1066,8 @@ mod sampling_tests {
drop(executor);
let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap();
let mut restored = metal::Executor::open(model, 32_768, false).unwrap();
assert!(restored.load_checkpoint(&checkpoint).unwrap());
let mut restored = metal::Executor::open(model, 32_768, false, 0).unwrap();
assert!(restored.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
assert_eq!(restored.position(), tokens.len() as u32);
assert_eq!(restored.tokens(), tokens);
assert_eq!(restored.checkpoint_tag(), [7; 32]);
@@ -1872,6 +1978,7 @@ mod tests {
model.render_prompt("", "Hello", ReasoningMode::Direct),
[0, 128_803, 19_923, 128_804, 128_822]
);
assert_eq!(model.tokenizer.tokenize_rendered("｜DSML｜").len(), 1);
assert!(model.is_stop_token_for_reasoning(128_822, ReasoningMode::Direct));
assert!(!model.is_stop_token_for_reasoning(128_822, ReasoningMode::High));
let think_start = *model
@@ -1884,18 +1991,21 @@ mod tests {
&[
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
@@ -1908,27 +2018,34 @@ mod tests {
]
);
assert_eq!(
model.render_continuation("Hello", ReasoningMode::Direct),
model.render_continuation("Hello", ReasoningMode::Direct, false),
[1, 128_803, 19_923, 128_804, 128_822]
);
assert_eq!(
model.render_continuation("Hello", ReasoningMode::Direct, true),
[128_803, 19_923, 128_804, 128_822]
);
assert_eq!(
model.render_conversation(
"",
&[
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: false,
skip_previous_eos: false,
reasoning: Some("Hello".into()),
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),