5. mcVS GPU索引的支持

本章描述GPU相关索引在mcVS上的一些参数和设定。默认使用Python接口作为示例介绍。

5.1. cuvs.neighbors.brute_force

brute_force 通过暴力搜索的方式精确查询最近邻向量，搜索效率较低，可以提供100%的召回, 适用于数据量较小的场景。

索引构建和索引搜索示例如下：

index = brute_force.build(dataset, metric='sqeuclidean', resources=None)
distances, neighbors = brute_force.search(index, queries, k)

5.2. cuvs.neighbors.ivf_flat

ivf_flat 是一种倒排文件索引 (IVF) 算法，在最近邻搜索中，数据被划分到多个聚类中，并在聚类的子集中执行暴力搜索。通常适用于数据量较大且不需精确召回的场景。

在 ivf_flat 索引的构建和搜索前，可以通过 IndexParams 和 SearchParams 来设定参数：

表 5.1 IndexParams 参数说明

参数名	说明	默认值
n_lists	划分索引的子聚类数量	sqrt(n)
metric	表示度量类型（metric type）的字符串	sqeuclidean
add_data_on_build	索引构建后是否应将训练数据添加到索引中	True
kmeans_n_iters	k-means 训练达到收敛前的最大迭代次数	20
kmeans_trainset_fraction	从原始数据集采样用于训练 k-means 聚类的点比例	0.5
adaptive_centers	现有的已训练质心是否应适应添加到索引中的新点	False
conservative_memory_allocation	是否分配更多内存并将旧列表复制到新的更大列表中	False

表 5.2 SearchParams 参数说明

参数名	说明	默认值
n_probes	对每个查询点要扫描的最近 IVF 列表数量	20

ivf_flat 索引构建示例如下，其中dataset是形状为形状为 (n_samples, dim) 的矩阵。

build_params = ivf_flat.IndexParams(n_lists=1024, metric='inner_product')
index = ivf_flat.build(build_params, dataset)

调用ivf_flat.search找到queries中每个查询的k个最近邻，索引搜索示例：

search_params = ivf_flat.SearchParams(n_probes=64)
distances, neighbors = ivf_flat.search(search_params, index, queries, k)

在IVF类索引中， n_lists 和 n_probes 是两个重要的参数，它们共同决定了索引的搜索性能与召回率之间的平衡：

• n_lists : 该参数控制着对数据集进行划分的聚类中心的数量。n_lists 的数量越大，索引的搜索效率越高，但是索引的构建时间也越长。

• n_probes : 该参数指定在搜索时需要检查多少个最有可能的聚类。n_probes 的数量越大，索引的搜索召回率越高，但是索引的搜索性能越低。

因此，在实际应用中，需要根据具体场景的需求来权衡这两个参数：

• 在追求高吞吐量的场景下，可适当减小 n_lists 或 n_probes ；

• 而在对召回率要求较高的场景下，则需要适当增大 n_probes，并可能需要相应调整 n_lists 来平衡性能。

5.3. cuvs.neighbors.ivf_pq

ivf_pq 同样是一种IVF算法，它是 ivf_flat 算法的扩展，数据首先被分成多个聚类，然后在每个聚类内部应用乘积量化（PQ）技术。其中，PQ是一种有损压缩方法，通过将向量进行压缩来减少内存占用。这使得 GPU 可以将更多向量存储在显存中，而将原始向量数据放到主内存中。但需要权衡的是，压缩程度越高，搜索精度（召回率）通常会相应下降。 在 ivf_pq 索引的构建和搜索前，可以通过 IndexParams 和 SearchParams 来设定参数。

表 5.3 IndexParams 参数说明

参数名	说明	默认值
n_lists	划分索引的子聚类数量	sqrt(n)
metric	表示度量类型（metric type）的字符串	sqeuclidean
pq_bits	使用 PQ 压缩后每个向量元素的位长，取值在4 到 8 之间的整数	8
pq_dim	使用 PQ 压缩后每个向量的维度，当为0时则自动选择最优值	0
codebook_kind	码本的创建方式	subspace
force_random_rotation	对输入数据和查询应用随机旋转矩阵	False
add_data_on_build	索引构建后是否应将训练数据添加到索引中	True
kmeans_n_iters	k-means 训练达到收敛前的最大迭代次数	20
kmeans_trainset_fraction	从原始数据集采样用于训练 k-means 聚类的点比例	0.5
adaptive_centers	现有的已训练质心是否应适应添加到索引中的新点	False
conservative_memory_allocation	是否分配更多内存并将旧列表复制到新的更大列表中	False

表 5.4 SearchParams 参数说明

参数名	说明	默认值
n_probes	对每个查询点要扫描的最近 IVF 列表数量	20
lut_dtype	用于存储 pq 查找表的数据类型	float32
internal_distance_dtype	搜索时计算的距离/相似度的存储数据类型	float32
preferred_smem_carveout	AP上的共享内存占比，默认值为 100%	1.0

ivf_pq 索引构建示例如下：

build_params = ivf_pq.IndexParams(n_lists=1024, metric='inner_product', pq_bits=8, pq_dim=0)
index = ivf_pq.build(build_params, dataset)

调用ivf_pq.search找到queries中每个查询的k个最近邻，索引搜索示例：

search_params = ivf_pq.SearchParams(n_probes=64, lut_dtype=np.float32)
distances, neighbors = ivf_pq.search(search_params, index, queries, k)

pq_dim 定义了乘积量化时每个子空间的维度，即原始向量被切分成的子向量数量。pq_bits 则指定了每个子空间量化时使用的比特数，即码本的大小。这两个参数共同影响着索引的性能表现：数值越大，通常意味着更高的搜索精度和召回率，但同时也伴随着更长的索引构建时间和更大的内存占用。

因此，在实际应用中，需要根据具体的业务需求（如对搜索速度的要求、可接受的内存开销、以及期望的召回率）来权衡选择 pq_dim 和 pq_bits 的值，以达到最佳的性能平衡点。

5.4. cuvs.neighbors.cagra

cagra 是一种基于图的最近邻搜索算法，专门为 GPU 加速而构建。CAGRA 通过首先构建训练点的 kNN 图，然后移除邻居之间的冗余路径来构建平面图的表示。该算法在索引构建和查询性能方面表现出色，适用于小批量和大批量搜索场景。

CAGRA 索引构建算法有两个基本步骤：

1. 构建 kNN 图

2. 从 kNN 图中剪枝冗余路径

CAGRA 可以使用 IVF-PQ 或 NN-Descent 策略来构建初始的 kNN 图，这些策略可以在索引构建期间的索引参数对象中选择。

在 cagra 索引的构建和搜索前，可以通过 IndexParams 和 SearchParams 来设定参数：

表 5.5 IndexParams 参数说明

参数名	说明	默认值
metric	表示度量类型的字符串	sqeuclidean
intermediate_graph_degree	初始 kNN 图的度数	128
graph_degree	最终 CAGRA 图的度数	64
build_algo	图构建算法，可选值：ivf_pq, nn_descent, iterative_cagra_search	ivf_pq
nn_descent_niter	NN-Descent 算法的最大迭代次数	20
compression	压缩参数，用于大数据集的向量压缩	None
ivf_pq_build_params	IVF-PQ 构建参数	None
ivf_pq_search_params	IVF-PQ 搜索参数	None
refinement_rate	细化率，用于提高搜索精度	1.0

表 5.6 SearchParams 参数说明

参数名	说明	默认值
max_queries	同时搜索的最大查询数（批量大小）	0（自动选择）
itopk_size	搜索期间保留的中间搜索结果数量	64
max_iterations	搜索迭代的最大次数	0（自动选择）
algo	搜索算法 (auto, single_cta, multi_cta)	auto
team_size	计算单个距离的线程数	0（自动选择）
search_width	每次迭代选择的起始图节点数	1
min_iterations	搜索迭代的最小次数	0
thread_block_size	CUDA 线程块大小	0（自动选择）
hashmap_mode	哈希表模式 (auto, small, hash)	auto
hashmap_min_bitlen	哈希表最小位数	0
hashmap_max_fill_rate	哈希表最大填充率	0.5

cagra 索引构建示例如下：

from cuvs.neighbors import cagra

build_params = cagra.IndexParams(
    metric='sqeuclidean',
    intermediate_graph_degree=128,
    graph_degree=64,
    build_algo='ivf_pq'
)
index = cagra.build(build_params, dataset)

调用 cagra.search 找到 queries 中每个查询的 k 个最近邻，索引搜索示例：

search_params = cagra.SearchParams(
    itopk_size=64,
    max_iterations=0,
    algo='auto'
)
distances, neighbors = cagra.search(search_params, index, queries, k)

在 CAGRA 索引中，intermediate_graph_degree 和 graph_degree 是两个重要的构建参数：

• intermediate_graph_degree：初始 kNN 图的度数，较大的值会增加初始图的连通性，但会增加内存使用和构建时间。

• graph_degree：最终 CAGRA 图的度数，所有顶点都将具有此度数。较大的图度数允许在搜索期间更多地探索搜索空间并提高召回率，但代价是搜索更多顶点。

在搜索参数中，itopk_size 是调整搜索精度和速度之间权衡的主要参数：

• itopk_size：搜索期间保留的中间搜索结果数量。该值需要 >= k。较高的值可提高搜索精度，但会降低搜索速度。