Accelerating Hierarchical Navigable Small World (HNSW)-RAG Vector Search with CUDA

How we moved the algorithm behind modern vector databases to the GPU step by step, and cut query time by 57.7%, without touching recall. Work done by: Miguel Gutierrez and Syaqui Rahmat Perdana If you use vector database search empowered on solutions like RAG, Recommendation Engines. There is a good chanche an algorithm (Hierarchical Navigable Small World )HNSW was doing the hard work behind. Here we improve CPUs implementation to use deep CUDA to improve the timeline. This proyect was done as a contest for the “High Performance Graph Data Analytics” Course contest at Politecnico Di Milano with Oracle. I did this proyect with my friend and teamate Syaqui Rahmat Perdana. The problem was : how much faster can HNSW get if we move it search to a GPU with CUDA ? This post is the story of optimization journey of widely used techniques for improving speedup. Five Techniques, one surprising failure and a final speedup of 57.7% per query. Plus an 84% reduction in raw search time and 33% faster index construction under the contest parameters. All the code is on GitHub . Introduction Hierarchical Navigable Small World ( HNSW), from Malkov & Yashunin’s 2020 paper [1], solves the nearest-neighbor problem: given a query vector, find the most similar vectors in a dataset of potentially millions. Comparing the query against every point works, but scales terribly. HNSW’s trick is to organize the data as a multi-layer graph : Each node is a data point; edges connect nearby points. The graph is stacked into layers : the top layers are sparse (few nodes, long-range links), and each layer down gets exponentially denser. There are 3 layers; as the layer increases, the number of nodes decreases exponentially. The nodes are connected to some close nodes at each layer. Search works like a ski lift: you enter at the sparse top layer, greedily hop toward the query, then drop down a layer and refine, repeating until you hit the dense bottom layer where the true nearest neighbors live. This mimics a skip list and gives you O(log N) search, which is why HNSW dominates the ANN benchmarks. Insertion follows the same idea: each new point gets a random maximum layer (drawn from an exponentially decaying distribution), then connects to its nearest neighbors on each layer it belongs to, capped at M connections. The steps start by adding the data to each layer with probability such that as the layer increase the number of nodes decrease exponentially. Then, for the next iterations it will check the topmost layer and then descent to the next layer and stop when maximum connections are met. The catch? Both search and insert are inherently sequential and branchy, a greedy walk through a graph, one hop at a time. That’s about the worst possible shape for a GPU, which wants thousands of identical operations running in lockstep. The interesting engineering question is finding the parallelism hiding inside the sequential walk. The Insight: Parallelize the Distances, Not the Walk Here’s the CPU inner loop of the layer search. For each candidate node, we walk its neighbors one at a time and compute a distance per neighbor: // Expand the current best candidate: look at each of its neighbors for (const auto neighbor : nearest_candidate_node.neighbors) { // Skip neighbors we've already evaluated in this search if (visited[neighbor.id]) continue; visited[neighbor.id] = true; // mark as seen // Fetch the neighbor's vector from the current layer (l_c) const auto& neighbor_node = layers[l_c][neighbor.id]; // ONE distance computation per loop iteration — this is the bottleneck const auto dist_from_neighbor = calc_dist(query, neighbor_node.data); ... } Each distance computation is independent of the others. That’s our parallelism. The greedy walk stays sequential on the host, but every time we expand a node, the batch of neighbor distances gets computed on the GPU simultaneously. The kernel itself is nothing exotic, one thread per vector, each computing a Euclidean distance: // __global__ = this function runs on the GPU, launched from the CPU __global__ void calculateDistances( const float* query, // the query vector (dim floats) const float* vectors, // all candidate vectors, packed back-to-back float* distances, // output: one distance per candidate int dim, // dimensionality of each vector int num_vectors // how many candidates are in this batch ) { // Each GPU thread gets a unique index — one thread per vector int idx = blockIdx.x * blockDim.x + threadIdx.x; // Threads beyond the batch size have nothing to do if (idx >= num_vectors) return; float distance = 0.0f; // Pointer to the start of *this thread's* vector in the packed array const float* vector = vectors + (idx * dim); // Sum of squared differences across all dimensions for (int i = 0; i < dim; i++) { float diff = vector[i] - query[i]; distance += diff * diff; } // Euclidean distance = sqrt of the sum; write to this thread's slot distances[idx] = sqrtf(distance); } On its own, this kernel barely moved the needle. And that’s the first lesson of GPU programming: the kernel is rarely the bottleneck — the memory traffic around it is. Everything that follows is about feeding this kernel efficiently. The Optimization Ladder 1. Batch Processing Instead of expanding one neighbor at a time, we collect unvisited neighbor IDs into a batch, ship the batch to the GPU, and compute all distances in one kernel launch: // Keep pulling candidates until the batch is full (or nothing is left) while (!candidates.empty() && batch_indices.size() < BATCH_SIZE) { ... // Collect ALL unvisited neighbors of the nearest candidate... for (const auto& neighbor : layers[l_c][nearest.id].neighbors) { if (!visited[neighbor.id]) { // ...into one list of IDs instead of processing them one by one batch_indices.push_back(neighbor.id); visited[neighbor.id] = true; } } } // batch_indices now goes to the GPU → one kernel launch computes // every distance in the batch simultaneously What used to be four sequential comparison rounds becomes a single parallel one: The sequential step does not process multiple nodes at each iteration The parallel process converts the 1,2,3,4 step in only one Sequential search evaluates neighbors one by one batching evaluates the whole neighborhood in a single round. This was the single biggest conceptual change, and the foundation everything else builds on. Result: from 3.48 ms down to ~1.69 ms per query, batching alone cut query time roughly in half. 2. Persistent GPU Memory Allocation Profiling showed we were paying cudaMalloc/cudaFree on every single search call, for the query buffer, the distances buffer, everything. So we allocated once at index construction and reused the buffers for the lifetime of the object: // Allocated ONCE at construction time, reused for every search: // GPU buffer for the incoming query vector cudaMalloc(&d_query_buffer, MAX_DIM * sizeof(float)); // GPU buffer where the kernel writes each batch's distances cudaMalloc(&d_distances_buffer, BATCH_SIZE * sizeof(float)); // The ENTIRE dataset, resident on the GPU for the object's lifetime — // searches send only indices, never the vectors themselves cudaMalloc(&d_all_vectors, total_vectors * vector_dim * sizeof(float)); That third line matters most: the entire dataset lives on the GPU permanently , so a search only ever needs to send the query and a small list of batch indices, not the vectors themselves. Host means the device traffic per query drops to almost nothing. Result: another −13.5% on top of batching. The cheapest memory transfer is the one you never make. 3. CUDA Streams, Our Instructive Failure Next, we tried the textbook trick: split each batch across 4 CUDA streams, so data transfer for one chunk overlaps with computation on another. In naive sequential operation, the GPU idles during transfers and PCIe idles during compute.The ideal: copies and kernels for different chunks overlap in time.On paper, beautiful. In practice: performance got worse, time per query went up 37%, erasing our previous gain. Why? Overlap only pays when the chunks are big enough to hide latency. Our per-batch workloads were small (a neighborhood of a graph node, not a giant matrix), so splitting them four ways mostly added stream-management and synchronization overhead — and we hadn’t tuned the number of streams for such small batches. It’s a classic case of applying a big-data optimization to a small-data inner loop. We kept the lesson and dropped the dependency on streams for the final hot path. 4. Pinned (Page-Locked) Host Memory By default, host memory is pageable, the OS can swap it out. Every host that maps a device copy from pageable memory silently goes through a staging step: pageable then pinned then GPU DRAM. Pageable transfers pay a hidden extra copy; pinned memory goes straight to DRAM. By allocating the query and result buffers with cudaMallocHost, we skip the staging copy and unlock true async transfers: // cudaMallocHost = pinned (page-locked) host memory: // the OS can never swap it out, so the GPU can DMA from it directly float* h_pinned_query; CUDA_CHECK(cudaMallocHost(&h_pinned_query, query.x.size() * sizeof(float))); // Copy the query into the pinned buffer (a cheap host-side memcpy) memcpy(h_pinned_query, query.x.data(), query.x.size() * sizeof(float)); // Async host→device transfer: no hidden staging copy, and the CPU // is free to keep working while the transfer is in flight CUDA_CHECK(cudaMemcpyAsync(d_query_buffer, h_pinned_query, query.x.size() * sizeof(float), cudaMemcpyHostToDevice, stream)); Result: −26.1% vs the streams version, bringing us back to the best time so far. 5. Everything Resident in CUDA The final version combined it all: batching, persistent allocations, pinned buffers, and the full dataset resident on the GPU shaving off a final 0.8%. Accelerating Insertion Too The contest’s primary target was search, but index construction was a secondary objective and since our insert already calls search_layer to find where each new point connects, every search optimization compounded for free. We added two insert-specific tricks: the dataset is copied to the GPU in chunks through pinned memory at build time (so the copy pipeline stays busy instead of one giant blocking transfer), and levels for a whole batch of nodes are pre-computed before insertion. The greedy graph-update logic itself stays sequential, HNSW insertion has real data dependencies between consecutive points, and respecting them is exactly what keeps recall intact. The Results We benchmarked on the SIFT small dataset (k=100, M=16, ef-construction=100, ef=100, n=1000, 100 repetitions per configuration — enough to get statistically meaningful comparisons across all five methods). Time per query at each rung of the optimization ladder. Note the streams bump. The progression tells the whole story: batching helps, killing redundant allocations helps a lot , streams backfire, pinned memory recovers it, and the final version lands at 1.47 ms/query vs 3.48 ms on CPU a -57.7% reduction. Under the contest’s required parameters, the raw numbers looked like this: And one more hypothesis confirmed: the more queries you run, the more the GPU pays off. With a single query the improvements were marginal, the parallelism only starts to shine at scale. Did We Break Recall? The contest’s hard constraint was accuracy. We computed Recall_CPU − Recall_GPU per query across 100 queries: The final implementation shows no significant recall difference from the CPU baseline. (The intermediate pinned-memory variant showed small per-query deviations we attribute to handling/numerical issues — worth flagging because it’s a reminder that memory optimizations can bite silently. The final version is clean.) What We’d Do Next Time limits (and Colab GPU quotas building the large index takes 5+ hours) left some threads dangling: Bigger datasets. Our hypothesis says GPU gains grow with scale; we’d like to prove it on the full SIFT benchmark. Sparse structures. We store the graph as a dense matrix; a CSR representation would slash memory and might improve locality. CPU + GPU together. The host sits mostly idle during search — multi-threading the graph walk while the GPU crunches distances could use the best of both worlds. Takeaways If you’re porting a graph algorithm to CUDA, our journey compresses to four rules: Don’t parallelize the algorithm; parallelize its inner arithmetic. The greedy walk stayed sequential — only the distance math went wide. Memory transfers dominate. After batching, every remaining gain came from allocating once and keeping data resident on the GPU — not from any kernel cleverness. Textbook optimizations have preconditions. CUDA streams are great — for workloads big enough to hide latency. Measure, don’t assume. Guard your accuracy metric from day one. A fast ANN index with degraded recall is just a bug with good benchmarks. Sources [1] — Yu A. Malkov and D. A. Yashunin. Efficient and Robust Approximate Nearest Neighbor Search Using Hierarchical Navigable Small World Graphs. IEEE Transactions on Pattern Analysis and Machine Intelligence, 42(4):824–836, April 2020. ISSN 1939–3539. doi: 10.1109/TPAMI.2018.2889473. [2] — Arailly Arailly. Arailly/hnsw: Implementation of HNSW. https://github.com/arailly/hnsw , 2024 [3] — Gaurav Jain. GauravJain28/Parallelized-and-Distributed-HNSW-Algorithm, March 2023. This work was done with Syauqi Rahmat Perdana for the High Performance Graph Data Analytics course at Politecnico di Milano, based on the HNSW paper by Malkov & Yashunin and the C++ implementation by Arailly. Thanks to Professors Ian Di Dio Lavore, Leonardo De Grandis, and Riccardo Strina for their supervision. Code: github.com/Syauqi99/hpgda_contest_MM Accelerating Hierarchical Navigable Small World (HNSW)-RAG Vector Search with CUDA was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.