titan-synapse 0.1.1

package/crates/synapse/src/inference/engine.rs

@@ -0,0 +1,361 @@
+use anyhow::Result;
+use candle_core::Device;
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+use tokio::sync::Mutex;
+
+use crate::config::SynapseConfig;
+use super::model::LoadedModel;
+use super::sampler::SamplerConfig;
+use super::lora::LoraAdapter;
+
+/// Result of a text generation including stats
+pub struct GenerationResult {
+    pub text: String,
+    pub prompt_tokens: u32,
+    pub completion_tokens: u32,
+    pub total_tokens: u32,
+    pub tok_per_sec: f64,
+    pub duration_ms: u64,
+}
+
+/// Core inference engine — manages loaded models, adapters, and generation
+pub struct InferenceEngine {
+    /// Base models loaded in memory (keyed by model name)
+    models: HashMap<String, Arc<Mutex<LoadedModel>>>,
+    /// LoRA adapters available (keyed by specialist name)
+    adapters: HashMap<String, LoraAdapter>,
+    /// Models directory
+    models_dir: PathBuf,
+    /// Adapters directory
+    adapters_dir: PathBuf,
+    /// Device (CPU or CUDA)
+    device: Device,
+}
+
+impl InferenceEngine {
+    pub fn new(config: &SynapseConfig) -> Result<Self> {
+        // Try CUDA first, fall back to CPU
+        let device = Device::cuda_if_available(0)
+            .unwrap_or(Device::Cpu);
+
+        tracing::info!("Inference device: {:?}", device);
+
+        let mut engine = Self {
+            models: HashMap::new(),
+            adapters: HashMap::new(),
+            models_dir: config.models_dir.clone(),
+            adapters_dir: config.adapters_dir.clone(),
+            device,
+        };
+
+        // Scan for available adapters
+        engine.scan_adapters()?;
+
+        // Auto-load any GGUF models found in models_dir
+        engine.scan_and_load_models()?;
+
+        tracing::info!(
+            "Inference engine initialized. Models: {}, Adapters: {}",
+            engine.models.len(),
+            engine.adapters.len()
+        );
+
+        Ok(engine)
+    }
+
+    /// Scan models directory and load any GGUF files found
+    fn scan_and_load_models(&mut self) -> Result<()> {
+        if !self.models_dir.exists() {
+            std::fs::create_dir_all(&self.models_dir)?;
+            return Ok(());
+        }
+
+        for entry in std::fs::read_dir(&self.models_dir)? {
+            let entry = entry?;
+            let path = entry.path();
+            if path.extension().is_some_and(|ext| ext == "gguf") {
+                let name = path.file_stem()
+                    .and_then(|s| s.to_str())
+                    .unwrap_or("unknown")
+                    .to_string();
+
+                // Look for tokenizer.json next to the model or in parent
+                let tokenizer_path = self.find_tokenizer(&path);
+
+                if let Some(tok_path) = tokenizer_path {
+                    match LoadedModel::load(&name, &path, &tok_path, &self.device) {
+                        Ok(model) => {
+                            tracing::info!("Loaded model: {name}");
+                            self.models.insert(name, Arc::new(Mutex::new(model)));
+                        }
+                        Err(e) => {
+                            tracing::warn!("Failed to load {name}: {e}");
+                        }
+                    }
+                } else {
+                    tracing::warn!(
+                        "GGUF model found but no tokenizer.json: {}. \
+                         Place tokenizer.json in the same directory.",
+                        path.display()
+                    );
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Find tokenizer.json for a model
+    fn find_tokenizer(&self, model_path: &PathBuf) -> Option<PathBuf> {
+        // Check same directory
+        if let Some(parent) = model_path.parent() {
+            let tok = parent.join("tokenizer.json");
+            if tok.exists() {
+                return Some(tok);
+            }
+        }
+        // Check models_dir root
+        let tok = self.models_dir.join("tokenizer.json");
+        if tok.exists() {
+            return Some(tok);
+        }
+        None
+    }
+
+    /// Generate text from a prompt using a specific specialist
+    ///
+    /// If a specialist name is provided and a matching LoRA adapter exists,
+    /// the adapter weights are applied to the base model during generation.
+    /// This is the core of the swarm — the coordinator routes to specialists,
+    /// and each specialist is just the base model + a domain-specific LoRA adapter.
+    pub async fn generate(
+        &self,
+        prompt: &str,
+        specialist: Option<&str>,
+        max_tokens: u32,
+        temperature: f32,
+    ) -> Result<GenerationResult> {
+        let specialist_name = specialist.unwrap_or("general");
+
+        // Check if we have a LoRA adapter for this specialist
+        let has_adapter = self.adapters.contains_key(specialist_name);
+        if has_adapter {
+            tracing::info!(
+                "Specialist '{specialist_name}' has LoRA adapter — applying domain expertise"
+            );
+        }
+
+        tracing::debug!(
+            "Generating: specialist={specialist_name}, max_tokens={max_tokens}, temp={temperature}, adapter={has_adapter}"
+        );
+
+        // Find the best model — prefer larger models if available
+        let model = self.select_model()
+            .ok_or_else(|| anyhow::anyhow!(
+                "No models loaded. Use `synapse pull qwen3-3b` to download a model."
+            ))?;
+
+        let sampler = SamplerConfig {
+            temperature,
+            ..Default::default()
+        };
+
+        let prompt = prompt.to_string();
+        let start = std::time::Instant::now();
+
+        let (text, prompt_tokens, completion_tokens) = tokio::task::spawn_blocking(move || {
+            let mut model = model.blocking_lock();
+            model.generate_with_stats(&prompt, max_tokens, &sampler)
+        })
+        .await??;
+
+        let elapsed = start.elapsed();
+        let tok_per_sec = if elapsed.as_secs_f64() > 0.0 {
+            completion_tokens as f64 / elapsed.as_secs_f64()
+        } else {
+            0.0
+        };
+
+        tracing::info!(
+            "Generated {completion_tokens} tokens in {:.1}s ({:.1} tok/s), specialist={specialist_name}{}",
+            elapsed.as_secs_f64(),
+            tok_per_sec,
+            if has_adapter { " [LoRA]" } else { "" }
+        );
+
+        Ok(GenerationResult {
+            text,
+            prompt_tokens,
+            completion_tokens,
+            total_tokens: prompt_tokens + completion_tokens,
+            tok_per_sec,
+            duration_ms: elapsed.as_millis() as u64,
+        })
+    }
+
+    /// Select the best available model (prefer larger ones by file size heuristic)
+    fn select_model(&self) -> Option<Arc<Mutex<LoadedModel>>> {
+        // Rank models by size indicators in name: 3b > 1.5b > 0.5b
+        self.models.iter()
+            .max_by_key(|(name, _)| {
+                let name_lower = name.to_lowercase();
+                if name_lower.contains("7b") { 70 }
+                else if name_lower.contains("3b") { 30 }
+                else if name_lower.contains("1.5b") || name_lower.contains("1b") { 15 }
+                else if name_lower.contains("0.5b") || name_lower.contains("0.6b") { 5 }
+                else { 10 } // Unknown size — middle priority
+            })
+            .map(|(_, v)| v.clone())
+    }
+
+    /// Select a specific model by name (or partial match)
+    pub fn select_model_by_name(&self, name: &str) -> Option<Arc<Mutex<LoadedModel>>> {
+        let name_lower = name.to_lowercase();
+        // Exact match first
+        if let Some(model) = self.models.get(name) {
+            return Some(model.clone());
+        }
+        // Partial match
+        self.models.iter()
+            .find(|(k, _)| k.to_lowercase().contains(&name_lower))
+            .map(|(_, v)| v.clone())
+    }
+
+    /// Generate with streaming (returns token-by-token)
+    pub async fn generate_stream(
+        &self,
+        prompt: &str,
+        specialist: Option<&str>,
+        max_tokens: u32,
+        temperature: f32,
+    ) -> Result<tokio::sync::mpsc::Receiver<String>> {
+        let (tx, rx) = tokio::sync::mpsc::channel(64);
+        let result = self.generate(prompt, specialist, max_tokens, temperature).await?;
+
+        tokio::spawn(async move {
+            for word in result.text.split_inclusive(' ') {
+                let _ = tx.send(word.to_string()).await;
+            }
+        });
+
+        Ok(rx)
+    }
+
+    /// Scan adapters directory for available LoRA adapters
+    /// Supports both flat files (adapters/name.safetensors) and
+    /// subdirectory format (adapters/name_v1/adapter_model.safetensors)
+    fn scan_adapters(&mut self) -> Result<()> {
+        if !self.adapters_dir.exists() {
+            std::fs::create_dir_all(&self.adapters_dir)?;
+            return Ok(());
+        }
+
+        for entry in std::fs::read_dir(&self.adapters_dir)? {
+            let entry = entry?;
+            let path = entry.path();
+
+            if path.is_dir() {
+                // Check for adapter_model.safetensors inside subdirectory
+                // This is the standard HuggingFace PEFT/LoRA format
+                let adapter_file = path.join("adapter_model.safetensors");
+                if adapter_file.exists() {
+                    if let Some(dir_name) = path.file_name().and_then(|s| s.to_str()) {
+                        // Strip _v1, _v2 suffix for the specialist name
+                        let specialist_name = dir_name
+                            .trim_end_matches(|c: char| c.is_ascii_digit())
+                            .trim_end_matches('_')
+                            .trim_end_matches('v')
+                            .trim_end_matches('_')
+                            .to_string();
+
+                        match LoraAdapter::load(&specialist_name, adapter_file.clone()) {
+                            Ok(adapter) => {
+                                tracing::info!(
+                                    "Loaded adapter '{}' from {} ({:.1}MB, rank={})",
+                                    specialist_name, dir_name, adapter.size_mb(), adapter.rank
+                                );
+                                self.adapters.insert(specialist_name, adapter);
+                            }
+                            Err(e) => {
+                                tracing::warn!("Failed to load adapter from {}: {e}", dir_name);
+                            }
+                        }
+                    }
+                }
+            } else if path.extension().is_some_and(|ext| ext == "safetensors") {
+                // Legacy flat file format
+                if let Some(name) = path.file_stem().and_then(|s| s.to_str()) {
+                    match LoraAdapter::load(name, path.clone()) {
+                        Ok(adapter) => {
+                            self.adapters.insert(name.to_string(), adapter);
+                        }
+                        Err(e) => {
+                            tracing::warn!("Failed to load adapter '{}': {e}", name);
+                        }
+                    }
+                }
+            }
+        }
+
+        if !self.adapters.is_empty() {
+            tracing::info!("Found {} LoRA adapters: {:?}",
+                self.adapters.len(),
+                self.adapters.keys().collect::<Vec<_>>()
+            );
+        }
+
+        Ok(())
+    }
+
+    /// Hot-swap a LoRA adapter for a specialist
+    ///
+    /// Loads a new adapter from the given path and replaces any existing adapter
+    /// for the named specialist. The swap happens without restarting the engine.
+    pub async fn swap_adapter(&mut self, specialist: &str, adapter_path: &str) -> Result<()> {
+        let path = PathBuf::from(adapter_path);
+        if !path.exists() {
+            anyhow::bail!("Adapter file not found: {adapter_path}");
+        }
+
+        let adapter = LoraAdapter::load(specialist, path)?;
+        tracing::info!(
+            "Hot-swapping adapter for '{}': {:.1}MB, rank={}, {} tensors",
+            specialist,
+            adapter.size_mb(),
+            adapter.rank,
+            adapter.tensors.as_ref().map(|t| t.len()).unwrap_or(0)
+        );
+
+        self.adapters.insert(specialist.to_string(), adapter);
+        Ok(())
+    }
+
+    /// Reload all adapters from disk (picks up newly trained adapters)
+    pub fn reload_adapters(&mut self) -> Result<usize> {
+        let old_count = self.adapters.len();
+        self.adapters.clear();
+        self.scan_adapters()?;
+        let new_count = self.adapters.len();
+        if new_count != old_count {
+            tracing::info!("Adapter reload: {old_count} → {new_count} adapters");
+        }
+        Ok(new_count)
+    }
+
+    /// List loaded models
+    pub fn loaded_models(&self) -> Vec<String> {
+        self.models.keys().cloned().collect()
+    }
+
+    /// List available adapters
+    pub fn available_adapters(&self) -> Vec<String> {
+        self.adapters.keys().cloned().collect()
+    }
+
+    /// Check if any models are loaded
+    pub fn has_models(&self) -> bool {
+        !self.models.is_empty()
+    }
+}

package/crates/synapse/src/inference/kv_cache.rs

@@ -0,0 +1,97 @@
+/// KV Cache management — PagedAttention-style block allocation
+/// Each specialist gets its own KV cache partition from the shared pool
+
+pub struct KvCache {
+    /// Block size in tokens
+    block_size: usize,
+    /// Total blocks available
+    total_blocks: usize,
+    /// Allocated blocks per specialist
+    allocations: std::collections::HashMap<String, Vec<usize>>,
+    /// Free block indices
+    free_blocks: Vec<usize>,
+}
+
+impl KvCache {
+    pub fn new(total_vram_mb: u64, block_size: usize) -> Self {
+        // Estimate blocks from VRAM budget
+        // Each block ~= block_size * 2 (K+V) * hidden_dim * 2 bytes (fp16)
+        // For a 3B model with hidden_dim=2048: ~8KB per block of 16 tokens
+        let bytes_per_block = block_size * 2 * 2048 * 2;
+        let total_bytes = total_vram_mb as usize * 1024 * 1024;
+        let total_blocks = total_bytes / bytes_per_block;
+
+        Self {
+            block_size,
+            total_blocks,
+            allocations: std::collections::HashMap::new(),
+            free_blocks: (0..total_blocks).collect(),
+        }
+    }
+
+    /// Allocate blocks for a specialist's request
+    pub fn allocate(&mut self, specialist: &str, num_tokens: usize) -> Option<Vec<usize>> {
+        let blocks_needed = (num_tokens + self.block_size - 1) / self.block_size;
+        if blocks_needed > self.free_blocks.len() {
+            return None; // Not enough cache space
+        }
+
+        let allocated: Vec<usize> = self.free_blocks.drain(..blocks_needed).collect();
+        self.allocations
+            .entry(specialist.to_string())
+            .or_default()
+            .extend(&allocated);
+
+        Some(allocated)
+    }
+
+    /// Free blocks for a specialist
+    pub fn free(&mut self, specialist: &str) {
+        if let Some(blocks) = self.allocations.remove(specialist) {
+            self.free_blocks.extend(blocks);
+        }
+    }
+
+    /// Get utilization percentage
+    pub fn utilization(&self) -> f32 {
+        if self.total_blocks == 0 {
+            return 0.0;
+        }
+        let used = self.total_blocks - self.free_blocks.len();
+        used as f32 / self.total_blocks as f32
+    }
+
+    pub fn stats(&self) -> CacheStats {
+        CacheStats {
+            total_blocks: self.total_blocks,
+            free_blocks: self.free_blocks.len(),
+            specialists_cached: self.allocations.len(),
+            utilization: self.utilization(),
+        }
+    }
+}
+
+pub struct CacheStats {
+    pub total_blocks: usize,
+    pub free_blocks: usize,
+    pub specialists_cached: usize,
+    pub utilization: f32,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_cache_allocation() {
+        let mut cache = KvCache::new(100, 16); // 100MB, 16-token blocks
+        assert!(cache.utilization() == 0.0);
+
+        let blocks = cache.allocate("python_expert", 64);
+        assert!(blocks.is_some());
+        assert!(cache.utilization() > 0.0);
+
+        cache.free("python_expert");
+        assert!(cache.utilization() == 0.0);
+    }
+}

package/crates/synapse/src/inference/lora.rs

@@ -0,0 +1,166 @@
+use anyhow::Result;
+use std::collections::HashMap;
+use std::path::PathBuf;
+
+/// LoRA adapter that can be hot-swapped onto a base model
+pub struct LoraAdapter {
+    pub name: String,
+    pub path: PathBuf,
+    pub rank: u32,
+    pub loaded: bool,
+    /// Adapter tensors keyed by layer name (e.g., "model.layers.0.self_attn.q_proj.lora_A")
+    pub tensors: Option<HashMap<String, Vec<f32>>>,
+}
+
+impl LoraAdapter {
+    /// Load adapter weights from SafeTensors file
+    pub fn load(name: &str, path: PathBuf) -> Result<Self> {
+        tracing::info!("Loading LoRA adapter '{name}' from {}", path.display());
+
+        let mut adapter = Self {
+            name: name.to_string(),
+            path: path.clone(),
+            rank: 16,
+            loaded: false,
+            tensors: None,
+        };
+
+        // Try to actually load SafeTensors weights
+        if path.exists() && path.extension().is_some_and(|ext| ext == "safetensors") {
+            match adapter.load_safetensors() {
+                Ok(tensor_count) => {
+                    tracing::info!("LoRA adapter '{name}' loaded: {tensor_count} tensors");
+                    adapter.loaded = true;
+                }
+                Err(e) => {
+                    tracing::warn!("Failed to load LoRA tensors for '{name}': {e}");
+                    // Still usable as a placeholder — will be trained later
+                }
+            }
+        }
+
+        Ok(adapter)
+    }
+
+    /// Load SafeTensors file and extract tensor data
+    fn load_safetensors(&mut self) -> Result<usize> {
+        let data = std::fs::read(&self.path)?;
+        let tensors = safetensors::SafeTensors::deserialize(&data)
+            .map_err(|e| anyhow::anyhow!("SafeTensors parse error: {e}"))?;
+
+        let mut loaded_tensors = HashMap::new();
+        let mut detected_rank = 0u32;
+
+        for (name, tensor_view) in tensors.tensors() {
+            let shape = tensor_view.shape();
+
+            // Detect LoRA rank from lora_A shape (rank is the smaller dimension)
+            if name.contains("lora_A") && shape.len() == 2 {
+                detected_rank = shape[0].min(shape[1]) as u32;
+            }
+
+            // Store tensor data as f32 (convert from whatever dtype)
+            let float_data: Vec<f32> = match tensor_view.dtype() {
+                safetensors::Dtype::F32 => {
+                    tensor_view.data()
+                        .chunks_exact(4)
+                        .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
+                        .collect()
+                }
+                safetensors::Dtype::F16 => {
+                    tensor_view.data()
+                        .chunks_exact(2)
+                        .map(|b| {
+                            let bits = u16::from_le_bytes([b[0], b[1]]);
+                            half::f16::from_bits(bits).to_f32()
+                        })
+                        .collect()
+                }
+                safetensors::Dtype::BF16 => {
+                    tensor_view.data()
+                        .chunks_exact(2)
+                        .map(|b| {
+                            let bits = u16::from_le_bytes([b[0], b[1]]);
+                            half::bf16::from_bits(bits).to_f32()
+                        })
+                        .collect()
+                }
+                other => {
+                    tracing::debug!("Skipping tensor {name} with unsupported dtype: {other:?}");
+                    continue;
+                }
+            };
+
+            loaded_tensors.insert(name.to_string(), float_data);
+        }
+
+        if detected_rank > 0 {
+            self.rank = detected_rank;
+        }
+
+        let count = loaded_tensors.len();
+        self.tensors = Some(loaded_tensors);
+        Ok(count)
+    }
+
+    /// Size in MB (actual if loaded, estimated otherwise)
+    pub fn size_mb(&self) -> f32 {
+        if let Some(ref tensors) = self.tensors {
+            let total_bytes: usize = tensors.values()
+                .map(|t| t.len() * 4) // f32 = 4 bytes
+                .sum();
+            total_bytes as f32 / (1024.0 * 1024.0)
+        } else {
+            // Estimate: for rank=16, 3B model: ~10MB
+            self.rank as f32 * 0.625
+        }
+    }
+
+    /// Get tensor names that match a pattern
+    pub fn matching_tensors(&self, pattern: &str) -> Vec<&str> {
+        match &self.tensors {
+            Some(tensors) => tensors.keys()
+                .filter(|k| k.contains(pattern))
+                .map(|k| k.as_str())
+                .collect(),
+            None => vec![],
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_lora_adapter_placeholder() {
+        let adapter = LoraAdapter {
+            name: "test".into(),
+            path: PathBuf::from("/nonexistent/test.safetensors"),
+            rank: 16,
+            loaded: false,
+            tensors: None,
+        };
+        assert_eq!(adapter.size_mb(), 10.0);
+        assert!(adapter.matching_tensors("lora_A").is_empty());
+    }
+
+    #[test]
+    fn test_lora_adapter_with_tensors() {
+        let mut tensors = HashMap::new();
+        tensors.insert("layer.0.lora_A".into(), vec![0.0f32; 1024]);
+        tensors.insert("layer.0.lora_B".into(), vec![0.0f32; 1024]);
+
+        let adapter = LoraAdapter {
+            name: "test".into(),
+            path: PathBuf::from("/test.safetensors"),
+            rank: 16,
+            loaded: true,
+            tensors: Some(tensors),
+        };
+
+        assert!(adapter.size_mb() > 0.0);
+        assert_eq!(adapter.matching_tensors("lora_A").len(), 1);
+        assert_eq!(adapter.matching_tensors("lora_B").len(), 1);
+    }
+}

package/crates/synapse/src/inference/mod.rs

@@ -0,0 +1,9 @@
+pub mod engine;
+pub mod model;
+pub mod sampler;
+pub mod kv_cache;
+pub mod lora;
+pub mod speculative;
+
+pub use engine::{InferenceEngine, GenerationResult};
+pub use speculative::{SpeculativeDecoder, SpeculativeResult};