Initial ForgeLM import

crates/inference-core/src/engine.rs

@@ -0,0 +1,311 @@
+//! 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;
+            }
+            (self.lib.set_warmup)(self.ctx.ptr, true);
+            // 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);
+            (self.lib.set_warmup)(self.ctx.ptr, false);
+            if ret != 0 {
+                warn!(ret, "Warmup decode returned non-zero (non-fatal)");
+            }
+
+            let memory = (self.lib.get_memory)(self.ctx.ptr);
+            (self.lib.memory_clear)(memory, false);
+        }
+
+        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,
+            self.config.n_batch,
+            prompt,
+            sampler_cfg,
+        )
+    }
+}
+
+// SAFETY: we are the sole owner of the pointers; no concurrent access.
+unsafe impl Send for InferenceEngine {}

crates/inference-core/src/lib.rs

@@ -0,0 +1,11 @@
+// inference-core: engine abstraction over llama-sys
+
+pub mod engine;
+pub mod profile;
+pub mod sampler;
+pub mod session;
+
+pub use engine::{EngineConfig, InferenceEngine};
+pub use profile::RuntimeProfile;
+pub use sampler::SamplerConfig;
+pub use session::GenerationSession;

crates/inference-core/src/profile.rs

@@ -0,0 +1,59 @@
+/// Predefined runtime profiles matching the plan targets.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum RuntimeProfile {
+    /// Qwen3.5-4B – maximum TPS, text-only, flash attention, moderate context.
+    Turbo4B,
+    /// Qwen3.5-9B – balanced quality/speed, tuned KV cache.
+    Balanced9B,
+    /// CPU-only fallback for machines without a GPU.
+    CpuFallback,
+}
+
+impl RuntimeProfile {
+    /// Returns sensible defaults for each profile.
+    pub fn defaults(&self) -> ProfileDefaults {
+        match self {
+            Self::Turbo4B => ProfileDefaults {
+                n_gpu_layers: 999,
+                n_ctx: 4096,
+                n_batch: 512,
+                n_ubatch: 128,
+                flash_attn: true,
+                n_threads: 8,
+            },
+            Self::Balanced9B => ProfileDefaults {
+                n_gpu_layers: 999,
+                n_ctx: 2048,
+                n_batch: 256,
+                n_ubatch: 64,
+                flash_attn: true,
+                n_threads: 8,
+            },
+            Self::CpuFallback => ProfileDefaults {
+                n_gpu_layers: 0,
+                n_ctx: 2048,
+                n_batch: 128,
+                n_ubatch: 32,
+                flash_attn: false,
+                n_threads: num_cpus(),
+            },
+        }
+    }
+}
+
+fn num_cpus() -> i32 {
+    // Simple fallback; replace with `num_cpus` crate if needed.
+    std::thread::available_parallelism()
+        .map(|n| n.get() as i32)
+        .unwrap_or(4)
+}
+
+#[derive(Debug, Clone)]
+pub struct ProfileDefaults {
+    pub n_gpu_layers: i32,
+    pub n_ctx: u32,
+    pub n_batch: u32,
+    pub n_ubatch: u32,
+    pub flash_attn: bool,
+    pub n_threads: i32,
+}

crates/inference-core/src/sampler.rs

@@ -0,0 +1,37 @@
+/// Sampling hyper-parameters for a single generation request.
+#[derive(Debug, Clone)]
+pub struct SamplerConfig {
+    pub temperature: f32,
+    pub top_p: f32,
+    pub top_k: i32,
+    pub repeat_penalty: f32,
+    pub max_tokens: u32,
+    pub stop_sequences: Vec<String>,
+    pub seed: u32,
+}
+
+impl Default for SamplerConfig {
+    fn default() -> Self {
+        Self {
+            temperature: 0.7,
+            top_p: 0.9,
+            top_k: 40,
+            repeat_penalty: 1.1,
+            max_tokens: 512,
+            stop_sequences: vec![],
+            seed: 0,
+        }
+    }
+}
+
+impl SamplerConfig {
+    pub fn from_request(temp: f32, top_p: f32, max_tokens: u32, stop: Vec<String>) -> Self {
+        Self {
+            temperature: temp,
+            top_p,
+            max_tokens,
+            stop_sequences: stop,
+            ..Default::default()
+        }
+    }
+}

crates/inference-core/src/session.rs

@@ -0,0 +1,476 @@
+//! GenerationSession: single-request token-by-token generation loop.
+//!
+//! Flow:
+//!   1. Tokenize the prompt (including system + user messages).
+//!   2. Prefill: feed all prompt tokens through llama_decode() in one batch.
+//!   3. Generation loop: sample → decode → yield token text → check stop.
+//!
+//! The session borrows raw pointers from InferenceEngine and must not outlive it.
+//! In Phase 7+ this will be refactored to use a shared context pool.
+
+use std::ffi::{CStr, CString};
+use std::ops::Range;
+use std::sync::Arc;
+use std::time::Instant;
+
+use anyhow::{bail, Context, Result};
+use tracing::{debug, info};
+
+use llama_sys::{llama_context, llama_token, llama_vocab, LlamaLib};
+
+use crate::sampler::SamplerConfig;
+
+// ─── GenerationSession ─────────────────────────────────────────────────────
+
+pub struct GenerationSession {
+    lib: Arc<LlamaLib>,
+    ctx: *mut llama_context,
+    vocab: *const llama_vocab,
+    sampler: *mut llama_sys::llama_sampler,
+    stop_sequences: Vec<String>,
+
+    // generation state
+    n_decoded: i32, // number of tokens decoded so far (prefill + gen)
+    prompt_tokens: u32,
+    pub max_tokens: u32,
+    pub tokens_generated: u32,
+    pub finished: bool,
+
+    // timing
+    pub prefill_ms: Option<f64>,
+    gen_start: Option<Instant>,
+    pub tokens_per_sec: f64,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct PrefillChunkPlan {
+    range: Range<usize>,
+    logits_index: Option<usize>,
+}
+
+// SAFETY: pointers come from InferenceEngine which is Send.
+unsafe impl Send for GenerationSession {}
+
+impl GenerationSession {
+    /// Tokenize prompt and run prefill decode. Returns a session ready to generate.
+    pub fn new(
+        _lib: Arc<LlamaLib>,
+        _ctx: *mut llama_context,
+        _sampler_cfg: SamplerConfig,
+    ) -> Result<Self> {
+        // Get vocab from context's model (via a workaround: we pass vocab separately)
+        // For now we'll get it indirectly — see Engine::new_session which passes vocab.
+        // This constructor signature will be refined; use the vocab variant below.
+        bail!("Use GenerationSession::with_vocab() instead")
+    }
+
+    /// Full constructor called by InferenceEngine::new_session().
+    pub fn with_vocab(
+        lib: Arc<LlamaLib>,
+        ctx: *mut llama_context,
+        vocab: *const llama_vocab,
+        n_batch: u32,
+        prompt: &str,
+        sampler_cfg: SamplerConfig,
+    ) -> Result<Self> {
+        // ── 1. Tokenize prompt ─────────────────────────────────────────────
+        let prompt_c = CString::new(prompt).context("Invalid UTF-8 in prompt")?;
+
+        // First call: get token count
+        let ret = unsafe {
+            (lib.tokenize)(
+                vocab,
+                prompt_c.as_ptr(),
+                prompt.len() as i32,
+                std::ptr::null_mut(),
+                0,
+                true, // add_special (BOS)
+                true, // parse_special
+            )
+        };
+
+        let n_tokens_needed = if ret < 0 { -ret } else { ret };
+
+        if n_tokens_needed == 0 {
+            bail!("llama_tokenize failed to return token count");
+        }
+
+        let mut tokens: Vec<llama_token> = vec![0; n_tokens_needed as usize];
+        let n_tokens = unsafe {
+            (lib.tokenize)(
+                vocab,
+                prompt_c.as_ptr(),
+                prompt.len() as i32,
+                tokens.as_mut_ptr(),
+                tokens.len() as i32,
+                true,
+                true,
+            )
+        };
+
+        if n_tokens < 0 {
+            bail!(
+                "llama_tokenize (second call) failed (returned {})",
+                n_tokens
+            );
+        }
+        tokens.truncate(n_tokens as usize);
+
+        let n_prompt_tokens = tokens.len() as i32;
+        debug!(n_tokens = n_prompt_tokens, "Prompt tokenised");
+
+        // ── 2. Build sampler chain ─────────────────────────────────────────
+        let mut sc_params = unsafe { (lib.sampler_chain_default_params)() };
+        sc_params.no_perf = true;
+        let sampler_chain = unsafe { (lib.sampler_chain_init)(sc_params) };
+        if sampler_chain.is_null() {
+            bail!("sampler_chain_init returned NULL");
+        }
+
+        unsafe {
+            if sampler_cfg.temperature <= 0.0 || sampler_cfg.temperature < 1e-6 {
+                // Greedy
+                let s = (lib.sampler_init_greedy)();
+                (lib.sampler_chain_add)(sampler_chain, s);
+            } else {
+                // top-k → temp → top-p → dist (for randomness)
+                let s_topk = (lib.sampler_init_top_k)(sampler_cfg.top_k);
+                (lib.sampler_chain_add)(sampler_chain, s_topk);
+
+                let s_temp = (lib.sampler_init_temp)(sampler_cfg.temperature);
+                (lib.sampler_chain_add)(sampler_chain, s_temp);
+
+                let s_topp = (lib.sampler_init_top_p)(sampler_cfg.top_p, 1);
+                (lib.sampler_chain_add)(sampler_chain, s_topp);
+
+                let s_dist = (lib.sampler_init_dist)(sampler_cfg.seed);
+                (lib.sampler_chain_add)(sampler_chain, s_dist);
+            }
+        }
+
+        // ── 3. Prefill decode ──────────────────────────────────────────────
+        let prefill_start = Instant::now();
+
+        if let Err(err) = decode_prefill(&lib, ctx, &tokens, n_batch) {
+            unsafe { (lib.sampler_free)(sampler_chain) };
+            return Err(err);
+        }
+
+        let prefill_ms = prefill_start.elapsed().as_secs_f64() * 1000.0;
+        let prefill_tps = n_prompt_tokens as f64 / (prefill_ms / 1000.0);
+        info!(
+            "Prefill finished in {:.2} ms ({:.2} t/s)",
+            prefill_ms, prefill_tps
+        );
+
+        // create session state
+        Ok(Self {
+            lib,
+            ctx,
+            vocab,
+            sampler: sampler_chain, // use sampler_chain here
+            stop_sequences: sampler_cfg.stop_sequences.clone(),
+            n_decoded: n_prompt_tokens,
+            prompt_tokens: n_prompt_tokens as u32,
+            max_tokens: sampler_cfg.max_tokens,
+            tokens_generated: 0,
+            finished: false,
+            prefill_ms: Some(prefill_ms),
+            gen_start: None,
+            tokens_per_sec: 0.0,
+        })
+    }
+
+    /// Sample and decode the next token. Returns `None` when generation ends.
+    ///
+    /// Each call advances the context by one token.
+    pub fn next_token(&mut self) -> Option<String> {
+        if self.finished {
+            return None;
+        }
+        if self.tokens_generated >= self.max_tokens {
+            self.finished = true;
+            self.finalize_timing();
+            return None;
+        }
+
+        // Start timing on first generated token
+        if self.gen_start.is_none() {
+            self.gen_start = Some(Instant::now());
+        }
+
+        unsafe {
+            // Sample the next token
+            let new_token = (self.lib.sampler_sample)(self.sampler, self.ctx, -1);
+
+            // Check end-of-generation
+            if (self.lib.vocab_is_eog)(self.vocab, new_token) {
+                self.finished = true;
+                self.finalize_timing();
+                return None;
+            }
+
+            // Decode (advance context with the new token)
+            let mut token_arr = [new_token];
+            let batch = (self.lib.batch_get_one)(token_arr.as_mut_ptr(), 1);
+            let ret = (self.lib.decode)(self.ctx, batch);
+            if ret != 0 {
+                debug!(ret, "llama_decode (gen step) error — stopping");
+                self.finished = true;
+                self.finalize_timing();
+                return None;
+            }
+
+            self.n_decoded += 1;
+            self.tokens_generated += 1;
+
+            // Detokenize
+            let token_text = self.detokenize_token(new_token);
+
+            // Check stop sequences
+            // (simple check: if generated text so far ends with a stop seq)
+            // For streaming we check per-token — good enough for Phase 2.
+            for stop in &self.stop_sequences {
+                if token_text.trim_end() == stop.as_str() {
+                    self.finished = true;
+                    self.finalize_timing();
+                    return None;
+                }
+            }
+
+            Some(token_text)
+        }
+    }
+
+    fn finalize_timing(&mut self) {
+        if let Some(start) = self.gen_start.take() {
+            let elapsed = start.elapsed().as_secs_f64();
+            if elapsed > 0.0 {
+                self.tokens_per_sec = self.tokens_generated as f64 / elapsed;
+            }
+        }
+        info!(
+            tokens_generated = self.tokens_generated,
+            tps = format!("{:.1}", self.tokens_per_sec),
+            prefill_ms = self
+                .prefill_ms
+                .map(|v| format!("{:.1}", v))
+                .unwrap_or_default(),
+            "Generation complete"
+        );
+    }
+
+    pub fn tokens_generated(&self) -> u32 {
+        self.tokens_generated
+    }
+
+    pub fn prompt_tokens(&self) -> u32 {
+        self.prompt_tokens
+    }
+
+    pub fn tps(&self) -> f64 {
+        self.tokens_per_sec
+    }
+
+    pub fn prefill_ms(&self) -> Option<f64> {
+        self.prefill_ms
+    }
+
+    fn detokenize_token(&self, token: llama_token) -> String {
+        if let Some(token_to_piece) = self.lib.token_to_piece {
+            return decode_token_piece_with(|buf, len| unsafe {
+                token_to_piece(self.vocab, token, buf, len, 0, true)
+            });
+        }
+
+        let text_ptr = unsafe { (self.lib.vocab_get_text)(self.vocab, token) };
+        if text_ptr.is_null() {
+            String::new()
+        } else {
+            unsafe { CStr::from_ptr(text_ptr) }
+                .to_string_lossy()
+                .into_owned()
+        }
+    }
+}
+
+fn build_prefill_plan(token_count: usize, max_batch_tokens: usize) -> Vec<PrefillChunkPlan> {
+    let chunk_size = max_batch_tokens.max(1);
+    let mut plans = Vec::new();
+    let mut start = 0usize;
+
+    while start < token_count {
+        let end = (start + chunk_size).min(token_count);
+        let logits_index = if end == token_count {
+            Some(end - start - 1)
+        } else {
+            None
+        };
+        plans.push(PrefillChunkPlan {
+            range: start..end,
+            logits_index,
+        });
+        start = end;
+    }
+
+    plans
+}
+
+fn decode_prefill(
+    lib: &LlamaLib,
+    ctx: *mut llama_context,
+    tokens: &[llama_token],
+    max_batch_tokens: u32,
+) -> Result<()> {
+    for plan in build_prefill_plan(tokens.len(), max_batch_tokens as usize) {
+        let chunk = &tokens[plan.range.clone()];
+        let mut batch = unsafe { (lib.batch_init)(chunk.len() as i32, 0, 1) };
+        batch.n_tokens = chunk.len() as i32;
+
+        let ret = unsafe {
+            for (i, token) in chunk.iter().copied().enumerate() {
+                *batch.token.add(i) = token;
+                *batch.pos.add(i) = (plan.range.start + i) as i32;
+                *batch.n_seq_id.add(i) = 1;
+
+                let seq_ids = *batch.seq_id.add(i);
+                *seq_ids.add(0) = 0;
+
+                *batch.logits.add(i) = if plan.logits_index == Some(i) { 1 } else { 0 };
+            }
+
+            let decode_batch = std::ptr::read(&batch);
+            (lib.decode)(ctx, decode_batch)
+        };
+
+        unsafe { (lib.batch_free)(batch) };
+
+        if ret != 0 {
+            bail!(
+                "llama_decode (prefill) returned {} while decoding prompt chunk {:?}",
+                ret,
+                plan.range
+            );
+        }
+    }
+
+    Ok(())
+}
+
+fn decode_token_piece_with<F>(mut call: F) -> String
+where
+    F: FnMut(*mut std::ffi::c_char, i32) -> i32,
+{
+    let mut capacity = 64usize;
+
+    loop {
+        let mut buffer = vec![0u8; capacity];
+        let written = call(
+            buffer.as_mut_ptr() as *mut std::ffi::c_char,
+            capacity as i32,
+        );
+
+        if written < 0 {
+            capacity = (-written) as usize;
+            continue;
+        }
+
+        let written = written as usize;
+        return String::from_utf8_lossy(&buffer[..written]).into_owned();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{build_prefill_plan, decode_token_piece_with, PrefillChunkPlan};
+
+    #[test]
+    fn retries_when_piece_buffer_is_too_small() {
+        let mut calls = 0;
+
+        let decoded = decode_token_piece_with(|buf, len| {
+            calls += 1;
+
+            if len < 17 {
+                return -17;
+            }
+
+            let bytes = b" Quantum computer";
+            unsafe {
+                std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
+            }
+            bytes.len() as i32
+        });
+
+        assert_eq!(decoded, " Quantum computer");
+        assert_eq!(calls, 1);
+    }
+
+    #[test]
+    fn decodes_piece_after_resizing_buffer() {
+        let mut calls = 0;
+
+        let decoded = decode_token_piece_with(|buf, len| {
+            calls += 1;
+
+            if calls == 1 {
+                return -18;
+            }
+
+            assert!(len >= 18);
+            let bytes = b" Quantum computer";
+            unsafe {
+                std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
+            }
+            bytes.len() as i32
+        });
+
+        assert_eq!(decoded, " Quantum computer");
+        assert_eq!(calls, 2);
+    }
+
+    #[test]
+    fn prefill_plan_requests_logits_only_for_last_prompt_token_at_batch_limit() {
+        let plan = build_prefill_plan(128, 128);
+
+        assert_eq!(
+            plan,
+            vec![PrefillChunkPlan {
+                range: 0..128,
+                logits_index: Some(127),
+            }]
+        );
+    }
+
+    #[test]
+    fn prefill_plan_splits_large_prompt_and_only_requests_final_logits() {
+        let plan = build_prefill_plan(300, 128);
+
+        assert_eq!(
+            plan,
+            vec![
+                PrefillChunkPlan {
+                    range: 0..128,
+                    logits_index: None,
+                },
+                PrefillChunkPlan {
+                    range: 128..256,
+                    logits_index: None,
+                },
+                PrefillChunkPlan {
+                    range: 256..300,
+                    logits_index: Some(43),
+                },
+            ]
+        );
+    }
+}
+
+impl Drop for GenerationSession {
+    fn drop(&mut self) {
+        if !self.sampler.is_null() {
+            unsafe { (self.lib.sampler_free)(self.sampler) };
+            self.sampler = std::ptr::null_mut();
+        }
+    }
+}