跳到主要内容

cuVS使用指南

用法说明​

在GPU索引配置中，可以通过 use_cuvs=True 参数启用 cuVS（CUDA向量搜索，CUDA Vector Search）的实现。需要注意的是，如果在编译Faiss源码时设置了 FAISS_ENABLE_CUVS=ON，或者使用了 faiss-gpu-cuvs 的conda包，则对于支持的索引类型，该标志会自动设为True。如果设置为 use_cuvs=False，则会使用传统的Faiss GPU实现。

构建cuVS GPU索引​

RMM（Rapids内存管理器）​

cuVS内部使用 RMM（Rapids Memory Manager），用于自定义设备端（GPU）和主机端（CPU）内存分配模式。当启用cuVS时，Faiss的GPU资源管理也会配置为使用RMM。下面的示例演示如何构建一个初始大小为1 GiB的最佳适应池（best fit pool）分配器。该内存池使用 rmm::mr::cuda_memory_resource 作为“上游”资源，由其完成实际的设备分配。

C++ 示例​

#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/per_device_resource.hpp>
// 为当前设备设置一个1 GiB的内存池资源，所有分配操作将共用此内存池。
rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
    rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull);
rmm::mr::set_current_device_resource(&pool_mr);

Python 示例​

import rmm
pool = rmm.mr.PoolMemoryResource(rmm.mr.CudaMemoryResource(),
                                 initial_pool_size=2 ** 30)
# 为当前设备设置RMM资源
rmm.mr.set_per_device_resource(pool)
# 或为特定设备设置RMM资源（如设备编号0）
rmm.mr.set_per_device_resource(0, pool)

在启用了cuVS的构建环境下，Faiss的 StandardGpuResources 对象会使用当前用户设置的 RMM 设备资源来进行GPU内存分配。

提示

RMM的Python接口不是Faiss的直接依赖，需要用户自行安装：

conda install -c rapidsai -c conda-forge rmm

设置这样的分配器通常能够获得比默认CUDA内存资源更优的检索性能。下一步是在索引配置中将 use_cuvs 标志设为True。

C++ 示例​

faiss::gpu::StandardGpuResources res;
faiss::gpu::GpuIndexIVFFlatConfig config;
config.use_cuvs = true;
faiss::gpu::GpuIndexIVFFlat index_gpu =
    faiss::gpu::GpuIndexIVFFlat(res, d, nlist, faiss::METRIC_L2, config);

Python 示例​

res = faiss.StandardGpuResources()
co = faiss.GpuIndexIVFFlatConfig()
co.use_cuvs = True
index_gpu = faiss.GpuIndexIVFFlat(res, ncols, nlist, faiss.METRIC_L2, co)

设置完配置后，可以像平时一样，使用常规Faiss GPU的API（如 add 增加向量，search 检索）对该索引对象进行操作。这些操作会在底层自动应用到cuVS索引上。

index_gpu.add(xb)
D, I = index.search(xq, k)

克隆CPU索引到GPU（cuVS）​

同样可以通过 use_cuvs 参数，将一个CPU索引直接克隆为cuVS索引。

C++ 示例​

faiss::gpu::GpuClonerOptions co;
co.use_cuvs = true;
index_gpu = faiss.GpuIndexIVFFlat(res, ncols, nlist, faiss::METRIC_L2, co)
index_gpu = faiss::gpu::index_cpu_to_gpu(
                res,
                0,
                &index_cpu,
                &co);

Python 示例​

co = faiss.GpuClonerOptions()
co.use_cuvs = True
index_gpu = faiss.index_cpu_to_gpu(res, 0, index_cpu, co)

CAGRA + HNSW融合用法​

除了CAGRA（高效图结构索引，Compressed Approximate Graph）索引可在GPU上加速检索外，还可以通过克隆GpuIndexCagra对象，来加速HNSW（分层导航小世界图，可用于高效近邻搜索）CPU索引的构建。这通过 IndexHNSWCagra 类实现。在GPU上利用更快（CAGRA）的索引构建来初始化HNSW索引的底层结构。

IndexHNSWCagra 有一个参数 base_level_only：

  • 如果 base_level_only=True，只用CAGRA图初始化HNSW索引的底层，此时得到的HNSW索引不可变，之后不能再添加新向量。
  • 如果 base_level_only=False，除了用CAGRA初始化底层外，还能像普通HNSW一样继续用添加API增加新向量。

C++ 示例​

faiss::gpu::GpuIndexCagra gpu_cagra_index(res, d);
// 针对CAGRA索引，train阶段会用所有向量构建图结构。
gpu_cagra_index.train(n, xb);
faiss::IndexHNSWCagra* cpu_hnsw_index;
// 若 base_level_only = false，则HNSW还可以添加新向量。
cpu_hnsw_index->base_level_only = false;
// 用CAGRA图初始化HNSW索引。
gpu_cagra_index.copyTo(cpu_hnsw_index);
// 此时，HNSW索引已包含所有CAGRA中的向量，也可以继续 add 新向量。
cpu_hnsw_index->add(n, newVecs);
important

通过CAGRA加速HNSW索引的构建，一方面利用了GPU强大的并行计算能力，提高构建效率；另一方面也让后续索引扩展变得灵活。对于需要大量相似检索和高维数据场景（如推荐系统、图片检索等），这种用法十分高效。