拓十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

亿级向量检索提速:cuVS+Elasticsearch GPU索引构建实战

亿级向量检索提速:cuVS+Elasticsearch GPU索引构建实战

向量检索这个事,做到一定规模之后,CPU 就真的扛不住了。我们线上有一批 1.38 亿条 384 维的 embedding 召回池,之前用 Elasticsearch 原生的 HNSW 建索引,一跑就是四五个小时,而且 JVM 堆内存和 GC 压力大得让人头疼;每天索引更新的调度窗口根本排不开。后来把 NVIDIA cuVS 引入检索链路,用 GPU 来做向量索引的构建和近邻搜索,整批 1.38 亿向量从原始数据到可查询的索引文件,10 分钟出头就能跑完,检索端延迟还比原来低了不少。

这篇文章就是这次改造的完整复盘。里面会讲清楚为什么选 cuVS + Elasticsearch 这套组合、索引参数到底怎么推算才不踩坑、完整搭建流程怎么走,以及我在实际环境中碰到的各种编译、显存、召回率问题。适合正在做推荐召回、多模态搜索、RAG 向量库,或者被亿级向量索引构建折磨得睡不着觉的工程师参考。不管你是第一次听到 cuVS,还是已经跑过 GPU 向量检索,应该都能从里面找到点能直接抄作业的东西。

1. 方案选型:为什么把 GPU 放到 ES 检索链路里

1.1 数据规模带来的第一个坎

先说一个最直接的感受:1.38 亿条 384 维 float32 向量,光原始数据就有 211.9GB,这个数字很多团队第一眼看到是没概念的。放到 Elasticsearch 里,Lucene 的 HNSW 索引构建需要不断计算邻居图、写段、合并段,不仅慢,而且内存占用远不止原始数据的大小。HNSW 的图结构本身会有额外开销,加上 JVM 堆、段合并的临时空间,跑批建索引的时候一台 256GB 内存的机器都会显得捉襟见肘。我们当时第一次全量构建,主节点直接拉起了 4 个多小时的任务,中间还因为堆内存告警被 kill 过两次。

这时候就需要重新思考一个问题:ES 在整个链路里的职责到底是什么。它适合做文档存储、过滤、聚合和最终打分,但在“亿级向量全量建索引”这件事上,纯 CPU 的近似最近邻算法确实到了瓶颈。GPU 的并行度天然适合这类计算密集型任务,k-means 聚类、向量距离计算、PQ 编码这些操作,在 GPU 上可以成量级地加速。于是方向就很明确了:把向量索引的构建和搜索从 ES 的 JVM 里拆出来,交给 GPU 侧的向量搜索库,ES 只负责它最擅长的文档管理和查询路由。

1.2 cuVS 是什么,有哪些算法

cuVS 是 NVIDIA 开源的 GPU 向量搜索库,归属于 RAPIDS 生态,里面实现了几种主流算法:IVF-Flat、IVF-PQ、CAGRA,也有一部分 HNSW 的支持。你可以把它理解成一个专门在 GPU 上做最近邻搜索的工具箱,输入一组向量,它帮你建索引,然后给你一个支持近似搜索的接口。算法选型上,IVF-PQ 是我们这次的首选,原因是它在建索引速度和内存占用上最均衡。IVF-PQ 做的事情可以理解成两步:先把向量空间用倒排索引切分成若干个区域,然后对每个区域内做乘积量化,把高维向量压缩成一小段码字。压缩之后,1.38 亿条 384 维向量,索引文件只有几个 GB,A100 80G 的显存完全放得下。

CAGRA 我也测过一轮,它是 cuVS 里性能更极端的图算法,搜索延迟可以做到更低,但建索引时对显存的要求更高,数据量上亿之后训练阶段的中间结果很容易把显存撑爆。如果你的数据量在几千万级别,且对查询延迟极度敏感,CAGRA 值得试;但像我这种 1.38 亿全量池,IVF-PQ 更稳。HNSW 在 cuVS 里也有 GPU 实现,但它的构建复杂度摆在那里,GPU 加速效果没有 IVF 系明显,所以我们没有走那条路。

1.3 架构设计:存储与检索分层

整套架构最后落地是这样分的:Elasticsearch 负责管理文档元数据、原始向量字段、过滤条件,以及最终的精确重排;cuVS 负责离线构建大规模近邻索引,在线接受查询向量,返回候选集的 doc_id。数据流上,离线阶段先用 GPU 把全量向量构建成 IVF-PQ 索引文件,同时把 1.38 亿文档的元数据批量导入 ES;查询阶段,线上请求的 query 向量先送进 cuVS 搜索接口,取回 Top-K 候选 doc_id,再去 ES 里做一次 mget 把文档捞出来,最后可以用原始向量做精确距离重排。

这样做的好处是很明显的。第一,ES 的 JVM 堆不再背“全量向量索引”这个重担,集群稳定性直线上升。第二,索引构建和更新完全独立,cuVS 重建索引不影响 ES 对外服务。第三,整套链路是插拔式的,哪天数据量再翻一倍,我可以只加 GPU 节点和调整 IVF-PQ 参数,不用推翻 ES 集群。如果你用的是 OpenSearch,它自带的 k-NN 插件支持 FAISS 思路类似,但 ES 这边没有原生 GPU 索引插件,自己接 cuVS 反而灵活。

2. 数据准备与索引参数推演:别拍脑袋选 nlist 和 M

2.1 先看清楚你的数据形态

任何向量索引方案,第一步不是写代码,而是先搞清楚你的数据长什么样。我们这批数据是文本和图像的混合 embedding,统一归一到 384 维,并且做了 L2 归一化。为什么要归一化很关键:如果你打算用余弦相似度,那归一化之后的内积就等价于余弦距离,这样在 cuVS 里可以直接用内积作为距离度量,少一层向量长度修正的开销。

不同维度下,1.38 亿条向量的数据体量完全不同,直接影响显存和参数选择。下面的表是我当时算的账:

向量维度1.38 亿条原始数据大小IVF-PQ 压缩后索引大小(M=24, nbits=8)
128 维约 70.6GB约 1.3GB
384 维约 211.9GB约 3.9GB
768 维约 423.9GB约 7.7GB

所以如果你的 embedding 是 768 维甚至 1536 维,压缩后的索引还能接受,但训练和编码的中间过程对显存的要求就高很多,这时就要考虑分块构建。我们最终选择在 A100 80G 上全量跑 384 维数据,显存刚好覆盖。

2.2 IVF-PQ 关键参数:nlist、M、nbits

IVF-PQ 有三个参数是必须想清楚的:nlist 是倒排索引的聚类中心数量,M 是把向量切成多少个子向量段,nbits 是每个子向量的量化位数。先说 nlist,经验法则是取数据量的平方根附近的值,sqrt(1.38 亿) 约等于 11747,所以取 16384 这个 2 的幂比较好。nlist 越大,每个桶里的向量越少,搜索时定位越准,但训练聚类的时间也更长,还可能导致部分桶过小反而影响召回。如果你数据是 1 亿级别,16384 是一个很稳的起点。

M 的选择相对灵活,前提是能被向量维度整除。384 维可以取 M=16 或 M=24,M=16 时每个子向量是 24 维,压缩后每个向量 20 字节;M=24 时每个子向量是 16 维,压缩后 28 字节。压缩率越高,索引越小,但量化误差越大,召回率会有损失。我们最终在召回率达标的前提下选了 M=24,单条压缩 28 字节,总索引约 3.9GB,这个量级在显存里非常舒服。nbits 默认 8 就够,表示每个子向量的码本里有 256 个中心点,再往上提收益很小反而显著增加训练耗时。

2.3 显存不够怎么办

如果你手里的卡不是 80G,或者向量维度更高,别慌,cuVS 支持分批构建。关键思路是分离“训练”和“编码”两个过程。训练阶段只需要一部分有代表性的样本,比如 200 万条随机采样就足够 k-means 收敛,200 万 × 384 维 × 4 字节也就 3GB 左右;编码阶段可以按照 500 万条一个 chunk 分批进行,每批算完结果写盘,最后拼接成完整的索引文件。这样即便数据总量几百 GB,构建过程的显存峰值也能控制在 20GB 左右。

还有一个实用技巧是索引加载时用内存映射而不是一次性全量塞进显存。cuVS 支持把索引文件映射到内存,搜索时按需读取对应的倒排桶,这样单张卡能承载的索引规模可以远大于显存容量。代价是查询延迟会有所上升,因为多了一层磁盘/内存读取,但对“索引不常驻显存”的场景来说,这是一个很值得用的降级方案。实测下来,索引从显存模式切到映射模式,单查询的 P99 延迟大约从 3ms 涨到 8ms 左右,但能支撑更大规模的数据。

3. 从零搭建 GPU 加速向量检索管线的完整实操

3.1 环境清单与安装

这一步看着简单,坑其实不少。先说硬件:我们用了两台 GPU 节点,每台 8 张 A100 80G,但实际构建单索引只需要一张卡,多卡主要用于并行处理不同数据分片。软件层面需要 NVIDIA 驱动支持 CUDA 12.0 以上,然后安装 cuVS 的 Python 包。安装命令很简单:

# 建议用独立的 conda 环境或 docker 镜像,避免污染线上环境 conda create -n cuvs-env python=3.10 -y conda activate cuvs-env pip install cuvs-cu12 nvidia-smi # 确认驱动和 GPU 状态正常

这里特别提醒一句:cuVS 的版本和 CUDA 版本必须匹配,装完第一件事是跑一遍官方自带的 smoke test。另外如果你之前装过 pytorch 的 GPU 版本,不要理所当然地认为 cuVS 一定能直接用,两个包的 CUDA runtime 依赖可能不一样。我们一开始就在一台装过老旧 CUDA 10 环境的机器上翻车了,最后换到干净镜像一次性通过。

3.2 用 cuVS 构建 IVF-PQ 索引的核心代码

cuVS 的 Python API 非常简明,核心逻辑就三段:加载向量集、构建索引、保存索引文件。我贴一段我们实际在用的简化代码:

import numpy as np import cuvs from cuvs.neighbors import ivf_pq # 加载向量,假设是 (N, 384) 的 float32 数组 vectors = np.fromfile("all_vectors.bin", dtype=np.float32).reshape(-1, 384) # 索引配置 nlist = 16384 M = 24 nbits = 8 # 构建 IVF-PQ 索引 index = ivf_pq.build( vectors, metric="inner_product", n_list=nlist, m=M, n_bits=nbits, kmeans_n_iters=20, # k-means 迭代次数,20 次足够收敛 train_size=2000000, # 训练样本数 ) # 保存索引 ivf_pq.save(index, "cuvs_138m.index")

这段代码跑完后,磁盘上会出现一个索引文件,文件名我们自己带上了版本号。构建过程中你会发现 GPU 利用率很高,而 CPU 几乎闲置,这说明计算确实被卸载到显卡上了。如果你的数据量太大一次 build 扛不住,可以改用先 train 再分批 add 的写法,cuVS 支持增量添加编码后的向量,训练完成后每个 chunk 走add接口进去,最后统一 save。需要说明的是,不同 cuVS 版本对函数签名略有调整,实际以你安装版本的官方文档为准,但整体流程就是这三步。

搜索端的代码同样很直接:

from cuvs.neighbors import ivf_pq index = ivf_pq.load("cuvs_138m.index") distances, indices = ivf_pq.search(index, queries, k=200, n_probes=64)

这里的queries是待检索的 query 向量 batch,n_probes是搜索时检查多少个倒排桶,值越大召回越好,延迟也越高。返回的indices就是我们需要的 doc_id 候选集,下一步拿着它去 ES 捞文档。

3.3 向量数据导入 ES

虽然主检索走 cuVS,但 ES 里还是要存一份原始向量,主要用于两件事:一是精确重排时计算真实距离,二是排查线上问题时能直接捞出来看。因为这批数据不会频繁更新,我们直接把索引的 refresh 关掉,用 bulk 批量写入,导入过程相对可控。

Mapping 配置如下:

PUT /vector_items { "mappings": { "properties": { "doc_id": { "type": "keyword" }, "title": { "type": "text" }, "embedding": { "type": "dense_vector", "dims": 384, "index": true, "similarity": "cosine", "index_options": { "type": "hnsw", "m": 32, "ef_construction": 200 } } } } }

这里有个容易纠结的点:ES 里既然已经有 HNSW 索引了,是不是可以直接用它做检索?如果你数据量在百万级,是的,直接查就行;但到了亿级,ES HNSW 的建索引时间和内存开销都会成为瓶颈,所以我们在 ES 里保留 dense_vector 字段但不会把主查询压给它。bulk 导入时把refresh_interval设为 -1,每批 5000 条,并发 8 个线程,1.38 亿文档大约 30 到 40 分钟写完。这个时间是可以接受的,而且和 GPU 建索引并行跑,整体调度窗口并不吃亏。

3.4 检索链路:GPU 粗排 + ES 精确重排

查询链路的设计决定用户体验。我们线上的流程是:query 向量先进 cuVS,一次取回 Top 200 的候选 doc_id;然后带着这批 doc_id 去 ES 做一次 multi-get,把文档取回来;最后用 ES 字段里的原始向量做一次精确的余弦相似度计算,重排出 Top 50 结果。为什么要做精确重排?因为 IVF-PQ 是近似搜索,量化误差会导致局部顺序不够准,但它的召回候选通常已经包含了真正相关的文档,精确重排一小批候选的成本很低,收益却很明显。

精确重排可以直接在应用层用 numpy 做,也可以把候选列表和目标向量传给 ES 的 script_score 查询,让 ES 顺便把分算了。数据量在 200 条候选这个级别,两种方案延迟差异不大,应用层做的好处是逻辑更可控。端到端来看,如果用户在 ES 上还有大量 filter 条件,比如按类目限制,那么更优的做法是先把 filter 后的 doc_id 集合传给 cuVS 搜索接口做限制,或者反过来先用 cuVS 粗排再在 ES 里 filter。这个顺序要根据你的过滤率来测,我们的场景过滤率不高,所以先 cuVS 再 ES filter,延迟表现最好。

4. 性能实测与参数调优:1.38 亿向量到底多快

4.1 索引构建实测数据

直接上我们在一张 A100 80G 上跑出来的数据。全量 1.38 亿条 384 维向量,IVF-PQ 参数 nlist=16384、M=24、nbits=8,完整构建过程分为训练、编码、写索引三段,总耗时约 9 分 40 秒。其中 k-means 训练用了 2 分 15 秒,分批编码和写入用了 7 分出头。注意这里包含从原始二进制文件读取数据的时间,如果数据已经在显存里,还能再快个几十秒。

对比之前的 CPU 方案:用 Lucene HNSW 在 32 核高主频机器上建同样一批数据的索引,实测耗时 4 小时 35 分钟。这已经不是“快多少倍”的问题了,而是从“每天只能半夜跑一次”变成“随时可以全量重建”。对我们这种需要频繁刷新 embedding 模型的场景来说,这个差异直接决定了能不能做到日更甚至小时级更新。如果你用的是 CAGRA 且数据量在千万级,构建时间还能压到 3 分钟以内,但显存需求会明显变高。

4.2 检索延迟与召回率

检索端的表现,我用两个指标说明:延迟和召回率。延迟方面,GPU 侧的 IVF-PQ 搜索在 n_probes=32 时,单查询平均 1.2ms,P99 在 3ms 左右;加上 ES 的 mget 和精确重排,完整端到端 P99 大约在 20 到 25ms。这个延迟水平和一台中等配置的 ES 集群直接查亿级 HNSW 差不多,但关键在于索引构建的瓶颈解决了。召回率方面,我们在一个 10 万条人工标注的验证集上测了 Recall@10,n_probes=64 的时候可以达到 94.6%,效果非常接近全量精确检索。

不同 n_probes 的权衡比较明显:

n_probesRecall@10GPU 单查询 P99 延迟
1688.3%约 1.8ms
3292.1%约 3.2ms
6494.6%约 5.6ms
12896.2%约 9.4ms

所以不要迷信单一参数。如果你的业务对召回更敏感,比如搜索场景,建议 n_probes 开到 64 甚至 128;如果是推荐召回这样的上游环节,下游还有粗排模型兜底,n_probes=32 就够。实际做法是拿线上真实 query 样本跑一遍 AB,别在验证集上自嗨。

4.3 索引不常驻显存怎么处理

前面提到过内存映射索引的降级方案,这里再补充一组实测数据。把索引从显存模式切换成映射模式后,同样的 n_probes=64,单查询 P99 从 5.6ms 涨到 11.3ms,QPS 大概掉了三成。这个代价在某些场景能接受,但如果你有比较充足的 GPU 资源,还是尽量让索引留在显存里。映射模式真正适合的是那些查询量不大、但索引规模特别大的冷启动场景。

显存不足时的另一个实用思路是缩小 nlist。比如 nlist 从 16384 降到 4096,索引文件变化不大,但搜索时要扫描每个桶的候选数量会变多,反而可能让延迟升高。不要为了省显存盲目调小 nlist,它和 n_probes 是配合使用的。显存如果真的不够,优先考虑换小一点的 M 值,比如 M=16,索引体积能再降三分之一左右。

5. 常见问题排查与避坑记录

5.1 CUDA 环境和 cuVS 版本不匹配

这个坑几乎每个人都踩过。cuVS 的 pip 包分为cuvs-cu12和cuvs-cu11等不同版本,如果驱动支持的 CUDA 版本和包不一致,import 阶段不报错,但一执行 build 就崩,报错信息往往是CUDA driver version is insufficient或者直接段错误。排查方法很直接:先nvidia-smi看驱动版本,再用python -c "import cuvs; print(cuvs.__version__)"确认包版本。还遇到过 GPU 架构太老导致的no kernel image is available,这个基本没法通过换包解决,只能换卡。我的建议是直接用官方 docker 镜像,比如nvcr.io/nvidia/rapidsai/base系列,省掉一半以上的环境折腾时间。

5.2 召回率掉得很厉害

召回率不达标时,先别急着调 n_probes。我最常见的一个原因是向量没有归一化但用了内积度量,这会导致不同长度的向量之间有虚假的“高相似度”,召回结果自然乱掉。确认数据都做了归一化之后,再看 nlist 和训练样本数。训练样本太少会导致 k-means 聚类中心没有代表性,经验是训练样本量至少是 nlist 的 100 到 200 倍,我们取了 200 万,对 16384 个中心来说绰绰有余。最后调整 n_probes 和 M 的组合,如果 M 过大导致压缩误差明显,召回率会卡在一个平台期上不去,这时适当减小 M,召回率会有一次明显回升。

5.3 ES 导入慢和 mget 瓶颈

ES 导入慢多数是 refresh 和段合并造成的。导入前把refresh_interval设为 -1,导入结束恢复默认值,这个方法可以让写入吞吐量翻一倍以上。另外 bulk 批量大小和并发线程数需要调,5000 条一批、8 到 16 个并发线程在我们这边表现最好,再往上加并发会导致 ES 端 HTTP 连接池打满,性能反而下降。查询阶段如果发现 mget 慢,先看一眼是不是候选 doc_id 没有走主键而是走了其它字段,把字段类型设为 keyword 且只建必要的 doc_values,能省不少时间。如果一次候选 200 条 mget 就要 10ms 以上,可以考虑把候选数先降到 100,配合精确重排效果差别不大。

5.4 数据更新和增量索引怎么做

全量重建虽然只要 10 分钟,但也不能每次模型升级都无脑重建。我们目前的流程是:新 embedding 模型上线前,先离线抽一批验证集评估召回,确认指标不降再触发全量重建。重建过程中,ES 的老索引继续服务线上,新索引构建完成后切换查询入口,回滚操作也很简单,把配置里的索引版本号指回去就行。如果以后数据规模达到十亿级别,就需要考虑索引分片:把数据按 ID hash 分成多个分片,每个分片单独建 IVF-PQ 索引,查询时并行搜索所有分片再合并结果。这个方案我们还没上线,但已经在测试环境验证过流程,思路是把一台 GPU 承载的索引按水平拆到多卡上,单卡显存压力会小很多。

最后再分享一个小技巧:给每次构建的 cuVS 索引文件打上 build_id,并把索引路径放到 ES 的自定义 metadata 里。这样线上出问题的时候,你能精确知道当前查询跑的是哪一批数据构建的索引,排查“为什么召回结果和预期不一致”这种问题会快很多。当你的数据更新频率越来越高,版本管理就不再是锦上添花,而是必需品。

返回列表