⚠️ 此页面为自动翻译,翻译可能不完美。
blog-post

Manticore中更快的K近邻搜索:2次遍历HNSW、批量距离计算和AVX-512

author image

TL;DR: 对HNSW搜索引擎的三项改进使高k值下的K近邻吞吐量提升最高达29%,在并发负载下提升超过20%。无需API变更、无需重建索引、无需配置。只需更快的搜索。

Manticore中更快的K近邻搜索

Manticore的K近邻搜索基于hnswlib ,这是一个开源的HNSW实现。历史上,我们的大部分K近邻工作集中在自定义距离函数上,例如用于二进制量化的方法,而不是hnswlib的核心搜索循环。我们还添加了诸如使用ACORN-1预过滤提前终止 等功能,但主搜索循环保持不变:hnswlib仍然访问邻居、计算距离,并以相同方式维护候选集。

这些改进更进一步,修改了hnswlib的核心搜索循环本身——重构了如何遍历邻居、如何调用距离函数以及如何与CPU的内存层次结构交互。结合列式库中新的AVX-512距离实现,这些改进针对三个性能瓶颈:低效的内存访问模式、冗余的数据加载和间接函数调用开销。

编译时距离函数特化

此前,距离函数是一个存储在HNSW索引中的运行时函数指针,并为每个候选项调用。对于大规模搜索预算,这意味着每个查询可能有大量的间接调用。间接调用会阻止编译器将距离函数内联到搜索循环中,并产生分支预测开销。

新代码使用C++模板在编译时解析距离函数。当搜索开始时,一个单一的switch语句根据距离度量和量化设置选择正确的模板特化。从那时起,整个内循环——邻居遍历、距离计算、候选集更新——作为一个单体函数运行,距离计算完全内联。现在编译器可以跨距离计算边界优化寄存器分配、指令调度和循环展开。

两遍邻居处理

HNSW算法通过访问节点并计算其邻居的距离来探索图。在原始实现中,每个邻居在单次遍历中被处理:检查是否已访问,获取其向量数据,计算距离,更新候选集。这意味着在数据需要之前,内存预取提示几乎没有时间生效。

新实现将其拆分为两遍。第一遍遍历当前节点的所有邻居,跳过已访问的邻居,并将未访问的邻居收集到一个小的批处理数组中。当每个邻居被添加到批处理时,会为其向量数据发出预取提示。第二遍遍历该批处理并计算距离。当第二遍到达每个向量时,第一遍的预取已经将数据加载到缓存中。

第二遍遍历的是候选ID的紧凑顺序数组,而不是图结构本身。底层的向量加载仍然是分散的,但数据已经被提前预取。

对于未过滤的查询(K近邻搜索没有WHERE子句),新代码还采用了一条快速路径,完全消除了每个候选的过滤检查。

批量距离计算

两遍结构在两个方面有所帮助:它为预取提供了更多时间,也使批处理变得容易。一旦第二遍有了紧凑的候选列表,它就可以一次处理两个候选,而不是逐个处理。

在评分两个候选时,查询向量在每个SIMD迭代中加载一次,并用于两个距离计算,消除了冗余加载。

这减少了重复的查询端加载,并允许评分循环成对处理候选,对于奇数余数有回退方案。为内积、L2距离及其二进制量化变体提供了批量2函数。

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 数据集(1M向量,1536维,余弦距离)上运行,使用AMD Ryzen 7 9700X(Zen 5,8个物理核心/16个逻辑核心)。所有数据使用1位二进制量化,禁用过采样和重评分。对于多线程运行,吞吐量报告为每线程平均查询每秒:每个工作者运行自己的查询批次,其QPS独立测量,最终数字是跨工作者的平均值。每个结果是6次独立运行的平均值。还禁用了提前终止,以隔离这些优化对原始HNSW遍历的影响。

Zen 5 被选中是因为它支持原生 512 位数据路径的 AVX-512,避免了某些较旧的英特尔处理器与拆分 512 执行行为和重大的 AVX-512 降频相关的性能问题。这有助于将这些变化的算法影响与 CPU 特定的 AVX-512 限速行为隔离开来。

仅算法改进

第一张图表通过将新的 AVX2 构建与之前的 AVX2 构建进行比较,隔离了算法更改(两遍处理、批量距离、编译时调度)的影响。两个构建使用相同的 SIMD 指令集,因此差异纯粹来自新的代码结构。

AVX2(新)与 AVX2(旧)吞吐量提升对比

在单线程情况下,从 k=10 的 +3% 到 k=1000 的 +24%,随着距离计算在搜索工作负载中占主导地位,增益稳步增长。随着更多线程竞争内存带宽,每个线程的增益缩小:在 4 或 8 个线程时为 +9-10%,在 16 个线程时仅为 +2-5%。

16 线程情况是超线程(每个物理核心运行两个线程)。距离计算是内存密集型的,当两个线程共享一个核心的 L1/L2 缓存时,预取和批量处理的优势部分被共享资源竞争所抵消。算法改进仍然有帮助,但提升空间缩小了。

SIMD 宽度优势(AVX-512 与 AVX2)

第二张图表通过将 AVX-512 构建与新的 AVX2 构建进行比较,隔离了 AVX-512 的影响(两者都具有相同的算法改进)。

AVX-512 与 AVX2(新)吞吐量差异

在 k=10 时,无论线程数如何,AVX-512 都比 AVX2 稍慢(约 -2%)。这是 AVX-512 特有的:仅算法改进不会显示这种回归,因此这不是每个查询的统一开销。从 k=30 开始,AVX-512 在每个线程数上都超越了 AVX-2。

有趣的是,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 会看到的效果。

AVX-512(新)与 AVX2(旧)吞吐量提升对比

单线程曲线从 k=10 的 +0.5% 上升到 k=1000 的 +29%。多线程曲线在 k=1000 时都达到 +22-24%。改进在不同线程数上广泛分布——算法和 SIMD 改进在不同并发级别上组合不同,但综合结果在中等至高 k 值时始终显著。

为什么增益随 k 增大而增长

所有三张图表都显示了相同的形状:低 k 时改进较小,高 k 时改进较大。原因是低 k 查询在图遍历(访问节点、检查已访问位、弹出候选集)上花费更多时间——这些工作与图结构相关,而不是 k。随着 k 增大,有效搜索预算按比例增长,查询在距离计算上花费更多时间。优化针对距离计算及其周围的循环,因此它们的增益与距离计算在工作量中所占比例成正比。

这对您意味着什么

这些改进不需要您采取任何行动。它们已在最近的 Manticore Search 27.1.5 版本 中提供;没有 API 更改,没有新的配置选项,也不需要重建索引。

这些增益与 KNN 提前终止 叠加:提前终止减少了每个查询的距离计算次数,而这些优化使每次计算更快。

最大的改进出现在以下情况:

  • 高维向量(每次距离计算涉及更多算术运算,更多 SIMD 优势)
  • 大的 k 值(总距离计算次数更多,更多批量处理和缓存优化的机会)
  • 带有过采样的查询(过采样乘以有效 k,使查询进入增益最大的范围)

进一步阅读

安装Manticore Search

安装Manticore Search