一起来认识 GPU-Cagra 索引！

Zilliz RDS

发布于 2024-07-10 19:19:47

870

发布于 2024-07-10 19:19:47

文章被收录于专栏：Reinvent Data Science

01.Context about GPU-based Index

向量搜索是一个非常典型计算密集型的场景。作为世界上最快的向量数据库，Milvus的向量检索引擎Knowhere占用了超过整个系统80%的计算资源。而当我们谈论到高性能的计算，GPU总是一个绕不开的话题，在向量领域也不例外。

通过跟Nvidia的合作，Milvus成为了世界上第一个支持使用GPU加速的向量数据库。基于Nvidia的向量检索库RAFT，Milvus在2.3版本正式支持了GPU索引，并且以此为基础支持了Nvidia的推荐系统Merlin。

Milvus在2.3版本支持了两种索引，IVFFLAT 和 IVFPQ，并在测试中表现出了远超目前最流行的HNSW的性能。然而，如何在小批量查询中获得性能提升，如何让基于GPU的索引更有性价比等问题让让我们持续探索新的解决方案。

Graph-based的向量搜索算法依靠强大的性能在近年来取代了IVF-based的算法成为了主流。同样在GPU上，工程师们也在积极探索可行的GPU图算法的实现，包括GGNN，SONG等。但是由于Graph-based算法的计算方式，它很难被直接搬到GPU上，这些实现方法都没能在小batch搜索下达到预想的性能。

NVIDIA在最近推出了基于GPU的图索引CAGRA，同时帮助Milvus在最新的2.4版本完成了相关支持。

02.Performance

性能是GPU索引的关键。We evaluated the Milvus' performance through the open-source vector database benchmark tool and compared the performance between CPU-base HNSW, GPU-based IVFFLAT and CAGRA.

为了让benchmark的结果更有现实指导意义，所有测试都在AWS可以获取的host上进行，其中GPU采用了Tesla T4和A10G。测试机的价格区间都在可比较的范围内，如下图。

所有性能测试结果基于top100@98%recall，Client为AWS m6i.4xlarge(16C64G)。The evaluation results are shown in the charts below.

这是在GPU最为不擅长的的小batch(1)场景下Milvus展现出来的性能。在两个数据集中，CAGRA都展现出了接近10倍的性能。

上图为使用更大的搜索batch(10, 100)时，Milvus的性能结果。相比较HNSW，CAGRA展现出来了几十倍的性能提升。

CAGRA除了在向量搜索上表现强悍，在构建索引上同样令人印象深刻，在GPU的帮助下它可以达到十倍左右的索引构建速度。

03.Cagra Introduction

CAGRA BUILD阶段

相比于CPU上的 HNSW 的构建算法，CAGRA放弃了插入-更新的迭代思路。CPU上的 HNSW 构建流程上并发空间很难充分利用GPU的并行能力。为此，CAGRA采用了全新的建图方式。

Cagra首先使用IVFPQ或者NN-DESCENT来构建一个原始图，原始图中，每一个节点的邻居节点的个数degree较多，CAGRA在原始图的基础上，再对所有的邻接边进行重要性排序，剪掉不重要的邻边。然后再进行正反图合并，CAGRA的build 流程如下图所示：

1.原始图的 build algo

IVFPQ

在 IVFPQ模式下，CAGRA会在数据集上先 build 一个IVFPQ索引，得益于PQ索引的量化特点，并不会明显占用太多的显存，然后使用IVFPQ索引对数据集中的每一个点执行 search 任务，将IVFPQ索引找到的近似的最近邻居作为邻接点，从而完成原始图的构建。

NN-DESCENT

CAGRA 进行的传统 NN-DESCENT 建图过程如下：

从数据集 v 中随机选择 k 个点作为初始邻接表 B[v]。
对邻接表 B[v] 取逆，得到反向邻接表 R[v]，将 B 和 R 合并得到 H[v]。
对数据集中的任意节点 v，根据 H[v] 找到所有邻居的邻居，并选取最近的 k 个节点作为其邻居。
重复步骤 2-3，直到 B 不再变化或达到迭代要求。

NN-Descent相比HNSW建图流程，容易并行化，任务和任务之间，需要交互的数据较少。在 Milvus 的实践中，NN-DESCENT原始图构建算法在 GPU上，可以极大的加快 CAGRA邻接图的 build 效率。当然，相比 IVF PQ模式，NN-DESCENT在原始图的质量上要稍差一些。

2.CAGRA剪枝

CAGRA图的剪枝策略，主要基于以下两条准则：

基于通路的重要性排序，并非传统的基于距离的排序方式，基于距离的重要性排序，不一定有利于图的联通性。
反向图合并，对你重要的人，可能你对他也重要。

CAGRA 会对其额外的边进行修剪。在初始图阶段，每个节点的相邻边根据距离具有不同的权重 w。CAGRA 使用

max(w_{x \to z},w_{z \to y}) < w_{x \to y}

在上图中，对节点 x 的邻边进行剪枝，可以发现根据以上的判别式，C点的可绕行通路为 2，B D的可绕行通路为 1，因此，C B或者 C D的重要性就低。

在对正向图进行基于路径的修剪后，对所有的边取反，然后分别从正向图和反向图中各取 1/2 的边进行合并，生成最终的 CAGRA 图。

CAGRA SEARCH

CAGRA的search过程主要围绕着一块固定部分有序的优先队列进行。

CAGRA的搜索如上图所示，CAGRA使用一个顺序内存缓冲区，包括一个内部的 top-M 列表（通常在其他算法中称为优先队列）和它的候选列表，如图6顶部所示。内部 top-M 列表的长度为 M(≥ k)，候选列表的长度为 p × d，其中 p 是在每次迭代中遍历的图的源节点数量，d 是 CAGRA 图的度。每个缓冲区元素是一个键/值对，包含节点索引和节点与查询之间的相应距离。

随机抽样（初始化步骤）：我们使用伪随机数生成器选择 p × d 个均匀随机的节点索引，并计算每个节点与查询之间的距离。结果存储在候选列表中。我们将内部的 top-M 列表设置为虚拟条目（填充 FLT_MAX），在下一次的排序中，可以不影响排序完成后的头部结果，FLT_MAX必然排序到末尾。
内部 top-M 列表更新：我们从整个缓冲区中挑选出具有最小距离的 top-M 个节点，并将结果存储在内部的 top-M 列表中。
候选列表索引更新（图遍历步骤）：我们选择内部 top-M 列表中顶部 p 个节点的所有邻居节点，过滤掉曾经作为父节点的节点（通过hash 表过滤）。把这些节点存储在候选列表中。此步骤不计算候选列表中每个节点与查询之间的距离。
距离计算：仅当节点第一次出现在查询的候选列表中时，我们才计算它与查询向量之间的距离。此条件分支修剪了不必要的计算，因为如果在之前的迭代中已经计算过距离，则不需要重新计算。如果一个节点已经出现在候选列表中且距离已经计算过，则：