1918 lines
64 KiB
Rust
1918 lines
64 KiB
Rust
mod gguf;
|
||
#[cfg(any(target_os = "macos", test))]
|
||
mod kvstore;
|
||
#[cfg(target_os = "macos")]
|
||
mod metal;
|
||
mod tokenizer;
|
||
mod validation;
|
||
|
||
#[cfg(target_os = "macos")]
|
||
use crate::metrics::{KvLookup, Metrics};
|
||
use crate::model::ModelChoice;
|
||
#[cfg(target_os = "macos")]
|
||
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};
|
||
#[cfg(target_os = "macos")]
|
||
use kvstore::{KvStore, StoreReason};
|
||
use sha2::{Digest, Sha256};
|
||
#[cfg(target_os = "macos")]
|
||
use std::collections::HashMap;
|
||
use std::path::Path;
|
||
#[cfg(target_os = "macos")]
|
||
use std::path::PathBuf;
|
||
#[cfg(target_os = "macos")]
|
||
use std::sync::Arc;
|
||
#[cfg(target_os = "macos")]
|
||
use std::sync::atomic::{AtomicBool, Ordering};
|
||
#[cfg(target_os = "macos")]
|
||
use std::time::Instant;
|
||
use tokenizer::Tokenizer;
|
||
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;
|
||
|
||
#[cfg(target_os = "macos")]
|
||
pub(crate) use kvstore::sweep_unreachable as sweep_transient_cache;
|
||
#[cfg(target_os = "macos")]
|
||
pub(crate) use metal::configure_sources as configure_metal_sources;
|
||
|
||
const DENSE: &[u32] = &[Q8_0, Q4_K, Q4_0];
|
||
const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K];
|
||
const PLAIN: &[u32] = &[F16, F32];
|
||
const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0];
|
||
|
||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||
enum ModelFamily {
|
||
DeepSeek,
|
||
Glm,
|
||
}
|
||
|
||
#[derive(Clone, Copy)]
|
||
struct Shape {
|
||
model: ModelChoice,
|
||
family: ModelFamily,
|
||
layers: u32,
|
||
embd: u64,
|
||
vocab: u64,
|
||
heads: u64,
|
||
head_kv: u64,
|
||
head_dim: u64,
|
||
value_dim: u64,
|
||
rot: u64,
|
||
out_groups: u64,
|
||
lora_q: u64,
|
||
lora_o: u64,
|
||
experts: u64,
|
||
experts_used: u64,
|
||
expert_shared: u64,
|
||
ff_expert: u64,
|
||
ff_dense: u64,
|
||
hash_layers: u32,
|
||
sliding_window: u64,
|
||
indexer_heads: u64,
|
||
indexer_head_dim: u64,
|
||
indexer_top_k: u64,
|
||
hc: u64,
|
||
hc_sinkhorn: u64,
|
||
nextn: u32,
|
||
leading_dense: u32,
|
||
kv_lora: u64,
|
||
key_mla: u64,
|
||
value_mla: u64,
|
||
rms_epsilon: f32,
|
||
hc_epsilon: f32,
|
||
expert_weight_scale: f32,
|
||
swiglu_clamp: f32,
|
||
rope_base: f32,
|
||
rope_scale: f32,
|
||
rope_beta_fast: f32,
|
||
rope_beta_slow: f32,
|
||
compress_rope_base: f32,
|
||
original_context: u64,
|
||
}
|
||
|
||
const FLASH: Shape = Shape {
|
||
model: ModelChoice::DeepSeekV4Flash,
|
||
family: ModelFamily::DeepSeek,
|
||
layers: 43,
|
||
embd: 4096,
|
||
vocab: 129_280,
|
||
heads: 64,
|
||
head_kv: 1,
|
||
head_dim: 512,
|
||
value_dim: 512,
|
||
rot: 64,
|
||
out_groups: 8,
|
||
lora_q: 1024,
|
||
lora_o: 1024,
|
||
experts: 256,
|
||
experts_used: 6,
|
||
expert_shared: 1,
|
||
ff_expert: 2048,
|
||
ff_dense: 0,
|
||
hash_layers: 3,
|
||
sliding_window: 128,
|
||
indexer_heads: 64,
|
||
indexer_head_dim: 128,
|
||
indexer_top_k: 512,
|
||
hc: 4,
|
||
hc_sinkhorn: 20,
|
||
nextn: 0,
|
||
leading_dense: 0,
|
||
kv_lora: 0,
|
||
key_mla: 0,
|
||
value_mla: 0,
|
||
rms_epsilon: 1.0e-6,
|
||
hc_epsilon: 1.0e-6,
|
||
expert_weight_scale: 1.5,
|
||
swiglu_clamp: 10.0,
|
||
rope_base: 10_000.0,
|
||
rope_scale: 16.0,
|
||
rope_beta_fast: 32.0,
|
||
rope_beta_slow: 1.0,
|
||
compress_rope_base: 160_000.0,
|
||
original_context: 65_536,
|
||
};
|
||
|
||
const PRO: Shape = Shape {
|
||
model: ModelChoice::DeepSeekV4Pro,
|
||
layers: 61,
|
||
embd: 7168,
|
||
heads: 128,
|
||
out_groups: 16,
|
||
lora_q: 1536,
|
||
experts: 384,
|
||
ff_expert: 3072,
|
||
indexer_top_k: 1024,
|
||
expert_weight_scale: 2.5,
|
||
..FLASH
|
||
};
|
||
|
||
const GLM: Shape = Shape {
|
||
model: ModelChoice::Glm52,
|
||
family: ModelFamily::Glm,
|
||
layers: 79,
|
||
embd: 6144,
|
||
vocab: 154_880,
|
||
heads: 64,
|
||
head_kv: 1,
|
||
head_dim: 576,
|
||
value_dim: 512,
|
||
rot: 64,
|
||
out_groups: 0,
|
||
lora_q: 2048,
|
||
lora_o: 0,
|
||
experts: 256,
|
||
experts_used: 8,
|
||
expert_shared: 1,
|
||
ff_expert: 2048,
|
||
ff_dense: 12_288,
|
||
hash_layers: 0,
|
||
sliding_window: 0,
|
||
indexer_heads: 32,
|
||
indexer_head_dim: 128,
|
||
indexer_top_k: 2048,
|
||
hc: 0,
|
||
hc_sinkhorn: 0,
|
||
nextn: 1,
|
||
leading_dense: 3,
|
||
kv_lora: 512,
|
||
key_mla: 256,
|
||
value_mla: 256,
|
||
rms_epsilon: 1.0e-5,
|
||
hc_epsilon: 0.0,
|
||
expert_weight_scale: 2.5,
|
||
swiglu_clamp: 0.0,
|
||
rope_base: 8_000_000.0,
|
||
rope_scale: 1.0,
|
||
rope_beta_fast: 0.0,
|
||
rope_beta_slow: 0.0,
|
||
compress_rope_base: 0.0,
|
||
original_context: 1_048_576,
|
||
};
|
||
|
||
pub(crate) struct Model {
|
||
main: Gguf,
|
||
support: Option<Gguf>,
|
||
support_kind: Option<SupportKind>,
|
||
shape: Shape,
|
||
tokenizer: Tokenizer,
|
||
}
|
||
|
||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||
pub(crate) struct ModelSummary {
|
||
pub(crate) model: ModelChoice,
|
||
pub(crate) mapped_bytes: u64,
|
||
pub(crate) tensor_count: usize,
|
||
pub(crate) vocabulary_size: usize,
|
||
pub(crate) support_loaded: bool,
|
||
}
|
||
|
||
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)?;
|
||
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)
|
||
}
|
||
|
||
fn open_main(path: &Path, expected: ModelChoice) -> Result<Self, String> {
|
||
let main = Gguf::open(path)?;
|
||
let shape = validate_main(&main, expected)?;
|
||
let tokenizer = Tokenizer::load(&main, shape.family)?;
|
||
if tokenizer.vocab_size() != shape.vocab as usize {
|
||
return Err(format!(
|
||
"tokenizer has {} entries, expected {}",
|
||
tokenizer.vocab_size(),
|
||
shape.vocab
|
||
));
|
||
}
|
||
Ok(Self {
|
||
main,
|
||
support: None,
|
||
support_kind: None,
|
||
shape,
|
||
tokenizer,
|
||
})
|
||
}
|
||
|
||
pub(crate) fn summary(&self) -> ModelSummary {
|
||
ModelSummary {
|
||
model: self.shape.model,
|
||
mapped_bytes: self.main.len() + self.support.as_ref().map_or(0, Gguf::len),
|
||
tensor_count: self.main.tensors.len()
|
||
+ self
|
||
.support
|
||
.as_ref()
|
||
.map_or(0, |support| support.tensors.len()),
|
||
vocabulary_size: self.tokenizer.vocab_size(),
|
||
support_loaded: self.support.is_some(),
|
||
}
|
||
}
|
||
|
||
fn checkpoint_identity(&self) -> [u8; 32] {
|
||
let mut hash = Sha256::new();
|
||
hash.update(b"DS4Server model checkpoint identity v1");
|
||
hash.update(self.main.checkpoint_identity());
|
||
if let Some(support) = &self.support {
|
||
hash.update(support.checkpoint_identity());
|
||
}
|
||
hash.finalize().into()
|
||
}
|
||
|
||
pub(crate) fn tokenize(&self, text: &str) -> Vec<i32> {
|
||
self.tokenizer.tokenize(text)
|
||
}
|
||
|
||
pub(crate) fn render_prompt(
|
||
&self,
|
||
system: &str,
|
||
prompt: &str,
|
||
reasoning: ReasoningMode,
|
||
) -> Vec<i32> {
|
||
self.tokenizer.encode_chat(system, prompt, reasoning)
|
||
}
|
||
|
||
fn render_conversation(
|
||
&self,
|
||
system: &str,
|
||
messages: &[ChatTurn],
|
||
reasoning: ReasoningMode,
|
||
) -> Vec<i32> {
|
||
self.tokenizer
|
||
.encode_conversation(system, messages, reasoning)
|
||
}
|
||
|
||
fn render_history(
|
||
&self,
|
||
system: &str,
|
||
messages: &[ChatTurn],
|
||
reasoning: ReasoningMode,
|
||
) -> Vec<i32> {
|
||
self.tokenizer.encode_history(system, messages, 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>> {
|
||
self.tokenizer.token_bytes(token)
|
||
}
|
||
|
||
pub(crate) fn eos_token(&self) -> i32 {
|
||
self.tokenizer.eos()
|
||
}
|
||
|
||
pub(crate) fn is_stop_token(&self, token: i32) -> bool {
|
||
self.tokenizer.is_stop(token)
|
||
}
|
||
|
||
pub(crate) fn is_think_start_token(&self, token: i32) -> bool {
|
||
self.tokenizer.is_think_start(token)
|
||
}
|
||
|
||
pub(crate) fn is_think_end_token(&self, token: i32) -> bool {
|
||
self.tokenizer.is_think_end(token)
|
||
}
|
||
|
||
pub(crate) fn is_stop_token_for_reasoning(&self, token: i32, reasoning: ReasoningMode) -> bool {
|
||
self.is_stop_token(token)
|
||
|| (reasoning == ReasoningMode::Direct
|
||
&& (self.is_think_start_token(token) || self.is_think_end_token(token)))
|
||
}
|
||
|
||
pub(crate) fn tensor_data(&self, name: &str) -> Result<&[u8], String> {
|
||
self.main.tensor_data(name)
|
||
}
|
||
}
|
||
|
||
#[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`.
|
||
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")]
|
||
impl Drop for Generator {
|
||
fn drop(&mut self) {
|
||
// Match ds4: every resident session graph must die before Metal cleanup.
|
||
self.resident_sessions.clear();
|
||
}
|
||
}
|
||
|
||
#[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,
|
||
pub(crate) tool: bool,
|
||
pub(crate) system: bool,
|
||
pub(crate) skip_previous_eos: bool,
|
||
pub(crate) reasoning: Option<String>,
|
||
pub(crate) reasoning_complete: bool,
|
||
pub(crate) content: String,
|
||
}
|
||
|
||
pub(crate) struct GenerationOutput {
|
||
pub(crate) message: ChatTurn,
|
||
pub(crate) prompt_tokens: u32,
|
||
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,
|
||
}
|
||
|
||
struct CheckpointSelection {
|
||
found: bool,
|
||
incompatible: bool,
|
||
}
|
||
|
||
fn checkpoint_rebuild_activity(incompatible: bool) -> &'static str {
|
||
if incompatible {
|
||
"Rebuilding context: the checkpoint belongs to a different model or model configuration."
|
||
} else {
|
||
"Rebuilding context: the saved history or generation settings changed."
|
||
}
|
||
}
|
||
|
||
#[cfg(target_os = "macos")]
|
||
pub(crate) struct CompactionOutput {
|
||
pub(crate) summary: String,
|
||
pub(crate) tail_start: usize,
|
||
pub(crate) context_tokens: u32,
|
||
pub(crate) checkpoint: PathBuf,
|
||
}
|
||
|
||
#[cfg(target_os = "macos")]
|
||
impl Generator {
|
||
pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
|
||
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(),
|
||
settings.diagnostics.expert_profile_path.as_deref(),
|
||
)?;
|
||
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_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,
|
||
);
|
||
Ok(Self {
|
||
executor,
|
||
_simulated_memory: simulated_memory,
|
||
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::<usize>().ok())
|
||
.filter(|limit| *limit > 0)
|
||
.unwrap_or(1)
|
||
.saturating_sub(1),
|
||
})
|
||
}
|
||
|
||
pub(crate) fn summary(&self) -> ModelSummary {
|
||
self.executor.model().summary()
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub(crate) fn generate(
|
||
&mut self,
|
||
checkpoint: &Path,
|
||
bootstrap: Option<&Path>,
|
||
messages: &[ChatTurn],
|
||
settings: &TurnSettings,
|
||
cancelled: &AtomicBool,
|
||
mut emit: impl FnMut(bool, String),
|
||
mut progress: impl FnMut(u32, u32, Option<f32>),
|
||
mut phase: impl FnMut(&'static str),
|
||
) -> Result<GenerationOutput, String> {
|
||
let checkpoint_present = checkpoint.is_file();
|
||
let selected = self.select_checkpoint(checkpoint, |tag| {
|
||
checkpoint_matches_prefix(
|
||
tag,
|
||
&settings.system_prompt,
|
||
settings.reasoning_mode,
|
||
messages,
|
||
)
|
||
})?;
|
||
if checkpoint_present && !selected.found {
|
||
phase(checkpoint_rebuild_activity(selected.incompatible));
|
||
}
|
||
if !selected.found
|
||
&& let Some(directory) = bootstrap
|
||
{
|
||
self.prepare_bootstrap(directory, settings, cancelled, &mut progress, &mut phase)?;
|
||
}
|
||
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 =
|
||
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)
|
||
} else {
|
||
[0; 32]
|
||
},
|
||
)?;
|
||
output.checkpoint_bytes = std::fs::metadata(checkpoint)
|
||
.map(|item| item.len())
|
||
.unwrap_or(0);
|
||
self.checkpoint = Some(checkpoint.to_owned());
|
||
self.resident_active = Some(checkpoint.to_owned());
|
||
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_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,
|
||
);
|
||
}
|
||
|
||
pub(crate) fn generate_transient(
|
||
&mut self,
|
||
directory: &Path,
|
||
messages: &[ChatTurn],
|
||
settings: &TurnSettings,
|
||
cancelled: &AtomicBool,
|
||
mut emit: impl FnMut(bool, String),
|
||
mut progress: impl FnMut(u32, u32, Option<f32>),
|
||
) -> Result<GenerationOutput, String> {
|
||
let history = messages
|
||
.split_last()
|
||
.map_or(messages, |(_, history)| history);
|
||
let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
|
||
let history_key =
|
||
conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
|
||
let history_tag: [u8; 32] = Sha256::digest(&history_key).into();
|
||
let previous_checkpoint =
|
||
self.restore_cached_prefix(directory, &store, &history_key, history_tag)?;
|
||
|
||
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
|
||
.as_deref()
|
||
.and_then(|path| std::fs::metadata(path).ok())
|
||
.map(|item| item.len());
|
||
if prompt_complete {
|
||
let completed_key =
|
||
conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed);
|
||
let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into();
|
||
if !settings.kv_cache.stores(
|
||
self.executor.position(),
|
||
history.is_empty(),
|
||
self.last_store_tokens,
|
||
) {
|
||
// A gated store still leaves the finished conversation in the
|
||
// 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);
|
||
self.save_checkpoint(&completed_checkpoint, completed_tag)?;
|
||
let retained = store.record(
|
||
&completed_checkpoint,
|
||
&completed_key,
|
||
completed_tag,
|
||
self.executor.position(),
|
||
self.executor.context(),
|
||
if history.is_empty() {
|
||
StoreReason::Cold
|
||
} else {
|
||
StoreReason::Continued
|
||
},
|
||
)?;
|
||
output.checkpoint_bytes = retained
|
||
.then(|| std::fs::metadata(&completed_checkpoint).ok())
|
||
.flatten()
|
||
.map_or(0, |item| item.len());
|
||
if retained {
|
||
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)
|
||
}
|
||
|
||
fn prepare_bootstrap(
|
||
&mut self,
|
||
directory: &Path,
|
||
settings: &TurnSettings,
|
||
cancelled: &AtomicBool,
|
||
progress: &mut impl FnMut(u32, u32, Option<f32>),
|
||
phase: &mut impl FnMut(&'static str),
|
||
) -> Result<(), String> {
|
||
let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &[]);
|
||
let tag: [u8; 32] = Sha256::digest(&key).into();
|
||
let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
|
||
self.restore_cached_prefix(directory, &store, &key, tag)?;
|
||
|
||
let tokens = self.executor.model().render_history(
|
||
&settings.system_prompt,
|
||
&[],
|
||
settings.reasoning_mode,
|
||
);
|
||
if tokens.len() >= self.executor.context() as usize {
|
||
return Err(format!(
|
||
"System prompt has {} tokens, but the configured context size is {} tokens",
|
||
tokens.len(),
|
||
self.executor.context()
|
||
));
|
||
}
|
||
let reused = self.executor.align_prompt(&tokens)?;
|
||
if reused == tokens.len() {
|
||
return Ok(());
|
||
}
|
||
|
||
phase("Updating system prompt cache…");
|
||
let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?;
|
||
if completed != tokens.len() - reused {
|
||
return Err("generation cancelled while updating the system prompt cache".into());
|
||
}
|
||
self.executor.note_checkpoint_tag(tag);
|
||
if !settings.kv_cache.stores(self.executor.position(), true, 0) {
|
||
self.checkpoint = None;
|
||
self.resident_active = Some(resident_key(directory, tag));
|
||
return Ok(());
|
||
}
|
||
|
||
let checkpoint = store.checkpoint_path(&key);
|
||
self.save_checkpoint(&checkpoint, tag)?;
|
||
let retained = store.record(
|
||
&checkpoint,
|
||
&key,
|
||
tag,
|
||
self.executor.position(),
|
||
self.executor.context(),
|
||
StoreReason::Cold,
|
||
)?;
|
||
self.checkpoint = retained.then_some(checkpoint.clone());
|
||
self.resident_active = Some(if retained {
|
||
checkpoint
|
||
} else {
|
||
resident_key(directory, tag)
|
||
});
|
||
Ok(())
|
||
}
|
||
|
||
fn restore_cached_prefix(
|
||
&mut self,
|
||
directory: &Path,
|
||
store: &KvStore,
|
||
key: &[u8],
|
||
tag: [u8; 32],
|
||
) -> Result<Option<PathBuf>, String> {
|
||
if self.executor.checkpoint_tag() == tag {
|
||
self.metrics.kv_lookup(KvLookup::MemoryHit);
|
||
return Ok(self.checkpoint.clone());
|
||
}
|
||
if let Some(entry) = store.find(key, self.executor.context()) {
|
||
if self
|
||
.select_checkpoint(&entry.checkpoint, |tag| tag == entry.tag)?
|
||
.found
|
||
{
|
||
store.touch(&entry)?;
|
||
self.last_store_tokens = entry.tokens;
|
||
return Ok(Some(entry.checkpoint));
|
||
}
|
||
store.discard(&entry);
|
||
} else {
|
||
let key = resident_key(directory, tag);
|
||
let restored = self.activate_resident(key)?;
|
||
if !restored || self.executor.checkpoint_tag() != tag {
|
||
self.executor.reset()?;
|
||
}
|
||
self.metrics.kv_lookup(KvLookup::Miss);
|
||
}
|
||
self.checkpoint = None;
|
||
self.last_store_tokens = 0;
|
||
Ok(None)
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub(crate) fn compact(
|
||
&mut self,
|
||
messages: &[ChatTurn],
|
||
settings: &TurnSettings,
|
||
rebuild_system_prompt: &str,
|
||
reason: &str,
|
||
checkpoint: &Path,
|
||
cancelled: &AtomicBool,
|
||
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;
|
||
let result = (|| {
|
||
phase("Compacting durable task state…");
|
||
let mut private_messages = messages.to_vec();
|
||
private_messages.push(ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: crate::compaction::summary_prompt(reason),
|
||
});
|
||
let mut private_settings = settings.clone();
|
||
private_settings.reasoning_mode = ReasoningMode::Direct;
|
||
let private_prompt = self.executor.model().render_conversation(
|
||
&private_settings.system_prompt,
|
||
&private_messages,
|
||
private_settings.reasoning_mode,
|
||
);
|
||
private_settings.max_generated_tokens = crate::compaction::summary_budget(
|
||
private_prompt.len().min(u32::MAX as usize) as u32,
|
||
self.executor.context(),
|
||
)
|
||
.ok_or_else(|| "not enough context left to request compaction summary".to_owned())?;
|
||
private_settings.temperature = 0.0;
|
||
private_settings.stops = vec![
|
||
"<｜DSML｜".into(),
|
||
"<DSML｜".into(),
|
||
"<tool_call>".into(),
|
||
"<think>".into(),
|
||
"</think>".into(),
|
||
];
|
||
let (output, prompt_complete) = self.generate_inner(
|
||
&private_messages,
|
||
&private_settings,
|
||
cancelled,
|
||
&mut |_, _| {},
|
||
&mut progress,
|
||
)?;
|
||
if cancelled.load(Ordering::Relaxed) || !prompt_complete {
|
||
return Err("context compaction interrupted".into());
|
||
}
|
||
let summary = crate::compaction::sanitize_summary(&output.message.content);
|
||
if summary.is_empty() {
|
||
return Err("context compaction produced an empty summary".into());
|
||
}
|
||
|
||
// The private request must never become the rebuilt session prefix.
|
||
phase("Rebuilding compacted context…");
|
||
self.executor.reset()?;
|
||
self.checkpoint = None;
|
||
|
||
let full = self.executor.model().render_conversation(
|
||
&settings.system_prompt,
|
||
messages,
|
||
settings.reasoning_mode,
|
||
);
|
||
let mut starts = Vec::with_capacity(messages.len());
|
||
for index in 0..messages.len() {
|
||
starts.push(
|
||
self.executor
|
||
.model()
|
||
.render_conversation(
|
||
&settings.system_prompt,
|
||
&messages[..index],
|
||
settings.reasoning_mode,
|
||
)
|
||
.len() as u32,
|
||
);
|
||
}
|
||
let start = crate::compaction::tail_start(
|
||
messages,
|
||
&starts,
|
||
full.len() as u32,
|
||
crate::compaction::tail_budget(self.executor.context()),
|
||
);
|
||
let tail = messages[start..].to_vec();
|
||
let rebuilt_system =
|
||
crate::compaction::summary_system_prompt(rebuild_system_prompt, Some(&summary));
|
||
let history_tokens = self.executor.model().render_history(
|
||
&rebuilt_system,
|
||
&tail,
|
||
settings.reasoning_mode,
|
||
);
|
||
if history_tokens.len() >= self.executor.context() as usize {
|
||
return Err("compacted context does not fit the configured context".into());
|
||
}
|
||
let context = self.executor.context();
|
||
let completed = self.executor.prefill(&history_tokens, |used| {
|
||
progress(used, context, None);
|
||
!cancelled.load(Ordering::Relaxed)
|
||
})?;
|
||
if completed != history_tokens.len() || cancelled.load(Ordering::Relaxed) {
|
||
return Err("context compaction interrupted during rebuild".into());
|
||
}
|
||
let tag = conversation_tag(&rebuilt_system, settings.reasoning_mode, &tail);
|
||
phase("Saving compacted context…");
|
||
self.save_checkpoint(checkpoint, tag)?;
|
||
Ok(CompactionOutput {
|
||
summary,
|
||
tail_start: start,
|
||
context_tokens: history_tokens.len() as u32,
|
||
checkpoint: checkpoint.to_owned(),
|
||
})
|
||
})();
|
||
if result.is_err() {
|
||
let _ = std::fs::remove_file(checkpoint);
|
||
self.executor.reset()?;
|
||
self.checkpoint = None;
|
||
} else {
|
||
self.checkpoint = Some(checkpoint.to_owned());
|
||
}
|
||
result
|
||
}
|
||
|
||
pub(crate) fn rendered_history_tokens(
|
||
&self,
|
||
messages: &[ChatTurn],
|
||
settings: &TurnSettings,
|
||
) -> Result<u32, String> {
|
||
u32::try_from(
|
||
self.executor
|
||
.model()
|
||
.render_history(&settings.system_prompt, messages, settings.reasoning_mode)
|
||
.len(),
|
||
)
|
||
.map_err(|_| "rendered conversation is too large".to_owned())
|
||
}
|
||
|
||
fn select_checkpoint(
|
||
&mut self,
|
||
checkpoint: &Path,
|
||
matches: impl Fn([u8; 32]) -> bool,
|
||
) -> Result<CheckpointSelection, String> {
|
||
let resident_hit = self.activate_resident(checkpoint.to_owned())?;
|
||
if !checkpoint.is_file() {
|
||
self.executor.reset()?;
|
||
self.checkpoint = None;
|
||
self.metrics.kv_lookup(KvLookup::Miss);
|
||
return Ok(CheckpointSelection {
|
||
found: false,
|
||
incompatible: false,
|
||
});
|
||
}
|
||
if resident_hit && matches(self.executor.checkpoint_tag()) {
|
||
self.checkpoint = Some(checkpoint.to_owned());
|
||
self.metrics.kv_lookup(KvLookup::MemoryHit);
|
||
return Ok(CheckpointSelection {
|
||
found: true,
|
||
incompatible: false,
|
||
});
|
||
}
|
||
if self.checkpoint.as_deref() == Some(checkpoint) {
|
||
let found = matches(self.executor.checkpoint_tag());
|
||
self.metrics.kv_lookup(if found {
|
||
KvLookup::MemoryHit
|
||
} else {
|
||
KvLookup::Miss
|
||
});
|
||
return Ok(CheckpointSelection {
|
||
found,
|
||
incompatible: false,
|
||
});
|
||
}
|
||
self.executor.reset()?;
|
||
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 found = matches!(loaded, Ok(true)) && matches(self.executor.checkpoint_tag());
|
||
let incompatible = loaded.is_err();
|
||
let lookup = match loaded {
|
||
Ok(true) if found => 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(CheckpointSelection {
|
||
found,
|
||
incompatible,
|
||
})
|
||
}
|
||
|
||
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) {
|
||
// DS4 defaults to one active session; opt into its batched-server
|
||
// behavior with DS4_RESIDENT_SESSIONS when memory permits.
|
||
if self.resident_limit != 0 {
|
||
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();
|
||
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 prefill_suffix(
|
||
&mut self,
|
||
tokens: &[i32],
|
||
reused: usize,
|
||
cancelled: &AtomicBool,
|
||
progress: &mut impl FnMut(u32, u32, Option<f32>),
|
||
) -> Result<usize, String> {
|
||
let suffix = &tokens[reused..];
|
||
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);
|
||
}
|
||
Ok(completed)
|
||
}
|
||
}
|
||
|
||
fn generate_inner(
|
||
&mut self,
|
||
messages: &[ChatTurn],
|
||
settings: &TurnSettings,
|
||
cancelled: &AtomicBool,
|
||
emit: &mut impl FnMut(bool, String),
|
||
progress: &mut impl FnMut(u32, u32, Option<f32>),
|
||
) -> Result<(GenerationOutput, bool), String> {
|
||
let tokens = match messages.split_last() {
|
||
Some((latest, history))
|
||
if latest.user
|
||
&& self.executor.checkpoint_tag()
|
||
== conversation_tag(
|
||
&settings.system_prompt,
|
||
settings.reasoning_mode,
|
||
history,
|
||
) =>
|
||
{
|
||
let mut tokens = self.executor.tokens().to_vec();
|
||
tokens.extend(self.executor.model().render_continuation(
|
||
&latest.content,
|
||
settings.reasoning_mode,
|
||
latest.skip_previous_eos,
|
||
));
|
||
tokens
|
||
}
|
||
_ => self.executor.model().render_conversation(
|
||
&settings.system_prompt,
|
||
messages,
|
||
settings.reasoning_mode,
|
||
),
|
||
};
|
||
if tokens.is_empty() {
|
||
return Err("the rendered prompt is empty".into());
|
||
}
|
||
let max_context = self.executor.context() as usize;
|
||
if tokens.len() >= max_context {
|
||
return Err(format!(
|
||
"Prompt has {} tokens, but the configured context size is {max_context} tokens",
|
||
tokens.len()
|
||
));
|
||
}
|
||
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,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: reasoning.then(String::new),
|
||
reasoning_complete: !reasoning,
|
||
content: String::new(),
|
||
};
|
||
let mut emitted_reasoning = 0;
|
||
let mut emitted_content = 0;
|
||
let mut pending_utf8 = Vec::new();
|
||
let prompt_tokens = tokens.len();
|
||
let suffix = &tokens[reused..];
|
||
let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?;
|
||
self.publish_execution_stats();
|
||
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;
|
||
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(
|
||
&mut generated,
|
||
&mut emitted_reasoning,
|
||
&mut emitted_content,
|
||
&settings.stops,
|
||
emit,
|
||
);
|
||
return Ok((
|
||
GenerationOutput {
|
||
message: generated,
|
||
prompt_tokens: prompt_tokens as u32,
|
||
cached_tokens: reused as u32,
|
||
completion_tokens: generated_tokens,
|
||
finish_reason: "stop",
|
||
previous_checkpoint_bytes: None,
|
||
checkpoint_bytes: 0,
|
||
},
|
||
true,
|
||
));
|
||
}
|
||
let token = sample(
|
||
self.executor.logits(),
|
||
settings.temperature,
|
||
settings.top_p,
|
||
settings.min_p,
|
||
settings.top_k,
|
||
&mut rng,
|
||
);
|
||
if self
|
||
.executor
|
||
.model()
|
||
.is_stop_token_for_reasoning(token, settings.reasoning_mode)
|
||
{
|
||
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
|
||
flush_generated(
|
||
&mut generated,
|
||
&mut emitted_reasoning,
|
||
&mut emitted_content,
|
||
&settings.stops,
|
||
emit,
|
||
);
|
||
return Ok((
|
||
GenerationOutput {
|
||
message: generated,
|
||
prompt_tokens: prompt_tokens as u32,
|
||
cached_tokens: reused as u32,
|
||
completion_tokens: generated_tokens,
|
||
finish_reason: "stop",
|
||
previous_checkpoint_bytes: None,
|
||
checkpoint_bytes: 0,
|
||
},
|
||
true,
|
||
));
|
||
}
|
||
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)
|
||
{
|
||
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
|
||
flush_generated(
|
||
&mut generated,
|
||
&mut emitted_reasoning,
|
||
&mut emitted_content,
|
||
&settings.stops,
|
||
emit,
|
||
);
|
||
return Ok((
|
||
GenerationOutput {
|
||
message: generated,
|
||
prompt_tokens: prompt_tokens as u32,
|
||
cached_tokens: reused as u32,
|
||
completion_tokens: generated_tokens,
|
||
finish_reason: "stop",
|
||
previous_checkpoint_bytes: None,
|
||
checkpoint_bytes: 0,
|
||
},
|
||
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),
|
||
),
|
||
);
|
||
}
|
||
}
|
||
append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
|
||
flush_generated(
|
||
&mut generated,
|
||
&mut emitted_reasoning,
|
||
&mut emitted_content,
|
||
&settings.stops,
|
||
emit,
|
||
);
|
||
Ok((
|
||
GenerationOutput {
|
||
message: generated,
|
||
prompt_tokens: prompt_tokens as u32,
|
||
cached_tokens: reused as u32,
|
||
completion_tokens: generated_tokens,
|
||
finish_reason: "length",
|
||
previous_checkpoint_bytes: None,
|
||
checkpoint_bytes: 0,
|
||
},
|
||
true,
|
||
))
|
||
}
|
||
}
|
||
|
||
#[cfg(any(target_os = "macos", test))]
|
||
fn append_generated_bytes(
|
||
generated: &mut ChatTurn,
|
||
reasoning: bool,
|
||
pending: &mut Vec<u8>,
|
||
bytes: &[u8],
|
||
final_flush: bool,
|
||
) {
|
||
let text = if reasoning {
|
||
generated.reasoning.get_or_insert_default()
|
||
} else {
|
||
&mut generated.content
|
||
};
|
||
pending.extend_from_slice(bytes);
|
||
loop {
|
||
match std::str::from_utf8(pending) {
|
||
Ok(valid) => {
|
||
text.push_str(valid);
|
||
pending.clear();
|
||
return;
|
||
}
|
||
Err(error) => {
|
||
let valid = error.valid_up_to();
|
||
text.push_str(std::str::from_utf8(&pending[..valid]).unwrap());
|
||
pending.drain(..valid);
|
||
match error.error_len() {
|
||
Some(length) => {
|
||
text.push('\u{fffd}');
|
||
pending.drain(..length);
|
||
}
|
||
None if final_flush => {
|
||
text.push_str(&String::from_utf8_lossy(pending));
|
||
pending.clear();
|
||
return;
|
||
}
|
||
None => return,
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(target_os = "macos")]
|
||
fn flush_generated(
|
||
generated: &mut ChatTurn,
|
||
emitted_reasoning: &mut usize,
|
||
emitted_content: &mut usize,
|
||
stops: &[String],
|
||
emit: &mut impl FnMut(bool, String),
|
||
) {
|
||
if let Some(reasoning) = &mut generated.reasoning {
|
||
let _ = emit_safe_text(reasoning, emitted_reasoning, stops, true, true, emit);
|
||
}
|
||
let _ = emit_safe_text(
|
||
&mut generated.content,
|
||
emitted_content,
|
||
stops,
|
||
true,
|
||
false,
|
||
emit,
|
||
);
|
||
}
|
||
|
||
#[cfg(any(target_os = "macos", test))]
|
||
fn emit_safe_text(
|
||
text: &mut String,
|
||
emitted: &mut usize,
|
||
stops: &[String],
|
||
final_flush: bool,
|
||
reasoning: bool,
|
||
emit: &mut impl FnMut(bool, String),
|
||
) -> bool {
|
||
let stop = stops
|
||
.iter()
|
||
.filter_map(|stop| {
|
||
text[*emitted..]
|
||
.find(stop)
|
||
.map(|position| *emitted + position)
|
||
})
|
||
.min();
|
||
if let Some(stop) = stop {
|
||
if stop > *emitted {
|
||
emit(reasoning, text[*emitted..stop].to_owned());
|
||
}
|
||
text.truncate(stop);
|
||
*emitted = stop;
|
||
return true;
|
||
}
|
||
let hold = if final_flush {
|
||
0
|
||
} else {
|
||
stops
|
||
.iter()
|
||
.map(|stop| stop.len().saturating_sub(1))
|
||
.max()
|
||
.unwrap_or(0)
|
||
};
|
||
let mut safe = text.len().saturating_sub(hold);
|
||
while safe > *emitted && !text.is_char_boundary(safe) {
|
||
safe -= 1;
|
||
}
|
||
if safe > *emitted {
|
||
emit(reasoning, text[*emitted..safe].to_owned());
|
||
*emitted = safe;
|
||
}
|
||
false
|
||
}
|
||
|
||
#[cfg(any(target_os = "macos", test))]
|
||
fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> Vec<u8> {
|
||
fn text(output: &mut Vec<u8>, value: &str) {
|
||
output.extend_from_slice(&(value.len() as u64).to_le_bytes());
|
||
output.extend_from_slice(value.as_bytes());
|
||
}
|
||
|
||
let mut output = b"DS4Server chat checkpoint v4".to_vec();
|
||
text(&mut output, system);
|
||
output.push(match reasoning {
|
||
ReasoningMode::Direct => 0,
|
||
ReasoningMode::High => 1,
|
||
ReasoningMode::Max => 2,
|
||
});
|
||
for message in messages {
|
||
output.push(u8::from(message.user));
|
||
output.push(u8::from(message.tool));
|
||
output.push(u8::from(message.system));
|
||
output.push(u8::from(message.skip_previous_eos));
|
||
match &message.reasoning {
|
||
Some(reasoning) => {
|
||
output.push(1);
|
||
text(&mut output, reasoning);
|
||
}
|
||
None => output.push(0),
|
||
}
|
||
output.push(u8::from(message.reasoning_complete));
|
||
text(&mut output, &message.content);
|
||
}
|
||
output
|
||
}
|
||
|
||
#[cfg(any(target_os = "macos", test))]
|
||
fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
|
||
Sha256::digest(conversation_key(system, reasoning, messages)).into()
|
||
}
|
||
|
||
#[cfg(any(target_os = "macos", test))]
|
||
fn checkpoint_matches_prefix(
|
||
checkpoint: [u8; 32],
|
||
system: &str,
|
||
reasoning: ReasoningMode,
|
||
messages: &[ChatTurn],
|
||
) -> bool {
|
||
// ponytail: appended control messages are few; carry incremental hashes if
|
||
// scanning a genuinely changed, very long history becomes measurable.
|
||
(0..messages.len())
|
||
.rev()
|
||
.any(|end| checkpoint == conversation_tag(system, reasoning, &messages[..end]))
|
||
}
|
||
|
||
#[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(any(target_os = "macos", test))]
|
||
fn sample(
|
||
logits: &[f32],
|
||
temperature: f32,
|
||
top_p: f32,
|
||
min_p: f32,
|
||
top_k: i32,
|
||
rng: &mut Rng,
|
||
) -> i32 {
|
||
if temperature <= 0.0 {
|
||
return logits
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|a, b| a.1.total_cmp(b.1))
|
||
.map_or(0, |(index, _)| index as i32);
|
||
}
|
||
let maximum = logits
|
||
.iter()
|
||
.copied()
|
||
.filter(|value| value.is_finite())
|
||
.fold(f32::NEG_INFINITY, f32::max);
|
||
if !maximum.is_finite() {
|
||
return 0;
|
||
}
|
||
let top_p = if top_p <= 0.0 || top_p > 1.0 {
|
||
1.0
|
||
} else {
|
||
top_p
|
||
};
|
||
let min_p = min_p.max(0.0);
|
||
let mut probabilities: Vec<(usize, f32)> = logits
|
||
.iter()
|
||
.enumerate()
|
||
.filter(|(_, logit)| logit.is_finite())
|
||
.map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp()))
|
||
.filter(|(_, probability)| *probability >= min_p)
|
||
.collect();
|
||
if probabilities.is_empty() {
|
||
return logits
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|a, b| a.1.total_cmp(b.1))
|
||
.map_or(0, |(index, _)| index as i32);
|
||
}
|
||
if top_p < 1.0 || top_k > 0 {
|
||
probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
|
||
if top_k > 0 {
|
||
probabilities.truncate(probabilities.len().min(top_k as usize));
|
||
}
|
||
}
|
||
if top_p < 1.0 {
|
||
let total: f32 = probabilities
|
||
.iter()
|
||
.map(|(_, probability)| probability)
|
||
.sum();
|
||
let mut kept = 0.0;
|
||
let count = probabilities
|
||
.iter()
|
||
.position(|(_, probability)| {
|
||
kept += *probability;
|
||
kept / total >= top_p
|
||
})
|
||
.map_or(probabilities.len(), |index| index + 1);
|
||
probabilities.truncate(count);
|
||
}
|
||
let kept_total: f32 = probabilities.iter().map(|(_, p)| p).sum();
|
||
let mut choice = rng.unit() * kept_total;
|
||
for (token, probability) in &probabilities {
|
||
choice -= probability;
|
||
if choice <= 0.0 {
|
||
return *token as i32;
|
||
}
|
||
}
|
||
probabilities.last().map_or(0, |(token, _)| *token as i32)
|
||
}
|
||
|
||
#[cfg(any(target_os = "macos", test))]
|
||
struct Rng(u64);
|
||
|
||
#[cfg(any(target_os = "macos", test))]
|
||
impl Rng {
|
||
fn new(seed: u64) -> Self {
|
||
Self(seed.max(1))
|
||
}
|
||
|
||
fn unit(&mut self) -> f32 {
|
||
let mut value = self.0;
|
||
if value == 0 {
|
||
value = 0x9e37_79b9_7f4a_7c15;
|
||
}
|
||
value ^= value >> 12;
|
||
value ^= value << 25;
|
||
value ^= value >> 27;
|
||
self.0 = value;
|
||
let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d);
|
||
((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod sampling_tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn zero_temperature_is_greedy() {
|
||
let mut rng = Rng::new(1);
|
||
assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, 0, &mut rng), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn top_k_and_stream_stops_are_applied_before_output() {
|
||
let mut rng = Rng::new(1);
|
||
assert_eq!(sample(&[1.0, 4.0, 2.0], 1.0, 1.0, 0.0, 1, &mut rng), 1);
|
||
|
||
let mut text = "hello STOP hidden".to_owned();
|
||
let mut emitted = 0;
|
||
let mut chunks = Vec::new();
|
||
assert!(emit_safe_text(
|
||
&mut text,
|
||
&mut emitted,
|
||
&["STOP".into()],
|
||
false,
|
||
false,
|
||
&mut |_, chunk| chunks.push(chunk),
|
||
));
|
||
assert_eq!(text, "hello ");
|
||
assert_eq!(chunks, ["hello "]);
|
||
}
|
||
|
||
#[test]
|
||
fn split_utf8_token_bytes_are_joined_before_decoding() {
|
||
let mut generated = ChatTurn {
|
||
user: false,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: String::new(),
|
||
};
|
||
let mut pending = Vec::new();
|
||
|
||
append_generated_bytes(&mut generated, false, &mut pending, &[0xf0, 0x9f], false);
|
||
assert!(generated.content.is_empty());
|
||
append_generated_bytes(&mut generated, false, &mut pending, &[0x98, 0x8a], false);
|
||
|
||
assert_eq!(generated.content, "😊");
|
||
assert!(pending.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn checkpoint_tag_covers_the_canonical_chat_state() {
|
||
let mut messages = vec![ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
}];
|
||
let tag = conversation_tag("System", ReasoningMode::High, &messages);
|
||
assert_eq!(
|
||
tag,
|
||
conversation_tag("System", ReasoningMode::High, &messages)
|
||
);
|
||
assert_ne!(
|
||
tag,
|
||
conversation_tag("Changed", ReasoningMode::High, &messages)
|
||
);
|
||
assert_ne!(
|
||
tag,
|
||
conversation_tag("System", ReasoningMode::Direct, &messages)
|
||
);
|
||
messages[0].system = true;
|
||
assert_ne!(
|
||
tag,
|
||
conversation_tag("System", ReasoningMode::High, &messages)
|
||
);
|
||
messages[0].system = false;
|
||
messages[0].skip_previous_eos = true;
|
||
assert_ne!(
|
||
tag,
|
||
conversation_tag("System", ReasoningMode::High, &messages)
|
||
);
|
||
let prefix = conversation_key("System", ReasoningMode::High, &messages);
|
||
messages.push(ChatTurn {
|
||
user: false,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: Some("because".into()),
|
||
reasoning_complete: true,
|
||
content: "Hi".into(),
|
||
});
|
||
assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix));
|
||
}
|
||
|
||
#[test]
|
||
fn checkpoint_tag_accepts_an_unchanged_prefix_before_reminders() {
|
||
let mut messages = vec![ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Question".into(),
|
||
}];
|
||
let checkpoint = conversation_tag("System", ReasoningMode::High, &messages);
|
||
messages.extend([
|
||
ChatTurn {
|
||
user: false,
|
||
tool: true,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Tool result".into(),
|
||
},
|
||
ChatTurn {
|
||
user: false,
|
||
tool: false,
|
||
system: true,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "System prompt reminder".into(),
|
||
},
|
||
]);
|
||
|
||
assert!(checkpoint_matches_prefix(
|
||
checkpoint,
|
||
"System",
|
||
ReasoningMode::High,
|
||
&messages,
|
||
));
|
||
assert!(!checkpoint_matches_prefix(
|
||
checkpoint,
|
||
"Changed",
|
||
ReasoningMode::High,
|
||
&messages,
|
||
));
|
||
}
|
||
|
||
#[test]
|
||
fn bootstrap_key_is_the_prefix_before_dynamic_session_context() {
|
||
let system = "System\n\nProject instructions from AGENTS.md:\n\nkeep this";
|
||
let bootstrap = conversation_key(system, ReasoningMode::High, &[]);
|
||
let messages = vec![
|
||
ChatTurn {
|
||
user: false,
|
||
tool: false,
|
||
system: true,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "current date and time".into(),
|
||
},
|
||
ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "hello".into(),
|
||
},
|
||
];
|
||
|
||
assert!(conversation_key(system, ReasoningMode::High, &messages).starts_with(&bootstrap));
|
||
}
|
||
|
||
#[test]
|
||
fn checkpoint_rebuilds_explain_compatibility_and_history_misses() {
|
||
assert!(checkpoint_rebuild_activity(true).contains("different model"));
|
||
assert!(checkpoint_rebuild_activity(false).contains("history"));
|
||
}
|
||
|
||
#[test]
|
||
#[cfg(target_os = "macos")]
|
||
#[ignore = "requires the 80 GiB Flash checkpoint and Apple Metal"]
|
||
fn metal_executes_real_flash_token() {
|
||
configure_metal_sources().unwrap();
|
||
let path = crate::model::engine_artifacts(
|
||
ModelChoice::DeepSeekV4Flash,
|
||
false,
|
||
false,
|
||
&crate::app::models_path(),
|
||
)
|
||
.model;
|
||
let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap();
|
||
let tokens = model.render_prompt(
|
||
"You are a helpful assistant",
|
||
"Hello",
|
||
ReasoningMode::Direct,
|
||
);
|
||
assert_eq!(tokens.len(), 10);
|
||
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();
|
||
}
|
||
assert!(executor.logits().iter().all(|logit| logit.is_finite()));
|
||
let argmax = executor
|
||
.logits()
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|a, b| a.1.total_cmp(b.1))
|
||
.unwrap();
|
||
eprintln!(
|
||
"Rust logits: argmax={} value={} logit0={}",
|
||
argmax.0,
|
||
argmax.1,
|
||
executor.logits()[0]
|
||
);
|
||
assert_eq!(argmax.0, 19_923);
|
||
assert!((executor.logits()[0] - -7.675_424).abs() < 0.1);
|
||
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], &mut |_| {})
|
||
.unwrap();
|
||
executor.eval(next).unwrap();
|
||
let continued_logit = executor.logits()[0];
|
||
let continued_argmax = executor
|
||
.logits()
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|a, b| a.1.total_cmp(b.1))
|
||
.unwrap()
|
||
.0;
|
||
drop(executor);
|
||
|
||
let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).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]);
|
||
restored.eval(next).unwrap();
|
||
assert_eq!(
|
||
restored
|
||
.logits()
|
||
.iter()
|
||
.enumerate()
|
||
.max_by(|a, b| a.1.total_cmp(b.1))
|
||
.unwrap()
|
||
.0,
|
||
continued_argmax
|
||
);
|
||
assert!((restored.logits()[0] - continued_logit).abs() < 1.0e-5);
|
||
std::fs::remove_file(checkpoint).unwrap();
|
||
}
|
||
}
|