feat: first cut at actual token generation and model loading

src/engine.rs

@@ -1,12 +1,19 @@
 mod gguf;
+#[cfg(target_os = "macos")]
+mod metal;
 mod tokenizer;
 
 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 std::sync::atomic::{AtomicBool, Ordering};
 use tokenizer::Tokenizer;
 
+#[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];
@@ -236,6 +243,16 @@ impl Model {
         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)
+    }
+
     pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
         self.tokenizer.token_bytes(token)
     }
@@ -253,6 +270,230 @@ impl Model {
     }
 }
 
+#[cfg(target_os = "macos")]
+pub(crate) struct Generator {
+    executor: metal::Executor,
+}
+
+#[derive(Clone)]
+pub(crate) struct ChatTurn {
+    pub(crate) user: bool,
+    pub(crate) reasoning: bool,
+    pub(crate) content: String,
+}
+
+#[cfg(target_os = "macos")]
+impl Generator {
+    pub(crate) fn open(settings: &EngineSettings) -> 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,
+        )?;
+        Ok(Self { executor })
+    }
+
+    pub(crate) fn generate(
+        &mut self,
+        messages: &[ChatTurn],
+        settings: &TurnSettings,
+        cancelled: &AtomicBool,
+        mut emit: impl FnMut(String),
+    ) -> Result<(), String> {
+        self.executor.reset()?;
+        let 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!(
+                "the conversation uses {} tokens; the current Rust attention port supports fewer than {max_context}",
+                tokens.len()
+            ));
+        }
+        let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
+        for token in tokens {
+            if cancelled.load(Ordering::Relaxed) {
+                return Ok(());
+            }
+            self.executor.eval(token)?;
+        }
+        for _ in 0..settings
+            .max_generated_tokens
+            .max(0)
+            .min((max_context - self.executor.position() as usize) as i32)
+        {
+            if cancelled.load(Ordering::Relaxed) {
+                return Ok(());
+            }
+            let token = sample(
+                self.executor.logits(),
+                settings.temperature,
+                settings.top_p,
+                settings.min_p,
+                &mut rng,
+            );
+            if self.executor.model().is_stop_token(token) {
+                return Ok(());
+            }
+            if let Some(bytes) = self.executor.model().token_bytes(token) {
+                emit(String::from_utf8_lossy(&bytes).into_owned());
+            }
+            self.executor.eval(token)?;
+        }
+        Ok(())
+    }
+}
+
+#[cfg(target_os = "macos")]
+fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, 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 {
+        probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
+        let total: f32 = logits
+            .iter()
+            .filter(|logit| logit.is_finite())
+            .map(|logit| ((*logit - maximum) / temperature).exp())
+            .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, &mut rng), 1);
+    }
+
+    #[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, 128, false).unwrap();
+        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);
+    }
+}
+
 pub(crate) fn validate_model_artifact(
     path: &Path,
     expected: ModelChoice,
@@ -1144,6 +1385,32 @@ mod tests {
             model.render_prompt("", "Hello", ReasoningMode::Direct),
             [0, 128_803, 19_923, 128_804, 128_822]
         );
+        assert_eq!(
+            model.render_conversation(
+                "",
+                &[
+                    ChatTurn {
+                        user: true,
+                        reasoning: false,
+                        content: "Hello".into(),
+                    },
+                    ChatTurn {
+                        user: false,
+                        reasoning: false,
+                        content: "Hello".into(),
+                    },
+                    ChatTurn {
+                        user: true,
+                        reasoning: false,
+                        content: "Hello".into(),
+                    },
+                ],
+                ReasoningMode::Direct,
+            ),
+            [
+                0, 128_803, 19_923, 128_804, 128_822, 19_923, 128_803, 19_923, 128_804, 128_822,
+            ]
+        );
     }
 
     #[test]