使用NVIDIA cuVS优化向量搜索:从索引构建到实时检索
使用NVIDIA cuVS优化向量搜索:从索引构建到实时检索
在AI驱动的搜索应用(如RAG、推荐系统和异常检测)中,高效的向量搜索至关重要。然而,高性能的索引、低延迟的检索和无缝的可扩展性一直是开发者面临的挑战。NVIDIA cuVS(CUDA Vector Search)的出现,为开发者和数据科学家带来了GPU加速的向量搜索和聚类功能,彻底改变了这一现状。
本文将深入探讨NVIDIA cuVS的最新特性,展示其如何实现更快的索引构建、实时更新和高效的大规模搜索,并提供丰富的代码示例,帮助你快速上手。
挑战:向量搜索的性能瓶颈
向量搜索的核心在于两个阶段。第一个阶段是索引构建,需要将海量向量数据组织成优化的数据结构,以便快速检索。这个过程通常计算密集且耗时。第二个阶段是查询/检索,在索引中查找与查询向量最相似的向量。对于实时应用,延迟必须尽可能低。
随着数据集规模的爆炸式增长,传统的基于CPU的解决方案在索引构建速度和查询吞吐量方面都面临着巨大的压力。特别是在处理数百万甚至数十亿规模的向量数据时,CPU的串行计算特性成为了明显的瓶颈。
解决方案:NVIDIA cuVS
NVIDIA cuVS通过利用GPU的大规模并行计算能力,为向量搜索提供了前所未有的加速。最新版本引入了优化的索引算法、扩展的语言支持,并与FAISS、Milvus、Weaviate、Lucene等主流向量数据库和搜索引擎深度集成。cuVS支持多种先进的索引算法,包括CAGRA(GPU原生图索引)、IVF-PQ(倒排文件-乘积量化)、HNSW(分层可导航小世界图)和DiskANN/Vamana(超大规模图算法)。
1. GPU上的闪电般索引构建
cuVS可以将索引构建速度提升高达40倍。这得益于其对多种先进索引算法的GPU原生实现。在实际应用中,Google Cloud AlloyDB使用cuVS实现了HNSW索引构建比pgvector快9倍的性能,Weaviate使用CAGRA索引构建时间缩短8倍,而Apache Lucene在GPU上的索引构建速度提升了40倍。
代码示例:使用cuVS构建IVF-PQ索引
下面的Python代码展示了如何使用cuVS构建一个IVF-PQ索引。IVF-PQ是一种非常流行的索引类型,它通过聚类(IVF)和量化(PQ)来平衡搜索速度和内存使用。
import numpy as np
from cuvs.neighbors import ivf_pq# 1. 准备数据集
# 在实际应用中,这些向量通常来自深度学习模型的嵌入层
n_samples = 1_000_000 # 一百万个样本向量
n_features = 128 # 每个向量的维度
dtype = np.float32 # 使用32位浮点数以获得更好的GPU性能# 生成随机数据集作为演示
# 使用固定的随机种子以确保结果可复现
np.random.seed(42)
dataset = np.random.rand(n_samples, n_features).astype(dtype)# 2. 配置IVF-PQ索引参数
# IVF-PQ是一种非常流行的近似最近邻搜索算法,它结合了倒排文件(IVF)和乘积量化(PQ)
# IVF通过聚类将数据集划分为多个分区(n_lists),搜索时只需检查少数几个分区,从而加速搜索
# PQ通过将高维向量分解为多个低维子向量并对其进行量化,来压缩向量,减少内存使用
index_params = ivf_pq.IndexParams(n_lists=1024, # IVF部分的聚类中心数量。通常是数据集大小的平方根附近,需要调优。M=8, # PQ部分的子量化器数量。维度越高,M值可以越大。bits_per_code=8, # 每个子编码的位数。8位意味着每个子空间有256个中心点。metric="L2Expanded", # 距离度量。L2Expanded是欧氏距离的平方,对于GPU计算更高效。add_data_on_build=True # 在构建索引的同时添加数据。如果为False,则需要单独调用add方法。
)# 3. 在GPU上构建索引
# cuVS会自动利用GPU的并行计算能力来执行聚类和量化,速度远超CPU
print("开始在GPU上构建IVF-PQ索引...")
index = ivf_pq.build(index_params, dataset)
print("索引构建完成!")# 4. 准备查询向量
n_queries = 10
np.random.seed(10)
queries = np.random.rand(n_queries, n_features).astype(dtype)# 5. 在GPU上执行搜索
# 搜索时,我们不必检查所有1024个分区,而是只检查与查询向量最接近的n_probes个分区
search_params = ivf_pq.SearchParams(n_probes=32 # 要探索的聚类分区数量。n_probes越大,召回率越高,但速度越慢。
)# 搜索每个查询向量的k个最近邻
k = 10
neighbors, distances = ivf_pq.search(search_params, index, queries, k=k)print(f"为 {n_queries} 个查询向量找到的前 {k} 个最近邻的索引是:")
print(neighbors)
print(f"\n对应的距离是:")
print(distances)
2. CPU-GPU互操作性:灵活部署
cuVS最受欢迎的特性之一是其CPU-GPU互操作性。你可以在强大的GPU上快速构建索引,然后将索引部署到成本更低的CPU服务器上进行查询。这对于那些索引构建时间是主要瓶颈,但查询负载相对较低的场景非常理想。
上图展示了cuVS索引如何在GPU和CPU之间灵活部署。从多模态结构化和非结构化数据开始,通过NeMo Retriever嵌入和特征工程,在GPU上进行高速索引构建,最后可以选择在GPU或CPU上执行搜索,并将结果存储到数据库中。
代码示例:构建并保存FAISS兼容索引
cuVS可以与Meta的FAISS库无缝协作,构建一个索引,然后将其序列化为FAISS格式,以便在任何支持FAISS的CPU环境中使用。
import numpy as np
from cuvs.neighbors import ivf_pq
from cuvs.neighbors.faiss import from_faiss, to_faiss
import faiss
import os# (继续使用上一个示例中构建的cuVS索引 'index' 和查询向量 'queries')# 1. 将cuVS GPU索引转换为FAISS CPU索引
# 这个过程会将GPU上的索引结构和数据传输到CPU内存中,并转换为FAISS兼容的格式
print("将cuVS索引转换为FAISS CPU索引...")
faiss_index_cpu = to_faiss(index)# 2. 保存FAISS索引到磁盘
# 序列化后的索引可以轻松地在不同机器、不同进程之间共享
index_path = "my_faiss_index.bin"
faiss.write_index(faiss_index_cpu, index_path)
print(f"FAISS索引已保存到 {index_path}")# --- 模拟在另一个纯CPU环境中 --- ## 3. 加载FAISS索引
# 即使没有GPU,也可以使用FAISS库加载并使用这个索引
print("\n在CPU环境中加载FAISS索引...")
loaded_faiss_index = faiss.read_index(index_path)# 4. 在CPU上执行搜索
# FAISS的CPU搜索接口与GPU版本非常相似
k = 10
distances_cpu, neighbors_cpu = loaded_faiss_index.search(queries, k=k)print(f"在CPU上为 {len(queries)} 个查询找到的前 {k} 个最近邻的索引是:")
print(neighbors_cpu)# 清理文件
os.remove(index_path)
这种互操作性极大地提高了部署的灵活性。FAISS库使用cuVS在CPU上加速索引构建12倍,GPU索引加速8倍以上。你可以根据实际需求,选择最合适的硬件配置来平衡性能和成本。
代码示例:使用CAGRA构建GPU原生图索引
CAGRA是NVIDIA专门为GPU设计的图索引算法,它在GPU上的性能表现尤为出色。
import numpy as np
from cuvs.neighbors import cagra# 准备数据集
n_samples = 500_000
n_features = 128
dataset = np.random.rand(n_samples, n_features).astype(np.float32)# 配置CAGRA索引参数
# CAGRA使用图结构来表示向量之间的近邻关系
index_params = cagra.IndexParams(metric="L2", # 使用L2距离graph_degree=64, # 图中每个节点的平均度数。度数越高,召回率越高,但内存占用也越大。intermediate_graph_degree=128 # 构建过程中使用的中间图度数
)# 在GPU上构建CAGRA索引
print("开始构建CAGRA索引...")
index = cagra.build(index_params, dataset)
print("CAGRA索引构建完成!")# 准备查询
queries = np.random.rand(100, n_features).astype(np.float32)# 配置搜索参数
search_params = cagra.SearchParams(max_queries=100, # 批处理的最大查询数量itopk_size=64, # 内部top-k缓冲区大小search_width=1, # 搜索宽度,影响搜索的广度max_iterations=0 # 最大迭代次数,0表示自动确定
)# 执行搜索
neighbors, distances = cagra.search(search_params, index, queries, k=10)print(f"使用CAGRA找到的最近邻索引:")
print(neighbors[:5]) # 显示前5个查询的结果
3. 量化:节省内存,提升性能
cuVS支持多种量化技术,如二进制量化和标量量化,可将向量的内存占用分别减少4倍和32倍,同时带来4倍到20倍的性能提升。量化技术通过降低向量表示的精度来换取更小的内存占用和更快的计算速度,在许多实际应用中,这种精度损失是可以接受的。
代码示例:使用二进制量化
二进制量化将每个浮点数向量转换为一个紧凑的二进制向量,极大地减少了内存占用,并允许使用高效的汉明距离进行计算。
import numpy as np
from cuvs.neighbors import ivf_pq, cagra
from cuvs.quantize import binary_quantize# (继续使用第一个示例中的 'dataset' 和 'queries')# 1. 对数据集进行二进制量化
# 二进制量化将每个浮点数向量转换为一个紧凑的二进制向量(通常是256位或512位)
# 这极大地减少了内存占用,并允许使用高效的汉明距离进行计算
print("对数据集进行二进制量化...")
# `binary_quantize` 返回量化后的二进制向量和原始向量的范数(用于某些距离计算)
binary_dataset = binary_quantize(dataset)
print(f"原始数据类型: {dataset.dtype}, 量化后数据类型: {binary_dataset.dtype}")
print(f"原始数据形状: {dataset.shape}, 量化后数据形状: {binary_dataset.shape}")# 2. 使用CAGRA直接在二进制量化向量上构建索引
# cuVS的CAGRA算法支持直接在二进制量化数据上构建图索引
index_params_bq = cagra.IndexParams(metric="Hamming", # 对于二进制向量,使用汉明距离graph_degree=32
)
print("\n在二进制量化数据上构建CAGRA索引...")
index_bq = cagra.build(index_params_bq, binary_dataset)# 3. 对查询向量进行同样的量化
queries_bq = binary_quantize(queries)# 4. 使用二进制量化索引进行搜索
search_params_bq = cagra.SearchParams()
neighbors_bq, distances_bq = cagra.search(search_params_bq, index_bq, queries_bq, k=10)print("\n使用二进制量化索引找到的最近邻:")
print(neighbors_bq[:3]) # 显示前3个查询的结果# 5. 计算内存节省
original_memory = dataset.nbytes
quantized_memory = binary_dataset.nbytes
print(f"\n内存占用对比:")
print(f"原始数据: {original_memory / 1024 / 1024:.2f} MB")
print(f"量化后数据: {quantized_memory / 1024 / 1024:.2f} MB")
print(f"压缩比: {original_memory / quantized_memory:.2f}x")
代码示例:标量量化
标量量化通过将32位浮点数转换为8位整数来实现更高的压缩率。
import numpy as np
from cuvs.neighbors import ivf_pq
from cuvs.quantize import scalar_quantize# 准备数据集
dataset = np.random.rand(100_000, 128).astype(np.float32)# 1. 执行标量量化
# 标量量化将每个浮点数独立地映射到一个较小的整数范围
print("执行标量量化...")
quantized_dataset, scale, offset = scalar_quantize(dataset,bits=8 # 使用8位整数表示,可以选择4、8或16位
)print(f"量化后数据类型: {quantized_dataset.dtype}")
print(f"缩放因子: {scale}")
print(f"偏移量: {offset}")# 2. 在量化数据上构建索引
# 注意:某些索引类型可能需要特定的量化支持
index_params = ivf_pq.IndexParams(n_lists=256,metric="L2Expanded"
)# 反量化用于索引构建(在实际应用中,某些索引可以直接使用量化数据)
dequantized_dataset = quantized_dataset.astype(np.float32) * scale + offset
index = ivf_pq.build(index_params, dequantized_dataset)print("\n标量量化索引构建完成!")# 3. 计算压缩比
print(f"\n压缩比: {dataset.nbytes / quantized_dataset.nbytes:.2f}x")
4. 高吞吐量搜索与动态批处理
对于广告投放、金融交易等需要高吞吐量和低延迟的在线服务,cuVS提供了动态批处理API,可将延迟降低多达10倍。它能智能地将传入的单个查询请求组合成批次,以充分利用GPU的并行处理能力。CAGRA持久搜索功能可以使高容量搜索吞吐量提升8倍以上。
代码示例:批处理搜索优化
import numpy as np
from cuvs.neighbors import cagra
import time# 准备索引和数据
dataset = np.random.rand(1_000_000, 128).astype(np.float32)
index_params = cagra.IndexParams(metric="L2", graph_degree=64)
index = cagra.build(index_params, dataset)# 准备大批量查询
n_queries = 10000
queries = np.random.rand(n_queries, 128).astype(np.float32)# 方法1:逐个查询(效率低)
print("方法1:逐个查询...")
start_time = time.time()
results_single = []
for i in range(100): # 仅测试100个查询以节省时间single_query = queries[i:i+1]neighbors, distances = cagra.search(cagra.SearchParams(), index, single_query, k=10)results_single.append(neighbors)
single_query_time = time.time() - start_time
print(f"逐个查询100次耗时: {single_query_time:.4f} 秒")# 方法2:批处理查询(高效)
print("\n方法2:批处理查询...")
start_time = time.time()
batch_queries = queries[:100]
neighbors_batch, distances_batch = cagra.search(cagra.SearchParams(),index,batch_queries,k=10
)
batch_query_time = time.time() - start_time
print(f"批处理100次查询耗时: {batch_query_time:.4f} 秒")
print(f"加速比: {single_query_time / batch_query_time:.2f}x")
代码示例:CAGRA预过滤能力
cuVS的预过滤功能允许你在搜索时排除某些向量,即使99%的向量被排除,仍能实现高召回率。
import numpy as np
from cuvs.neighbors import cagra# 准备数据集和索引
n_samples = 100_000
n_features = 128
dataset = np.random.rand(n_samples, n_features).astype(np.float32)index_params = cagra.IndexParams(metric="L2", graph_degree=64)
index = cagra.build(index_params, dataset)# 创建一个过滤掩码
# 假设我们只想在前10%的向量中搜索(例如,基于某些业务规则)
filter_mask = np.zeros(n_samples, dtype=bool)
filter_mask[:n_samples // 10] = True # 只有前10%的向量可以被返回print(f"过滤后可搜索的向量数量: {filter_mask.sum()} / {n_samples}")# 准备查询
queries = np.random.rand(10, n_features).astype(np.float32)# 使用预过滤进行搜索
search_params = cagra.SearchParams()
# 注意:具体的过滤API可能因版本而异,这里展示概念
# neighbors, distances = cagra.search_with_filter(
# search_params, index, queries, k=10, filter_mask=filter_mask
# )print("预过滤搜索完成!即使大部分向量被排除,仍能保持高召回率。")
5. 扩展的语言支持
除了Python,cuVS还提供了对Rust、Go和Java的API支持,使更多开发者能够利用GPU加速的向量搜索。这些语言绑定提供了与Python API相似的功能和性能,让不同技术栈的团队都能享受到GPU加速的好处。
代码示例:Rust API
// 这是一个概念性的Rust代码片段,展示了cuvs-rs绑定的基本用法
// 你需要将 `cuvs-rs` 添加到你的 Cargo.toml 依赖中// 引入cuVS Rust绑定的关键模块
use cuvs::neighbors::ivf_pq::{IndexParams, SearchParams, build, search};
use cuvs::error::CuVsError;fn main() -> Result<(), CuVsError> {// 定义向量的维度和数据集大小const DIM: usize = 128;const N_SAMPLES: usize = 1_000_000;const N_QUERIES: usize = 10;// 1. 生成模拟数据集// 在实际应用中,你会从文件或嵌入模型加载数据println!("生成数据集...");let dataset: Vec<f32> = (0..N_SAMPLES * DIM).map(|_| rand::random::<f32>()).collect();// 2. 创建IVF-PQ索引参数let index_params = IndexParams::new().n_lists(1024).metric("L2").pq_dim(8).pq_bits(8);// 3. 在GPU上构建索引println!("开始在GPU上构建索引...");let index = build(&index_params, &dataset, DIM)?;println!("索引构建完成!");// 4. 生成查询向量let queries: Vec<f32> = (0..N_QUERIES * DIM).map(|_| rand::random::<f32>()).collect();// 5. 执行搜索println!("执行搜索...");let search_params = SearchParams::new().n_probes(32);let (neighbors, distances) = search(&search_params, &index, &queries, 10)?;println!("搜索完成!");println!("找到的最近邻索引 (前5个): {:?}", &neighbors[..50]);Ok(())
}
代码示例:Go API(概念性)
package mainimport ("fmt""github.com/rapidsai/cuvs-go/neighbors/ivfpq"
)func main() {// 1. 准备数据集// 在实际应用中,从文件或数据库加载向量数据nSamples := 1000000nFeatures := 128dataset := make([]float32, nSamples*nFeatures)// ... 填充数据集 ...// 2. 配置索引参数indexParams := ivfpq.IndexParams{NLists: 1024,Metric: "L2",PQDim: 8,PQBits: 8,AddDataOnBuild: true,}// 3. 构建索引fmt.Println("开始构建索引...")index, err := ivfpq.Build(indexParams, dataset, nFeatures)if err != nil {panic(err)}fmt.Println("索引构建完成!")// 4. 准备查询nQueries := 10queries := make([]float32, nQueries*nFeatures)// ... 填充查询数据 ...// 5. 执行搜索searchParams := ivfpq.SearchParams{NProbes: 32,}neighbors, distances, err := ivfpq.Search(searchParams, index, queries, 10)if err != nil {panic(err)}fmt.Printf("找到 %d 个查询的最近邻\n", nQueries)fmt.Println("最近邻索引:", neighbors[:10])
}
生态系统与集成
cuVS的强大之处不仅在于其自身的功能,还在于其与整个AI生态系统的深度集成。无论是开源向量数据库Milvus、Weaviate,还是搜索引擎巨头Apache Lucene、Elasticsearch、OpenSearch,都在积极集成cuVS以提供GPU加速能力。这种广泛的生态系统支持使得开发者可以在现有的技术栈中无缝引入GPU加速,而无需进行大规模的架构重构。
集成伙伴 | 加速效果 | 应用场景 |
---|---|---|
Meta FAISS | CPU索引构建加速12倍,GPU索引加速8倍 | 通用向量搜索,研究原型 |
Milvus with DDN | 使用CAGRA索引构建加速22倍 | 大规模向量数据库 |
Weaviate | 使用CAGRA索引构建时间缩短8倍 | 知识图谱,语义搜索 |
Apache Lucene | GPU加速索引构建40倍 | 全文搜索,企业搜索 |
Apache Solr | 端到端索引构建加速6倍 | 企业搜索平台 |
Google Cloud AlloyDB | HNSW索引构建比pgvector快9倍 | 云数据库服务 |
Oracle Database 23ai | 端到端索引构建加速5倍 | 企业级数据库 |
OpenSearch 3.0 | 索引构建时间加快9.4倍 | 日志分析,搜索引擎 |
Elasticsearch | 通过插件支持GPU加速 | 企业搜索,日志分析 |
这些集成不仅提供了显著的性能提升,还保持了与现有API和工作流的兼容性,使得迁移成本降到最低。
实际应用场景
cuVS在多个领域都展现出了卓越的性能。在检索增强生成(RAG)应用中,cuVS可以快速检索相关文档片段,为大语言模型提供上下文信息。在推荐系统中,cuVS能够实时计算用户和物品之间的相似度,提供个性化推荐。在探索性数据分析中,cuVS与RAPIDS cuML结合,通过UMAP等降维算法和nn-descent等聚类算法,实现了迭代式的近实时数据探索工作流。
在单细胞基因组学领域,rapids-singlecell库利用cuVS和cuML实现了突破性的性能提升,使得研究人员能够分析更大规模的细胞数据集。在主题建模方面,BERTopic库通过集成cuVS,使得大规模主题建模工作流变得可行。此外,Microsoft Advertising正在探索将cuVS的CAGRA搜索算法集成到其广告投放管道中,以提高吞吐量和降低延迟。
核心算法深入:nn-descent与kNN图构建
cuVS还提供了nn-descent算法,这是一种用于构建k近邻图的高效算法。kNN图在数据挖掘、流形学习和聚类技术中都有广泛应用。nn-descent算法支持核外(out-of-core)构建,这意味着即使数据集大小超过GPU内存,也能够构建kNN图。
代码示例:使用nn-descent构建kNN图
import numpy as np
from cuvs.neighbors import nn_descent# 准备大规模数据集
n_samples = 500_000
n_features = 128
dataset = np.random.rand(n_samples, n_features).astype(np.float32)# 配置nn-descent参数
# nn-descent通过迭代改进来构建近似kNN图
params = nn_descent.IndexParams(metric="L2", # 距离度量graph_degree=32, # 每个节点的邻居数量(k值)intermediate_graph_degree=64, # 构建过程中的中间图度数max_iterations=20, # 最大迭代次数termination_threshold=0.001 # 收敛阈值
)# 构建kNN图
print("开始构建kNN图...")
knn_graph = nn_descent.build(params, dataset)
print("kNN图构建完成!")# knn_graph包含每个点的k个最近邻的索引和距离
print(f"图的形状: {knn_graph.shape}")
print(f"第一个节点的邻居: {knn_graph[0]}")# 这个kNN图可以用于后续的聚类、降维等任务
性能调优与最佳实践
为了充分发挥cuVS的性能,需要注意以下几个方面的调优。
索引参数调优:不同的索引类型有不同的参数需要调优。对于IVF-PQ索引,n_lists
通常设置为数据集大小的平方根附近,M
和bits_per_code
需要根据维度和精度要求进行平衡。对于CAGRA索引,graph_degree
越高,召回率越高,但内存占用也越大。
搜索参数调优:搜索参数直接影响召回率和延迟的权衡。对于IVF-PQ,增加n_probes
可以提高召回率但会降低速度。对于CAGRA,search_width
和max_iterations
控制搜索的广度和深度。
批处理大小优化:GPU在处理大批量查询时性能最佳。尽可能将多个查询组合成批次,以充分利用GPU的并行能力。cuVS Bench工具可以帮助你在特定硬件上找到最优的批处理大小。
内存管理:对于超大规模数据集,考虑使用量化技术来减少内存占用。二进制量化和标量量化可以在保持可接受精度的同时,显著减少内存需求。
混合部署策略:利用CPU-GPU互操作性,在GPU上快速构建索引,然后根据查询负载和成本考虑,选择在GPU或CPU上执行搜索。
开始使用NVIDIA cuVS
NVIDIA cuVS正在重新定义GPU加速的向量搜索,通过持续的技术创新、扩展的语言支持和深度的生态集成,为开发者提供了前所未有的性能和灵活性。无论你是希望增强现有向量数据库的搜索能力,还是构建定制的AI检索系统,cuVS都为你提供了将性能推向新高度所需的速度、灵活性和易用性。
准备好开始了吗?访问**rapidsai/cuvs GitHub仓库,获取完整的端到端示例和自动调优指南。你也可以使用cuVS Bench**在GPU和CPU环境中对ANN搜索实现进行基准测试和比较。cuVS支持通过pip、conda等多种方式安装,并提供了详细的文档和示例代码,帮助你快速上手。
推荐阅读
- 【NVIDIA开发者】博客文章翻译-CUDA-X Data Science加速模型训练
- 【NVIDIA开发者】博客文章翻译-使用NVIDIA Dynamo减少KV Cache瓶颈
- 【NVIDIA开发者】博客文章翻译-AI智能体开发工具安全性
- 【NVIDIA开发者】博客文章翻译-Isaac Lab与Newton物理引擎四足机器人训练