Manticore中更快的KNN搜索：两遍HNSW、批量距离计算和AVX-512支持

# Manticore 中的更快 KNN 搜索：双遍 HNSW、批量距离计算与 AVX-512 来源：https://medium.com/@s_nikolaev/faster-knn-search-in-manticore-2-pass-hnsw-batched-distances-and-avx-512-b85604647aab Sergey Nikolaev (https://medium.com/@s_nikolaev?source=post_page---byline--b85604647aab---------------------------------------) 按回车键或点击可查看完整尺寸图片 **TL;DR：** 对 HNSW 搜索引擎的三项改动，使得 KNN 吞吐量在高 k 值时提升高达 29%，并发负载下增益超过 20%。无需 API 变更、无需重建索引、无需配置。只是搜索更快了。 ## Manticore 中更快的 KNN 搜索 Manticore 的 KNN 搜索构建于 `hnswlib`（https://github.com/nmslib/hnswlib）之上，这是一个开源的 HNSW 实现。历史上，我们大部分的 KNN 工作集中在自定义距离函数上，例如用于二值量化的那些，而非聚焦于 `hnswlib` 的核心搜索循环。我们还添加了诸如使用 ACORN-1（https://manual.manticoresearch.com/Searching/KNN#Filtering-strategies:-prefilter-vs.-postfilter）进行预过滤以及提前终止（https://manual.manticoresearch.com/Searching/KNN#Filtering-strategies:-prefilter-vs.-postfilter）等功能，但主要搜索循环保持不变：hnswlib 仍然以相同的方式访问邻居、计算距离并维护其候选集。 这次的改动更进一步，对 hnswlib 的核心搜索循环本身进行了修改——重构了遍历邻居的方式、调用距离函数的方式以及与 CPU 内存层次结构的交互方式。结合列式库中新增的 AVX-512 距离实现，这些改动针对三个开销来源：低效的内存访问模式、冗余的数据加载以及间接函数调用的开销。 ## 编译时距离函数特化 此前，距离函数是一个运行时函数指针，存储在 HNSW 索引中，并为每个候选者调用。对于大的搜索预算，这意味着每次查询会涉及大量间接调用。间接调用阻止编译器将距离函数内联到搜索循环中，并且会产生分支预测开销。 新代码使用 C++ 模板在编译时解析距离函数。当搜索开始时，一个单一的 switch 语句根据距离度量和量化设置选择正确的模板特化。此后，整个内层循环——邻居遍历、距离计算、候选集更新——作为一个单一的函数运行，距离计算被完全内联。现在编译器可以在距离计算边界上优化寄存器分配、指令调度和循环展开。 ## 双遍邻居处理 HNSW 算法通过访问节点并计算与其邻居的距离来探索图。在原始实现中，每个邻居在一个单遍中处理：检查是否已访问，获取其向量数据，计算距离，更新候选集。这意味着在需要数据之前，内存预取提示几乎没有时间生效。 新的实现将其拆分为两个遍。第1遍遍历当前节点的所有邻居，跳过已访问的邻居，并将未访问的邻居收集到一个小的批次数组中。当每个邻居被添加到批次中时，会为其向量数据发出预取提示。第2遍遍历批次并计算距离。到第2遍到达每个向量时，第1遍的预取已经将数据带入缓存。 第2遍遍历的是一个紧凑的候选 ID 顺序数组，而不是图结构本身。底层的向量加载仍然是分散的，但数据已经被提前预取。 对于未过滤的查询（KNN 搜索上没有 WHERE 子句），新代码还采用了一条快速路径，完全消除了每个候选者的过滤检查。 ## 批量距离计算 双遍结构在两个方向上有帮助：它为预取争取了更多时间，并且使批量化变得容易。一旦第2遍拥有一个紧凑的候选列表，它就可以一次对两个候选者进行评分，而不是逐个进行。 在对两个候选者进行评分时，查询向量在每个 SIMD 迭代中加载一次，并重复用于两个距离计算，从而消除了冗余加载。 这减少了重复的查询端加载，并让评分循环能够成对处理候选者，对奇数余量有回退处理。批量-2 函数针对内积、L2 及其二值量化变体提供。 ## AVX-512 支持 新的 AVX-512 距离代码每次迭代处理 16 个浮点数，而 AVX2 为 8 个。对于内积和 L2 距离，核心循环使用融合乘加（`_mm512_fmadd_ps`），该指令将乘法和累加合并为一条指令。对于二值量化向量，AVX-512 VPOPCNTDQ 扩展加速了距离计算中使用的位计数操作。 Manticore 现在提供三个库变体：基础构建、AVX2 构建和 AVX-512 构建。在启动时，Manticore 检测 CPU 的能力并自动加载相应的库。无需配置。 ## 基准测试结果 以下基准测试在 `dbpedia-openai-1M-1536-angular`（https://storage.googleapis.com/ann-filtered-benchmark/datasets/dbpedia_openai_1M.tgz）数据集上运行（100 万个向量，1536 维，余弦距离），使用 AMD Ryzen 7 9700X（Zen 5，8 个物理核心 / 16 个逻辑核心）。所有数据使用 1 位二值量化，禁用了过采样和重新评分。对于多线程运行，吞吐量报告为每线程平均每秒查询数：每个工作线程运行自己的查询批次，独立测量其 QPS，最终数字是所有工作线程的平均值。每个结果是 6 次独立运行的平均值。同时禁用了提前终止，以隔离这些优化对原始 HNSW 遍历的影响。 选择 Zen 5 是因为它支持带有本地 512 位数据路径的 AVX-512，避免了与一些较旧 Intel 处理器相关的分离 512 执行行为和严重的 AVX-512 降频问题。这有助于将算法效果与 CPU 特定的 AVX-512 节流行为隔离开来。 ## 纯算法改进 第一个图表通过将新的 AVX2 构建与之前的 AVX2 构建进行比较，隔离了算法改动（双遍处理、批量距离、编译时分派）的效果。两个构建使用相同的 SIMD 指令集，因此差异完全来自新的代码结构。 按回车键或点击可查看完整尺寸图片 在单线程上，增益从 k=10 时的 +3% 稳步增长到 k=1000 时的 +24%，因为距离计算成为搜索工作负载的主导。随着更多线程竞争内存带宽，每线程增益缩小：在 4 或 8 个线程时为 +9–10%，在 16 个线程时仅为 +2–5%。 16 线程的情况是 SMT（每个物理核心运行两个线程）。距离计算受内存限制，因此当两个线程共享核心的 L1/L2 缓存时，预取和批量的优势部分被共享资源争用所抵消。算法改进仍然有帮助，但余量缩小。 ## SIMD 宽度收益（AVX-512 与 AVX2） 第二个图表通过比较 AVX-512 构建与新的 AVX2 构建（两者共享相同的算法改进）隔离了 AVX-512 的效果。 按回车键或点击可查看完整尺寸图片 在 k=10 时，AVX-512 比 AVX2 稍慢（约 -2%），与线程数无关。这是 AVX-512 特有的：仅算法改进并未表现出此退化，因此它不是统一的每查询开销。从 k=30 开始，AVX-512 在所有线程数下都领先。 有趣的是，AVX-512 的收益随线程数增长。尽管此基准测试禁用了过采样，但默认的 Manticore KNN 查询使用 `LIMIT 20`，而默认的 `oversampling=3.0`（将量化搜索后的有效 HNSW 搜索预算乘以重新评分）使得内部 k 变为 60。在 k=60 时，AVX-512 相对于新 AVX2 在单线程上为 +1.2%，4 线程为 +2.6%，8 线程为 +3.4%，16 线程为 +6.5%。 ## 综合改进（AVX-512 与旧代码） 第三个图表展示了累积效果：带有所有新代码的 AVX-512 与之前的 AVX2 构建相比较。这是如果用户从之前的 Manticore 版本升级到新版本，且其 CPU 支持 AVX-512 时将会看到的效果。 按回车键或点击可查看完整尺寸图片 单线程曲线从 k=10 时的 +0.5% 上升到 k=1000 时的 +29%。多线程曲线在 k=1000 时都达到 +22–24%。改进广泛分布在各线程数上——算法和 SIMD 增益在不同并发级别下以不同方式组成，但组合结果在中高 k 值时始终很大。 ## 为什么增益随 k 增长 所有三个图表都呈现相同的形状：低 k 时改进很小，高 k 时改进很大。原因是低 k 查询将更大部分时间花在图遍历上（访问节点、检查已访问位、弹出候选集）——这些工作与图结构规模相关，而非 k。随着 k 增长，有效搜索预算按比例增长，查询在距离计算上花费更多时间。这些优化针对的是距离计算及其周围的循环，因此它们的收益随距离计算所代表的工作份额而缩放。 ## 这对你意味着什么 这些改进无需任何操作。它们在最近的 Manticore Search 27.1.5 版本（https://manticoresearch.com/blog/manticore-search-27-1-5/）中可用；没有 API 变更，没有新的配置选项，也不需要重建索引。 这些增益与 KNN 提前终止（https://manual.manticoresearch.com/Searching/KNN#Early-termination）叠加：提前终止减少了每次查询的距离计算次数，而这些优化使得每次计算更快。 最大的改进出现在以下情况： - **高维向量**（每次距离计算更多的算术运算，更多的 SIMD 收益） - **大的 k 值**（更多的总距离计算，更多的批量和缓存优化机会） - **带有过采样的查询**（过采样乘以有效 k，将查询推入增益最大的范围） ## 进一步阅读 - KNN 提前终止文档（https://manual.manticoresearch.com/Searching/KNN#Early-termination）——Manticore 如何检测 HNSW 收敛并提前停止 - KNN 过滤文档（https://manual.manticoresearch.com/Searching/KNN#Filtering-strategies:-prefilter-vs.-postfilter）——预过滤如何与 ACORN-1 配合工作 - KNN 向量搜索参考（https://manual.manticoresearch.com/Searching/KNN）——完整语法、参数、量化、重新评分