利用 NVIDIA cuVS 加速 Faiss 中的 GPU 向量搜索

随着企业收集的非结构化数据日益增多，并更加广泛地应用大语言模型（LLM），对更高效、更具可扩展性系统的需要也愈发迫切。以检索增强生成（RAG）为代表的高级信息检索工具，在处理海量数据——有时达到TB甚至PB级别——时，可能耗时数小时乃至数天。

与此同时，广告推荐系统等在线搜索应用难以仅依靠CPU实现实时响应。为满足实时性需求，往往需要动用数千个CPU，导致基础设施成本大幅上升。

本文将探讨如何利用 NVIDIA cuVS 和 Meta Faiss 库应对密集向量在相似性搜索与聚类中的挑战。cuVS 借助 GPU 加速，显著提升了搜索索引的构建速度以及实际查询的执行效率。通过该方案，可实现更快速、更经济且更高性能的向量检索，同时保持 CPU 与 GPU 之间的良好兼容性。

具体而言，本文内容包括：

• 融合 cuVS 与 Faiss 的优势
• cuVS 在向量搜索性能提升中的作用及其适用场景
• GPU 加速的反向文件索引（IVF）与基于图的索引性能分析
• 通过基准测试与 Python 代码示例，展示如何利用 Faiss 构建并搜索由 cuVS 加速的索引

集成 cuVS 和 Faiss 有哪些优势？

无论是每秒查询数百万个向量、处理大型多模态嵌入，还是利用 GPU 构建大规模索引，与 Faiss 的 cuVS 集成都能显著提升性能与灵活性。

cuVS 使您能够：

• 在 GPU 上实现高达 12 倍的索引构建速度提升，同时保持 95% 的召回率
• 搜索延迟最高可降低至原来的八分之一，召回率仍达 95%
• 支持在 GPU 和 CPU 环境间灵活迁移索引，满足多样化的部署需求

Faiss 中的 GPU 加速

Faiss 是一个广泛应用于研究和生产环境的向量搜索库。它既可独立使用，也支持与 PyTorch 集成，同时还可嵌入到 RocksDB、OpenSearch 和 Milvus 等向量数据库中。

Faiss 于 2018 年率先提供 GPU 支持，并在此后持续发展。在 NeurIPS 2021 big-ann-benchmarks 竞赛中，NVIDIA 凭借其 GPU 加速算法取得了领先成绩。这些技术随后被贡献至 Faiss，并已集成到开源的 cuVS 库中运行。

自 Faiss v1.10.0 版本起，用户可选择使用 cuVS 来增强反向文件索引算法，包括 IVF-PQ、IVF-Flat、Flat（即暴力搜索），以及 CAGRA（基于 CUDA 的图结构近似最近邻算法）——一种专为 GPU 从头构建的高性能图索引技术。

轻松实现 CPU 与 GPU 的互操作性

使用 cuVS 加速 Faiss 中的 GPU 索引 显著提升了 CPU 与 GPU 之间的互操作性。借助 Faiss，用户可以在 GPU 上构建索引，随后将其部署到 CPU 上运行。这使得用户既能享受 GPU 带来的索引构建加速，又能保留原有的 CPU 搜索架构。整个过程可在 Faiss 库中无缝完成。

例如，分层可导航小世界（HNSW）索引在 CPU 上的构建过程非常缓慢，尤其在处理大规模数据时，可能需要数小时甚至数天。相比之下，CAGRA 索引的构建速度可提升多达 12 倍。此外，CAGRA 图结构可以在 Faiss 中转换为 HNSW 格式，随后部署到 CPU 上用于搜索。

使用 cuVS 对 Faiss 进行基准测试

我们对两个数据集进行了性能基准测试，比较了启用与不启用 cuVS 的 Faiss 表现。

1. Deep100M： Deep1B 数据集的一个包含 100M 个向量的子集，每个向量具有 96 个维度。
2. OpenAI 文本嵌入: 基于文本嵌入模型 ada-002 生成的 500 万个向量，每个向量具有 1536 个维度。

测试在 NVIDIA H100 Tensor Core GPU 和英特尔至强 Platinum 8480CL CPU 上进行，测量对象为：

• 索引构建耗时
• 单次查询延迟（在线搜索）
• 大规模批量吞吐量（离线搜索）

随着非结构化数据快速增长，持续提升索引构建性能变得至关重要。然而，若不兼顾搜索性能和生成模型的质量，单纯衡量索引构建时间便失去了意义。为此，该团队设计了一套专属的索引构建基准测试方法。更多详细信息，请参阅 cuVS 文档。

除了考虑搜索性能和质量外，还需将模型与性能较优的参数配置进行对比，这一过程通过帕累托曲线实现，以确保比较的公平性。在95%召回率的条件下，对不同索引的延迟和吞吐量加速效果进行了比较。

IVF：cuVS 与经典版 Faiss GPU 的对比

我们首先对 IVF 索引 IVF-Flat 和 IVF-PQ 进行了基准测试，比较了 Faiss 的经典 GPU 实现与支持 cuVS 的新 Faiss 版本。

• 构建速度：采用 cuVS 后，IVF-PQ 和 IVF-Flat 的构建速度提升了 4.7 倍（图 1）。
• 搜索延迟：IVF-PQ 的搜索延迟最高降低至原来的 1/8，IVF-Flat 的搜索延迟降低了 90%（图 1）。
• 吞吐量：cuVS 在两个数据集上将 IVF-PQ 的大批量搜索吞吐量提升了 3 倍（图 2），同时在 IVF-Flat 上保持了相近的性能表现，展现出处理大规模离线搜索任务的强大能力。

在线延迟

图 1a 和 1b 展示了各 IVF 指数变体在在线搜索延迟和索引构建时间方面的表现。与传统的 Faiss 相比，cuVS 能够实现更快的索引构建速度，并在两个数据集上显著降低搜索延迟。

批量 (离线) 吞吐量

图2展示了IVF指数变体的批量吞吐量。cuVS能够提升批处理性能，在图像和文本嵌入任务中实现更高的每秒查询处理量。

这些改进得益于更优的 GPU 集群（例如，平衡 k-means）、更广泛的参数支持（例如，IVF-PQ 中增加的子量化器）以及代码层面的优化。

基于图的索引：cuVS CAGRA 与 Faiss HNSW（CPU）的对比

CAGRA 是一种针对 GPU 优化的固定度数平面图索引，在性能上相较于基于 CPU 的 HNSW 具有显著优势，包括：

• 构建时间: CAGRA 的构建速度相比之前可提升达 12.3 倍（图 3）
• 延迟: 在 Deep100M 数据集上，在线搜索速度最高可提升至原来的 4.7 倍（图 3）
• 吞吐量: 在离线搜索场景中，CAGRA 可将图像数据的吞吐量提升 18 倍，文本嵌入的吞吐量提升超过 8 倍（图 4），尤其适用于需要在低延迟条件下处理大规模推理任务的工作负载

cuVS 可将 CAGRA 图直接转换为 HWSW 图，从而在 GPU 上实现更高效的图构建，同时利用 CPU 进行搜索，其速度和质量可与之媲美。

在线延迟

图 3a 和 3b 展示了 GPU CAGRA 与 CPU HNSW 在线延迟和索引构建时间的对比。与 Deep100M 数据集上基于 CPU 的 HNSW 相比，CAGRA 能显著加快索引构建速度，并降低在线查询延迟，搜索性能最高可提升 4.7 倍。

批量 (离线) 吞吐量

图4展示了GPU上的CAGRA与CPU上的HNSW在批量吞吐量方面的对比情况。在批量查询场景下，CAGRA表现出较高的吞吐性能，每秒可处理数百万次查询，并且在两个数据集上的表现均优于基于CPU的HNSW。

如何在 Faiss 中开始使用 cuVS

本节简要介绍在 cuVS 支持下安装 Faiss 的步骤，并提供使用 Python 创建和搜索索引的代码示例。

安装

您可以选择使用 cuVS或预构建的 Conda 软件包来构建 Faiss。

# Conda install (CUDA 12.4)
conda install -c rapidsai -c conda-forge -c nvidia pytorch::faiss-gpu-cuvs 
'cuda-version>=12.0,<=12.9'

或者，您可以使用以下命令安装启用了 cuVS 的 Faiss 软件包的最新预发布版本：

conda install -c rapidsai -c rapidsai-nightly -c conda-forge -c nvidia 
pytorch/label/nightly::faiss-gpu-cuvs 'cuda-version>=12.0,<=12.9'

内存管理

使用以下代码段通过 RMM 启用 GPU 显存池化（推荐）。该方法有助于提升性能。

import rmm
pool = rmm.mr.PoolMemoryResource(
    rmm.mr.CudaMemoryResource(),
    initial_pool_size=2**30
)
rmm.mr.set_current_device_resource(pool)

使用 cuVS 构建 IVFPQ 索引

借助 faiss-gpu-cuvs 软件包，cuVS 可自动应用于受支持的索引类型，用户无需修改代码即可享受其带来的性能提升。以下展示了使用 cuVS 后端创建 IVFPQ 索引的示例：

import faiss
import numpy as np

np.random.seed(1234)
xb = np.random.random((1000000, 96)).astype('float32')
xq = np.random.random((10000, 96)).astype('float32')
xt = np.random.random((100000, 96)).astype('float32')

res = faiss.StandardGpuResources()
# Disable the default temporary memory allocation since an RMM pool resource has already been set.
res.noTempMemory()

# Case 1: Creating cuVS GPU index
config = faiss.GpuIndexIVFPQConfig()
config.interleavedLayout = True
index_gpu = faiss.GpuIndexIVFPQ(res, 96, 1024, 96, 6, faiss.METRIC_L2, config) # expanded parameter set with cuVS (bits per code = 6).
index_gpu.train(xt)
index_gpu.add(xb)

# Case 2: Cloning a CPU index to a cuVS GPU index
quantizer = faiss.IndexFlatL2(96)
index_cpu = faiss.IndexIVFPQ(quantizer,96, 1024, 96, 8, faiss.METRIC_L2)
index_cpu.train(xt)
co = faiss.GpuClonerOptions()
index_gpu = faiss.index_cpu_to_gpu(res, 0, index_cpu, co)
# The cuVS index now uses the trained quantizer as it's IVF centroids.
assert(index_gpu.is_trained)
index_gpu.add(xb)
k = 10
D, I = index_gpu.search(xq, k)

构建 cuVS CAGRA 索引

以下示例展示了如何结合 Faiss 和 cuVS 加速功能来构建并查询 CAGRA 索引。

import faiss
import numpy as np

# Step 1: Create the CAGRA index config
config = faiss.GpuIndexCagraConfig()
config.graph_degree = 32
config.intermediate_graph_degree = 64

# Step 2: Initialize the CAGRA index
res = faiss.StandardGpuResources()
gpu_cagra_index = faiss.GpuIndexCagra(res, 96, faiss.METRIC_L2, config)

# Step 3: Add the 1M vectors to the index
n = 1000000
data = np.random.random((n, 96)).astype('float32')
gpu_cagra_index.train(data)

# Step 4: Search the index for top 10 neighbors for each query.
xq = np.random.random((10000, 96)).astype('float32')
D, I = gpu_cagra_index.search(xq,10)

将 CAGRA 转换为 HNSW (用于 CPU 搜索)

可通过新的 faiss.IndexHNSWCagra CPU 类别将 CAGRA 索引自动转换为 HNSW 格式，从而实现 GPU 加速的索引构建，随后支持基于 CPU 的搜索。

# Create the HNSW index object for vectors with 96 dimensions.
M = 16
cpu_hnsw_index = faiss.IndexHNSWCagra(96, M, faiss.METRIC_L2)
cpu_hnsw_index.base_level_only=False

# Initializes the HNSW base layer with the CAGRA graph. 
gpu_cagra_index.copyTo(cpu_hnsw_index)

# Add new vectors to the hierarchy.
newVecs = np.random.random((100000, 96)).astype('float32')
cpu_hnsw_index.add(newVecs)

有关完整的代码示例，请参考 Faiss cuVS 笔记本。

从向量中获取更多信息

将 NVIDIA cuVS 集成到 Faiss 后，能够显著提升近似最近邻（ANN）搜索的速度与可扩展性。无论您使用的是倒排文件（IVF）索引，还是基于图的方法，cuVS 与 Faiss 的集成均能带来性能的明显优化。

• 索引构建速度显著提升：在 GPU 上实现最高达 12 倍的加速
• 搜索延迟大幅降低：实时搜索性能提升最高达 4.7 倍
• 无缝支持 CPU 与 GPU 互操作：可在 GPU 上构建索引、CPU 上执行搜索，也可反向操作

该团队还推出了 CAGRA，一种专为 GPU 设计的高性能图索引，在构建时间和吞吐量方面均优于传统的基于 CPU 的 HNSW。更值得一提的是，CAGRA 图可以转换为 HNSW 格式，从而支持基于 CPU 的高效推理，实现 GPU 加速构建与 CPU 灵活部署的优势结合，适用于混合场景下的应用需求。

无论您需要扩展搜索基础设施以每秒处理数百万次查询，还是希望快速尝试新的嵌入模型，将 Faiss 与 cuVS 集成都能为您提供有力工具，助力更高效、更智能地迭代，并自信地完成部署。

准备好了吗？安装 faiss-gpu-cuvs 软件包并查看 示例 Notebook。