chore: initialize autoresearch baseline

crates/inference-core/src/engine.rs

@@ -0,0 +1,305 @@
+//! InferenceEngine: loads a GGUF model via llama-sys FFI and manages the
+//! lifetime of `llama_model` + `llama_context`.
+//!
+//! Design decisions:
+//! - Single model, single context per engine instance (Phase 2 design).
+//! - `LlamaLib` is loaded once and stored in an `Arc` so future session spawning
+//!   can share it without re-loading the dynamic library.
+//! - All raw pointers are wrapped in RAII guards that call the matching
+//!   `llama_free` / `llama_model_free` on drop.
+//! - The engine is `Send` but NOT `Sync` – callers must coordinate access
+//!   (the scheduler in a later phase will own this exclusively).
+
+use std::ffi::CString;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use anyhow::{bail, Context, Result};
+use tracing::{info, warn};
+
+use llama_sys::{
+    llama_context, llama_model, LlamaLib, LLAMA_FLASH_ATTN_TYPE_DISABLED,
+    LLAMA_FLASH_ATTN_TYPE_ENABLED,
+};
+
+use crate::profile::RuntimeProfile;
+use crate::sampler::SamplerConfig;
+use crate::session::GenerationSession;
+
+// ─── Configuration ─────────────────────────────────────────────────────────
+
+/// Configuration for the inference engine.
+#[derive(Debug, Clone)]
+pub struct EngineConfig {
+    /// Path to the GGUF model file.
+    pub model_path: PathBuf,
+    /// Runtime performance profile (used to fill in defaults).
+    pub profile: RuntimeProfile,
+    /// How many GPU layers to offload (999 = all).
+    pub n_gpu_layers: i32,
+    /// Context window size in tokens.
+    pub n_ctx: u32,
+    /// Physical batch size (prompt processing).
+    pub n_batch: u32,
+    /// Micro-batch size (generation step).
+    pub n_ubatch: u32,
+    /// Number of CPU threads for generation.
+    pub n_threads: i32,
+    /// Number of CPU threads for batch/prefill.
+    pub n_threads_batch: i32,
+    /// Enable flash attention.
+    pub flash_attn: bool,
+    /// Use mmap for model loading (false = load fully into RAM/VRAM).
+    pub use_mmap: bool,
+    /// RNG seed for sampling (0 = random).
+    pub seed: u32,
+}
+
+impl Default for EngineConfig {
+    fn default() -> Self {
+        let prof = RuntimeProfile::Turbo4B;
+        let d = prof.defaults();
+        Self {
+            model_path: PathBuf::new(),
+            profile: prof,
+            n_gpu_layers: d.n_gpu_layers,
+            n_ctx: d.n_ctx,
+            n_batch: d.n_batch,
+            n_ubatch: d.n_ubatch,
+            n_threads: d.n_threads,
+            n_threads_batch: d.n_threads,
+            flash_attn: d.flash_attn,
+            use_mmap: false, // --no-mmap recommended for CUDA
+            seed: 0,
+        }
+    }
+}
+
+impl EngineConfig {
+    /// Build from a profile, applying its defaults.
+    pub fn from_profile(profile: RuntimeProfile, model_path: PathBuf) -> Self {
+        let d = profile.defaults();
+        Self {
+            model_path,
+            profile,
+            n_gpu_layers: d.n_gpu_layers,
+            n_ctx: d.n_ctx,
+            n_batch: d.n_batch,
+            n_ubatch: d.n_ubatch,
+            n_threads: d.n_threads,
+            n_threads_batch: d.n_threads,
+            flash_attn: d.flash_attn,
+            use_mmap: false,
+            seed: 0,
+        }
+    }
+}
+
+// ─── RAII wrappers ─────────────────────────────────────────────────────────
+
+/// Owns a `*mut llama_model` and frees it on drop.
+struct OwnedModel {
+    ptr: *mut llama_model,
+    lib: Arc<LlamaLib>,
+}
+
+impl Drop for OwnedModel {
+    fn drop(&mut self) {
+        if !self.ptr.is_null() {
+            unsafe { (self.lib.model_free)(self.ptr) };
+        }
+    }
+}
+
+// SAFETY: llama_model is reference-counted internally; we are the only owner.
+unsafe impl Send for OwnedModel {}
+
+/// Owns a `*mut llama_context` and frees it on drop.
+pub(crate) struct OwnedContext {
+    pub ptr: *mut llama_context,
+    pub lib: Arc<LlamaLib>,
+}
+
+impl Drop for OwnedContext {
+    fn drop(&mut self) {
+        if !self.ptr.is_null() {
+            unsafe { (self.lib.free)(self.ptr) };
+        }
+    }
+}
+
+// SAFETY: we enforce single-threaded access through InferenceEngine.
+unsafe impl Send for OwnedContext {}
+
+// ─── Engine ────────────────────────────────────────────────────────────────
+
+/// Core inference engine.
+///
+/// - Owns a loaded model and its associated context.
+/// - Provides `new_session()` to start a streaming generation.
+/// - NOT `Sync` – the scheduler must ensure single-threaded use.
+pub struct InferenceEngine {
+    pub config: EngineConfig,
+    pub(crate) lib: Arc<LlamaLib>,
+    _model: OwnedModel,
+    pub(crate) ctx: OwnedContext,
+}
+
+impl InferenceEngine {
+    /// Load a GGUF model and initialise a llama_context.
+    ///
+    /// Blocks until the model is fully loaded.
+    pub fn load(config: EngineConfig) -> Result<Self> {
+        if config.model_path.as_os_str().is_empty() {
+            bail!("EngineConfig.model_path is empty");
+        }
+        if !config.model_path.exists() {
+            bail!("Model file not found: {}", config.model_path.display());
+        }
+
+        info!(
+            model = %config.model_path.display(),
+            profile = ?config.profile,
+            n_gpu_layers = config.n_gpu_layers,
+            n_ctx = config.n_ctx,
+            n_batch = config.n_batch,
+            n_ubatch = config.n_ubatch,
+            flash_attn = config.flash_attn,
+            use_mmap = config.use_mmap,
+            "Loading model via llama-sys"
+        );
+
+        // 1. Load dynamic library
+        let lib = Arc::new(LlamaLib::load_auto().context("Failed to load llama.dll")?);
+
+        // 2. Backend init
+        unsafe {
+            (lib.backend_load_all)();
+            (lib.backend_init)();
+        }
+        info!("llama backend initialised");
+
+        // 3. Model params
+        let mut mparams = unsafe { (lib.model_default_params)() };
+        mparams.n_gpu_layers = config.n_gpu_layers;
+        mparams.use_mmap = config.use_mmap;
+
+        // 4. Load model from file
+        let model_path_c = CString::new(config.model_path.to_str().context("Non-UTF8 model path")?)
+            .context("CString::new failed")?;
+
+        let model_ptr = unsafe { (lib.model_load_from_file)(model_path_c.as_ptr(), mparams) };
+        if model_ptr.is_null() {
+            bail!("llama_model_load_from_file returned NULL – check model path and VRAM");
+        }
+
+        // Log model info
+        {
+            let mut desc = vec![0u8; 256];
+            unsafe {
+                (lib.model_desc)(model_ptr, desc.as_mut_ptr() as *mut _, desc.len());
+            }
+            let desc_str = std::ffi::CStr::from_bytes_until_nul(&desc)
+                .map(|s| s.to_string_lossy().into_owned())
+                .unwrap_or_else(|_| "<unknown>".into());
+            let n_params = unsafe { (lib.model_n_params)(model_ptr) };
+            let is_hybrid = unsafe { (lib.model_is_hybrid)(model_ptr) };
+            let is_recurrent = unsafe { (lib.model_is_recurrent)(model_ptr) };
+            info!(
+                desc = %desc_str,
+                params_b = n_params as f64 / 1e9,
+                is_hybrid,
+                is_recurrent,
+                "Model loaded"
+            );
+        }
+
+        let model = OwnedModel {
+            ptr: model_ptr,
+            lib: Arc::clone(&lib),
+        };
+
+        // 5. Context params
+        let mut cparams = unsafe { (lib.context_default_params)() };
+        cparams.n_ctx = config.n_ctx;
+        cparams.n_batch = config.n_batch;
+        cparams.n_ubatch = config.n_ubatch;
+        cparams.n_threads = config.n_threads;
+        cparams.n_threads_batch = config.n_threads_batch;
+        cparams.flash_attn_type = if config.flash_attn {
+            LLAMA_FLASH_ATTN_TYPE_ENABLED
+        } else {
+            LLAMA_FLASH_ATTN_TYPE_DISABLED
+        };
+        cparams.offload_kqv = true;
+        cparams.no_perf = false;
+
+        // 6. Init context
+        let ctx_ptr = unsafe { (lib.init_from_model)(model_ptr, cparams) };
+        if ctx_ptr.is_null() {
+            bail!("llama_init_from_model returned NULL – insufficient VRAM or bad params");
+        }
+
+        let actual_ctx = unsafe { (lib.n_ctx)(ctx_ptr) };
+        info!(n_ctx = actual_ctx, "Context initialised");
+
+        let ctx = OwnedContext {
+            ptr: ctx_ptr,
+            lib: Arc::clone(&lib),
+        };
+
+        Ok(Self {
+            config,
+            lib,
+            _model: model,
+            ctx,
+        })
+    }
+
+    /// Warm up: run a single dummy decode to pre-allocate CUDA buffers.
+    ///
+    /// Call once after `load()` before accepting real requests.
+    pub fn warmup(&mut self) -> Result<()> {
+        info!("Running engine warmup (dummy decode)…");
+
+        unsafe {
+            // Mark context as warmup so llama.cpp skips perf counters
+            if let Ok(f) = std::panic::catch_unwind(|| {}) {
+                let _ = f;
+            }
+            // set_warmup may not exist in all builds; call decode with 1 token
+            let dummy_token: llama_sys::llama_token = 1; // BOS-ish token
+            let mut tokens = vec![dummy_token];
+            let batch = (self.lib.batch_get_one)(tokens.as_mut_ptr(), 1);
+            let ret = (self.lib.decode)(self.ctx.ptr, batch);
+            if ret != 0 {
+                warn!(ret, "Warmup decode returned non-zero (non-fatal)");
+            }
+        }
+
+        info!("Warmup complete");
+        Ok(())
+    }
+
+    /// Create a new generation session for a single request.
+    ///
+    /// `prompt` is the fully-formatted string (system + user turns already joined).
+    /// The engine must not be accessed concurrently while the session is active.
+    pub fn new_session(
+        &mut self,
+        prompt: &str,
+        sampler_cfg: SamplerConfig,
+    ) -> Result<GenerationSession> {
+        let vocab = unsafe { (self.lib.model_get_vocab)(self._model.ptr) };
+        GenerationSession::with_vocab(
+            Arc::clone(&self.lib),
+            self.ctx.ptr,
+            vocab,
+            prompt,
+            sampler_cfg,
+        )
+    }
+}
+
+// SAFETY: we are the sole owner of the pointers; no concurrent access.
+unsafe impl Send for InferenceEngine {}