Home · ASPLOS 2024 · Machine Learning Systems

JUNO: Optimizing High-Dimensional Approximate Nearest Neighbour Search with Sparsity-Aware Algorithm and Ray-Tracing Core Mapping

Zihan Liu, Wentao Ni, Jingwen Leng, Yu Feng, Cong Guo, Quan Chen and 3 others

Conference
ASPLOS 2024
Published
2024
Pages
549-565
Citations
20

Vector search — finding the closest matches to a query among billions of high-dimensional vectors, the engine behind recommendations, retrieval, and RAG — leans on product quantization, which compresses vectors and looks up approximate distances, but it wastes 90-99% of its time computing distances to codebook centroids that never matter (only about a quarter of them encode the actual top results); Juno (Shanghai Jiao Tong and Rochester) cuts that waste with a sparsity-aware algorithm that scores only the centroids near the query, and — the clever twist — realizes that "which centroids fall within a distance threshold" is exactly a ray-sphere intersection test, so it repurposes the GPU's ray-tracing cores (built for graphics) to do the search in logarithmic time, deriving each distance for free from the ray's hit time and pipelining with the tensor cores — for 2.2-8.5x higher search throughput at the same accuracy.

The argument in one glance

  1. Problem

    Vector search is everywhere, and its dominant method wastes most of its work

    Approximate nearest neighbour (ANN) search — finding the vectors closest to a query in a huge high-dimensional dataset — is the engine behind recommendation systems, retrieval, and the vector databases that power retrieval-augmented generation. Exact search is hopeless in high dimensions (the curse of dimensionality), so the dominant method is PRODUCT QUANTIZATION (PQ): compress each vector by splitting it into sub-vectors and quantizing each to one of a small set of codebook centroids, then approximate a query's distance to every data point by looking up and summing precomputed sub-vector distances. It is fast and it is what frameworks like FAISS use. But it is deeply wasteful: it computes distances to ALL the codebook centroids for every query, when only a small fraction of them ever encode that query's actual nearest neighbours — so the overwhelming majority of the distance work is thrown away.

  2. What people did

    Run product quantization faster on the GPU

    The standard way to make PQ faster is to throw more parallel hardware at it: run it on a GPU, using the many CUDA cores to compute the pairwise centroid distances and accumulate the lookup tables in parallel. FAISS does exactly this, and it is the state of the art. Build the lookup tables for all centroids across all sub-spaces at once, sum them for every candidate point, and let the GPU's throughput carry the day. More cores, more parallelism, faster search.

  3. Why it fell short

    Parallelizing wasted work is still wasted work

    Running PQ on a GPU parallelizes the computation but does not remove the waste — and the waste is enormous. Profiling FAISS shows that for a query's top-100 results, only about 25% of the codebook centroids are actually used to encode those neighbours; in many sub-spaces, ALL of the top-100 points are encoded by just ONE centroid. Yet PQ dutifully computes and accumulates distances to 100% of the centroids, so 90% to 99.9% of the time in the lookup-table and distance-accumulation stages is spent on centroids that never matter to the answer. A faster engine that does the wrong work faster is still doing the wrong work. And even the useful work suffers on a GPU, because gathering the relevant centroids involves irregular memory accesses that GPUs handle poorly. The real problem is not throughput but that the algorithm computes far too much, in a pattern the hardware does not like.

  4. This paper's idea

    Skip the useless distances, and cast the search as ray tracing

    Juno (Shanghai Jiao Tong University and University of Rochester — Liu, Ni, Leng, Feng, Zhu, and colleagues) attacks both the wasted work and the hardware mismatch. First, a SPARSITY- and LOCALITY-AWARE algorithm: because the useful centroids are few and cluster NEAR the query, Juno sets a per-sub-space distance THRESHOLD (predicted by a small regression model trained offline on data density) and scores only the centroids inside it — skipping about 70% of the distance calculations. Second, the twist: "which centroids lie within a distance threshold of the query" is exactly a RAY-SPHERE INTERSECTION test, so Juno maps the search onto the GPU's RAY-TRACING CORES — fixed-function hardware built for graphics — which find the intersecting spheres via a bounding-volume tree in logarithmic time, deriving each distance for free from the ray's hit time and never touching global memory. The distance accumulation is pipelined onto the tensor cores. The result is 2.2-8.5x higher search throughput at the same accuracy.

Why this problem is worth solving

Vector search sits under a huge and growing set of applications — recommendation, semantic retrieval, and the databases that feed retrieval-augmented generation — so a large constant-factor speedup at the same accuracy is broadly valuable. The dominant algorithm’s cost, and its waste, both live in

the standard way to make nearest-neighbour search tractable in high dimensions. Instead of storing full vectors and computing exact distances, PQ compresses: it splits each D-dimensional vector into sub-vectors, and for each sub-space it trains a small codebook of E representative CENTROIDS, encoding each sub-vector as the nearest centroid. To answer a query, it computes the query’s distance to every centroid in every sub-space once (building a lookup table), and then approximates the distance to any data point by summing the looked-up sub-vector distances — turning an expensive full-vector distance into a few table lookups and an add. This is why PQ is fast and ubiquitous. But the lookup-table construction computes the query’s distance to ALL E centroids per sub-space, and the accumulation sums over all candidates, so its cost scales with the full codebook regardless of which centroids actually matter
. The algorithm treats every centroid as equally worth measuring, which is where the opportunity hides.

The measurement Juno makes is that the centroids are not equally worth measuring at all — the useful ones are

the empirical structure of which codebook centroids matter to a query’s result. Profiling shows that of all the centroids PQ scores, only about a quarter are ever used to encode a query’s true top-100 neighbours (SPARSITY) — in many sub-spaces a single centroid encodes all of them — and, crucially, those useful centroids are the ones spatially CLOSEST to the query (LOCALITY): selecting the roughly 25% of centroids nearest the query captures essentially all the ones that matter. So the answer depends on a small, nearby subset, while PQ spends 90-99.9% of its distance work on the far, irrelevant majority. This is exactly the kind of structure a well-designed search should exploit — compute distances only for the nearby centroids and skip the rest
. Turning that observation into speed, though, requires two things: an algorithm that identifies the nearby centroids cheaply per query, and hardware that can do that neighbourhood search efficiently despite the irregular memory access it implies. Juno’s insight is that the second requirement is already met by hardware sitting unused on the GPU — the ray-tracing cores — once you see that finding nearby centroids is the same problem those cores were built to solve.

What people did before, and where it broke

  • CPU product quantization (FAISS-CPU). The classic PQ pipeline on a CPU. Correct and widely used, but limited parallelism, and it still computes distances to all centroids — the algorithmic waste is present regardless of the hardware.
  • GPU product quantization (FAISS-GPU). Parallelize the lookup-table construction and accumulation across many CUDA cores. Faster, and the state of the art, but it parallelizes the WASTE — 90-99.9% of the distance work is on irrelevant centroids — and the useful gather is irregular memory access the GPU handles poorly.
  • Graph-based ANN (HNSW). Navigate a proximity graph to the nearest neighbours, often more accurate per candidate. Effective and orthogonal, but it is a different index; the PQ distance computation Juno optimizes is a component many systems (including graph-based ones) still rely on.
  • Fixed-threshold pruning. Prune centroids beyond a static distance cutoff. A step in the right direction, but a single fixed threshold is wrong across queries and data regions of different density — too loose wastes work, too tight loses recall.
  • The gap. No approach both removed the distance work on irrelevant centroids with a query-adaptive threshold AND executed the resulting neighbourhood search on hardware suited to it. Juno does both, by mapping the threshold search onto the GPU’s ray-tracing cores.

How Juno works

Prune the far centroids, cast the search as rays, run it on RT cores, pipeline the sum

  1. Diagnose the waste in product quantization new

    PQ computes distances to every codebook centroid, but only about 25% ever encode a query's true top results, so 90-99.9% of the distance work is thrown away. Juno targets exactly this waste rather than parallelizing it — the goal is to compute distances only for the centroids that matter.

  2. Exploit sparsity and locality with a query-adaptive threshold new

    The useful centroids are few and cluster near the query, so Juno sets a per-sub-space distance THRESHOLD and scores only the centroids inside it. The threshold is predicted per query by a small polynomial regression trained offline on the data density around the query (a denser region needs a wider threshold to capture the top neighbours), skipping about 70% of distance calculations.

  3. Recast the threshold search as ray-sphere intersection new

    Finding "which centroids lie within the threshold distance of the query" is geometrically identical to a ray-sphere intersection: place a sphere of radius = threshold at each centroid, cast a ray from the query, and a HIT means the centroid is within range. Because distance is symmetric, Juno pre-builds the centroid spheres offline and casts query rays at them at runtime, so no scene is rebuilt per query.

  4. Run it on the GPU ray-tracing cores new

    The RT cores are fixed-function hardware built for graphics that compute ray-geometry intersections through a bounding-volume-hierarchy tree in logarithmic time — exactly the neighbourhood search Juno needs. Spheres for different sub-spaces are stacked along a depth axis and the ray length capped so a query only tests its own sub-space, turning a wasteful linear scan into a log-time tree traversal on otherwise-idle hardware.

  5. Derive each distance from the ray hit time, not a memory fetch new

    Gathering a centroid's coordinates from global memory is the irregular access GPUs hate. Instead Juno computes the distance directly from the ray's intersection TIME (a value the RT hardware produces as a byproduct): the hit time encodes how far along the ray the sphere was struck, from which the query-to-centroid distance follows by a closed-form formula — no memory fetch at all, and inner-product search is handled by adjusting the sphere radii offline.

  6. Pipeline the accumulation onto the tensor cores new

    While the RT cores build the per-sub-space distances, the tensor cores sum them across sub-spaces as a matrix multiply (a distance matrix times an all-ones vector). Splitting the GPU's streaming multiprocessors roughly 90:10 between RT and tensor work and overlapping the two stages via CUDA MPS adds another 44-50%, keeping both specialized units busy.

Evaluated on SIFT (128 dimensions), DEEP (96 dimensions), and TTI (200 dimensions, inner-product) datasets from 1M to 100M points, against FAISS-GPU, across low/medium/high recall settings, on GPUs with and without RT cores (A40, RTX 4090, A100).

The idea

Do not do the work that does not matter. The core waste in PQ is that it measures the query’s distance to every codebook centroid, when only about a quarter of them ever encode the true nearest neighbours and those are the ones closest to the query. Juno’s first move is simply to stop computing the useless distances: predict a per-query distance threshold and score only the centroids inside it, cutting roughly 70% of the distance calculations. This is the spend- effort-only-where-it-counts discipline the site has been sharpening — the sparse-tiling page filling tiles by real occupancy instead of processing empty coordinate space, the edge-inference page spending its budget on the parameters that matter — and here it is the recognition that a fast engine doing wasted work is still wasted: before optimizing how the distances are computed, compute far fewer of them. The query-adaptive threshold (denser data needs a wider net) is what makes the pruning safe for recall.

Find the problem your idle hardware was already built to solve. The standout move is recognizing an isomorphism: “which centroids lie within a distance threshold of the query” is exactly “which spheres does a ray intersect” — the ray-sphere intersection test that a GPU’s RAY-TRACING cores perform, in logarithmic time via a bounding-volume tree, as their entire reason for existing in graphics. So Juno maps nearest-neighbour search onto graphics hardware, and gets a log-time tree search on units that would otherwise sit idle during an ANN workload. This is the map-your-problem-onto-a-solved-one discipline the site rewards — the heterogeneous-serving page casting scheduling as max-flow, the TLB page borrowing a hashing data structure — with an extra twist that the “solved problem” comes with DEDICATED SILICON: the payoff is not just a good algorithm but free acceleration from hardware built for something else entirely. Recognizing that your problem, suitably reframed, is one that existing hardware already accelerates is a rare and powerful lever.

Compute the answer from a byproduct instead of paying for a fetch. The elegant detail that makes the mapping fast is refusing to read the centroid coordinates from memory. On a GPU, gathering scattered centroid data is the irregular access that kills performance — so instead Juno derives the query-to-centroid distance from the ray’s intersection TIME, a value the RT hardware already produces when it reports a hit. The hit time says how far along the ray the sphere was struck, and a closed-form relation turns that into the distance, with no memory access at all. This is the use-the- byproduct, avoid-the-memory-roundtrip discipline the site keeps meeting — the persistent-memory and task-parallel pages dissolving data movement by passing references or keeping data on-chip — here sharpened to computing the needed quantity from a value the pipeline was going to generate anyway, turning a memory-bound gather into pure arithmetic.

Turning nearest-neighbour search into ray tracing

From scoring every centroid to a log-time ray query

Answering a query with product-quantization ANN search. Numbers from Juno's evaluation.

Step 1 of 5

How PQ answers a query

Product quantization compresses each vector into codebook centroids per sub-space, then answers a query by computing its distance to every centroid (building lookup tables) and summing the looked-up distances across sub-spaces for each candidate. This turns an expensive full-vector distance into table lookups and adds — fast, and the standard approach.

the lookup-table cost scales with the full codebook size E per sub-space, for every query
FAISS runs this on the GPU's CUDA cores in parallel

Why it works

  • PQ computes distances to all centroids but only ~25% matter → Juno prunes to the nearby centroids via a query-adaptive threshold → about 70% of distance calculations are skipped before any hardware is involved.
  • “Centroids within a distance threshold” is exactly a ray-sphere intersection → Juno maps the search onto the GPU’s ray-tracing cores → a log-time bounding-volume tree search replaces a linear scan, on otherwise-idle hardware.
  • Gathering centroid coordinates is the irregular memory access GPUs handle poorly → Juno derives each distance from the ray’s hit time by a closed-form formula → the gather becomes pure arithmetic with no memory fetch.
  • The RT cores do the filtering while the tensor cores can do the summation → pipelining the two with a roughly 90:10 SM split → both specialized units stay busy, adding 44-50%.
  • The useful centroids are sparse and locally clustered → the threshold-and-ray approach captures them while ignoring the far majority → 2.2-8.5x over FAISS-GPU at matched recall, and 2.5x from the algorithm alone even without RT cores.

What it costs

  • No accuracy guarantee. Like PQ-based ANN generally, Juno gives no theoretical recall guarantee and relies on the quality of the offline clustering, so it can miss true neighbours (false negatives); the density-based threshold is a heuristic.
  • Bound by the RT cores. The speedup depends on RT-core performance, so on a GPU without them (an A100) high-quality search degrades to only the algorithmic sparsity gain (about 2.5x for lower-quality search), and faster RT cores are needed for the top results.
  • Optimizes only the PQ component. Juno accelerates the product-quantization distance computation; graph-based refinements (HNSW) are orthogonal and not integrated, so systems that lean on those get less of the benefit.
  • Single-GPU scale. Evaluated up to 100M points on one GPU; billion-scale datasets would need chunking, partitioning, or other storage techniques the paper leaves to future work.
  • The aggressive approximation trades recall. The hit-count ranking (skip distance computation, rank by how many spheres a ray hits) buys extra speed only at moderate recall, introducing false positives and negatives unacceptable for high-recall applications.

The evaluation

SIFT (128 dimensions, L2), DEEP (96 dimensions, L2), and TTI (200 dimensions, inner-product) datasets from 1M to 100M points, against the FAISS-GPU product-quantization baseline, across low, medium, and high recall targets, on GPUs with RT cores (A40, RTX 4090) and without (A100), measuring search throughput at matched recall.

Speeding up product-quantization vector search

PropertyFAISS-GPU (parallel PQ)Sparsity-aware only (no RT cores)Juno (sparsity + RT cores)
Centroids scored per query
Wasted distance work
Where the search runs
Distance obtained fromComputed distances
Throughput vs FAISS-GPUReference
Juno's win is two-layered: the sparsity-aware threshold removes ~70% of the distance work (2.5x on its own, even without RT cores), and mapping the resulting neighbourhood search onto the GPU's idle ray-tracing cores — reading each distance from the ray hit time and pipelining the sum onto the tensor cores — multiplies it to 2.2-8.5x at matched accuracy.Underlined values explain themselves when clicked.

Check yourself

Check your understanding

  1. 1. Running product-quantization search on a GPU (FAISS) parallelizes it across many cores, yet Juno argues this misses the real problem. Predict why more parallelism is not the fix.

  2. 2. Juno's standout move is to run nearest-neighbour search on the GPU's ray-tracing cores, which were designed for graphics. Predict what makes this possible and powerful.

  3. 3. When a ray hits a centroid's sphere, Juno computes the query-to-centroid distance from the ray's intersection TIME rather than fetching the centroid's coordinates from memory. Predict why this detail matters so much on a GPU.

  4. 4. On a GPU without ray-tracing cores (an A100), Juno's high-quality speedup largely vanishes, yet its sparsity algorithm alone still delivers about 2.5x. Predict what this separation reveals about the design.

Sources

Related papers on Machine Learning Systems

See all 212 →