Ascend-RaBitQ：异构NPU-CPU加速1位量化的十亿级相似性搜索
Ascend-RaBitQ: Heterogeneous NPU-CPU Acceleration of Billion-Scale Similarity Search with 1-bit Quantization
摘要
向量相似性搜索是现代人工智能系统的关键组成部分，但由于计算开销和内存带宽限制，传统的基于 CPU 的实现面临着数十亿规模语料库的基本可扩展性瓶颈。虽然神经处理单元 (NPU) 提供高几个数量级的计算密度，但由于基本的硬件不匹配，现有的 CPU/GPU 优化的 1 位 RaBitQ 量化实现无法直接移植到 NPU 架构，并且同质设计范例难以同时平衡精度、内存占用和性能。本文介绍了 Ascend-RaBitQ，这是第一个用于十亿级矢量搜索的异构 NPU-CPU 优化 IVF-RaBitQ 系统，其建立在将粗排序 (NPU) 与精细排序 (CPU) 解耦的核心洞察之上，允许每个阶段利用其最佳硬件，打破了长期存在的精度-内存-性能权衡。我们提出了一个三阶段异构执行路径，包括 AI Core 加速的 1 位量化向量粗排序、设备上 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 实现提高了 11.7 倍，比数学上等效的 CPU 基准提高了两个数量级以上，同时在分布式多 NPU 系统上展示了令人鼓舞的可扩展性。