小红书向量检索团队获得CIKM 2024最佳应用研究论文奖！

在CIKM 2024大会上，小红书向量检索团队的论文《A Real-Time Adaptive Multi-Stream GPU System for Online Approximate Nearest Neighborhood Search》荣获应用研究赛道最佳论文奖。论文提出了一种实时多流GPU在线近似最近邻检索系统（RTAMS-GANNS），通过动态向量插入算法和多流并发机制，解决了GPU向量检索在高QPS场景下的实时向量插入难题，相关技术已成功应用于搜索与推荐系统。CIKM是信息和知识管理领域的顶级国际会议，由ACM和SIGIR主办，本届应用研究赛道吸引了来自工业界和学术界400余篇投稿，316篇论文进入评审程序，最终有103篇论文被录用，仅一篇论文获评最佳应用研究论文奖。

论文链接：

https://dl.acm.org/doi/10.1145/3627673.3680054

1.1 背景和局限

近似最近邻搜索（Approximate Nearest Neighbor Search, ANNS）已在现代搜索、推荐系统及大模型增强生成（Retrieval-Augmented Generation, RAG）等领域获得广泛应用。其核心目标是通过近似算法高效查找Top-K结果，降低计算成本，使得在海量数据中快速检索成为可能。典型的ANNS算法包括聚类、量化、树结构和图结构等方法，这些方法虽牺牲了一部分召回率，但显著减少了距离计算带来的计算负担。然而，基于CPU的ANNS系统在处理大规模数据时，常面临高计算负载和内存带宽达到瓶颈的问题，为此，研究逐渐转向利用GPU的高算力和显存带宽来提升检索性能。

Faiss和Raft是率先引入GPU用于ANNS的系统，其开创性地将检索任务从CPU转移至GPU，实现了显著的性能提升。这些系统的实现细节已公开，且在工业界获得广泛应用。然而，这些系统多用于离线场景，难以满足在线场景中频繁实时向量插入的需求。实时更新对于搜索和推荐系统尤为关键，因为只有确保最新数据被纳入检索，才能提供更精确且高时效性的搜索结果。现有系统在应对这一需求时，存在两方面的主要局限。

首先，现有系统的向量插入显存管理不够高效。在典型的基于倒排索引的ANNS算法（如IVFFLAT和IVFPQ）中，向量数据通常被划分到多个聚类中心下，检索时仅计算与查询相关的聚类列表，获取最终的Top-K结果。但在新向量插入时，现有系统通常需将新增向量列表从离线段复制到新段，这个过程不仅增加了显存消耗，还导致GPU的显存分配和复制需要频繁触发内核启动，消耗大量资源。值得说明的是，本文仅讨论基于倒排索引的算法，因为对于基于图的算法即使在CPU系统中实现并发实时更新上仍具有相当大的技术挑战。

其次，现有系统多采用单流串行执行模式。尽管在任务依赖性较强的场景下单流模式可以最大化吞吐量，但在GPU任务独立、负载多样化的场景中，单流模式未必能够充分发挥GPU的并行性能。例如，当执行实时向量插入任务时，大量数据复制往往导致长时间等待，阻碍正常检索的执行，尤其在处理热门数据（长向量列表）时，单次实时插入可能需耗费数十毫秒，从而对实时检索带来严重影响。

1.2 贡献

为克服上述不足，本文提出了一种面向GPU的在线实时多流ANNS系统——RTAMS-GANNS，其创新点包括以下三个方面：

1) 设计了一种基于显存块的动态算法用于高效实时向量插入，该方法在避免频繁显存分配和复制的同时，支持新向量无缝整合，并基于原地重排优化了显存的高效利用。

2) 引入多流并行执行模式，通过构建流缓存的动态资源池，在GPU流的并行特性下实现高性能的向量检索与实时插入任务的并发执行，避免了串行执行带来的额外等待，消除了向量扩展带来的阻塞问题。

3）通过在多种基准数据集及工业级数据集上的广泛实验，RTAMS-GANNS验证了其在不同QPS场景中的优越性能，延迟降低达40%-80%。相关系统现已成功部署在真实的工业搜索和推荐系统中，日均为数亿用户提供服务，显著提升了数据实时更新与检索效率，为搜索推荐类应用的用户体验带来了显著改善。

2.1 GPU加速近似最近邻检索

对于近似最近邻检索（ANNS）中的GPU加速方法，Faiss和Raft率先将向量检索系统从CPU迁移至GPU，开启了GPU加速ANNS的新时代。基于此基础，后续研究在算法优化和系统架构上进一步推进。例如，RobustiQ优化了GPU上的量化算法，而其他研究SONG、CAGRA则增强了GPU图算法。在系统优化方面，一些研究探索了CPU-GPU异构调度、数据分配及周期性同步机制，但这些系统依然依赖串行的离线架构，无法有效解决在线场景中多任务特别是实时向量插入的需求。本研究首次提出了一种并行架构的GPU ANNS系统，专注于在线场景中的高频实时向量插入，从系统架构层面确保动态环境下检索结果的时效性和准确性，解决了现有方法中的瓶颈。

2.2 CPU/GPU在线向量插入方法

对于在CPU系统中的实时向量更新方法，由于CPU中向量插入较为直接，例如SPFresh在倒排索引算法中采用增量再平衡技术，优化了向量分区以适应数据分布的变化。然而，该方法主要解决CPU场景下的id列表平衡的数据问题，和我们在GPU上提出的方法想改进执行机制的出发点并不同也无法直接使用，因为GPU插入操作涉及昂贵的显存分配开销、非对齐显存访问和阻塞执行等问题。同样，FreshDiskANN在图算法中提出了双通道的流合并算法，用于将内存中的索引与SSD索引合并，这个方法主要关注于超大数据场景下的更新，而不是如何使更新机制更适合于SSD，而且由于SSD和GPU之间的重大差别，这方法不能直接应用于GPU。此外，一些研究探索了CPU和GPU的异构硬件结合，在CPU与GPU数据交换时加载最新向量，但未根本上解决实时插入问题，仍面临高开销的复制和显存重新分配。

3.1 基于显存块的实时向量插入算法

本节主要介绍我们提出的基于显存块的动态向量插入算法。系统的整体显存布局如上图(b)所示。不同于传统方法中将所有向量放置于连续显存，我们将每个显存块 m 用作链表的节点，连接每个簇c中的向量id列表l。

在基于IVF的算法中，相同簇的显存块均以相同颜色表示，如上图(b)所示。为了减少处理中的显存分配开销，我们设计了一种显存分配方法。

该方法由中央显存池P（几乎占据整个GPU显存）和最小单位显存块m组成。中央显存池P预先按显存块划分，并使用从0到|P|的标号。最新分配显存块的当前位置用curP表示。每当需要分配一个显存块时，我们通过atomicAdd对curP进行操作，通过逻辑定位显存块的位置而无需实际进行额外的显存分配：

对于新插入向量的显存块，其向量id列表l定义为m'：

显存块m包含三个主要部分：头部mh、核心id列表ml和核心向量mv。头部mh存储重要信息，如前一和后一头部地址（prev_header和next_header）及块信息（例如向量容量、大小等等）。核心id列表ml保存属于该簇显存块的向量id，而核心向量mv中的每个向量一起排列或按32维交错排列(Faiss/Raft模式，由GPU warp线程数决定)，以实现coalesced读指令：

3.2 基于流缓存的多流执行机制

在本节中，我们提出了一种并行执行机制，使实时向量插入Kernel能够与GPU上的常规检索排序Kernel并行运行。如上图所示，以往系统通常以串行方式执行任务，这是GPU应用中的常见做法，因为串行执行可以最大化处理器的吞吐量。然而，在在线近似最近邻检索（ANNS）中，单次请求的计算量可能并不足以充分利用GPU的性能，而串行执行实时向量插入的Kernel也可能会阻塞正常的检索排序Kernel。因此，并行系统更适合在线系统。

通过并行执行这些Kernel，我们不仅大幅减少了执行时间（上图中标注为Time Saved），还使实时向量插入更加即时。为实现并行，我们在CUDA中采用多流并行技术，这需要领域知识来设计Kernel的执行机制。所有提交给GPU的Kernel都应在特定流中运行。对于搜索Kernel，每个Kernel将运行在不同的流上，而流资源可以通过资源池复用。对于插入Kernel，则应在同一数据流上运行。因此，最终结构是为搜索Kernel分配多个流，而插入Kernel使用单一流。

在设计时还需要考虑以下两点：

1）系统必须避免在线的显存分配和释放，以防止出现全局阻塞并退化为串行执行机制。

2）系统中的每个Kernel不应占用过多资源，否则会导致阻塞。

对于第一个考虑点，我们借鉴了TCMalloc 中的方法，设计了一个双层的基于流的资源池。首先，为每个批次搜索请求分配一个专用流，该流拥有一个独立的小显存分配（足够支持在线IVF算法的临时使用）。当资源需求较大时，则会从中央显存池中请求一个更大的资源分配，在使用后返还。对于实时向量插入Kernel，我们在设计阶段考虑了并行性和GPU显存的无锁分配（3.1节），因此在运行过程中不会遇到阻塞问题，能够与检索Kernel平滑并行执行。

针对第二个考虑点，由于多个Kernel会在流多处理器（SM）上同时调度，确保单个Kernel不会独占所有计算资源至关重要，因为当线程数量较大时，GPU上每个块的线程数会受到限制从而退化成串行执行。因此，我们会限制每个Kernel的Block数量和线程数量，以防止过度资源占用，尽管这可能会增加计算需求高的在线请求的延迟，不过这种操作在在线场景中也是较为常见和可以接受的。

3.3 部署细节

本节介绍实时多流GPU系统在实际单索引场景中的部署细节。该系统已在T4/A10 GPU服务器上运行六个月以上，成功集成到工业界的搜索和推荐系统中，服务超过1亿日活跃用户。系统由离线和在线两部分组成，分别负责预处理数据和实时向量插入任务。离线部分专注于预处理数据，在线部分则处理实时向量插入。在线模块通过多个CPU线程协调，并从Kafka获取消息，采用动态批处理策略：每秒或累计到128个向量的倍数后触发批处理，单次最大批次为1024向量，之后将批次发送到GPU。在搜索Kernel执行前，我们从资源池中分配32个独立资源，每个资源预留50MB缓存显存。若临时需求超出该值，则从中央显存池额外获取，每次分配200MB。为确保高QPS环境下资源的公平分配，我们设计了无锁队列机制，当所有资源耗尽时拒绝请求。此外，我们为向量插入任务分配了专用流以优化处理。中央显存池P跨搜索和插入两个任务使用。为便于管理，采用高地址给搜索、低地址给插入。资源利用率超过90%时触发警报机制。在Kernel执行期间，我们将每个显存块配置为容纳1024个向量，线程数限制在1024以下，避免多线程冲突。目前，向量插入和重排Kernel仅激活一个块以防资源竞争，这种并行度足以满足当前需求。

4.1 不同QPS下性能比较

实验结果显示，RTAMS-GANNS（所提出方法）在小数据集和大数据集上均显著降低了延迟。首先，RTAMS-GANNS在低QPS_insertion下可减少至多40%的延迟，受益于动态向量插入算法和多流并行机制。请求一到达即被处理，减少至少1毫秒延迟。其次，随着QPS增大，RTAMS-GANNS延迟增长率最低，在高QPS下减少延迟超过80%，尤其是在大数据集上表现出几乎线性增长。相比之下，其他系统在QPS_search超过5000、QPS_insertion超过1000时频繁超时（延迟超过20ms），而RTAMS-GANNS的动态插入算法可在多流并行框架下稳定运行。四个系统中，CPU系统因性能不足导致延迟最高，但在DSSMRT40M上QPS_search较低时呈线性增长，使其在QPS_insertion = 2000时优于Faiss。这一现象符合实践经验，即实时插入操作对CPU系统的正常检索性能影响不大。此外，实验发现Faiss的插入和检索性能均落后于Raft和RTAMS-GANNS。原因在于Faiss的add操作需将向量列表复制至CPU再返回GPU，对大数据集而言效率极低，同时搜索速度较慢。

4.2 重排

其次，我们特别关注重排机制是否会阻塞后续的向量插入，以及是否能够优化在线搜索的性能。我们模拟了一个场景，在同一列表中持续插入向量，并观察在不同重排阈值设置下重排前后延迟的波动，以及重排后的性能优化程度。我们将请求参数固定为QPS_search = 5000和QPS_insertion = 2000。在实验中设置的阈值特别大，通常不会触发到如此程度。我们可以看到，重排时间在50毫秒以内，这对于实时插入是非常理想的。这意味着一旦重排完成，系统会立即更新在此期间等待的向量。此外，重排后我们观察到显存读取和跳转的性能在某种程度上得到了优化，因为向量被重新分组。这导致两个指标的延迟均略微改善了约0.1毫秒。虽然这一改进看起来微小，但仍然具有重要意义。

4.3 显存块大小参数比较

最后，我们将比较不同显存块大小的影响。较大的显存块能够容纳更多的向量，但这也意味着可能会有大量显存闲置。相反，较小的显存块大小可能导致频繁的显存块分配，而在搜索过程中会出现更多的头部跳转，从而影响检索性能。我们将参数同样固定为QPS_search = 5000和QPS_insertion = 2000，专门观察显存块大小变化对性能的影响。我们发现，随着显存块大小的增加，延迟会降低。然而，超过1024时，性能提升的幅度逐渐减小，保留大于1024的显存块会导致收益递减。此外，如果ivf的聚类数量为4000，保留1024的显存块在大数据集下需要至少1GB的填充显存。由于GPU显存是宝贵的资源，保留超过1GB的填充显存会导致显著的浪费。此外，如果某些列表经常更新新向量，那么其他部分显存将永远不会被利用。

在本文中，我们提出了一种实时多流GPU在线近似最近邻检索系统（RTAMS-GANNS）。我们观察到，当前的GPU近似最近邻检索系统主要集中于离线场景，而忽视了在线场景中对高频向量插入的需求。现有系统面临着效率低下甚至全局检索阻塞的问题，这主要是由于需要新的显存分配和在单一流内操作的瓶颈。我们提出的系统通过引入两个创新组件有效地克服了这些挑战：基于显存块的动态向量插入算法和具有流缓存的多流执行架构。实验结果表明，我们的系统能够有效处理不同数据集中的不同QPS水平，将延迟降低高达40%到80%。此外，实验还显示，动态重排机制能够聚合新插入的向量而不影响检索，从而进一步减少延迟。我们还详细介绍了我们的行业部署经验，相关系统在搜索和推荐应用中每日处理超过一亿用户的请求。我们的工作展示了在实际部署场景中稳定、可靠和创新的解决方案。

该论文部分技术已于今年3月申请专利公开

1.https://patents.google.com/patent/CN118096496A

2.https://patents.google.com/patent/CN118350979A