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.
3.7 kB · 129 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130use crate::backend_rust::RustBackend;use crate::backend_swift::SwiftBackend;use crate::config::Config;use crate::dto::{EmbedRequest, EmbedResponse};use crate::errors::OrchestratorError;use crate::tokenizer::{round_up, TokenCounter};
/// Build a sub-request containing only the texts at `indices`, preserving order.fn sub_request(req: &EmbedRequest, indices: &[usize]) -> EmbedRequest { EmbedRequest { texts: indices.iter().map(|&i| req.texts[i].clone()).collect(), input_type: req.input_type, size: req.size.clone(), }}
/// Partition indices into (short, long) by the threshold.pub fn partition(counts: &[usize], threshold: usize) -> (Vec<usize>, Vec<usize>) { let mut short_idx = Vec::new(); let mut long_idx = Vec::new(); for (i, &c) in counts.iter().enumerate() { if c <= threshold { short_idx.push(i); } else { long_idx.push(i); } } (short_idx, long_idx)}
pub async fn dispatch( req: EmbedRequest, tokenizer: &TokenCounter, swift: &SwiftBackend, rust: &RustBackend, cfg: &Config,) -> Result<EmbedResponse, OrchestratorError> { if req.texts.is_empty() { return Err(OrchestratorError::EmptyTexts); } if let Some(s) = req.size.as_deref() { if s != "0.6B" { return Err(OrchestratorError::UnsupportedSize); } }
// 1. Wrap + count each text. let counts: Vec<usize> = req .texts .iter() .map(|t| tokenizer.count_wrapped(t, req.input_type)) .collect::<Result<_, _>>()?;
// 2. Validate model max. for (i, &c) in counts.iter().enumerate() { if c > cfg.model_max_tokens { return Err(OrchestratorError::TextTooLong { index: i, tokens: c, limit: cfg.model_max_tokens, }); } }
// 3. Partition. let (short_idx, long_idx) = partition(&counts, cfg.swift_threshold);
// 4. Dispatch in parallel. let short_fut = async { if short_idx.is_empty() { return Ok::<Vec<Vec<f32>>, OrchestratorError>(vec![]); } let sub = sub_request(&req, &short_idx); Ok(swift.embed(&sub).await?.vectors) }; let long_fut = async { if long_idx.is_empty() { return Ok::<Vec<Vec<f32>>, OrchestratorError>(vec![]); } let max_tok = long_idx.iter().map(|i| counts[*i]).max().unwrap(); let bucket = round_up(max_tok, cfg.rust_bucket_granularity).clamp(1024, cfg.model_max_tokens); let sub = sub_request(&req, &long_idx); Ok(rust.run_one(bucket, &sub).await?.vectors) }; let (shorts, longs) = tokio::try_join!(short_fut, long_fut)?;
// 5. Merge back in original order. let mut out: Vec<Vec<f32>> = vec![Vec::new(); req.texts.len()]; let mut dim = 0usize; for (slot, v) in short_idx.iter().zip(shorts) { if dim == 0 { dim = v.len(); } out[*slot] = v; } for (slot, v) in long_idx.iter().zip(longs) { if dim == 0 { dim = v.len(); } out[*slot] = v; }
Ok(EmbedResponse { vectors: out, dim, model: "qwen3-embedding-0.6B".into(), })}
#[cfg(test)]mod tests { use super::*;
#[test] fn partition_basic() { let counts = vec![25, 30, 8200, 28, 15040]; let (s, l) = partition(&counts, 8192); assert_eq!(s, vec![0, 1, 3]); assert_eq!(l, vec![2, 4]); }
#[test] fn partition_all_short() { let (s, l) = partition(&[1, 2, 3], 8192); assert_eq!(s, vec![0, 1, 2]); assert!(l.is_empty()); }}