建實戰(zhàn))
向量檢索這個事做到一定規(guī)模之后CPU 就真的扛不住了。我們線上有一批 1.38 億條 384 維的 embedding 召回池之前用 Elasticsearch 原生的 HNSW 建索引一跑就是四五個小時而且 JVM 堆內(nèi)存和 GC 壓力大得讓人頭疼每天索引更新的調(diào)度窗口根本排不開。后來把 NVIDIA cuVS 引入檢索鏈路用 GPU 來做向量索引的構(gòu)建和近鄰搜索整批 1.38 億向量從原始數(shù)據(jù)到可查詢的索引文件10 分鐘出頭就能跑完檢索端延遲還比原來低了不少。這篇文章就是這次改造的完整復(fù)盤。里面會講清楚為什么選 cuVS Elasticsearch 這套組合、索引參數(shù)到底怎么推算才不踩坑、完整搭建流程怎么走以及我在實際環(huán)境中碰到的各種編譯、顯存、召回率問題。適合正在做推薦召回、多模態(tài)搜索、RAG 向量庫或者被億級向量索引構(gòu)建折磨得睡不著覺的工程師參考。不管你是第一次聽到 cuVS還是已經(jīng)跑過 GPU 向量檢索應(yīng)該都能從里面找到點能直接抄作業(yè)的東西。1. 方案選型為什么把 GPU 放到 ES 檢索鏈路里1.1 數(shù)據(jù)規(guī)模帶來的第一個坎先說一個最直接的感受1.38 億條 384 維 float32 向量光原始數(shù)據(jù)就有 211.9GB這個數(shù)字很多團隊第一眼看到是沒概念的。放到 Elasticsearch 里L(fēng)ucene 的 HNSW 索引構(gòu)建需要不斷計算鄰居圖、寫段、合并段不僅慢而且內(nèi)存占用遠(yuǎn)不止原始數(shù)據(jù)的大小。HNSW 的圖結(jié)構(gòu)本身會有額外開銷加上 JVM 堆、段合并的臨時空間跑批建索引的時候一臺 256GB 內(nèi)存的機器都會顯得捉襟見肘。我們當(dāng)時第一次全量構(gòu)建主節(jié)點直接拉起了 4 個多小時的任務(wù)中間還因為堆內(nèi)存告警被 kill 過兩次。這時候就需要重新思考一個問題ES 在整個鏈路里的職責(zé)到底是什么。它適合做文檔存儲、過濾、聚合和最終打分但在“億級向量全量建索引”這件事上純 CPU 的近似最近鄰算法確實到了瓶頸。GPU 的并行度天然適合這類計算密集型任務(wù)k-means 聚類、向量距離計算、PQ 編碼這些操作在 GPU 上可以成量級地加速。于是方向就很明確了把向量索引的構(gòu)建和搜索從 ES 的 JVM 里拆出來交給 GPU 側(cè)的向量搜索庫ES 只負(fù)責(zé)它最擅長的文檔管理和查詢路由。1.2 cuVS 是什么有哪些算法cuVS 是 NVIDIA 開源的 GPU 向量搜索庫歸屬于 RAPIDS 生態(tài)里面實現(xiàn)了幾種主流算法IVF-Flat、IVF-PQ、CAGRA也有一部分 HNSW 的支持。你可以把它理解成一個專門在 GPU 上做最近鄰搜索的工具箱輸入一組向量它幫你建索引然后給你一個支持近似搜索的接口。算法選型上IVF-PQ 是我們這次的首選原因是它在建索引速度和內(nèi)存占用上最均衡。IVF-PQ 做的事情可以理解成兩步先把向量空間用倒排索引切分成若干個區(qū)域然后對每個區(qū)域內(nèi)做乘積量化把高維向量壓縮成一小段碼字。壓縮之后1.38 億條 384 維向量索引文件只有幾個 GBA100 80G 的顯存完全放得下。CAGRA 我也測過一輪它是 cuVS 里性能更極端的圖算法搜索延遲可以做到更低但建索引時對顯存的要求更高數(shù)據(jù)量上億之后訓(xùn)練階段的中間結(jié)果很容易把顯存撐爆。如果你的數(shù)據(jù)量在幾千萬級別且對查詢延遲極度敏感CAGRA 值得試但像我這種 1.38 億全量池IVF-PQ 更穩(wěn)。HNSW 在 cuVS 里也有 GPU 實現(xiàn)但它的構(gòu)建復(fù)雜度擺在那里GPU 加速效果沒有 IVF 系明顯所以我們沒有走那條路。1.3 架構(gòu)設(shè)計存儲與檢索分層整套架構(gòu)最后落地是這樣分的Elasticsearch 負(fù)責(zé)管理文檔元數(shù)據(jù)、原始向量字段、過濾條件以及最終的精確重排cuVS 負(fù)責(zé)離線構(gòu)建大規(guī)模近鄰索引在線接受查詢向量返回候選集的 doc_id。數(shù)據(jù)流上離線階段先用 GPU 把全量向量構(gòu)建成 IVF-PQ 索引文件同時把 1.38 億文檔的元數(shù)據(jù)批量導(dǎo)入 ES查詢階段線上請求的 query 向量先送進 cuVS 搜索接口取回 Top-K 候選 doc_id再去 ES 里做一次 mget 把文檔撈出來最后可以用原始向量做精確距離重排。這樣做的好處是很明顯的。第一ES 的 JVM 堆不再背“全量向量索引”這個重?fù)?dān)集群穩(wěn)定性直線上升。第二索引構(gòu)建和更新完全獨立cuVS 重建索引不影響 ES 對外服務(wù)。第三整套鏈路是插拔式的哪天數(shù)據(jù)量再翻一倍我可以只加 GPU 節(jié)點和調(diào)整 IVF-PQ 參數(shù)不用推翻 ES 集群。如果你用的是 OpenSearch它自帶的 k-NN 插件支持 FAISS 思路類似但 ES 這邊沒有原生 GPU 索引插件自己接 cuVS 反而靈活。2. 數(shù)據(jù)準(zhǔn)備與索引參數(shù)推演別拍腦袋選 nlist 和 M2.1 先看清楚你的數(shù)據(jù)形態(tài)任何向量索引方案第一步不是寫代碼而是先搞清楚你的數(shù)據(jù)長什么樣。我們這批數(shù)據(jù)是文本和圖像的混合 embedding統(tǒng)一歸一到 384 維并且做了 L2 歸一化。為什么要歸一化很關(guān)鍵如果你打算用余弦相似度那歸一化之后的內(nèi)積就等價于余弦距離這樣在 cuVS 里可以直接用內(nèi)積作為距離度量少一層向量長度修正的開銷。不同維度下1.38 億條向量的數(shù)據(jù)體量完全不同直接影響顯存和參數(shù)選擇。下面的表是我當(dāng)時算的賬向量維度1.38 億條原始數(shù)據(jù)大小IVF-PQ 壓縮后索引大小M24, nbits8128 維約 70.6GB約 1.3GB384 維約 211.9GB約 3.9GB768 維約 423.9GB約 7.7GB所以如果你的 embedding 是 768 維甚至 1536 維壓縮后的索引還能接受但訓(xùn)練和編碼的中間過程對顯存的要求就高很多這時就要考慮分塊構(gòu)建。我們最終選擇在 A100 80G 上全量跑 384 維數(shù)據(jù)顯存剛好覆蓋。2.2 IVF-PQ 關(guān)鍵參數(shù)nlist、M、nbitsIVF-PQ 有三個參數(shù)是必須想清楚的nlist 是倒排索引的聚類中心數(shù)量M 是把向量切成多少個子向量段nbits 是每個子向量的量化位數(shù)。先說 nlist經(jīng)驗法則是取數(shù)據(jù)量的平方根附近的值sqrt(1.38 億) 約等于 11747所以取 16384 這個 2 的冪比較好。nlist 越大每個桶里的向量越少搜索時定位越準(zhǔn)但訓(xùn)練聚類的時間也更長還可能導(dǎo)致部分桶過小反而影響召回。如果你數(shù)據(jù)是 1 億級別16384 是一個很穩(wěn)的起點。M 的選擇相對靈活前提是能被向量維度整除。384 維可以取 M16 或 M24M16 時每個子向量是 24 維壓縮后每個向量 20 字節(jié)M24 時每個子向量是 16 維壓縮后 28 字節(jié)。壓縮率越高索引越小但量化誤差越大召回率會有損失。我們最終在召回率達標(biāo)的前提下選了 M24單條壓縮 28 字節(jié)總索引約 3.9GB這個量級在顯存里非常舒服。nbits 默認(rèn) 8 就夠表示每個子向量的碼本里有 256 個中心點再往上提收益很小反而顯著增加訓(xùn)練耗時。2.3 顯存不夠怎么辦如果你手里的卡不是 80G或者向量維度更高別慌cuVS 支持分批構(gòu)建。關(guān)鍵思路是分離“訓(xùn)練”和“編碼”兩個過程。訓(xùn)練階段只需要一部分有代表性的樣本比如 200 萬條隨機采樣就足夠 k-means 收斂200 萬 × 384 維 × 4 字節(jié)也就 3GB 左右編碼階段可以按照 500 萬條一個 chunk 分批進行每批算完結(jié)果寫盤最后拼接成完整的索引文件。這樣即便數(shù)據(jù)總量幾百 GB構(gòu)建過程的顯存峰值也能控制在 20GB 左右。還有一個實用技巧是索引加載時用內(nèi)存映射而不是一次性全量塞進顯存。cuVS 支持把索引文件映射到內(nèi)存搜索時按需讀取對應(yīng)的倒排桶這樣單張卡能承載的索引規(guī)?？梢赃h(yuǎn)大于顯存容量。代價是查詢延遲會有所上升因為多了一層磁盤/內(nèi)存讀取但對“索引不常駐顯存”的場景來說這是一個很值得用的降級方案。實測下來索引從顯存模式切到映射模式單查詢的 P99 延遲大約從 3ms 漲到 8ms 左右但能支撐更大規(guī)模的數(shù)據(jù)。3. 從零搭建 GPU 加速向量檢索管線的完整實操3.1 環(huán)境清單與安裝這一步看著簡單坑其實不少。先說硬件我們用了兩臺 GPU 節(jié)點每臺 8 張 A100 80G但實際構(gòu)建單索引只需要一張卡多卡主要用于并行處理不同數(shù)據(jù)分片。軟件層面需要 NVIDIA 驅(qū)動支持 CUDA 12.0 以上然后安裝 cuVS 的 Python 包。安裝命令很簡單# 建議用獨立的 conda 環(huán)境或 docker 鏡像避免污染線上環(huán)境 conda create -n cuvs-env python3.10 -y conda activate cuvs-env pip install cuvs-cu12 nvidia-smi # 確認(rèn)驅(qū)動和 GPU 狀態(tài)正常這里特別提醒一句cuVS 的版本和 CUDA 版本必須匹配裝完第一件事是跑一遍官方自帶的 smoke test。另外如果你之前裝過 pytorch 的 GPU 版本不要理所當(dāng)然地認(rèn)為 cuVS 一定能直接用兩個包的 CUDA runtime 依賴可能不一樣。我們一開始就在一臺裝過老舊 CUDA 10 環(huán)境的機器上翻車了最后換到干凈鏡像一次性通過。3.2 用 cuVS 構(gòu)建 IVF-PQ 索引的核心代碼cuVS 的 Python API 非常簡明核心邏輯就三段加載向量集、構(gòu)建索引、保存索引文件。我貼一段我們實際在用的簡化代碼import numpy as np import cuvs from cuvs.neighbors import ivf_pq # 加載向量假設(shè)是 (N, 384) 的 float32 數(shù)組 vectors np.fromfile(all_vectors.bin, dtypenp.float32).reshape(-1, 384) # 索引配置 nlist 16384 M 24 nbits 8 # 構(gòu)建 IVF-PQ 索引 index ivf_pq.build( vectors, metricinner_product, n_listnlist, mM, n_bitsnbits, kmeans_n_iters20, # k-means 迭代次數(shù)20 次足夠收斂 train_size2000000, # 訓(xùn)練樣本數(shù) ) # 保存索引 ivf_pq.save(index, cuvs_138m.index)這段代碼跑完后磁盤上會出現(xiàn)一個索引文件文件名我們自己帶上了版本號。構(gòu)建過程中你會發(fā)現(xiàn) GPU 利用率很高而 CPU 幾乎閑置這說明計算確實被卸載到顯卡上了。如果你的數(shù)據(jù)量太大一次 build 扛不住可以改用先 train 再分批 add 的寫法cuVS 支持增量添加編碼后的向量訓(xùn)練完成后每個 chunk 走add接口進去最后統(tǒng)一 save。需要說明的是不同 cuVS 版本對函數(shù)簽名略有調(diào)整實際以你安裝版本的官方文檔為準(zhǔn)但整體流程就是這三步。搜索端的代碼同樣很直接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 向量數(shù)據(jù)導(dǎo)入 ES雖然主檢索走 cuVS但 ES 里還是要存一份原始向量主要用于兩件事一是精確重排時計算真實距離二是排查線上問題時能直接撈出來看。因為這批數(shù)據(jù)不會頻繁更新我們直接把索引的 refresh 關(guān)掉用 bulk 批量寫入導(dǎo)入過程相對可控。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 } } } } }這里有個容易糾結(jié)的點ES 里既然已經(jīng)有 HNSW 索引了是不是可以直接用它做檢索如果你數(shù)據(jù)量在百萬級是的直接查就行但到了億級ES HNSW 的建索引時間和內(nèi)存開銷都會成為瓶頸所以我們在 ES 里保留 dense_vector 字段但不會把主查詢壓給它。bulk 導(dǎo)入時把refresh_interval設(shè)為 -1每批 5000 條并發(fā) 8 個線程1.38 億文檔大約 30 到 40 分鐘寫完。這個時間是可以接受的而且和 GPU 建索引并行跑整體調(diào)度窗口并不吃虧。3.4 檢索鏈路GPU 粗排 ES 精確重排查詢鏈路的設(shè)計決定用戶體驗。我們線上的流程是query 向量先進 cuVS一次取回 Top 200 的候選 doc_id然后帶著這批 doc_id 去 ES 做一次 multi-get把文檔取回來最后用 ES 字段里的原始向量做一次精確的余弦相似度計算重排出 Top 50 結(jié)果。為什么要做精確重排因為 IVF-PQ 是近似搜索量化誤差會導(dǎo)致局部順序不夠準(zhǔn)但它的召回候選通常已經(jīng)包含了真正相關(guān)的文檔精確重排一小批候選的成本很低收益卻很明顯。精確重排可以直接在應(yīng)用層用 numpy 做也可以把候選列表和目標(biāo)向量傳給 ES 的 script_score 查詢讓 ES 順便把分算了。數(shù)據(jù)量在 200 條候選這個級別兩種方案延遲差異不大應(yīng)用層做的好處是邏輯更可控。端到端來看如果用戶在 ES 上還有大量 filter 條件比如按類目限制那么更優(yōu)的做法是先把 filter 后的 doc_id 集合傳給 cuVS 搜索接口做限制或者反過來先用 cuVS 粗排再在 ES 里 filter。這個順序要根據(jù)你的過濾率來測我們的場景過濾率不高所以先 cuVS 再 ES filter延遲表現(xiàn)最好。4. 性能實測與參數(shù)調(diào)優(yōu)1.38 億向量到底多快4.1 索引構(gòu)建實測數(shù)據(jù)直接上我們在一張 A100 80G 上跑出來的數(shù)據(jù)。全量 1.38 億條 384 維向量IVF-PQ 參數(shù) nlist16384、M24、nbits8完整構(gòu)建過程分為訓(xùn)練、編碼、寫索引三段總耗時約 9 分 40 秒。其中 k-means 訓(xùn)練用了 2 分 15 秒分批編碼和寫入用了 7 分出頭。注意這里包含從原始二進制文件讀取數(shù)據(jù)的時間如果數(shù)據(jù)已經(jīng)在顯存里還能再快個幾十秒。對比之前的 CPU 方案用 Lucene HNSW 在 32 核高主頻機器上建同樣一批數(shù)據(jù)的索引實測耗時 4 小時 35 分鐘。這已經(jīng)不是“快多少倍”的問題了而是從“每天只能半夜跑一次”變成“隨時可以全量重建”。對我們這種需要頻繁刷新 embedding 模型的場景來說這個差異直接決定了能不能做到日更甚至小時級更新。如果你用的是 CAGRA 且數(shù)據(jù)量在千萬級構(gòu)建時間還能壓到 3 分鐘以內(nèi)但顯存需求會明顯變高。4.2 檢索延遲與召回率檢索端的表現(xiàn)我用兩個指標(biāo)說明延遲和召回率。延遲方面GPU 側(cè)的 IVF-PQ 搜索在 n_probes32 時單查詢平均 1.2msP99 在 3ms 左右加上 ES 的 mget 和精確重排完整端到端 P99 大約在 20 到 25ms。這個延遲水平和一臺中等配置的 ES 集群直接查億級 HNSW 差不多但關(guān)鍵在于索引構(gòu)建的瓶頸解決了。召回率方面我們在一個 10 萬條人工標(biāo)注的驗證集上測了 Recall10n_probes64 的時候可以達到 94.6%效果非常接近全量精確檢索。不同 n_probes 的權(quán)衡比較明顯n_probesRecall10GPU 單查詢 P99 延遲1688.3%約 1.8ms3292.1%約 3.2ms6494.6%約 5.6ms12896.2%約 9.4ms所以不要迷信單一參數(shù)。如果你的業(yè)務(wù)對召回更敏感比如搜索場景建議 n_probes 開到 64 甚至 128如果是推薦召回這樣的上游環(huán)節(jié)下游還有粗排模型兜底n_probes32 就夠。實際做法是拿線上真實 query 樣本跑一遍 AB別在驗證集上自嗨。4.3 索引不常駐顯存怎么處理前面提到過內(nèi)存映射索引的降級方案這里再補充一組實測數(shù)據(jù)。把索引從顯存模式切換成映射模式后同樣的 n_probes64單查詢 P99 從 5.6ms 漲到 11.3msQPS 大概掉了三成。這個代價在某些場景能接受但如果你有比較充足的 GPU 資源還是盡量讓索引留在顯存里。映射模式真正適合的是那些查詢量不大、但索引規(guī)模特別大的冷啟動場景。顯存不足時的另一個實用思路是縮小 nlist。比如 nlist 從 16384 降到 4096索引文件變化不大但搜索時要掃描每個桶的候選數(shù)量會變多反而可能讓延遲升高。不要為了省顯存盲目調(diào)小 nlist它和 n_probes 是配合使用的。顯存如果真的不夠優(yōu)先考慮換小一點的 M 值比如 M16索引體積能再降三分之一左右。5. 常見問題排查與避坑記錄5.1 CUDA 環(huán)境和 cuVS 版本不匹配這個坑幾乎每個人都踩過。cuVS 的 pip 包分為cuvs-cu12和cuvs-cu11等不同版本如果驅(qū)動支持的 CUDA 版本和包不一致import 階段不報錯但一執(zhí)行 build 就崩報錯信息往往是CUDA driver version is insufficient或者直接段錯誤。排查方法很直接先nvidia-smi看驅(qū)動版本再用python -c import cuvs; print(cuvs.__version__)確認(rèn)包版本。還遇到過 GPU 架構(gòu)太老導(dǎo)致的no kernel image is available這個基本沒法通過換包解決只能換卡。我的建議是直接用官方 docker 鏡像比如nvcr.io/nvidia/rapidsai/base系列省掉一半以上的環(huán)境折騰時間。5.2 召回率掉得很厲害召回率不達標(biāo)時先別急著調(diào) n_probes。我最常見的一個原因是向量沒有歸一化但用了內(nèi)積度量這會導(dǎo)致不同長度的向量之間有虛假的“高相似度”召回結(jié)果自然亂掉。確認(rèn)數(shù)據(jù)都做了歸一化之后再看 nlist 和訓(xùn)練樣本數(shù)。訓(xùn)練樣本太少會導(dǎo)致 k-means 聚類中心沒有代表性經(jīng)驗是訓(xùn)練樣本量至少是 nlist 的 100 到 200 倍我們?nèi)×?200 萬對 16384 個中心來說綽綽有余。最后調(diào)整 n_probes 和 M 的組合如果 M 過大導(dǎo)致壓縮誤差明顯召回率會卡在一個平臺期上不去這時適當(dāng)減小 M召回率會有一次明顯回升。5.3 ES 導(dǎo)入慢和 mget 瓶頸ES 導(dǎo)入慢多數(shù)是 refresh 和段合并造成的。導(dǎo)入前把refresh_interval設(shè)為 -1導(dǎo)入結(jié)束恢復(fù)默認(rèn)值這個方法可以讓寫入吞吐量翻一倍以上。另外 bulk 批量大小和并發(fā)線程數(shù)需要調(diào)5000 條一批、8 到 16 個并發(fā)線程在我們這邊表現(xiàn)最好再往上加并發(fā)會導(dǎo)致 ES 端 HTTP 連接池打滿性能反而下降。查詢階段如果發(fā)現(xiàn) mget 慢先看一眼是不是候選 doc_id 沒有走主鍵而是走了其它字段把字段類型設(shè)為 keyword 且只建必要的 doc_values能省不少時間。如果一次候選 200 條 mget 就要 10ms 以上可以考慮把候選數(shù)先降到 100配合精確重排效果差別不大。5.4 數(shù)據(jù)更新和增量索引怎么做全量重建雖然只要 10 分鐘但也不能每次模型升級都無腦重建。我們目前的流程是新 embedding 模型上線前先離線抽一批驗證集評估召回確認(rèn)指標(biāo)不降再觸發(fā)全量重建。重建過程中ES 的老索引繼續(xù)服務(wù)線上新索引構(gòu)建完成后切換查詢?nèi)肟诨貪L操作也很簡單把配置里的索引版本號指回去就行。如果以后數(shù)據(jù)規(guī)模達到十億級別就需要考慮索引分片把數(shù)據(jù)按 ID hash 分成多個分片每個分片單獨建 IVF-PQ 索引查詢時并行搜索所有分片再合并結(jié)果。這個方案我們還沒上線但已經(jīng)在測試環(huán)境驗證過流程思路是把一臺 GPU 承載的索引按水平拆到多卡上單卡顯存壓力會小很多。最后再分享一個小技巧給每次構(gòu)建的 cuVS 索引文件打上 build_id并把索引路徑放到 ES 的自定義 metadata 里。這樣線上出問題的時候你能精確知道當(dāng)前查詢跑的是哪一批數(shù)據(jù)構(gòu)建的索引排查“為什么召回結(jié)果和預(yù)期不一致”這種問題會快很多。當(dāng)你的數(shù)據(jù)更新頻率越來越高版本管理就不再是錦上添花而是必需品。 億級向量檢索提速：cuVS+Elasticsearch GPU索引構(gòu)建實戰(zhàn)
向量檢索這個事做到一定規(guī)模之后CPU 就真的扛不住了。我們線上有一批 1.38 億條 384 維的 embedding 召回池之前用 Elasticsearch 原生的 HNSW 建索引一跑就是四五個小時而且 JVM 堆內(nèi)存和 GC 壓力大得讓人頭疼每天索引更新的調(diào)度窗口根本排不開。后來把 NVIDIA cuVS 引入檢索鏈路用 GPU 來做向量索引的構(gòu)建和近鄰搜索整批 1.38 億向量從原始數(shù)據(jù)到可查詢的索引文件10 分鐘出頭就能跑完檢索端延遲還比原來低了不少。這篇文章就是這次改造的完整復(fù)盤。里面會講清楚為什么選 cuVS Elasticsearch 這套組合、索引參數(shù)到底怎么推算才不踩坑、完整搭建流程怎么走以及我在實際環(huán)境中碰到的各種編譯、顯存、召回率問題。適合正在做推薦召回、多模態(tài)搜索、RAG 向量庫或者被億級向量索引構(gòu)建折磨得睡不著覺的工程師參考。不管你是第一次聽到 cuVS還是已經(jīng)跑過 GPU 向量檢索應(yīng)該都能從里面找到點能直接抄作業(yè)的東西。1. 方案選型為什么把 GPU 放到 ES 檢索鏈路里1.1 數(shù)據(jù)規(guī)模帶來的第一個坎先說一個最直接的感受1.38 億條 384 維 float32 向量光原始數(shù)據(jù)就有 211.9GB這個數(shù)字很多團隊第一眼看到是沒概念的。放到 Elasticsearch 里L(fēng)ucene 的 HNSW 索引構(gòu)建需要不斷計算鄰居圖、寫段、合并段不僅慢而且內(nèi)存占用遠(yuǎn)不止原始數(shù)據(jù)的大小。HNSW 的圖結(jié)構(gòu)本身會有額外開銷加上 JVM 堆、段合并的臨時空間跑批建索引的時候一臺 256GB 內(nèi)存的機器都會顯得捉襟見肘。我們當(dāng)時第一次全量構(gòu)建主節(jié)點直接拉起了 4 個多小時的任務(wù)中間還因為堆內(nèi)存告警被 kill 過兩次。這時候就需要重新思考一個問題ES 在整個鏈路里的職責(zé)到底是什么。它適合做文檔存儲、過濾、聚合和最終打分但在“億級向量全量建索引”這件事上純 CPU 的近似最近鄰算法確實到了瓶頸。GPU 的并行度天然適合這類計算密集型任務(wù)k-means 聚類、向量距離計算、PQ 編碼這些操作在 GPU 上可以成量級地加速。于是方向就很明確了把向量索引的構(gòu)建和搜索從 ES 的 JVM 里拆出來交給 GPU 側(cè)的向量搜索庫ES 只負(fù)責(zé)它最擅長的文檔管理和查詢路由。1.2 cuVS 是什么有哪些算法cuVS 是 NVIDIA 開源的 GPU 向量搜索庫歸屬于 RAPIDS 生態(tài)里面實現(xiàn)了幾種主流算法IVF-Flat、IVF-PQ、CAGRA也有一部分 HNSW 的支持。你可以把它理解成一個專門在 GPU 上做最近鄰搜索的工具箱輸入一組向量它幫你建索引然后給你一個支持近似搜索的接口。算法選型上IVF-PQ 是我們這次的首選原因是它在建索引速度和內(nèi)存占用上最均衡。IVF-PQ 做的事情可以理解成兩步先把向量空間用倒排索引切分成若干個區(qū)域然后對每個區(qū)域內(nèi)做乘積量化把高維向量壓縮成一小段碼字。壓縮之后1.38 億條 384 維向量索引文件只有幾個 GBA100 80G 的顯存完全放得下。CAGRA 我也測過一輪它是 cuVS 里性能更極端的圖算法搜索延遲可以做到更低但建索引時對顯存的要求更高數(shù)據(jù)量上億之后訓(xùn)練階段的中間結(jié)果很容易把顯存撐爆。如果你的數(shù)據(jù)量在幾千萬級別且對查詢延遲極度敏感CAGRA 值得試但像我這種 1.38 億全量池IVF-PQ 更穩(wěn)。HNSW 在 cuVS 里也有 GPU 實現(xiàn)但它的構(gòu)建復(fù)雜度擺在那里GPU 加速效果沒有 IVF 系明顯所以我們沒有走那條路。1.3 架構(gòu)設(shè)計存儲與檢索分層整套架構(gòu)最后落地是這樣分的Elasticsearch 負(fù)責(zé)管理文檔元數(shù)據(jù)、原始向量字段、過濾條件以及最終的精確重排cuVS 負(fù)責(zé)離線構(gòu)建大規(guī)模近鄰索引在線接受查詢向量返回候選集的 doc_id。數(shù)據(jù)流上離線階段先用 GPU 把全量向量構(gòu)建成 IVF-PQ 索引文件同時把 1.38 億文檔的元數(shù)據(jù)批量導(dǎo)入 ES查詢階段線上請求的 query 向量先送進 cuVS 搜索接口取回 Top-K 候選 doc_id再去 ES 里做一次 mget 把文檔撈出來最后可以用原始向量做精確距離重排。這樣做的好處是很明顯的。第一ES 的 JVM 堆不再背“全量向量索引”這個重?fù)?dān)集群穩(wěn)定性直線上升。第二索引構(gòu)建和更新完全獨立cuVS 重建索引不影響 ES 對外服務(wù)。第三整套鏈路是插拔式的哪天數(shù)據(jù)量再翻一倍我可以只加 GPU 節(jié)點和調(diào)整 IVF-PQ 參數(shù)不用推翻 ES 集群。如果你用的是 OpenSearch它自帶的 k-NN 插件支持 FAISS 思路類似但 ES 這邊沒有原生 GPU 索引插件自己接 cuVS 反而靈活。2. 數(shù)據(jù)準(zhǔn)備與索引參數(shù)推演別拍腦袋選 nlist 和 M2.1 先看清楚你的數(shù)據(jù)形態(tài)任何向量索引方案第一步不是寫代碼而是先搞清楚你的數(shù)據(jù)長什么樣。我們這批數(shù)據(jù)是文本和圖像的混合 embedding統(tǒng)一歸一到 384 維并且做了 L2 歸一化。為什么要歸一化很關(guān)鍵如果你打算用余弦相似度那歸一化之后的內(nèi)積就等價于余弦距離這樣在 cuVS 里可以直接用內(nèi)積作為距離度量少一層向量長度修正的開銷。不同維度下1.38 億條向量的數(shù)據(jù)體量完全不同直接影響顯存和參數(shù)選擇。下面的表是我當(dāng)時算的賬向量維度1.38 億條原始數(shù)據(jù)大小IVF-PQ 壓縮后索引大小M24, nbits8128 維約 70.6GB約 1.3GB384 維約 211.9GB約 3.9GB768 維約 423.9GB約 7.7GB所以如果你的 embedding 是 768 維甚至 1536 維壓縮后的索引還能接受但訓(xùn)練和編碼的中間過程對顯存的要求就高很多這時就要考慮分塊構(gòu)建。我們最終選擇在 A100 80G 上全量跑 384 維數(shù)據(jù)顯存剛好覆蓋。2.2 IVF-PQ 關(guān)鍵參數(shù)nlist、M、nbitsIVF-PQ 有三個參數(shù)是必須想清楚的nlist 是倒排索引的聚類中心數(shù)量M 是把向量切成多少個子向量段nbits 是每個子向量的量化位數(shù)。先說 nlist經(jīng)驗法則是取數(shù)據(jù)量的平方根附近的值sqrt(1.38 億) 約等于 11747所以取 16384 這個 2 的冪比較好。nlist 越大每個桶里的向量越少搜索時定位越準(zhǔn)但訓(xùn)練聚類的時間也更長還可能導(dǎo)致部分桶過小反而影響召回。如果你數(shù)據(jù)是 1 億級別16384 是一個很穩(wěn)的起點。M 的選擇相對靈活前提是能被向量維度整除。384 維可以取 M16 或 M24M16 時每個子向量是 24 維壓縮后每個向量 20 字節(jié)M24 時每個子向量是 16 維壓縮后 28 字節(jié)。壓縮率越高索引越小但量化誤差越大召回率會有損失。我們最終在召回率達標(biāo)的前提下選了 M24單條壓縮 28 字節(jié)總索引約 3.9GB這個量級在顯存里非常舒服。nbits 默認(rèn) 8 就夠表示每個子向量的碼本里有 256 個中心點再往上提收益很小反而顯著增加訓(xùn)練耗時。2.3 顯存不夠怎么辦如果你手里的卡不是 80G或者向量維度更高別慌cuVS 支持分批構(gòu)建。關(guān)鍵思路是分離“訓(xùn)練”和“編碼”兩個過程。訓(xùn)練階段只需要一部分有代表性的樣本比如 200 萬條隨機采樣就足夠 k-means 收斂200 萬 × 384 維 × 4 字節(jié)也就 3GB 左右編碼階段可以按照 500 萬條一個 chunk 分批進行每批算完結(jié)果寫盤最后拼接成完整的索引文件。這樣即便數(shù)據(jù)總量幾百 GB構(gòu)建過程的顯存峰值也能控制在 20GB 左右。還有一個實用技巧是索引加載時用內(nèi)存映射而不是一次性全量塞進顯存。cuVS 支持把索引文件映射到內(nèi)存搜索時按需讀取對應(yīng)的倒排桶這樣單張卡能承載的索引規(guī)?？梢赃h(yuǎn)大于顯存容量。代價是查詢延遲會有所上升因為多了一層磁盤/內(nèi)存讀取但對“索引不常駐顯存”的場景來說這是一個很值得用的降級方案。實測下來索引從顯存模式切到映射模式單查詢的 P99 延遲大約從 3ms 漲到 8ms 左右但能支撐更大規(guī)模的數(shù)據(jù)。3. 從零搭建 GPU 加速向量檢索管線的完整實操3.1 環(huán)境清單與安裝這一步看著簡單坑其實不少。先說硬件我們用了兩臺 GPU 節(jié)點每臺 8 張 A100 80G但實際構(gòu)建單索引只需要一張卡多卡主要用于并行處理不同數(shù)據(jù)分片。軟件層面需要 NVIDIA 驅(qū)動支持 CUDA 12.0 以上然后安裝 cuVS 的 Python 包。安裝命令很簡單# 建議用獨立的 conda 環(huán)境或 docker 鏡像避免污染線上環(huán)境 conda create -n cuvs-env python3.10 -y conda activate cuvs-env pip install cuvs-cu12 nvidia-smi # 確認(rèn)驅(qū)動和 GPU 狀態(tài)正常這里特別提醒一句cuVS 的版本和 CUDA 版本必須匹配裝完第一件事是跑一遍官方自帶的 smoke test。另外如果你之前裝過 pytorch 的 GPU 版本不要理所當(dāng)然地認(rèn)為 cuVS 一定能直接用兩個包的 CUDA runtime 依賴可能不一樣。我們一開始就在一臺裝過老舊 CUDA 10 環(huán)境的機器上翻車了最后換到干凈鏡像一次性通過。3.2 用 cuVS 構(gòu)建 IVF-PQ 索引的核心代碼cuVS 的 Python API 非常簡明核心邏輯就三段加載向量集、構(gòu)建索引、保存索引文件。我貼一段我們實際在用的簡化代碼import numpy as np import cuvs from cuvs.neighbors import ivf_pq # 加載向量假設(shè)是 (N, 384) 的 float32 數(shù)組 vectors np.fromfile(all_vectors.bin, dtypenp.float32).reshape(-1, 384) # 索引配置 nlist 16384 M 24 nbits 8 # 構(gòu)建 IVF-PQ 索引 index ivf_pq.build( vectors, metricinner_product, n_listnlist, mM, n_bitsnbits, kmeans_n_iters20, # k-means 迭代次數(shù)20 次足夠收斂 train_size2000000, # 訓(xùn)練樣本數(shù) ) # 保存索引 ivf_pq.save(index, cuvs_138m.index)這段代碼跑完后磁盤上會出現(xiàn)一個索引文件文件名我們自己帶上了版本號。構(gòu)建過程中你會發(fā)現(xiàn) GPU 利用率很高而 CPU 幾乎閑置這說明計算確實被卸載到顯卡上了。如果你的數(shù)據(jù)量太大一次 build 扛不住可以改用先 train 再分批 add 的寫法cuVS 支持增量添加編碼后的向量訓(xùn)練完成后每個 chunk 走add接口進去最后統(tǒng)一 save。需要說明的是不同 cuVS 版本對函數(shù)簽名略有調(diào)整實際以你安裝版本的官方文檔為準(zhǔn)但整體流程就是這三步。搜索端的代碼同樣很直接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 向量數(shù)據(jù)導(dǎo)入 ES雖然主檢索走 cuVS但 ES 里還是要存一份原始向量主要用于兩件事一是精確重排時計算真實距離二是排查線上問題時能直接撈出來看。因為這批數(shù)據(jù)不會頻繁更新我們直接把索引的 refresh 關(guān)掉用 bulk 批量寫入導(dǎo)入過程相對可控。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 } } } } }這里有個容易糾結(jié)的點ES 里既然已經(jīng)有 HNSW 索引了是不是可以直接用它做檢索如果你數(shù)據(jù)量在百萬級是的直接查就行但到了億級ES HNSW 的建索引時間和內(nèi)存開銷都會成為瓶頸所以我們在 ES 里保留 dense_vector 字段但不會把主查詢壓給它。bulk 導(dǎo)入時把refresh_interval設(shè)為 -1每批 5000 條并發(fā) 8 個線程1.38 億文檔大約 30 到 40 分鐘寫完。這個時間是可以接受的而且和 GPU 建索引并行跑整體調(diào)度窗口并不吃虧。3.4 檢索鏈路GPU 粗排 ES 精確重排查詢鏈路的設(shè)計決定用戶體驗。我們線上的流程是query 向量先進 cuVS一次取回 Top 200 的候選 doc_id然后帶著這批 doc_id 去 ES 做一次 multi-get把文檔取回來最后用 ES 字段里的原始向量做一次精確的余弦相似度計算重排出 Top 50 結(jié)果。為什么要做精確重排因為 IVF-PQ 是近似搜索量化誤差會導(dǎo)致局部順序不夠準(zhǔn)但它的召回候選通常已經(jīng)包含了真正相關(guān)的文檔精確重排一小批候選的成本很低收益卻很明顯。精確重排可以直接在應(yīng)用層用 numpy 做也可以把候選列表和目標(biāo)向量傳給 ES 的 script_score 查詢讓 ES 順便把分算了。數(shù)據(jù)量在 200 條候選這個級別兩種方案延遲差異不大應(yīng)用層做的好處是邏輯更可控。端到端來看如果用戶在 ES 上還有大量 filter 條件比如按類目限制那么更優(yōu)的做法是先把 filter 后的 doc_id 集合傳給 cuVS 搜索接口做限制或者反過來先用 cuVS 粗排再在 ES 里 filter。這個順序要根據(jù)你的過濾率來測我們的場景過濾率不高所以先 cuVS 再 ES filter延遲表現(xiàn)最好。4. 性能實測與參數(shù)調(diào)優(yōu)1.38 億向量到底多快4.1 索引構(gòu)建實測數(shù)據(jù)直接上我們在一張 A100 80G 上跑出來的數(shù)據(jù)。全量 1.38 億條 384 維向量IVF-PQ 參數(shù) nlist16384、M24、nbits8完整構(gòu)建過程分為訓(xùn)練、編碼、寫索引三段總耗時約 9 分 40 秒。其中 k-means 訓(xùn)練用了 2 分 15 秒分批編碼和寫入用了 7 分出頭。注意這里包含從原始二進制文件讀取數(shù)據(jù)的時間如果數(shù)據(jù)已經(jīng)在顯存里還能再快個幾十秒。對比之前的 CPU 方案用 Lucene HNSW 在 32 核高主頻機器上建同樣一批數(shù)據(jù)的索引實測耗時 4 小時 35 分鐘。這已經(jīng)不是“快多少倍”的問題了而是從“每天只能半夜跑一次”變成“隨時可以全量重建”。對我們這種需要頻繁刷新 embedding 模型的場景來說這個差異直接決定了能不能做到日更甚至小時級更新。如果你用的是 CAGRA 且數(shù)據(jù)量在千萬級構(gòu)建時間還能壓到 3 分鐘以內(nèi)但顯存需求會明顯變高。4.2 檢索延遲與召回率檢索端的表現(xiàn)我用兩個指標(biāo)說明延遲和召回率。延遲方面GPU 側(cè)的 IVF-PQ 搜索在 n_probes32 時單查詢平均 1.2msP99 在 3ms 左右加上 ES 的 mget 和精確重排完整端到端 P99 大約在 20 到 25ms。這個延遲水平和一臺中等配置的 ES 集群直接查億級 HNSW 差不多但關(guān)鍵在于索引構(gòu)建的瓶頸解決了。召回率方面我們在一個 10 萬條人工標(biāo)注的驗證集上測了 Recall10n_probes64 的時候可以達到 94.6%效果非常接近全量精確檢索。不同 n_probes 的權(quán)衡比較明顯n_probesRecall10GPU 單查詢 P99 延遲1688.3%約 1.8ms3292.1%約 3.2ms6494.6%約 5.6ms12896.2%約 9.4ms所以不要迷信單一參數(shù)。如果你的業(yè)務(wù)對召回更敏感比如搜索場景建議 n_probes 開到 64 甚至 128如果是推薦召回這樣的上游環(huán)節(jié)下游還有粗排模型兜底n_probes32 就夠。實際做法是拿線上真實 query 樣本跑一遍 AB別在驗證集上自嗨。4.3 索引不常駐顯存怎么處理前面提到過內(nèi)存映射索引的降級方案這里再補充一組實測數(shù)據(jù)。把索引從顯存模式切換成映射模式后同樣的 n_probes64單查詢 P99 從 5.6ms 漲到 11.3msQPS 大概掉了三成。這個代價在某些場景能接受但如果你有比較充足的 GPU 資源還是盡量讓索引留在顯存里。映射模式真正適合的是那些查詢量不大、但索引規(guī)模特別大的冷啟動場景。顯存不足時的另一個實用思路是縮小 nlist。比如 nlist 從 16384 降到 4096索引文件變化不大但搜索時要掃描每個桶的候選數(shù)量會變多反而可能讓延遲升高。不要為了省顯存盲目調(diào)小 nlist它和 n_probes 是配合使用的。顯存如果真的不夠優(yōu)先考慮換小一點的 M 值比如 M16索引體積能再降三分之一左右。5. 常見問題排查與避坑記錄5.1 CUDA 環(huán)境和 cuVS 版本不匹配這個坑幾乎每個人都踩過。cuVS 的 pip 包分為cuvs-cu12和cuvs-cu11等不同版本如果驅(qū)動支持的 CUDA 版本和包不一致import 階段不報錯但一執(zhí)行 build 就崩報錯信息往往是CUDA driver version is insufficient或者直接段錯誤。排查方法很直接先nvidia-smi看驅(qū)動版本再用python -c import cuvs; print(cuvs.__version__)確認(rèn)包版本。還遇到過 GPU 架構(gòu)太老導(dǎo)致的no kernel image is available這個基本沒法通過換包解決只能換卡。我的建議是直接用官方 docker 鏡像比如nvcr.io/nvidia/rapidsai/base系列省掉一半以上的環(huán)境折騰時間。5.2 召回率掉得很厲害召回率不達標(biāo)時先別急著調(diào) n_probes。我最常見的一個原因是向量沒有歸一化但用了內(nèi)積度量這會導(dǎo)致不同長度的向量之間有虛假的“高相似度”召回結(jié)果自然亂掉。確認(rèn)數(shù)據(jù)都做了歸一化之后再看 nlist 和訓(xùn)練樣本數(shù)。訓(xùn)練樣本太少會導(dǎo)致 k-means 聚類中心沒有代表性經(jīng)驗是訓(xùn)練樣本量至少是 nlist 的 100 到 200 倍我們?nèi)×?200 萬對 16384 個中心來說綽綽有余。最后調(diào)整 n_probes 和 M 的組合如果 M 過大導(dǎo)致壓縮誤差明顯召回率會卡在一個平臺期上不去這時適當(dāng)減小 M召回率會有一次明顯回升。5.3 ES 導(dǎo)入慢和 mget 瓶頸ES 導(dǎo)入慢多數(shù)是 refresh 和段合并造成的。導(dǎo)入前把refresh_interval設(shè)為 -1導(dǎo)入結(jié)束恢復(fù)默認(rèn)值這個方法可以讓寫入吞吐量翻一倍以上。另外 bulk 批量大小和并發(fā)線程數(shù)需要調(diào)5000 條一批、8 到 16 個并發(fā)線程在我們這邊表現(xiàn)最好再往上加并發(fā)會導(dǎo)致 ES 端 HTTP 連接池打滿性能反而下降。查詢階段如果發(fā)現(xiàn) mget 慢先看一眼是不是候選 doc_id 沒有走主鍵而是走了其它字段把字段類型設(shè)為 keyword 且只建必要的 doc_values能省不少時間。如果一次候選 200 條 mget 就要 10ms 以上可以考慮把候選數(shù)先降到 100配合精確重排效果差別不大。5.4 數(shù)據(jù)更新和增量索引怎么做全量重建雖然只要 10 分鐘但也不能每次模型升級都無腦重建。我們目前的流程是新 embedding 模型上線前先離線抽一批驗證集評估召回確認(rèn)指標(biāo)不降再觸發(fā)全量重建。重建過程中ES 的老索引繼續(xù)服務(wù)線上新索引構(gòu)建完成后切換查詢?nèi)肟诨貪L操作也很簡單把配置里的索引版本號指回去就行。如果以后數(shù)據(jù)規(guī)模達到十億級別就需要考慮索引分片把數(shù)據(jù)按 ID hash 分成多個分片每個分片單獨建 IVF-PQ 索引查詢時并行搜索所有分片再合并結(jié)果。這個方案我們還沒上線但已經(jīng)在測試環(huán)境驗證過流程思路是把一臺 GPU 承載的索引按水平拆到多卡上單卡顯存壓力會小很多。最后再分享一個小技巧給每次構(gòu)建的 cuVS 索引文件打上 build_id并把索引路徑放到 ES 的自定義 metadata 里。這樣線上出問題的時候你能精確知道當(dāng)前查詢跑的是哪一批數(shù)據(jù)構(gòu)建的索引排查“為什么召回結(jié)果和預(yù)期不一致”這種問題會快很多。當(dāng)你的數(shù)據(jù)更新頻率越來越高版本管理就不再是錦上添花而是必需品。