full chat implemented

src/app.rs

@@ -457,6 +457,7 @@ enum GenerationCommand {
         engine: crate::settings::EngineSettings,
         turn: crate::settings::TurnSettings,
         messages: Vec<ChatTurn>,
+        checkpoint: PathBuf,
         idle_timeout: Duration,
         cancel: Arc<AtomicBool>,
     },
@@ -972,9 +973,23 @@ impl App {
                         Some("Stop the active generation before deleting its project.".into());
                     return Task::none();
                 }
+                let checkpoint_ids = self
+                    .projects
+                    .iter()
+                    .find(|item| item.project.id == project_id)
+                    .map(|item| {
+                        item.sessions
+                            .iter()
+                            .map(|session| session.id)
+                            .collect::<Vec<_>>()
+                    })
+                    .unwrap_or_default();
                 if let Some(database) = &mut self.database {
                     match database.delete_project(project_id) {
                         Ok(()) => {
+                            for session_id in checkpoint_ids {
+                                let _ = fs::remove_file(session_checkpoint_path(session_id));
+                            }
                             if self.selected_project == Some(project_id) {
                                 self.selected_project = None;
                                 self.selected_session = None;
@@ -1036,6 +1051,7 @@ impl App {
                 if let Some(database) = &mut self.database {
                     match database.delete_session(session_id) {
                         Ok(()) => {
+                            let _ = fs::remove_file(session_checkpoint_path(session_id));
                             if self.selected_session == Some(session_id) {
                                 self.selected_session = None;
                                 self.conversation.clear();
@@ -1465,6 +1481,7 @@ impl App {
                 engine: effective.engine,
                 turn: effective.turn,
                 messages,
+                checkpoint: session_checkpoint_path(session_id),
                 idle_timeout,
                 cancel: Arc::clone(&cancel),
             };
@@ -1594,6 +1611,7 @@ fn spawn_generation_worker() -> Result<GenerationWorker, String> {
                         engine,
                         turn,
                         messages,
+                        checkpoint,
                         idle_timeout: requested_timeout,
                         cancel,
                     }) => {
@@ -1614,6 +1632,7 @@ fn spawn_generation_worker() -> Result<GenerationWorker, String> {
                         }
                         if let Some((_, generator)) = &mut loaded {
                             let result = generator.generate(
+                                &checkpoint,
                                 &messages,
                                 &turn,
                                 &cancel,
@@ -1687,6 +1706,12 @@ fn models_path() -> PathBuf {
     application_support_path().join("models")
 }
 
+fn session_checkpoint_path(session_id: i32) -> PathBuf {
+    application_support_path()
+        .join("kv-cache")
+        .join(format!("{session_id}.bin"))
+}
+
 pub(crate) fn app_icon() -> window::Icon {
     let decoder = png::Decoder::new(std::io::Cursor::new(include_bytes!(
         "../assets/app-icon.png"

src/engine.rs

@@ -7,7 +7,8 @@ use crate::model::ModelChoice;
 use crate::settings::TurnSettings;
 use crate::settings::{EngineSettings, ReasoningMode};
 use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value};
-use std::path::Path;
+use sha2::{Digest, Sha256};
+use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
 use tokenizer::Tokenizer;
 
@@ -253,6 +254,10 @@ impl Model {
             .encode_conversation(system, messages, reasoning)
     }
 
+    fn render_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec<i32> {
+        self.tokenizer.encode_continuation(prompt, reasoning)
+    }
+
     pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
         self.tokenizer.token_bytes(token)
     }
@@ -287,6 +292,7 @@ impl Model {
 #[cfg(target_os = "macos")]
 pub(crate) struct Generator {
     executor: metal::Executor,
+    checkpoint: Option<PathBuf>,
 }
 
 #[derive(Clone)]
@@ -315,24 +321,81 @@ impl Generator {
             settings.context_tokens.max(1) as u32,
             settings.execution.quality,
         )?;
-        Ok(Self { executor })
+        Ok(Self {
+            executor,
+            checkpoint: None,
+        })
     }
 
     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),
     ) -> Result<(), String> {
+        self.select_checkpoint(checkpoint)?;
+        let (generated, prompt_complete) =
+            self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
+        let mut completed = messages.to_vec();
+        completed.push(generated);
+        self.executor.save_checkpoint(
+            checkpoint,
+            if prompt_complete {
+                conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
+            } else {
+                [0; 32]
+            },
+        )
+    }
+
+    fn select_checkpoint(&mut self, checkpoint: &Path) -> Result<(), String> {
+        if self.checkpoint.as_deref() == Some(checkpoint) {
+            return Ok(());
+        }
         self.executor.reset()?;
-        progress(self.executor.position(), self.executor.context());
-        let tokens = self.executor.model().render_conversation(
-            &settings.system_prompt,
-            messages,
-            settings.reasoning_mode,
-        );
+        if self.executor.load_checkpoint(checkpoint).is_err() {
+            self.executor.reset()?;
+            let _ = std::fs::remove_file(checkpoint);
+        }
+        self.checkpoint = Some(checkpoint.to_owned());
+        Ok(())
+    }
+
+    fn generate_inner(
+        &mut self,
+        messages: &[ChatTurn],
+        settings: &TurnSettings,
+        cancelled: &AtomicBool,
+        emit: &mut impl FnMut(bool, String),
+        progress: &mut impl FnMut(u32, u32),
+    ) -> Result<(ChatTurn, 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),
+                );
+                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());
         }
@@ -343,12 +406,20 @@ impl Generator {
                 tokens.len()
             ));
         }
+        let reused = self.executor.align_prompt(&tokens)?;
+        progress(self.executor.position(), self.executor.context());
         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,
+            reasoning: reasoning.then(String::new),
+            reasoning_complete: !reasoning,
+            content: String::new(),
+        };
         let prompt_tokens = tokens.len();
-        for (index, token) in tokens.into_iter().enumerate() {
+        for (index, token) in tokens.into_iter().enumerate().skip(reused) {
             if cancelled.load(Ordering::Relaxed) {
-                return Ok(());
+                return Ok((generated, false));
             }
             self.executor.eval(token)?;
             if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens {
@@ -361,7 +432,7 @@ impl Generator {
             .min((max_context - self.executor.position() as usize) as i32)
         {
             if cancelled.load(Ordering::Relaxed) {
-                return Ok(());
+                return Ok((generated, true));
             }
             let token = sample(
                 self.executor.logits(),
@@ -375,23 +446,65 @@ impl Generator {
                 .model()
                 .is_stop_token_for_reasoning(token, settings.reasoning_mode)
             {
-                return Ok(());
+                return Ok((generated, true));
             }
             if self.executor.model().is_think_start_token(token) {
                 reasoning = true;
+                generated.reasoning.get_or_insert_default();
             } else if self.executor.model().is_think_end_token(token) {
                 reasoning = false;
+                generated.reasoning_complete = true;
                 emit(false, String::new());
             } else if let Some(bytes) = self.executor.model().token_bytes(token) {
-                emit(reasoning, String::from_utf8_lossy(&bytes).into_owned());
+                let content = String::from_utf8_lossy(&bytes).into_owned();
+                if reasoning {
+                    generated
+                        .reasoning
+                        .get_or_insert_default()
+                        .push_str(&content);
+                } else {
+                    generated.reasoning_complete = true;
+                    generated.content.push_str(&content);
+                }
+                emit(reasoning, content);
             }
             self.executor.eval(token)?;
             progress(self.executor.position(), self.executor.context());
         }
-        Ok(())
+        Ok((generated, true))
     }
 }
 
+#[cfg(target_os = "macos")]
+fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
+    fn text(hasher: &mut Sha256, value: &str) {
+        hasher.update((value.len() as u64).to_le_bytes());
+        hasher.update(value.as_bytes());
+    }
+
+    let mut hasher = Sha256::new();
+    hasher.update(b"DS4Server chat checkpoint v1");
+    text(&mut hasher, system);
+    hasher.update([match reasoning {
+        ReasoningMode::Direct => 0,
+        ReasoningMode::High => 1,
+        ReasoningMode::Max => 2,
+    }]);
+    for message in messages {
+        hasher.update([u8::from(message.user)]);
+        match &message.reasoning {
+            Some(reasoning) => {
+                hasher.update([1]);
+                text(&mut hasher, reasoning);
+            }
+            None => hasher.update([0]),
+        }
+        hasher.update([u8::from(message.reasoning_complete)]);
+        text(&mut hasher, &message.content);
+    }
+    hasher.finalize().into()
+}
+
 #[cfg(target_os = "macos")]
 fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, rng: &mut Rng) -> i32 {
     if temperature <= 0.0 {
@@ -490,6 +603,29 @@ mod sampling_tests {
         assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, &mut rng), 1);
     }
 
+    #[test]
+    fn checkpoint_tag_covers_the_canonical_chat_state() {
+        let messages = [ChatTurn {
+            user: true,
+            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)
+        );
+    }
+
     #[test]
     #[ignore = "requires the 80 GiB Flash checkpoint and Apple Metal"]
     fn metal_executes_real_flash_token() {
@@ -506,7 +642,7 @@ mod sampling_tests {
         assert_eq!(tokens.len(), 10);
         let mut executor = metal::Executor::open(model, 32_768, false).unwrap();
         assert_eq!(executor.context(), 32_768);
-        for token in tokens {
+        for &token in &tokens {
             executor.eval(token).unwrap();
         }
         assert!(executor.logits().iter().all(|logit| logit.is_finite()));
@@ -524,6 +660,40 @@ mod sampling_tests {
         );
         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]).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).unwrap();
+        assert!(restored.load_checkpoint(&checkpoint).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();
     }
 }
 
@@ -1453,6 +1623,10 @@ mod tests {
                 0, 128_803, 19_923, 128_804, 128_822, 19_923, 1, 128_803, 19_923, 128_804, 128_822,
             ]
         );
+        assert_eq!(
+            model.render_continuation("Hello", ReasoningMode::Direct),
+            [1, 128_803, 19_923, 128_804, 128_822]
+        );
         assert_eq!(
             model.render_conversation(
                 "",

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
     }