GPU-Resident Top-K for Agentic RAG: I Built a CUDA Kernel So My Retrieval Step Would Stop Bouncing Off the GPU

A 343-line CUDA kernel that keeps Top-K retrieval on the GPU. According to the original article, keeping similarity search resident on device memory makes the retrieval step up to 8.6x faster than a CPU brute-force baseline, measured on a GTX 1080.

The Problem: Retrieval Round-Trips

An agentic RAG loop calls the retrieval step again and again. In a default pipeline, the query embedding is shipped from the GPU to Python, the CPU scores the N corpus rows and picks the best K, and the result is shipped back to the GPU. The article treats this repeated trip across the PCIe bus as a silent tax on the pipeline.

The author's answer is to keep the corpus on the device and run the whole retrieval step there.

Concretely, the baseline in the repository models that round trip step by step: the query is copied from the device to the host, a CPU reference implementation scores a host-resident corpus, and the result is copied back to the device. The treatment keeps the corpus on the device and only moves a small query vector in and the K results out.

Key Takeaway: in an agentic loop, the cost is not only the similarity math but the round trip across the PCIe bus. Keeping the corpus resident on the GPU removes that trip.

Architecture and Concepts

The project, CUDA-TopK-Retrieval, uploads the corpus to the GPU once (upload_corpus_rowmajor) and then serves each query with search_resident. On the device, the kernel scores the corpus, keeps a partial top-K per thread block, and merges those partial results into the final top-K.

The hot path in the repository has four parts:

• Row scoring kernel. One CUDA thread scores one corpus row, as a dot product or as cosine similarity with precomputed row norms, and writes the result to a dense scores[N] buffer.
• Partial Top-K per block. The rows are split into P disjoint ranges, with P = min(128, ceil(N / 4096)). Each block keeps a sorted local Top-K for its range, and ties are broken deterministically: the lower row index wins.
• Merge kernel. A single-thread pass merges the P sorted partial lists into the global Top-K, best first.
• CPU oracle. topk_reference scores the corpus by brute force with std::partial_sort and the same comparator. The GPU indices must match it before any benchmark number is trusted.

The public API is a small C++ class, GpuTopkEngine. initialize allocates device buffers for a maximum corpus size, dimension and K (K is capped at 256), upload_corpus_rowmajor copies the corpus once, upload_row_l2 copies the row norms for cosine mode, and search_resident serves one query. A convenience search call re-uploads the corpus every time, which is the cold-cache case the resident path is meant to avoid.

How the Benchmark Was Run

• Hardware. The committed example run records an NVIDIA GeForce GTX 1080, driver 535.309.01, CUDA 12.2, an Intel Core i7-8700K host and -O3 -march=native compiler flags. GPU clocks were not locked for that run (the repository ships scripts/lock_gpu_clocks.sh for that purpose).
• Sweep. 45 configurations: 5 corpus sizes (10,000, 50,000, 100,000, 500,000 and 1,000,000 vectors), 3 embedding dimensions (384, 768, 1024) and 3 values of K (8, 32, 100).
• Measurement. Timing uses synchronized CUDA events. The example CSV records 3 timed trials after 1 warm-up per configuration, so treat the means as indicative rather than statistically tight.
• Correctness first. The GPU result is compared with the CPU oracle before timings are recorded.

Benchmark Results

The headline number, 8.57x, is the best case of the sweep: K = 8, one million vectors, 1024 dimensions, where the CPU round trip takes about 978 ms per query and the GPU-resident path about 114 ms. The full picture depends strongly on K.

Two things stand out in the data. First, the GPU path loses at K = 32 for the smallest corpus (10,000 vectors) at 384 and 768 dimensions, and at K = 100 it is slower than the CPU baseline in 14 of the 15 configurations, down to 0.02x at 10,000 vectors and 384 dimensions (about 3.4 ms on the CPU against 155 ms on the GPU). Second, the speedup is largest at 1024 dimensions for every corpus size.

Key Takeaway: the 8.57x figure is real but it is the best case. The GPU-resident path is a clear win for small K and large dimensions, and a loss for large K.

Reading the Numbers

This section is my reading of the data, not a claim made by the author. ⚠️ TO VERIFY with a profiler before relying on it.

• The slowdown at larger K is consistent with how the partial Top-K stage is written. In the source, each block runs on a single thread and inserts candidates into a sorted array with a bubble pass, which grows quickly with K. The README says as much: it is "not yet a warp-specialized tournament selection kernel", chosen for auditability.
• The baseline is the author's own brute-force CPU oracle, not a tuned library or an approximate nearest neighbor index. The speedup therefore measures the value of staying on the device against a simple CPU path, not against a production vector database.
• Absolute timings are high for large corpora (more than 100 ms per query at one million vectors on the GPU path), because the scoring kernel maps one thread to one row and reads the row sequentially. The source itself notes that tiling the dimension across warps is a follow-up.

Limits Stated by the Repository

• The partial Top-K uses single-lane selection per block, and the repository calls the kernel an auditable first version.
• It is a similarity and Top-K slice, not a vector database: no compression, no metadata filtering and no multi-GPU sharding.
• Continuous integration on GitHub-hosted runners only checks compilation and the CPU oracle tests. GPU parity checks and profiler captures need local hardware.
• GPU and CPU scores are compared within a small floating-point tolerance, with ties resolved by index.

When It Applies

If your agent retrieves a handful of passages per step from a corpus that already lives on the GPU, avoiding the PCIe round trip is worth measuring. If you need dozens of results per step, or a filtered or compressed index, this kernel as published is not the right tool yet. In both cases, measure on your own corpus size, embedding dimension and K before adopting it.

Reproduce It

The README gives this build and test sequence (it does not compile llama.cpp; the GGML_CUDA flag is only kept for ergonomics):

cmake -S . -B build -DGGML_CUDA=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build -j
cd build && ctest --output-on-failure

The topk_demo and topk_bench executables then run a demo query and the benchmark sweep. Results from your own hardware belong next to the committed example in examples/example-run-results/.

Sources

• GPU-Resident Top-K for Agentic RAG: I Built a CUDA Kernel So My Retrieval Step Would Stop Bouncing Off the GPU (Towards Data Science)
• cuda-topk-retrieval (GitHub repository: README, kernel source, benchmark CSV and environment.json)