14× 更快的嵌入：我们如何重建Manticore中的ONNX路径

# 14倍更快的嵌入：我们如何重建Manticore中的ONNX路径 来源：https://manticoresearch.com/blog/onnx-embeddings-speedup/ 当我们发布自动嵌入（https://manticoresearch.com/blog/auto-embeddings/）——这个功能可以自动将任何文本列转换为向量，无需单独运行模型服务——时，最常见的反馈是关于速度的。之前的路径通过基于Candle（https://github.com/huggingface/candle）的SentenceTransformers运行，Candle是Hugging Face的纯Rust机器学习推理运行时，这浪费了大量CPU：无论我们如何输入，大多数工作负载的文档/秒都停留在低两位数，并且并发调用在单个模型会话上串行化。因此，我们花了几个星期重建Manticore运行ONNX模型的方式。新的ONNX Runtime后端已随Manticore Search 27.1.5（https://manticoresearch.com/blog/manticore-search-27-1-5/）发布。 ONNX（开放神经网络交换）是一种可移植的模型格式，大多数流行的开源嵌入模型——MiniLM、BGE、E5及同类——已经发布该格式。结果是，在相同硬件（平均廉价16核/32线程服务器）、相同模型、相同权重下，新后端**平均比之前的SentenceTransformers/Candle路径快约14倍**，涵盖完整的`threads × batch`工作负载网格——无论你运行1个客户端线程还是32个线程，这一优势都成立。旧路径在整个网格中保持在5–11文档/秒的范围内；新路径处于70–230文档/秒的区间。 本篇博文是工程日志：我们尝试了什么，什么让我们惊讶，我们丢弃了什么，以及最终设计的样貌。 ## 摘要 - **平均比之前的SentenceTransformers/Candle路径快约14倍**，在同一台机器（16核/32线程）上，同一模型、相同权重，涵盖完整的`threads × batch`工作负载网格（1/2/4/8/16/32线程 × 批处理大小1...128）。 - 在Manticore Search 27.1.5（https://manticoresearch.com/blog/manticore-search-27-1-5/）中发布，新的ONNX路径现在是任何带有`.onnx`文件的Hugging Face模型的默认快速路径。 - 在`all-MiniLM-L12-v2`上，旧的Candle路径在我们尝试的所有配置下都停留在**5–11文档/秒**。新的ONNX路径落在**70–230文档/秒**的范围——**无论你运行1个客户端线程还是32个线程，约14倍的差距保持一致**。 - 我们测试机上的单次插入延迟：**单个客户端约14毫秒**，在8路并发负载下约**56毫秒**——两者都远低于Candle达到的200+毫秒。 - **想要最大的批量摄入吞吐量？** 使用**高批处理大小**（32–128）在**单个客户端线程**上。新后端在调用内部进行并行化，因此客户端扇出只会增加协调开销——我们机器上的峰值是**1线程 + batch=64 时达到233文档/秒**。 - 最重要的两个变化：关闭**`intra_op_spinning`**，以及放弃在工作线程内对文档进行批处理。 - 无需修改面向用户的API。已经指向支持ONNX的`MODEL_NAME`的表会自动拾取新路径。将现有表切换到不同模型不是一行命令——Manticore不允许就地更改`FLOAT_VECTOR`字段上的`MODEL_NAME`——但你不必重建整个表：你可以添加一个包含新模型的新列，重建其嵌入，然后删除旧列。 ## 为什么这很重要 使用自动嵌入时，数据库本身在每次`INSERT`上运行模型。这意味着嵌入速度*就是*INSERT速度——你的摄入吞吐量取决于嵌入步骤能维持的速度。旧的SentenceTransformers/Candle路径未能充分释放性能。并发导致锁竞争，批处理调用由于填充开销而停滞，在调用之间运行时以阻止下一次调用从上次调用停止处继续的方式休眠线程。主要症状很简单：无论你如何输入，`top`都会显示机器远未达到满负载。整个范围——单行INSERT、128行批量INSERT、一个客户端线程、三十二个客户端线程——都停留在**5–11文档/秒**，因为你的输入方式无法让你获得更多CPU。 新的ONNX路径将下限提高了一个数量级，*并且*为用户提供了有意义的性能调优选项。单线程、单行INSERT现在达到**72文档/秒**——已经是旧Candle上限的约7倍。增加并发或批处理大小后，它会攀升到**130–230文档/秒**的范围，网格顶部在**单个客户端线程上，使用`--batch-size=64`时达到233文档/秒**。在整个`threads × batch`矩阵上平均，新路径是**旧路径的约14倍**。 ## 为什么是ONNX，而不是Candle Manticore的嵌入库已经支持一些后端一段时间了。Candle路径在正确性方面很好，且易于发布。但对于像MiniLM和BGE系列这样的小型编码器的生产推理，ONNX Runtime难以被超越： - ONNX Runtime（或**ORT**——微软官方的手动调优的C++推理引擎，用于ONNX模型）进行图融合、常量折叠、内核自动调优。 - HuggingFace上大多数流行的嵌入模型已经在它们的`onnx/`目录中发布了预融合的`model.onnx`。磁盘上的文件已经是ORT想要的形状。 - 在相同的`all-MiniLM-L12-v2`权重上，在CPU上，ONNX路径比Candle路径有明显的改进。相同的质量，每个文档的工作量少得多。 ORT会话使用一组小的配置选项创建： ```rust let session = ort::session::Session::builder()? .with_optimization_level(GraphOptimizationLevel::Level3)? .with_intra_threads(0)? // 让ORT选择（=所有核心） .with_intra_op_spinning(false)? // 不要在调用之间忙等 .with_flush_to_zero()? // 消除注意力softmax上的非规格化数 .with_approximate_gelu()? // 激活函数快约10%，无质量损失 .commit_from_file(&onnx_path)?; ``` 大多数这些选项是无可争议的，“当然要打开”的旋钮。有一个不是：`intra_op_spinning(false)`。我们稍后会回到它——这是整个分支中最大的单个收益，而且它与其说是ORT设置，不如说是一个负载形状的决策。 ## 并发模型——大多数读者会觉得新鲜的部分 如果你给一个Rust开发者“让ONNX变快”的任务而没有其他约束，他们会倾向于两种模式之一。我们都试过了。对于这个工作负载，它们都是错误的。 **模式1：一个共享的`Session`后面加一个`Mutex`**（`Mutex`是一个锁，一次只允许一个线程接触会话）。易于推理，易于正确实现。并发时吞吐量崩溃，因为每个调用者都在锁上串行化。对于CLI工具来说很好，对于服务许多并发INSERT的数据库来说很糟糕。 **模式2：会话池，每个CPU一个`Session`。**不再有锁竞争，但冷启动时间成倍增加，内存使用成倍增加，并且小输入为了落在某个会话上而支付分发成本。我们在一个开发分支中有这个功能的工作版本，但它从未真正交付过。 解锁设计的关键是大多数Rust ONNX包装器都搞错了的事情：**在Linux和macOS上，ORT的C `Run()` API是线程安全的。**你可以跨许多并发调用者共享一个`Session`而无需任何锁定。C++端已经序列化了需要序列化的部分；Rust API只是将其隐藏在借检查器规则后面，而这些规则与底层库实际允许的并不匹配。 因此，我们将会话包装在一个小的平台感知类型中： ```rust #[cfg(not(target_os = "windows"))] struct SessionWrapper { inner: std::cell::UnsafeCell<Session>, } #[cfg(not(target_os = "windows"))] unsafe impl Sync for SessionWrapper {} #[cfg(not(target_os = "windows"))] unsafe impl Send for SessionWrapper {} impl SessionWrapper { fn with_session<R>(&self, f: impl FnOnce(&mut Session) -> R) -> R { f(unsafe { &mut *self.inner.get() }) } } ``` 是的，这是`unsafe`的。我们将借检查器排除在外，因为底层库被文档证明在我们使用的访问模式下是安全的。这是一个有意的`unsafe`，附带一行解释，而不是一个容易出错的功能。 在Windows上，ORT的线程模型存在已知问题，因此我们使用`Mutex`序列化`Run()`。重要的是，锁在整个闭包期间持有，而不仅仅是调用`run()`期间——这就是我们修复在Windows上看到的竞态条件的关键：一个线程的`SessionOutputs`还在被读取时，另一个线程已经开始新的`run()`。闭包范围的锁定，而不是调用范围的锁定。 ## 自适应并行——我们走过的弯路 这部分工作花费的时间最长，因为每本教科书都说“为了让ONNX快，要批处理你的输入”。所以我们的第一次尝试遵循了教科书。我们以8、16、32个文档为一组进行分词，将它们填充到`max_len`，然后每个工作线程运行一次前向传播。吞吐量数字比通过相同会话逐个处理相同文本更低。我们再次运行。相同的结果。我们花了一段时间试图反驳它，然后才接受它。回滚提交`980b24b "Revert: perf(model): batch inference in worker threads"`是我们停止对抗并围绕分析器不断告诉我们的内容重建的时刻。 两件事导致了这一意外。 **填充税。**一批混合长度的文本将每一行填充到最长行的长度。然后模型执行的工作量与`batch_size * max_len * hidden_dim`成正比，无论批中实际内容有多少。真实文本输入的长度变化很大：一个典型的8个随机句子的批次可能有一个60个token的异常值和七个8个token的行。模型花费大部分周期将填充token乘以注意力权重。对于单文档批次，模型只做与该文档实际token数量成正比的工作。一旦输入长度变化是现实的，逐文档来看，“不批处理”比“批处理”更便宜。 **旋转。** ORT的算子内线程池默认在调度之间*旋转*——线程在一个紧密循环中忙等，等待下一块工作。每次会话调用一个大批量时，这是不可见的：线程总是忙于真实工作。当有许多并发小调用时，它变成一场灾难：每个工作线程的算子内池在调用之间被固定到100% CPU，没有CPU留给其他任何事情。我们在`top`中确切地看到了这种模式：每个核心100%，吞吐量*低于*不旋转的情况。这听起来不对，直到你记得系统的其余部分也需要CPU时间——分词器、HNSW构建、`searchd`的其余部分。打开`with_intra_op_spinning(false)`是一行改动，立即提高了吞吐量并同时降低了CPU使用率。 因此，最终形状与教科书配方相反： - **一个共享会话**，没有池。 - **每次推理调用一个文档**，工作线程内不批处理。 - **许多并发调用者**，按CPU数量缩放。 - **调用之间不旋转**——像一个礼貌的公民一样让出CPU。 ```rust fn predict_pipelined(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, _> { let bs = batch_size(); // 小输入——单次分词+推理，没有线程开销。 // 这是1文档INSERT走的路径。 if texts.len() <= bs { return Self::tokenize_and_infer(&self.session, &self.tokenizer, texts, ...); } // 大输入——在工作线程之间分割，每个线程通过共享会话一次处理一个文档。 // 这故意模仿ORT最喜欢的那种许多并发调用者的模式。 let num_workers = (texts.len() / bs).min(available_cpus()).max(1); let docs_per_worker = texts.len().div_ceil(num_workers); std::thread::scope(|s| { for worker_texts in texts.chunks(docs_per_worker) { s.spawn(move || { for text in worker_texts { Self::tokenize_and_infer(&session, &tokenizer, std::slice::from_ref(text), ...)?; } Ok::<(), Error>(()) }); } }); // ... } ``` 两分支设计是有意为之。一个1行的INSERT以`texts.len() == 1`进入，这`<= bs`，因此它采用**零线程生成、零通道发送、零协调开销**的快速路径。一个包含数千行的批量REPLACE INTO采用并行分支并获得吞吐量收益。廉价的情况保持廉价，昂贵的情况保持并行。 我们还启用了启动时的一次性并行分词（`TOKENIZERS_PARALLELISM=true`），并在BPE之前按字符数预先截断输入，这样一个100KB的文本块不会让CPU在分词器上固定一秒钟，然后模型才看到它。 ## 数据 所有运行在我们的标准基准机器上，使用`all-MiniLM-L12-v2-onnx`，每次运行1000个文档。使用manticore-load（https://manticoresearch.com/blog/manticore-load/）生成： ``` manticore-load --quiet --drop --batch-size=1 --threads=8 --total=1000 \ --init="CREATE TABLE t ( f text, v FLOAT_VECTOR KNN_TYPE='hnsw' HNSW_SIMILARITY='l2' MODEL_NAME='onnx-models/all-MiniLM-L12-v2-onnx' FROM='' )" \ --load="INSERT INTO t(f) VALUES('')" ``` 相同命令，`--batch-size=2`、`8`、`32`、`128`，全部在8个线程下： | `--batch-size` | 文档/秒 | 平均调用延迟 (ms) | 每文档延迟 (ms) | |----------------|---------|-------------------|-----------------| | 1 | 143 | 55.9 | 55.9 | | 2 | 131 | 141.6 | 70.8 | | 8 | 90 | 170.3 | 87.9 | | 32 | 146 | 1753.4 | 54.8 | | 128 | 147 | 7696.0 | 54.4 | 与Candle在相同8个线程下相比——Candle在所有批处理大小下都平坦地停留在**10文档/秒**——根据你选择的批处理大小，这相当于**9倍到15倍更多的文档每秒**。“平均调用延迟”列是一个完整`INSERT`语句返回的时间，不是每文档时间；除以批处理大小，每文档成本落在55–90毫秒的范围内。 如果你将表切换到1个客户端线程——事实证明这对批量加载是最优的配置——数字进一步攀升：在批处理大小1/2/8/32/64/128下分别为**72/76/93/175/233/222文档/秒**。整个网格中的峰值是**1线程 × batch=64 时的233文档/秒**，每文档延迟约为**4.3毫秒**。 ### 如何输入以获得最大吞吐量 如果你需要批量加载大量数据并希望获得最大文档/秒，配方很简单：发送大型`INSERT ... VALUES (...), (...), ...`语句（批处理大小32–128），从**单个客户端线程**发送，而不是从多个线程发送许多小插入。新后端已经在调用内部进行了并行化（参见上面的`predict_pipelined`代码），因此客户端扇出只会增加ORT已经正在做的事情之上的协调开销——这就是为什么1线程 × batch=64（233文档/秒）明显优于8线程 × batch=128（147文档/秒）。 如果你的工作负载天然是逐行的——网络请求、队列消费者、MCP服务器——只需使用`INSERT INTO`。单线程/单行的下限72文档/秒已经是旧Candle路径的约7倍，延迟足够低，这不再是你需要优化的层级。 ### 前后对比，覆盖整个网格 为了使前后对比具体化，我们还根据旧Candle/`trans`路径在相同机器、相同权重下扫描了整个`threads × batch`网格： ONNX vs Candle在不同线程数和批处理大小下的吞吐量和CPU利用率 *每个X轴刻度是`backend threads/batch-size`。左半部分（`trans ...`）是旧的Candle路径——无论有多少线程或批处理大小有多大，文档/秒在整个网格中停留在5–11，而CPU已被固定。右半部分（`onnx ...`）是新路径——文档/秒在整个扫描中高出一个数量级。在新路径内部：在小批处理时，增加客户端线程有帮助（1T/batch=1 = 72 → 8T/batch=1 = 143）；在大型批处理时，单个客户端线程获胜（1T/batch=64 = 233 是全局峰值）。* ONNX vs Candle效率：不同配置下每% CPU的文档/秒 *相同扫描，但将效率（每% CPU的文档/秒）与文档/秒一起绘制。在Candle（`trans`）侧，两条线都紧贴地面——机器在花费CPU却没有产生文档。在ONNX（`onnx`）侧，效率在1–2线程搭配中等批处理大小时最高，此时每% CPU购买最多的嵌入，即使我们将线程数增加到32，它仍然远高于旧路径。*