full chat implemented

src/engine/metal.rs

@@ -2,8 +2,15 @@ use super::gguf::{F16, Gguf, Q8_0, Tensor as GgufTensor};
 use super::{Model, ModelFamily};
 use std::env;
 use std::ffi::c_void;
+use std::fs::{self, File};
+use std::io::{Read, Write};
 use std::path::Path;
 use std::ptr::NonNull;
+use std::time::UNIX_EPOCH;
+
+const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4RKV01";
+const CHECKPOINT_VERSION: u32 = 1;
+const CHECKPOINT_IO_CHUNK: usize = 8 * 1024 * 1024;
 
 const SOURCES: [(&str, &str); 19] = [
     ("DS4_METAL_FLASH_ATTN_SOURCE", "flash_attn.metal"),
@@ -78,6 +85,12 @@ unsafe extern "C" {
         data: *mut c_void,
         bytes: u64,
     ) -> i32;
+    fn ds4_gpu_tensor_write(
+        tensor: *mut GpuTensor,
+        offset: u64,
+        data: *const c_void,
+        bytes: u64,
+    ) -> i32;
     fn ds4_gpu_tensor_copy(
         dst: *mut GpuTensor,
         dst_offset: u64,
@@ -562,6 +575,34 @@ impl Buffer {
         )
     }
 
+    fn read(&self, offset: u64, values: &mut [u8]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_read(
+                    self.raw(),
+                    offset,
+                    values.as_mut_ptr().cast(),
+                    values.len() as u64,
+                )
+            },
+            "reading a Metal buffer",
+        )
+    }
+
+    fn write(&self, offset: u64, values: &[u8]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_write(
+                    self.raw(),
+                    offset,
+                    values.as_ptr().cast(),
+                    values.len() as u64,
+                )
+            },
+            "restoring a Metal buffer",
+        )
+    }
+
     fn fill(&self, value: f32, count: u64) -> Result<(), String> {
         call(
             unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) },
@@ -960,6 +1001,10 @@ pub(super) struct Executor {
     weights: Weights,
     session: Session,
     logits: Vec<f32>,
+    tokens: Vec<i32>,
+    quality: bool,
+    checkpoint_tag: [u8; 32],
+    model_modified: (u64, u32),
     _context: Context,
     model: Model,
 }
@@ -969,10 +1014,20 @@ impl Executor {
         let weights = Weights::bind(&model)?;
         let context_handle = Context::open(&model, quality)?;
         let session = Session::new(&model, context)?;
+        let model_modified = fs::metadata(model.main.path())
+            .and_then(|metadata| metadata.modified())
+            .ok()
+            .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok())
+            .map(|duration| (duration.as_secs(), duration.subsec_nanos()))
+            .unwrap_or_default();
         Ok(Self {
             weights,
             session,
             logits: vec![0.0; model.shape.vocab as usize],
+            tokens: Vec::new(),
+            quality,
+            checkpoint_tag: [0; 32],
+            model_modified,
             _context: context_handle,
             model,
         })
@@ -993,6 +1048,7 @@ impl Executor {
         commands.finish()?;
         self.session.scratch.logits.read_f32(&mut self.logits)?;
         self.session.position += 1;
+        self.tokens.push(token);
         Ok(())
     }
 
@@ -1014,6 +1070,268 @@ impl Executor {
 
     pub(super) fn reset(&mut self) -> Result<(), String> {
         self.session = Session::new(&self.model, self.session.context)?;
+        self.tokens.clear();
+        self.checkpoint_tag = [0; 32];
+        Ok(())
+    }
+
+    pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
+        if !tokens.starts_with(&self.tokens) {
+            self.reset()?;
+        }
+        Ok(self.tokens.len())
+    }
+
+    pub(super) fn tokens(&self) -> &[i32] {
+        &self.tokens
+    }
+
+    pub(super) fn checkpoint_tag(&self) -> [u8; 32] {
+        self.checkpoint_tag
+    }
+
+    pub(super) fn save_checkpoint(&self, path: &Path, tag: [u8; 32]) -> Result<(), String> {
+        if let Some(parent) = path.parent() {
+            fs::create_dir_all(parent).map_err(|error| error.to_string())?;
+        }
+        let temporary = path.with_extension("tmp");
+        let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
+        self.write_checkpoint(&mut file, tag)?;
+        file.sync_all().map_err(|error| error.to_string())?;
+        fs::rename(&temporary, path).map_err(|error| error.to_string())
+    }
+
+    pub(super) fn load_checkpoint(&mut self, path: &Path) -> Result<bool, String> {
+        let mut file = match File::open(path) {
+            Ok(file) => file,
+            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
+            Err(error) => return Err(error.to_string()),
+        };
+        self.read_checkpoint(&mut file)?;
+        Ok(true)
+    }
+
+    fn write_checkpoint(&self, file: &mut File, tag: [u8; 32]) -> Result<(), String> {
+        let shape = self.model.shape;
+        file.write_all(CHECKPOINT_MAGIC)
+            .map_err(|error| error.to_string())?;
+        for value in [
+            CHECKPOINT_VERSION,
+            self.session.context,
+            self.session.raw_cap,
+            shape.layers,
+            shape.head_dim as u32,
+            shape.indexer_head_dim as u32,
+            shape.vocab as u32,
+            u32::from(self.quality),
+        ] {
+            write_u32(file, value)?;
+        }
+        write_u64(file, self.model.main.len())?;
+        write_u64(file, self.model_modified.0)?;
+        write_u32(file, self.model_modified.1)?;
+        for weight in [self.weights.token_embedding, self.weights.output] {
+            write_u64(file, weight.offset)?;
+            write_u64(file, weight.bytes)?;
+            write_u32(file, weight.kind)?;
+        }
+        file.write_all(&tag).map_err(|error| error.to_string())?;
+        let token_count = u32::try_from(self.tokens.len())
+            .map_err(|_| "KV checkpoint has too many tokens".to_owned())?;
+        let raw_live = token_count.min(self.session.raw_cap);
+        write_u32(file, token_count)?;
+        write_u32(file, raw_live)?;
+        for &token in &self.tokens {
+            write_u32(file, token as u32)?;
+        }
+        for &logit in &self.logits {
+            write_u32(file, logit.to_bits())?;
+        }
+        for layer in &self.session.layers {
+            write_u32(
+                file,
+                layer.compression.as_ref().map_or(0, |state| state.rows),
+            )?;
+        }
+        for layer in &self.session.layers {
+            write_u32(file, layer.indexer.as_ref().map_or(0, |state| state.rows))?;
+        }
+
+        let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
+        let raw_first = token_count - raw_live;
+        for layer in &self.session.layers {
+            for position in raw_first..token_count {
+                let physical = position % self.session.raw_cap;
+                write_buffer(
+                    file,
+                    &layer.raw_cache,
+                    u64::from(physical) * shape.head_dim * 4,
+                    shape.head_dim * 4,
+                    &mut chunk,
+                )?;
+            }
+            if let Some(state) = &layer.compression {
+                write_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.head_dim * 2,
+                    &mut chunk,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
+                write_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?;
+                write_buffer(file, &state.state_score, 0, bytes, &mut chunk)?;
+            }
+            if let Some(state) = &layer.indexer {
+                write_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.indexer_head_dim * 4,
+                    &mut chunk,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
+                write_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?;
+                write_buffer(file, &state.state_score, 0, bytes, &mut chunk)?;
+            }
+        }
+        Ok(())
+    }
+
+    fn read_checkpoint(&mut self, file: &mut File) -> Result<(), String> {
+        let mut magic = [0; 8];
+        file.read_exact(&mut magic)
+            .map_err(|error| error.to_string())?;
+        if &magic != CHECKPOINT_MAGIC {
+            return Err("KV checkpoint has an invalid signature".into());
+        }
+        let shape = self.model.shape;
+        let header = [
+            CHECKPOINT_VERSION,
+            self.session.context,
+            self.session.raw_cap,
+            shape.layers,
+            shape.head_dim as u32,
+            shape.indexer_head_dim as u32,
+            shape.vocab as u32,
+            u32::from(self.quality),
+        ];
+        for expected in header {
+            if read_u32(file)? != expected {
+                return Err("KV checkpoint does not match the current executor".into());
+            }
+        }
+        if read_u64(file)? != self.model.main.len() {
+            return Err("KV checkpoint was written for a different model".into());
+        }
+        if read_u64(file)? != self.model_modified.0 || read_u32(file)? != self.model_modified.1 {
+            return Err("KV checkpoint model file has changed".into());
+        }
+        for weight in [self.weights.token_embedding, self.weights.output] {
+            if read_u64(file)? != weight.offset
+                || read_u64(file)? != weight.bytes
+                || read_u32(file)? != weight.kind
+            {
+                return Err("KV checkpoint was written for a different model layout".into());
+            }
+        }
+        let mut checkpoint_tag = [0; 32];
+        file.read_exact(&mut checkpoint_tag)
+            .map_err(|error| error.to_string())?;
+        let token_count = read_u32(file)?;
+        let raw_live = read_u32(file)?;
+        if token_count > self.session.context || raw_live != token_count.min(self.session.raw_cap) {
+            return Err("KV checkpoint token count is invalid".into());
+        }
+        let mut tokens = Vec::with_capacity(token_count as usize);
+        for _ in 0..token_count {
+            let token = read_u32(file)?;
+            if u64::from(token) >= shape.vocab {
+                return Err("KV checkpoint contains an invalid token".into());
+            }
+            tokens.push(token as i32);
+        }
+        let mut logits = Vec::with_capacity(shape.vocab as usize);
+        for _ in 0..shape.vocab {
+            logits.push(f32::from_bits(read_u32(file)?));
+        }
+        let mut compressed_rows = Vec::with_capacity(shape.layers as usize);
+        let mut indexer_rows = Vec::with_capacity(shape.layers as usize);
+        for layer in 0..shape.layers {
+            let rows = read_u32(file)?;
+            let ratio = compression_ratio(layer);
+            if rows != token_count.checked_div(ratio).unwrap_or(0) {
+                return Err("KV checkpoint compressed row count is invalid".into());
+            }
+            compressed_rows.push(rows);
+        }
+        for layer in 0..shape.layers {
+            let rows = read_u32(file)?;
+            let expected = if compression_ratio(layer) == 4 {
+                token_count / 4
+            } else {
+                0
+            };
+            if rows != expected {
+                return Err("KV checkpoint indexer row count is invalid".into());
+            }
+            indexer_rows.push(rows);
+        }
+
+        self.reset()?;
+        let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
+        let raw_first = token_count - raw_live;
+        for (index, layer) in self.session.layers.iter_mut().enumerate() {
+            for position in raw_first..token_count {
+                let physical = position % self.session.raw_cap;
+                read_buffer(
+                    file,
+                    &layer.raw_cache,
+                    u64::from(physical) * shape.head_dim * 4,
+                    shape.head_dim * 4,
+                    &mut chunk,
+                )?;
+            }
+            if let Some(state) = &mut layer.compression {
+                state.rows = compressed_rows[index];
+                read_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.head_dim * 2,
+                    &mut chunk,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.head_dim);
+                read_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?;
+                read_buffer(file, &state.state_score, 0, bytes, &mut chunk)?;
+            }
+            if let Some(state) = &mut layer.indexer {
+                state.rows = indexer_rows[index];
+                read_buffer(
+                    file,
+                    &state.cache,
+                    0,
+                    u64::from(state.rows) * shape.indexer_head_dim * 4,
+                    &mut chunk,
+                )?;
+                let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim);
+                read_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?;
+                read_buffer(file, &state.state_score, 0, bytes, &mut chunk)?;
+            }
+        }
+        let mut trailing = [0];
+        if file
+            .read(&mut trailing)
+            .map_err(|error| error.to_string())?
+            != 0
+        {
+            self.reset()?;
+            return Err("KV checkpoint has trailing data".into());
+        }
+        self.session.position = token_count;
+        self.tokens = tokens;
+        self.logits = logits;
+        self.checkpoint_tag = checkpoint_tag;
         Ok(())
     }
 
@@ -1892,6 +2210,71 @@ fn compression_ratio(layer: u32) -> u32 {
     }
 }
 
+fn compressor_state_bytes(ratio: u32, head_dim: u64) -> u64 {
+    let coefficient = if ratio == 4 { 2 } else { 1 };
+    coefficient * head_dim * coefficient * u64::from(ratio) * 4
+}
+
+fn write_buffer(
+    file: &mut File,
+    buffer: &Buffer,
+    mut offset: u64,
+    mut bytes: u64,
+    chunk: &mut [u8],
+) -> Result<(), String> {
+    while bytes != 0 {
+        let length = bytes.min(chunk.len() as u64) as usize;
+        buffer.read(offset, &mut chunk[..length])?;
+        file.write_all(&chunk[..length])
+            .map_err(|error| error.to_string())?;
+        offset += length as u64;
+        bytes -= length as u64;
+    }
+    Ok(())
+}
+
+fn read_buffer(
+    file: &mut File,
+    buffer: &Buffer,
+    mut offset: u64,
+    mut bytes: u64,
+    chunk: &mut [u8],
+) -> Result<(), String> {
+    while bytes != 0 {
+        let length = bytes.min(chunk.len() as u64) as usize;
+        file.read_exact(&mut chunk[..length])
+            .map_err(|error| error.to_string())?;
+        buffer.write(offset, &chunk[..length])?;
+        offset += length as u64;
+        bytes -= length as u64;
+    }
+    Ok(())
+}
+
+fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
+    file.write_all(&value.to_le_bytes())
+        .map_err(|error| error.to_string())
+}
+
+fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
+    file.write_all(&value.to_le_bytes())
+        .map_err(|error| error.to_string())
+}
+
+fn read_u32(file: &mut File) -> Result<u32, String> {
+    let mut bytes = [0; 4];
+    file.read_exact(&mut bytes)
+        .map_err(|error| error.to_string())?;
+    Ok(u32::from_le_bytes(bytes))
+}
+
+fn read_u64(file: &mut File) -> Result<u64, String> {
+    let mut bytes = [0; 8];
+    file.read_exact(&mut bytes)
+        .map_err(|error| error.to_string())?;
+    Ok(u64::from_le_bytes(bytes))
+}
+
 fn encode_output(
     s: &Scratch,
     w: &Weights,

src/engine/tokenizer.rs

@@ -209,6 +209,24 @@ impl Tokenizer {
         output
     }
 
+    pub(super) fn encode_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec<i32> {
+        let mut output = Vec::new();
+        if self.family == ModelFamily::DeepSeek {
+            output.push(self.eos);
+        }
+        output.push(self.user);
+        output.extend(self.tokenize(prompt));
+        output.push(self.assistant);
+        if reasoning != ReasoningMode::Direct {
+            output.push(self.think_start);
+        } else if self.family == ModelFamily::Glm {
+            output.extend([self.think_start, self.think_end]);
+        } else {
+            output.push(self.think_end);
+        }
+        output
+    }
+
     pub(super) fn eos(&self) -> i32 {
         self.eos
     }