中文

Ascend-RaBitQ:面向十亿级相似性搜索的异构NPU-CPU加速1位量化

信息检索 2026-05-18 v1

摘要

向量相似性搜索是现代AI系统的关键组成部分,但传统的基于CPU的实现由于计算开销和内存带宽限制,在处理十亿级语料库时面临根本性的可扩展性瓶颈。虽然神经处理单元提供了数量级更高的计算密度,但现有的CPU/GPU优化的1位RaBitQ量化实现由于根本性的硬件不匹配,无法直接移植到NPU架构,并且同构设计范式难以同时平衡精度、内存占用和性能。本文提出了Ascend-RaBitQ,这是首个面向十亿级向量搜索的异构NPU-CPU优化的IVF-RaBitQ系统,其核心见解是将粗排序(NPU)与精排序(CPU)解耦,允许每个阶段利用其最优硬件,打破了长期存在的精度-内存-性能权衡。我们提出了一种三阶段异构流水线,包括在1位量化向量上进行AI Core加速的粗排序、设备端AI CPU Top-k处理,以及在主机CPU上对全精度向量进行精重排序。我们引入了四种NPU架构原生优化:用于并行距离计算的融合AIC-AIV算子、利用旋转正交性的计算流重构、打破查询边界的细粒度索引块级负载均衡,以及AI Core和AI CPU之间的NPU内流水线并行以掩盖Top-k延迟。在标准数据集上的评估表明,Ascend-RaBitQ的索引构建速度比CPU基线快3.0到62.8倍,吞吐量比最快的CPU IVF-RaBitQ实现提升高达4.6倍,比数学上等效的CPU基线提升超过100倍,同时在分布式多NPU系统上展示了令人鼓舞的可扩展性。

关键词

引用

@article{arxiv.2605.16007,
  title  = {Ascend-RaBitQ: Heterogeneous NPU-CPU Acceleration of Billion-Scale Similarity Search with 1-bit Quantization},
  author = {Fujun He and Chuyue Ye and Huaxiang Cai and Zetao Lv and Baolong Cui and Wenru Yan and Chao Zhan and Zigang Zhang and Hao Yi and Jie Xiang and Xiabing Li and Yuhang Gai and Ziyang Zhang and Pengfei Zheng and Yunfei Du},
  journal= {arXiv preprint arXiv:2605.16007},
  year   = {2026}
}