1361 lines
44 KiB
Rust
1361 lines
44 KiB
Rust
mod gguf;
|
||
#[cfg(target_os = "macos")]
|
||
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;
|
||
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};
|
||
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;
|
||
use tokenizer::Tokenizer;
|
||
use validation::{validate_dspark, validate_main};
|
||
|
||
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>,
|
||
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 let Some(path) = &settings.artifacts.mtp {
|
||
let support = Gguf::open(path)?;
|
||
validate_dspark(&support, &model.shape)?;
|
||
model.support = Some(support);
|
||
}
|
||
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,
|
||
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,
|
||
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>,
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
pub(crate) struct ChatTurn {
|
||
pub(crate) user: bool,
|
||
pub(crate) tool: 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,
|
||
}
|
||
|
||
#[cfg(target_os = "macos")]
|
||
pub(crate) struct CompactionOutput {
|
||
pub(crate) summary: String,
|
||
pub(crate) tail: Vec<ChatTurn>,
|
||
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> {
|
||
if settings.model != ModelChoice::DeepSeekV4Flash {
|
||
return Err("local generation currently supports DeepSeek V4 Flash only".into());
|
||
}
|
||
if settings.speculative.dspark || settings.ssd.enabled || settings.steering.file.is_some() {
|
||
return Err(
|
||
"DSpark, SSD streaming, and steering are not yet available in the Rust executor"
|
||
.into(),
|
||
);
|
||
}
|
||
let model = Model::open(settings)?;
|
||
let executor = metal::Executor::open(
|
||
model,
|
||
settings.context_tokens.max(1) as u32,
|
||
settings.execution.quality,
|
||
settings.execution.prefill_chunk,
|
||
)?;
|
||
Ok(Self {
|
||
executor,
|
||
checkpoint: None,
|
||
last_store_tokens: 0,
|
||
metrics,
|
||
})
|
||
}
|
||
|
||
pub(crate) fn summary(&self) -> ModelSummary {
|
||
self.executor.model().summary()
|
||
}
|
||
|
||
pub(crate) fn generate(
|
||
&mut self,
|
||
checkpoint: &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);
|
||
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());
|
||
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);
|
||
Ok(output)
|
||
}
|
||
|
||
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 mut previous_checkpoint = self.checkpoint.clone();
|
||
if self.executor.checkpoint_tag() != history_tag {
|
||
if let Some(entry) = store.find(&history_key, self.executor.context()) {
|
||
if self.select_checkpoint(&entry.checkpoint, entry.tag)? {
|
||
store.touch(&entry)?;
|
||
self.last_store_tokens = entry.tokens;
|
||
previous_checkpoint = Some(entry.checkpoint);
|
||
} else {
|
||
store.discard(&entry);
|
||
self.last_store_tokens = 0;
|
||
previous_checkpoint = None;
|
||
}
|
||
} else {
|
||
self.executor.reset()?;
|
||
self.checkpoint = None;
|
||
self.last_store_tokens = 0;
|
||
previous_checkpoint = None;
|
||
self.metrics.kv_lookup(KvLookup::Miss);
|
||
}
|
||
} else {
|
||
self.metrics.kv_lookup(KvLookup::MemoryHit);
|
||
}
|
||
|
||
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());
|
||
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;
|
||
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);
|
||
}
|
||
Ok(output)
|
||
}
|
||
|
||
pub(crate) fn compact(
|
||
&mut self,
|
||
messages: &[ChatTurn],
|
||
settings: &TurnSettings,
|
||
reason: &str,
|
||
checkpoint: &Path,
|
||
cancelled: &AtomicBool,
|
||
mut progress: impl FnMut(u32, u32, Option<f32>),
|
||
) -> Result<CompactionOutput, String> {
|
||
let _ = std::fs::remove_file(checkpoint);
|
||
self.executor.reset()?;
|
||
self.checkpoint = None;
|
||
let result = (|| {
|
||
let mut private_messages = messages.to_vec();
|
||
private_messages.push(ChatTurn {
|
||
user: true,
|
||
tool: 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;
|
||
private_settings.max_generated_tokens = crate::compaction::SUMMARY_MAX_TOKENS;
|
||
private_settings.temperature = 0.0;
|
||
private_settings.stops.extend([
|
||
"<｜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.
|
||
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(&settings.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);
|
||
self.save_checkpoint(checkpoint, tag)?;
|
||
Ok(CompactionOutput {
|
||
summary,
|
||
tail,
|
||
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
|
||
}
|
||
|
||
fn select_checkpoint(
|
||
&mut self,
|
||
checkpoint: &Path,
|
||
expected_tag: [u8; 32],
|
||
) -> Result<bool, String> {
|
||
if self.checkpoint.as_deref() == Some(checkpoint) {
|
||
if !checkpoint.is_file() {
|
||
self.executor.reset()?;
|
||
self.checkpoint = None;
|
||
self.metrics.kv_lookup(KvLookup::Miss);
|
||
return Ok(false);
|
||
}
|
||
let found = self.executor.checkpoint_tag() == expected_tag;
|
||
self.metrics.kv_lookup(if found {
|
||
KvLookup::MemoryHit
|
||
} else {
|
||
KvLookup::Miss
|
||
});
|
||
return Ok(found);
|
||
}
|
||
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)) && self.executor.checkpoint_tag() == expected_tag;
|
||
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(found)
|
||
}
|
||
|
||
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],
|
||
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,
|
||
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 = 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
|
||
.max_generated_tokens
|
||
.max(0)
|
||
.min((max_context - self.executor.position() as usize) as i32)
|
||
{
|
||
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,
|
||
));
|
||
}
|
||
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,
|
||
));
|
||
}
|
||
}
|
||
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(
|
||
&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(target_os = "macos")]
|
||
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(target_os = "macos")]
|
||
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(target_os = "macos")]
|
||
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 v3".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.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(target_os = "macos")]
|
||
fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
|
||
Sha256::digest(conversation_key(system, reasoning, messages)).into()
|
||
}
|
||
|
||
#[cfg(target_os = "macos")]
|
||
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(target_os = "macos")]
|
||
struct Rng(u64);
|
||
|
||
#[cfg(target_os = "macos")]
|
||
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(all(test, target_os = "macos"))]
|
||
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,
|
||
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,
|
||
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].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,
|
||
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]
|
||
#[ignore = "requires the 80 GiB Flash checkpoint and Apple Metal"]
|
||
fn metal_executes_real_flash_token() {
|
||
configure_metal_sources().unwrap();
|
||
let path = Path::new(env!("CARGO_MANIFEST_DIR")).join(
|
||
"../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf",
|
||
);
|
||
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();
|
||
}
|
||
}
|