09_vector_retrievalGPU 向量检索引擎/张磊 - #70
Open
zhang200210-neu wants to merge 18 commits into
Open
zhang200210-neu wants to merge 18 commits into
zhang200210-neu wants to merge 18 commits into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
GPU Vector Search Engine (vsearch)
一个面向 RAG / 推荐 / 多模态检索场景的 CUDA 向量检索演示引擎,支持:
exact:GPU 精确检索 baseline(全量距离 + 分块归并 Top-K),支持 L2 / inner product / cosine;
ivf_flat:IVF(倒排文件)+ Flat 向量距离;
ivf_pq:IVF + Product Quantization(ADC 距离表 + 压缩码本;默认对 top-256 候选做原始向量精排,设 pq_rerank = 0 可切换为纯 ADC);
fp16 / fp32 二进制向量库读取(内存统一按 fp32 计算);
索引保存 / 加载,避免每次查询重新建索引;
批量查询、K=1/10/50/100、CPU 参考实现、结果/性能/质量日志;
主机与 GPU 正确性测试。
程序可在四个 GPU 平台构建与运行:
平台 构建方式 验证情况
NVIDIA CMake + nvcc RTX 4090 上构建、主机/GPU 测试与 bench 全部通过
天数智芯 CoreX make -f Makefile.corex(clang -x ivcore) MR-V100(IX-ML 4.4.0)上全部通过
沐曦 MetaX make -f Makefile.maca(cucc + -lmcblas) MXC500(MACA 3.5.3)上全部通过
摩尔线程 MUSA make -f Makefile.musa(mcc -x musa + -lmublas) MUSA 5.1.0 上全部通过
四平台最近一次回归均通过(test_host + test_gpu 全部 PASS,端到端 bench 正常)。 四平台使用同一套检索算法与 kernel 实现,差异仅体现为构建脚本与少量编译期分支。 代表性结果见文末与 REPORT.md。