306 lines
10 KiB
Rust
306 lines
10 KiB
Rust
//! 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 = true;
|
||
|
||
// 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 {}
|