Something went wrong. Try again.
Routing embedding stack for Qwen3-Embedding-0.6B on Apple Silicon: orchestrator + MLX + llama.cpp backends
Something went wrong. Try again.
11 kB · 302 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303#![allow( clippy::cast_possible_wrap, clippy::cast_possible_truncation, clippy::cast_precision_loss, clippy::cast_sign_loss)]
use std::num::NonZeroU32;use std::path::PathBuf;use std::time::{Duration, Instant};
use anyhow::Result;use tracing::{debug, error, info, warn};
use llama_cpp_2::context::params::{LlamaContextParams, LlamaPoolingType};use llama_cpp_2::llama_backend::LlamaBackend;use llama_cpp_2::llama_batch::LlamaBatch;use llama_cpp_2::model::params::LlamaModelParams;use llama_cpp_2::model::{AddBos, LlamaModel};use llama_cpp_2::token::LlamaToken;
use crate::dto::InputKind;use crate::prompt;
// ---------------------------------------------------------------------------// Public message type — sent from axum handlers to the worker thread// ---------------------------------------------------------------------------
pub struct WorkRequest { pub texts: Vec<String>, pub input_kind: InputKind, pub reply: tokio::sync::oneshot::Sender<Result<Vec<Vec<f32>>>>,}
// ---------------------------------------------------------------------------// Worker configuration// ---------------------------------------------------------------------------
pub struct WorkerConfig { pub model_path: PathBuf, pub max_batch: usize, pub max_len: usize, pub wait_ms: u64,}
// ---------------------------------------------------------------------------// GPU worker thread entry point// ---------------------------------------------------------------------------
/// Runs on a dedicated `std::thread`. All llama model/context objects are/// created and used exclusively on this thread (`LlamaContext` is `!Send`).////// The `LlamaBackend` is initialized once on the main thread and leaked to a/// `&'static` reference shared across all worker threads.////// Sends `n_embd` (or an error) back through `ready_tx` so the main thread/// knows the embedding dimension and can mark the server ready.pub fn run_worker( backend: &'static LlamaBackend, config: WorkerConfig, rx: std::sync::mpsc::Receiver<WorkRequest>, ready_tx: tokio::sync::oneshot::Sender<Result<usize>>,) { let WorkerConfig { model_path, max_batch, max_len, wait_ms, } = config;
// ── 1. Load model ────────────────────────────────────────────────── let model_params = LlamaModelParams::default().with_n_gpu_layers(1000); info!(path = %model_path.display(), "Loading model…");
let model = match LlamaModel::load_from_file(backend, &model_path, &model_params) { Ok(m) => m, Err(e) => { let err = anyhow::anyhow!("LlamaModel::load_from_file: {e}"); error!("{err:#}"); let _ = ready_tx.send(Err(err)); return; } };
// ── 2. Resolve special tokens / dimensions ───────────────────────── let n_embd = model.n_embd() as usize; info!(n_embd, "Model loaded");
// ── 3. Create context ────────────────────────────────────────────── let total_ctx_tokens = max_batch * (max_len + 1); // +1 for appended EOS let num_cpus = std::thread::available_parallelism() .map(|n| n.get() as i32) .unwrap_or(4);
// n_batch: max logical tokens per decode call (we want this large enough // to accept a single long input in one decode()). // n_ubatch: max PHYSICAL tokens processed at once. Smaller = less compute // graph memory. 2048 keeps peak workspace under ~1 GB. let ubatch_cap: u32 = std::env::var("EMBED_UBATCH") .ok() .and_then(|s| s.parse().ok()) .unwrap_or(2048); let n_ubatch = (total_ctx_tokens as u32).min(ubatch_cap);
let ctx_params = LlamaContextParams::default() .with_n_ctx(NonZeroU32::new(total_ctx_tokens as u32)) .with_n_batch(total_ctx_tokens as u32) .with_n_ubatch(n_ubatch) .with_n_seq_max(max_batch as u32) .with_n_threads_batch(num_cpus) .with_embeddings(true) .with_pooling_type(LlamaPoolingType::Last); // CRITICAL for Qwen3
let mut ctx = match model.new_context(backend, ctx_params) { Ok(c) => c, Err(e) => { let err = anyhow::anyhow!("new_context: {e}"); error!("{err:#}"); let _ = ready_tx.send(Err(err)); return; } };
info!( n_ctx = ctx.n_ctx(), n_batch = ctx.n_batch(), "Context created" );
// ── 4. Signal readiness ──────────────────────────────────────────── let _ = ready_tx.send(Ok(n_embd));
// ── 5. Main loop — accumulate then fire ──────────────────────────── let wait_dur = Duration::from_millis(wait_ms);
loop { // Block for first request (exit on channel close) let first = match rx.recv() { Ok(req) => req, Err(_) => { info!("Channel closed, worker exiting"); break; } };
let mut pending: Vec<WorkRequest> = vec![first]; let mut total_texts: usize = pending[0].texts.len(); let deadline = Instant::now() + wait_dur;
// Accumulate more requests until batch is full or deadline expires while total_texts < max_batch { let remaining = deadline.saturating_duration_since(Instant::now()); if remaining.is_zero() { break; } match rx.recv_timeout(remaining) { Ok(req) => { total_texts += req.texts.len(); pending.push(req); } Err(std::sync::mpsc::RecvTimeoutError::Timeout) => break, Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break, } }
debug!( requests = pending.len(), texts = total_texts, "Processing batch" );
process_batch(&model, &mut ctx, &mut pending, max_len, n_embd); }}
// ---------------------------------------------------------------------------// Batch processing — where correctness lives// ---------------------------------------------------------------------------
fn process_batch( model: &LlamaModel, ctx: &mut llama_cpp_2::context::LlamaContext<'_>, pending: &mut Vec<WorkRequest>, max_len: usize, n_embd: usize,) { // ── 1. Flatten all texts, tracking request boundaries ────────────── struct SeqEntry { request_idx: usize, tokens: Vec<LlamaToken>, }
let mut sequences: Vec<SeqEntry> = Vec::new();
for (req_idx, req) in pending.iter().enumerate() { for text in &req.texts { // Prompt-wrap let wrapped = prompt::wrap(text, req.input_kind);
// Tokenize let mut tokens = match model.str_to_token(&wrapped, AddBos::Always) { Ok(t) => t, Err(e) => { warn!(text_len = text.len(), "Tokenization failed: {e}"); // Push empty so we still have a slot; embedding will be zeros Vec::new() } };
// Truncate to max_len tokens. // // NB: str_to_token(.., AddBos::Always) calls llama_tokenize with // `add_special=true`, which honors the GGUF's add_bos_token / // add_eos_token metadata. For Qwen3-Embedding-GGUF, add_eos_token // is true, so EOS is already appended by the tokenizer. // We must NOT push another one or last-token pooling reads the // wrong hidden state. if tokens.len() > max_len { tokens.truncate(max_len); }
sequences.push(SeqEntry { request_idx: req_idx, tokens, }); } }
// ── 2. Build batch ───────────────────────────────────────────────── let total_tokens: usize = sequences.iter().map(|s| s.tokens.len()).sum(); let n_sequences = sequences.len(); let seq_lens: Vec<usize> = sequences.iter().map(|s| s.tokens.len()).collect(); info!(n_sequences, total_tokens, ?seq_lens, "Batch token stats");
let mut batch = LlamaBatch::new(total_tokens, n_sequences as i32); for (seq_id, entry) in sequences.iter().enumerate() { if entry.tokens.is_empty() { continue; } if let Err(e) = batch.add_sequence(&entry.tokens, seq_id as i32, false) { error!(seq_id, "Failed to add sequence to batch: {e}"); } }
// ── 3. Decode ────────────────────────────────────────────────────── ctx.clear_kv_cache();
if let Err(e) = ctx.decode(&mut batch) { error!("decode() failed: {e}"); // Send errors to all pending requests for req in pending.drain(..) { let _ = req.reply.send(Err(anyhow::anyhow!("decode failed: {e}"))); } return; }
// ── 4. Extract embeddings & L2-normalize ─────────────────────────── let mut all_embeddings: Vec<Option<Vec<f32>>> = vec![None; sequences.len()];
for (seq_id, _entry) in sequences.iter().enumerate() { match ctx.embeddings_seq_ith(seq_id as i32) { Ok(raw) => { all_embeddings[seq_id] = Some(l2_normalize(raw)); } Err(e) => { warn!(seq_id, "embeddings_seq_ith failed: {e}"); // Leave as None; will become zero vector } } }
// ── 5. Fan out — group vectors by request boundary ───────────────── let mut results_per_request: Vec<Vec<Vec<f32>>> = vec![Vec::new(); pending.len()];
for (seq_id, entry) in sequences.iter().enumerate() { let embedding = all_embeddings[seq_id] .take() .unwrap_or_else(|| vec![0.0; n_embd]); results_per_request[entry.request_idx].push(embedding); }
for (req_idx, req) in pending.drain(..).enumerate() { let vectors = std::mem::take(&mut results_per_request[req_idx]); let _ = req.reply.send(Ok(vectors)); }}
// ---------------------------------------------------------------------------// Utility// ---------------------------------------------------------------------------
fn l2_normalize(input: &[f32]) -> Vec<f32> { let magnitude = input .iter() .fold(0.0_f32, |acc, &val| val.mul_add(val, acc)) .sqrt();
if magnitude == 0.0 { return input.to_vec(); } input.iter().map(|&val| val / magnitude).collect()}