亿级向量检索提速：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 维向量索引文件只有几个 GBA100 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 和 M2.1 先看清楚你的数据形态任何向量索引方案第一步不是写代码而是先搞清楚你的数据长什么样。我们这批数据是文本和图像的混合 embedding统一归一到 384 维并且做了 L2 归一化。为什么要归一化很关键如果你打算用余弦相似度那归一化之后的内积就等价于余弦距离这样在 cuVS 里可以直接用内积作为距离度量少一层向量长度修正的开销。不同维度下1.38 亿条向量的数据体量完全不同直接影响显存和参数选择。下面的表是我当时算的账向量维度1.38 亿条原始数据大小IVF-PQ 压缩后索引大小M24, nbits8128 维约 70.6GB约 1.3GB384 维约 211.9GB约 3.9GB768 维约 423.9GB约 7.7GB所以如果你的 embedding 是 768 维甚至 1536 维压缩后的索引还能接受但训练和编码的中间过程对显存的要求就高很多这时就要考虑分块构建。我们最终选择在 A100 80G 上全量跑 384 维数据显存刚好覆盖。2.2 IVF-PQ 关键参数nlist、M、nbitsIVF-PQ 有三个参数是必须想清楚的nlist 是倒排索引的聚类中心数量M 是把向量切成多少个子向量段nbits 是每个子向量的量化位数。先说 nlist经验法则是取数据量的平方根附近的值sqrt(1.38 亿) 约等于 11747所以取 16384 这个 2 的幂比较好。nlist 越大每个桶里的向量越少搜索时定位越准但训练聚类的时间也更长还可能导致部分桶过小反而影响召回。如果你数据是 1 亿级别16384 是一个很稳的起点。M 的选择相对灵活前提是能被向量维度整除。384 维可以取 M16 或 M24M16 时每个子向量是 24 维压缩后每个向量 20 字节M24 时每个子向量是 16 维压缩后 28 字节。压缩率越高索引越小但量化误差越大召回率会有损失。我们最终在召回率达标的前提下选了 M24单条压缩 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 python3.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, dtypenp.float32).reshape(-1, 384) # 索引配置 nlist 16384 M 24 nbits 8 # 构建 IVF-PQ 索引 index ivf_pq.build( vectors, metricinner_product, n_listnlist, mM, n_bitsnbits, kmeans_n_iters20, # k-means 迭代次数20 次足够收敛 train_size2000000, # 训练样本数 ) # 保存索引 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, k200, n_probes64)这里的queries是待检索的 query 向量 batchn_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 参数 nlist16384、M24、nbits8完整构建过程分为训练、编码、写索引三段总耗时约 9 分 40 秒。其中 k-means 训练用了 2 分 15 秒分批编码和写入用了 7 分出头。注意这里包含从原始二进制文件读取数据的时间如果数据已经在显存里还能再快个几十秒。对比之前的 CPU 方案用 Lucene HNSW 在 32 核高主频机器上建同样一批数据的索引实测耗时 4 小时 35 分钟。这已经不是“快多少倍”的问题了而是从“每天只能半夜跑一次”变成“随时可以全量重建”。对我们这种需要频繁刷新 embedding 模型的场景来说这个差异直接决定了能不能做到日更甚至小时级更新。如果你用的是 CAGRA 且数据量在千万级构建时间还能压到 3 分钟以内但显存需求会明显变高。4.2 检索延迟与召回率检索端的表现我用两个指标说明延迟和召回率。延迟方面GPU 侧的 IVF-PQ 搜索在 n_probes32 时单查询平均 1.2msP99 在 3ms 左右加上 ES 的 mget 和精确重排完整端到端 P99 大约在 20 到 25ms。这个延迟水平和一台中等配置的 ES 集群直接查亿级 HNSW 差不多但关键在于索引构建的瓶颈解决了。召回率方面我们在一个 10 万条人工标注的验证集上测了 Recall10n_probes64 的时候可以达到 94.6%效果非常接近全量精确检索。不同 n_probes 的权衡比较明显n_probesRecall10GPU 单查询 P99 延迟1688.3%约 1.8ms3292.1%约 3.2ms6494.6%约 5.6ms12896.2%约 9.4ms所以不要迷信单一参数。如果你的业务对召回更敏感比如搜索场景建议 n_probes 开到 64 甚至 128如果是推荐召回这样的上游环节下游还有粗排模型兜底n_probes32 就够。实际做法是拿线上真实 query 样本跑一遍 AB别在验证集上自嗨。4.3 索引不常驻显存怎么处理前面提到过内存映射索引的降级方案这里再补充一组实测数据。把索引从显存模式切换成映射模式后同样的 n_probes64单查询 P99 从 5.6ms 涨到 11.3msQPS 大概掉了三成。这个代价在某些场景能接受但如果你有比较充足的 GPU 资源还是尽量让索引留在显存里。映射模式真正适合的是那些查询量不大、但索引规模特别大的冷启动场景。显存不足时的另一个实用思路是缩小 nlist。比如 nlist 从 16384 降到 4096索引文件变化不大但搜索时要扫描每个桶的候选数量会变多反而可能让延迟升高。不要为了省显存盲目调小 nlist它和 n_probes 是配合使用的。显存如果真的不够优先考虑换小一点的 M 值比如 M16索引体积能再降三分之一左右。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 里。这样线上出问题的时候你能精确知道当前查询跑的是哪一批数据构建的索引排查“为什么召回结果和预期不一致”这种问题会快很多。当你的数据更新频率越来越高版本管理就不再是锦上添花而是必需品。