Billion-Scale Similarity Search with GPUs

Deep learning turns images, sentences, and videos into vectors — lists of numbers that capture meaning. Two similar images produce vectors that are close together in this high-dimensional space. Finding the $k$ nearest neighbors of a query vector $x$ among a database of vectors $[y_i]$ means solving:

$L = k\text{-}\arg\min_{i} \|x - y_i\|^2$

For a million vectors this is already expensive. For a billion, the naive approach — computing all distances and sorting — is completely impractical. Even on a GPU with teraflops of throughput, two bottlenecks remain: the raw volume of distance computations and the challenge of selecting the top-$k$ results efficiently.

A 128-dimensional vector takes 512 bytes as 32-bit floats. A billion of them needs 512 GB — far more than any GPU's memory. Product quantization (PQ) solves this by splitting each vector into $b$ sub-vectors and replacing each sub-vector with the index of its closest centroid from a small learned codebook of 256 entries. The result: each vector is encoded in just $b$ bytes (typically 8–64 bytes), a compression ratio of 8–64×.

Think of it like color quantization in a GIF image: instead of storing exact RGB values for every pixel, you build a 256-color palette and store just the palette index per pixel. PQ does the same thing, but independently for each slice of the vector.

Even with compressed vectors, scanning all of them is wasteful. The inverted file (IVF) adds a routing step: first, cluster the entire database into $|C_1|$ groups using k-means. At query time, find the $\tau$ closest cluster centers to the query (the multi-probe parameter), then search only the vectors in those clusters. If $|C_1| = \sqrt{n}$ and $\tau$ is small, you skip the vast majority of the database.

Picture a city-wide post office system. Instead of checking every mailbox in the city, you first identify which $\tau$ postal districts your letter most likely belongs to, then search only the mailboxes in those districts. More districts ($\tau$) means better recall at the cost of more scanning.

The GPU's power comes from thousands of threads working in parallel, but finding the top-$k$ smallest distances is inherently sequential — classic heaps are slow on GPUs because they have irregular memory access and high warp divergence.

WarpSelect solves this by keeping all intermediate state in the GPU register file — the fastest memory tier, with 10–100× the bandwidth of even shared memory. Each warp (32 threads) maintains a sorted "warp queue" of the $k$ smallest values seen so far, plus small per-thread queues that act as filters. New values are rejected instantly if they're larger than the per-thread queue head. Only when the thread queues fill up does a bitonic sorting network merge everything back into the warp queue.

The result: a single-pass algorithm that processes input at up to 55% of peak GPU memory bandwidth — close to the theoretical speed limit of simply reading the data.

A brilliant property of product quantization: you never need to decompress the vectors to compute distances. For a given query $x$ and a known coarse centroid $q_1(y)$, precompute a lookup table $T_j$ of size 256 for each of the $b$ sub-quantizers. Each entry stores the squared distance from the query's sub-vector to a centroid. Then, computing the approximate distance to any compressed vector is just $b$ table lookups and additions — no multiply-adds, no decompression.

This turns the problem from $n \times d$ floating-point operations (expensive) to $n \times b$ byte lookups (cheap and cache-friendly). On a GPU, these lookup tables fit perfectly in shared memory, the fast on-chip cache that all threads in a block can access.

The full search pipeline for a batch of queries works in three stages:

• Coarse quantization: find the $\tau$ closest cluster centers for each query using exact brute-force search (matrix multiplication via cuBLAS).

• List scanning: for each query, scan only the $\tau$ selected inverted lists. For each compressed vector in a list, compute its approximate distance via lookup tables in shared memory, and feed the result into WarpSelect running in registers.

• k-selection finalization: merge partial results across lists and tiles to produce the final top-$k$ for each query.

Tiling over queries ($t_q$ at a time) and multi-streaming keep GPU utilization high while bounding memory usage. The entire pipeline avoids unnecessary global memory round-trips — the critical insight that makes FAISS so fast.

When the index is too large for a single GPU, FAISS supports two complementary strategies:

• Replication: copy the entire index to $R$ GPUs. Each GPU handles $n_q/R$ queries against the full database. This gives near-linear speedup for large query batches.

• Sharding: split the database across $S$ GPUs, each holding $\ell/S$ vectors. Every query is sent to all shards, and partial results are merged with an additional round of k-selection. This lets you index datasets that exceed any single GPU's memory.

Both strategies compose: $S$ shards with $R$ replicas each use $S \times R$ GPUs total. In practice, 4 GPUs can handle a billion 128-dimensional vectors with room to spare.

FAISS is to vector search what cuBLAS is to matrix multiplication: the optimized building block that every higher-level system builds upon. Dense Passage Retrieval would not exist without it, and neither would the retrieval backends powering today's language model assistants.