PaperlayerView exact revision on arXiv

Source-linked AI summary

GPU-Native Approximate Nearest Neighbor Search with IVF-RaBitQ: Fast Index Build and Search

Jifan Shi, Jianyang Gao, James Xia, Tamás Béla Fehér, Cheng Long

arXiv:2602.23999v2cs.DBcs.DScs.IR

TL;DR

GPU ANNS must support high recall without imposing excessive build time, memory traffic, or storage costs. IVF-RaBitQ integrates IVF with GPU-native RaBitQ quantization, distance computation, and fused search. Across cuVS Bench datasets, it offers a strong recall-throughput-build-time-storage trade-off, including higher throughput and faster construction than CAGRA and IVF-PQ in the reported settings.

  • Problem

    GPU ANNS workloads need scalable methods that balance high recall and throughput with index construction time, memory traffic, and storage requirements.

  • Method

    IVF-RaBitQ integrates IVF partitioning with GPU-native RaBitQ quantization, GPU distance computation, fused search, and a GPU-oriented index layout.

  • Results

    Across cuVS Bench datasets, IVF-RaBitQ offers a strong trade-off across recall, throughput, build time, and storage, with 0.8×~8.2× speedup over CAGRA and 1.4×~58.3× over IVF-PQ at Recall≈0.95.

  • Takeaways & Limitations

    IVF-RaBitQ provides a practical GPU ANNS solution that avoids raw-vector reranking while maintaining a favorable efficiency–accuracy balance.

  • Takeaways & Limitations

    The GPU inner-product design assumes that the inner product between a 1-bit RaBitQ code and a query is independent of the code’s cluster.

Abstract

from arXiv · show

Approximate nearest neighbor search (ANNS) on GPUs is gaining increasing popularity for modern retrieval and recommendation workloads that operate over massive high-dimensional vectors. Graph-based indexes deliver high recall and throughput but incur heavy build-time and storage costs. In contrast, cluster-based methods build and scale efficiently yet often need many probes for high recall, straining memory bandwidth and compute. Aiming to simultaneously achieve fast index build, high-throughput search, high recall, and low storage requirement for GPUs, we present IVF-RaBitQ (GPU), a GPU-native ANNS solution that integrates the cluster-based method IVF with RaBitQ quantization into an efficient GPU index build/search pipeline. Specifically, for index build, we develop a scalable GPU-native RaBitQ quantization method that enables fast and accurate low-bit encoding at scale. For search, we develop GPU-native distance computation schemes for RaBitQ codes and a fused search kernel to achieve high throughput with high recall. With IVF-RaBitQ implemented and integrated into the NVIDIA cuVS Library, experiments on cuVS Bench across multiple datasets show that IVF-RaBitQ offers a strong performance frontier in recall, throughput, index build time, and storage footprint. For Recall approximately equal 0.95, IVF-RaBitQ achieves 3.0x higher QPS than the state-of-the-art graph-based method CAGRA, while also constructing indices 14.7x faster on average. Compared to the cluster-based method IVF-PQ, IVF-RaBitQ delivers on average over 4.5x higher throughput while avoiding accessing the raw vectors for reranking.

1 INTRODUCTION

GPU ANNS must balance recall, throughput, construction cost, and storage as workloads scale. IVF-RaBitQ combines IVF with GPU-native RaBitQ pipelines to target this trade-off, achieving strong benchmark performance against graph- and cluster-based baselines.

  • Motivation: Graph-based GPU methods provide strong recall and query performance, but cluster-based methods offer more regular computation patterns suited to GPU batching and high throughput.Cluster-based methods use dense centroid-distance computation and parallel scans within selected clusters.
  • Motivation: IVF-Flat is accurate but memory-bandwidth intensive, while IVF-PQ reduces traffic yet requires codebook training and may need reranking for high recall.IVF-PQ also provides approximate distance estimates without strict accuracy guarantees.
  • Motivation: RaBitQ avoids codebook training and provides tighter distance-estimation error bounds, but its rescaling-factor search and candidate evaluation pose GPU execution challenges.CPU implementations benefit from SIMD-oriented processing that does not directly map to GPUs.
  • Contributions: IVF-RaBitQ integrates IVF partitioning with GPU-native RaBitQ quantization to jointly target fast builds, high-throughput search, high recall, and compact storage.The design includes a GPU quantization pipeline and fused search components.
  • Evaluation: For Recall≈0.95, IVF-RaBitQ reaches 0.8×~8.2× speedup over CAGRA while reducing index construction time by 3.4×~55.3× and requiring less storage.Against cuVS IVF-PQ, it reaches 1.4×~58.3× higher query throughput and accuracy under similar build cost and storage budget.

2 PRELIMINARIES

The paper frames IVF search as cluster selection followed by local candidate scanning, then reviews RaBitQ’s rotated-code quantization and distance-estimation procedure. RaBitQ selects a codebook vector aligned with a rescaled, rotated normalized data vector.

  • IVF Search: IVF search first selects nearby centroids, then probes their inverted lists to compute candidate distances and retrieve approximate nearest neighbors.This two-stage organization separates coarse cluster selection from within-cluster search.
  • RaBitQ Quantization: RaBitQ normalizes a vector relative to its center, applies a random orthogonal rotation, and quantizes the resulting unit vector using a B-bit codebook.The vector norm and direction-related quantities are represented through the normalized and rotated formulation.
  • RaBitQ Quantization: Quantization chooses the codebook vector ˆx nearest to the rotated normalized vector, equivalently maximizing cosine similarity with the rotated data vector.Each codebook vector corresponds to a pre-rotation vector in the candidate set G.
  • RaBitQ Quantization: Rather than enumerate all 2^(B·D) codebook vectors, RaBitQ searches critical rescaling factors because the best code is closest to t·o′ for some t.Only D·2^(B−1) critical factors need consideration.
  • Distance Estimation: Distance estimation rotates the query with the same matrix and combines its dot product with the RaBitQ code with additional factors.The query remains higher precision while the data vector is represented by its compact code.

3.1 Overview

IVF-RaBitQ is an end-to-end GPU-native ANNS design that combines IVF partitioning, RaBitQ codes, GPU distance computation, fused local search, and a layout tailored to GPU access patterns.

  • 3.1 Overview: IVF-RaBitQ combines an IVF coarse partition with RaBitQ quantization and implements index construction and search as GPU-friendly pipelines.The design targets fast building, high-throughput search, high recall, and compact storage.
  • Index Build: During index construction, vectors are assigned to IVF clusters, normalized relative to centroids, rotated, and encoded into 1-bit codes plus remaining-bit ex-codes.This representation supports staged distance computation during search.
  • Search: During batched search, queries are transformed with the shared rotation, the closest n_probe clusters are selected, and query–cluster pairs are processed independently in parallel.Each local search returns top-K candidates from its probed inverted list.
  • GPU Index Layout: A flattened CSR-like layout stores clusters, 1-bit codes, ex-codes, and auxiliary factors in contiguous arrays, with interleaving tuned for the dominant filter stage.The layout is intended to improve GPU memory coalescing.

3.2 Index Build on GPU

The GPU index-build pipeline combines balanced clustering, orthogonal rotation, block-level vector processing, and parallel rescaling-factor search. Its coarse-to-fine search replaces sequential CPU-style enumeration with fixed GPU-parallel rounds.

  • 3.2 Index Build on GPU: Balanced K-means assigns vectors to IVF centroids, after which vectors and centroids are normalized and transformed by a shared orthogonal matrix without changing distances.Rotation is implemented with GEMM or a Hadamard-transform-based alternative.
  • Kernel Design: GPU quantization must avoid one-vector-per-launch processing because launch and scheduling overheads accumulate across millions or billions of vectors.The design therefore considers coarser kernel granularity and block-level processing.
  • Kernel Design: Mapping a thread block to each vector exposes intra-vector parallelism for exploring multiple critical scaling factors, improving compute efficiency and memory throughput.This mapping also helps maintain high occupancy during cluster-by-cluster processing.
  • Rescaling-Factor Search: CPU RaBitQ enumeration relies on sequential priority-queue state, motivating a GPU-friendly alternative for searching rescaling factors.The sequential dependency is poorly suited to massively parallel execution.
  • Rescaling-Factor Search: The proposed two-phase grid search first samples rescaling factors coarsely, then narrows around the best candidate and refines it in parallel to produce the quantization code.The approach is inspired by the approximately unimodal objective over rescaling factors.

3.3 Search Pipeline on GPU

The GPU search pipeline rotates queries, selects nearby IVF clusters, reorganizes query–cluster work for locality, and performs two-stage cluster-local candidate search before global top-K merging.

  • Cluster Selection: Query batches are rotated once, then compared with centroids to select the n_probe closest clusters using GPU matrix operations and top-K selection.The dominant query–centroid product is computed by one GEMM.
  • Workload Reorganization: Sorting query–cluster pairs by cluster identifier improves locality and makes large-batch GPU execution more regular.The reorganization groups work on the same or nearby clusters closer in execution time.
  • Cluster-local Search: Each probed cluster uses a two-stage estimator that filters candidates with 1-bit distances and refines survivors using additional RaBitQ bits.The first stage also uses a lower bound to prune vectors beyond the current K-th threshold.
  • Final Merge: Cluster-local candidates from all n_probe clusters are merged per query, followed by a GPU top-K operation producing the final approximate nearest neighbors.

3.4 Kernel and Data-Layout Co-design

IVF-RaBitQ co-designs GPU-native inner-product methods, fused search execution, and data layout to reduce conversion, synchronization, memory-traffic, and locality costs.

  • GPU Inner-Product Computation: GPU-native inner-product computation avoids expanding packed 1-bit codes into bytes, which would add conversions and underuse GPU compute and memory resources.The method targets asymmetric products between floating-point queries and compact 1-bit RaBitQ data codes.
  • GPU Inner-Product Computation: The LUT strategy precomputes query-dependent products for all 2^U binary patterns in each dimension block, then sums entries selected by each code block.LUTs are reusable across data vectors because the product is independent of cluster membership.
  • GPU Inner-Product Computation: Bitwise decomposition expresses the product as weighted binary inner products computed with AND and population-count operations, processing 32 dimensions per GPU register operation.Only the quantized query representation needs shared memory, reducing the footprint compared with LUT-based computation.
  • Fused Search Kernel: A fused cluster-local kernel combines all four search stages, propagating intermediate results through on-chip memory instead of repeated kernel launches and global-memory transfers.The design replaces kernel-level synchronization with block-level synchronization and reduces associated overheads.
  • GPU-Oriented Index Layout: The GPU-oriented layout stores cluster vectors contiguously and interleaves 1-bit codes to improve locality, memory coalescing, and cache efficiency.The layout is co-designed with cluster-local access patterns and the fused kernel.

4.1 Experimental Setup

The evaluation uses GPU and CPU platforms, cuVS Bench, multiple IVF-based baselines, and representative datasets, with search results summarized through time–accuracy trade-offs.

  • Experimental Platform: GPU experiments run on a single NVIDIA L40S with 48 GB device memory, while CPU experiments use two Intel Xeon Gold 6418H processors and 1 TB of DDR5 memory.Both systems run Ubuntu 22.04 LTS; the GPU setup uses CUDA 13.1.
  • Compared Methods: The comparison includes IVF-Flat, IVF-PQ with and without refinement, and CAGRA, with all indexes pre-loaded into GPU memory before search.IVF-PQ with refinement re-evaluates shortlisted candidates using raw vectors.
  • Benchmark: GPU benchmarks use cuVS Bench, which standardizes data loading, query batching, and performance measurement for GPU-resident ANN methods.The default batch size is 10^4.
  • Search Evaluation: Figure 3 reports representative-dataset time–accuracy trade-offs on a log scale, distinguishing IVF-RaBitQ’s LUT and Bitwise modes from IVF-PQ refinement settings.

4.2 Experimental Results

Across diverse datasets and operating conditions, IVF-RaBitQ combines strong search performance with fast construction and compact storage, including when raw vectors exceed a single GPU’s memory.

  • GPU ANNS Search Performance: At Recall=0.95, IVF-RaBitQ delivers 1.4×~10.0× higher speed than IVF-PQ (w/ref) and 2.4×~58.3× higher speed than IVF-PQ (wo/ref).Against CAGRA, IVF-RaBitQ (Bitwise) averages 3.3× higher speed at recall 0.95, although CAGRA is slightly better on low-dimensional SIFT.
  • Kernel and Data-Layout Ablation: The bitwise kernel remains lightweight and scalable as dimensionality increases, while the LUT kernel becomes bounded by shared-memory capacity.The relative performance depends on GPU memory bandwidth and shared-memory size.
  • Scalability with Query Batch Size: At larger batch sizes, IVF-RaBitQ significantly outperforms IVF-PQ (w/ref) and CAGRA in throughput while maintaining substantially lower latency.The comparison targets Recall@10 values of 0.90, 0.95, and 0.99; all methods have sub-millisecond latency at small batch sizes.
  • Index Build Time: 14.7× average faster index construction than CAGRA and 4.2× faster than IVF-PQ, under matched cluster-training and 8-bit-per-dimension compression settings.IVF-RaBitQ is slower only than IVF-Flat among the evaluated IVF methods.
  • Storage Footprint: Less than 25% of the storage required by raw-vector methods while achieving comparable recall, and similar storage to IVF-PQ (wo/ref) with substantially higher recall and throughput.The storage comparison includes raw vectors for CAGRA, IVF-Flat, and IVF-PQ (w/ref).
  • Datasets Exceeding GPU Memory: On DEEP-150M, IVF-RaBitQ uses about 18.0 GB without raw vectors, whereas IVF-PQ (w/ref) uses 73.3 GB and IVF-PQ (wo/ref) cannot reach Recall≥0.90 or ≥0.95.The dataset’s 57.6 GB of raw vectors exceeds the 48 GB memory of one L40S GPU; IVF-RaBitQ achieves 1.58× and 1.70× higher QPS than IVF-PQ (w/ref) at the reported recall thresholds.

5 RELATED WORK

GPU ANNS research includes graph-based methods that traverse proximity graphs and cluster-based methods that organize vectors into searchable partitions. Recent GPU work focuses on parallel graph construction, traversal, candidate management, and related acceleration strategies.

  • Graph-Based ANNS on GPU: Graph-based GPU ANNS constructs proximity graphs and searches them by iteratively expanding candidate frontiers from one or a few entry points.Recent designs parallelize graph construction and optimize traversal, priority queues, visited sets, and hash-table filtering.
  • Graph-Based ANNS on GPU: Representative systems include GGNN’s hierarchical construction, GGraphCon’s divide-and-conquer graph building, and CAGRA’s fixed-degree graph with GPU-oriented search optimizations.These methods target parallel construction, load balance, graph reachability, and visited-set overhead.
  • Other GPU ANNS Approaches: GPU ANNS methods also include tree- and IVF-based approaches, with FAISS integrating GPU top-K selection and IVF-PQ for compressed search.These lines of work demonstrate GPU acceleration for large-scale approximate nearest-neighbor workloads.

6 CONCLUSION

IVF-RaBitQ combines IVF partitioning with RaBitQ quantization in an end-to-end GPU-native ANNS framework, targeting fast builds, high-throughput search, high recall, and compact storage. Integrated into NVIDIA cuVS and cuVS Bench, it achieves a strong trade-off frontier across these metrics, while dynamic updates and transactional operations remain future extensions.

  • IVF-RaBitQ integrates IVF partitioning and RaBitQ quantization into an end-to-end GPU-native ANNS framework.Its design combines scalable GPU quantization, GPU-native inner-product computation, a fused cluster-local search kernel, and a GPU-oriented index layout.
  • The framework delivers a strong trade-off frontier across recall, throughput, index build time, and memory footprint, especially in high-recall regimes.The system is integrated into NVIDIA cuVS and the cuVS Bench pipeline for reproducible evaluation and fair comparison with existing GPU baselines.
  • Future work will explore dynamic index updates and transactional operations.
Loading 2602.23999v2…