From 0c1cc17bfa61453f8c3bb943a40cb76a3458db51 Mon Sep 17 00:00:00 2001 From: saagpatel Date: Sat, 11 Jul 2026 13:25:48 -0700 Subject: [PATCH] perf(llm): reuse embedding context across batch inputs --- examples/embed_batch_parity.rs | 69 ++++++++++++++++++++++ src/llm.rs | 105 +++++++++++++++++++++++---------- 2 files changed, 144 insertions(+), 30 deletions(-) create mode 100644 examples/embed_batch_parity.rs diff --git a/examples/embed_batch_parity.rs b/examples/embed_batch_parity.rs new file mode 100644 index 0000000..95e6ab0 --- /dev/null +++ b/examples/embed_batch_parity.rs @@ -0,0 +1,69 @@ +//! Compare shared-context document embeddings with fresh-context embeddings. +//! +//! Usage: `embed_batch_parity ` + +use std::path::PathBuf; +use std::time::Instant; + +use engraph::llm::{EmbedModel, LlamaEmbed}; + +fn main() -> anyhow::Result<()> { + let models_dir = PathBuf::from( + std::env::args() + .nth(1) + .expect("models_dir argument is required"), + ); + let config = engraph::config::Config::default(); + let mut embedder = LlamaEmbed::new(&models_dir, &config)?; + + let texts: Vec = (0..32) + .map(|index| { + format!( + "Note {index}: this chunk discusses retrieval pipelines, wikilink graphs, \ + and reciprocal rank fusion in a knowledge vault. Sample {index} varies \ + the content slightly to avoid identical sequences." + ) + }) + .collect(); + let text_refs: Vec<&str> = texts.iter().map(String::as_str).collect(); + + let started = Instant::now(); + let shared_context = embedder.embed_batch(&text_refs)?; + let shared_context_ms = started.elapsed().as_millis(); + + let started = Instant::now(); + let mut fresh_context = Vec::with_capacity(text_refs.len()); + for text in &text_refs { + fresh_context.extend(embedder.embed_batch(&[text])?); + } + let fresh_context_ms = started.elapsed().as_millis(); + + let mut max_abs_diff = 0.0_f32; + let mut vectors_beyond_tolerance = 0_usize; + for (shared, fresh) in shared_context.iter().zip(&fresh_context) { + assert_eq!(shared.len(), fresh.len()); + let difference = shared + .iter() + .zip(fresh) + .map(|(left, right)| (left - right).abs()) + .fold(0.0_f32, f32::max); + max_abs_diff = max_abs_diff.max(difference); + if difference > 1e-6 { + vectors_beyond_tolerance += 1; + } + } + + println!( + "{{\"n\":{},\"shared_context_ms\":{},\"fresh_context_ms\":{},\"max_abs_diff\":{:e},\"vectors_beyond_1e-6\":{}}}", + text_refs.len(), + shared_context_ms, + fresh_context_ms, + max_abs_diff, + vectors_beyond_tolerance + ); + + if vectors_beyond_tolerance > 0 { + anyhow::bail!("shared-context embeddings exceeded the 1e-6 parity tolerance"); + } + Ok(()) +} diff --git a/src/llm.rs b/src/llm.rs index 4c417c3..2fae219 100644 --- a/src/llm.rs +++ b/src/llm.rs @@ -527,32 +527,47 @@ impl FlexTokenizer { /// Load tokenizer for a model. Tries external tokenizer.json first, falls back to GGUF-embedded. fn load_tokenizer_for_model(uri: &HfModelUri, models_dir: &Path) -> Result { - // First try: external tokenizer.json from candidate repos. - if let Some(tok) = try_external_tokenizer(uri, models_dir) { + // 1. Already-cached external tokenizer.json (no network). + if let Some(tok) = try_external_tokenizer(uri, models_dir, false) { return Ok(FlexTokenizer::HuggingFace(Box::new(tok))); } - // Fallback: load tokenizer from GGUF file metadata. + // 2. Tokenizer embedded in the GGUF we already have (no network). + // The model file is guaranteed present here: callers run ensure_model() + // before loading the tokenizer. Trying the network first meant every + // process start paid seconds of doomed HTTP requests (404/401 on all + // candidate repos) before landing on this exact fallback. let model_path = uri.cache_path(models_dir); - if model_path.exists() { - tracing::info!( - "no external tokenizer found, loading from GGUF: {}", + if model_path.exists() + && let Ok(tok) = shimmytok::Tokenizer::from_gguf_file(&model_path) + { + tracing::debug!( + "loaded tokenizer from GGUF metadata: {}", model_path.display() ); - let tok = shimmytok::Tokenizer::from_gguf_file(&model_path) - .map_err(|e| anyhow::anyhow!("loading tokenizer from GGUF metadata: {e}"))?; return Ok(FlexTokenizer::Gguf(Box::new(tok))); } + // 3. Last resort: download tokenizer.json from candidate repos. + if let Some(tok) = try_external_tokenizer(uri, models_dir, true) { + return Ok(FlexTokenizer::HuggingFace(Box::new(tok))); + } + bail!( - "could not find tokenizer for model '{}': no external tokenizer.json \ - and GGUF file not yet downloaded", + "could not find tokenizer for model '{}': no cached tokenizer.json, \ + no GGUF-embedded tokenizer, and no candidate repo served one", uri.repo ) } -/// Try downloading tokenizer.json from candidate HuggingFace repos. -fn try_external_tokenizer(uri: &HfModelUri, models_dir: &Path) -> Option { +/// Try loading tokenizer.json from candidate HuggingFace repos. +/// With `download = false`, only already-cached files are considered — no +/// network I/O. With `download = true`, missing candidates are fetched. +fn try_external_tokenizer( + uri: &HfModelUri, + models_dir: &Path, + download: bool, +) -> Option { let mut candidates: Vec = vec![uri.repo.clone()]; // Non-GGUF variant: "org/model-GGUF" → "org/model" @@ -594,7 +609,8 @@ fn try_external_tokenizer(uri: &HfModelUri, models_dir: &Path) -> Option Result> { - // Tokenize using llama.cpp's built-in tokenizer. - // Use AddBos::Never because PromptFormat already adds for embeddinggemma. + /// Tokenize with llama.cpp's built-in tokenizer. + /// AddBos::Never because PromptFormat already adds for embeddinggemma. + fn tokenize(&self, text: &str) -> Result> { let tokens = self .model .str_to_token(text, AddBos::Never) @@ -711,28 +726,38 @@ impl LlamaEmbed { if tokens.is_empty() { bail!("tokenizer returned empty token sequence"); } + Ok(tokens) + } - // Create a context with embeddings enabled (per-call, since LlamaContext is !Send). - // n_ubatch must be >= n_tokens for the encoder, and n_ctx must fit all tokens. - let n_tokens = tokens.len() as u32; - let n_ctx = std::num::NonZeroU32::new(n_tokens.max(64) + 16); + /// Create an embeddings context that can encode sequences up to `max_tokens`. + /// n_ubatch must be >= n_tokens for the encoder, and n_ctx must fit all tokens. + fn make_context(&self, max_tokens: u32) -> Result> { + let n_ctx = std::num::NonZeroU32::new(max_tokens.max(64) + 16); let ctx_params = LlamaContextParams::default() .with_embeddings(true) .with_n_ctx(n_ctx) - .with_n_ubatch(n_tokens.max(512)) - .with_n_batch(n_tokens.max(512)); - let mut ctx = self - .model + .with_n_ubatch(max_tokens.max(512)) + .with_n_batch(max_tokens.max(512)); + self.model .new_context(llama_backend()?, ctx_params) - .map_err(|e| anyhow::anyhow!("creating embedding context: {e}"))?; + .map_err(|e| anyhow::anyhow!("creating embedding context: {e}")) + } + /// Encode one token sequence in `ctx` and return the truncated, + /// L2-normalized embedding. + fn encode_tokens( + &self, + ctx: &mut llama_cpp_2::context::LlamaContext<'_>, + tokens: &[llama_cpp_2::token::LlamaToken], + ) -> Result> { // Create batch and add tokens — mark all as outputs for embedding. let mut batch = LlamaBatch::new(tokens.len() + 16, 1); batch - .add_sequence(&tokens, 0, true) + .add_sequence(tokens, 0, true) .map_err(|e| anyhow::anyhow!("adding sequence to batch: {e}"))?; // Encode (compute embeddings). Use encode() for embedding models. + ctx.clear_kv_cache(); ctx.encode(&mut batch) .map_err(|e| anyhow::anyhow!("embedding encode failed: {e}"))?; @@ -759,18 +784,38 @@ impl LlamaEmbed { Ok(normalized) } + + /// Run embedding inference and return the truncated, L2-normalized embedding. + fn embed_text(&self, text: &str) -> Result> { + let tokens = self.tokenize(text)?; + // Context is created per-call (LlamaContext is !Send, so it can't be + // stored in this Send struct). + let mut ctx = self.make_context(tokens.len() as u32)?; + self.encode_tokens(&mut ctx, &tokens) + } } impl EmbedModel for LlamaEmbed { fn embed_batch(&mut self, texts: &[&str]) -> Result>> { - // Process texts sequentially — llama.cpp context is per-call. + if texts.is_empty() { + return Ok(Vec::new()); + } // Apply document prompt format for indexing (asymmetric models need this). - texts + // Tokenize everything first, then share ONE context sized for the longest + // sequence: context creation costs llama.cpp graph planning + Metal + // pipeline setup per call, which dominated indexing when paid per chunk. + let token_seqs: Vec> = texts .iter() .map(|t| { let formatted = self.prompt_format.format_document("", t); - self.embed_text(&formatted) + self.tokenize(&formatted) }) + .collect::>()?; + let max_tokens = token_seqs.iter().map(|t| t.len()).max().unwrap_or(0) as u32; + let mut ctx = self.make_context(max_tokens)?; + token_seqs + .iter() + .map(|tokens| self.encode_tokens(&mut ctx, tokens)) .collect() }