Key Takeaways<\/p>\n
The problem: in agentic RAG, every tool call that needs context fires a similarity search. A default pipeline ships the query embedding from the GPU to Python, lets the CPU score N corpus rows and pick the best K, then ships the answer back. That round-trip is the silent tax. The compute is fine; the travel is the bill. We all know, travel is never cheap, no matter where you want to go (pun intended!)<\/p>\n
The easy fix: upload the corpus to VRAM once, then keep the similarity scoring, the Top-K selection, and the merge step on the device. Only the tiny per-query embedding (D floats) and the K results travel across PCIe.<\/p>\n
The receipts: on the same 7-year-old GTX 1080 used in Parts 1 and 2, the GPU-resident path runs the retrieval hop up to 8.57\u00d7 faster than a CPU brute-force baseline. At K=8 it wins on all 15 sweep configurations (N \u2208 {10k, 50k, 100k, 500k, 1M}, D \u2208 {384, 768, 1024}) with speedups from 2.43\u00d7 to 8.57\u00d7. At K=32 it wins on 13 of 15 configs, peaking at 7.76\u00d7. At K=100 \u2014 where the V1 selector intentionally stays simple \u2014 the CPU wins on 14 of 15 configs. That last sentence is the honest part (Well, even if I had lied, you could have easily caught it).<\/p>\n
The kicker: the wins are not \u201cmagic kernel\u201d wins. They are \u201cwe stopped shipping the corpus back to host RAM for no reason\u201d wins. It is also exactly the kind of \u201cmeasure many candidates, report only the best K back to the consumer\u201d decision a 5G base station and your phone have been making every few milliseconds since CSI feedback became a thing.<\/p>\n
TL;DR: Default agentic RAG treats the GPU as a serving box and the retrieval as a Python concern. Every tool call ships the query embedding D\u2192H, lets the CPU compute N dot products, sort the candidates, pick the top K, and ship indices and scores H\u2192D. For an agent that calls a vector store ten times per reasoning step, that round-trip is the dominant cost \u2014 not the model, not the embedding, it is the travel. CUDA-TopK-Retrieval keeps the corpus resident on the device, runs scoring + per-block partial Top-K + a multi-way merge entirely on the GPU, and exposes a tiny C++ orchestrator API (upload_corpus_rowmajor once, search_resident per query). The host-touching bytes per query collapse to one D-length embedding up and 2K results down. On a GTX 1080, across a 45-config sweep, the GPU-resident path beats the CPU-round-trip baseline on all 15 K=8 configs (2.43\u00d7 to 8.57\u00d7, peaking at N=1M, D=1024) and on 13 of 15 K=32 configs (the two losses are both at the smallest N=10k for D=384 and D=768, where the round-trip itself is already cheap; big-N K=32 wins climb to 7.76\u00d7). At K=100 the V1 kernel deliberately stays simple \u2014 single-lane-per-block bubble sort with a serial merge \u2014 and the CPU wins on 14 of 15 configs; that ceiling is the article\u2019s honest punchline and a clean setup for Part 4.<\/p>\n
GitHub Repo: https:\/\/github.com\/AnubhabBanerjee\/cuda-topk-retrieval<\/a><\/p>\n (Quick confession before we start: I came at this from a 5G\/6G RAN engineering background. Beam selection at a base station looks shockingly close to RAG Top-K \u2014 the UE scores a codebook of candidate beams by received power and reports the best handful back over the air. There is a whole section on that below \u2014 section 8 \u2014 but it is also why this kernel exists in the shape it does.)<\/p>\n Architecture mental model \u2014 keep this open while you read.<\/p>\n agent.embed(query) \u2192 cudaMemcpy H\u2192D (D floats) \u2192 row_dot_scores_kernel \u2192 partial_topk_block_kernel (P blocks) \u2192 merge_partial_topk_kernel \u2192 cudaMemcpy D\u2192H (K indices + K scores)<\/p>\n Everything below is just commentary on one part of that line.<\/p>\n 1. A confession: every RAG step in your agent is a tiny PCIe road trip<\/p>\n In Part 2<\/a> of this series, we successfully isolated our LLM agent\u2019s inference loop, keeping token generation running hot and fast on the device. We designed a system which avoids stalling. But the moment we give that agent a tool to search an external knowledge base\u2014the core of any multi-hop Retrieval-Augmented Generation (RAG) pipeline\u2014we silently destroy all that hard-won performance and we hit the wall. If you have ever wired an \u201cagentic\u201d pipeline to a vector store through a Python retriever, here is what really happens on each tool call (with a little intentional dramatization):<\/p>\n You: \u201cAgent, find me the five chunks most relevant to \u2018how do I claim deduction under section 80C?\u2019\u201c<\/p>\n Agent: \u201cSure. Embedding the query on the GPU. \u2705\u201d<\/p>\n Agent: \u201cNow bouncing the query embedding back to host.\u201d<\/p>\n (cudaMemcpy D\u2192H, ~1,024 floats) Python retriever: \u201cGot it. NumPy loop. Dot product N times. argpartition. Top-5.\u201d<\/p>\n (CPU scores half a million corpus rows, one row at a time, while a 9 TFLOP GPU watches) Python retriever: \u201cDone. Here are the indices and scores.\u201d<\/p>\n Agent: \u201cCool. Bouncing them back to the GPU now.\u201d<\/p>\n (cudaMemcpy H\u2192D, 10 numbers) Agent: \u201cReady. What was the question again?\u201d<\/p>\n The agent has a perfectly good GPU. The corpus is sitting in 4 GB of the VRAM. The query embedding was already on the GPU \u2014 we just generated it there. And then, on every single retrieval hop, we ship the query back to the host, do brute-force similarity in NumPy \/ FAISS-on-CPU \/ a hand-rolled loop, and ship the answer back.<\/p>\n Your GPU\u2019s utility meter: spends most of the retrieval step idle. Your PCIe bus: gets a workout it did not sign up for. Your agent\u2019s tool-call latency: dominated by something that is neither the model nor the embedding. That is the joke.<\/p>\n That is also the dirty secret of every agentic RAG demo that scales past the toy \u201cten chunks in memory\u201d stage. The retrieval hop bounces off the GPU and back, every time, and the bigger the corpus, the worse the tax. At a million rows of 1024-dim embeddings, the round-trip alone \u2014 not even the scoring, yes, just the round-trip \u2014 eats most of the budget of the retrieval step itself.<\/p>\n CUDA-TopK-Retrieval is what happens when you decide the round-trip is optional and you would rather write 343 lines of CUDA than let the agent vacation through host RAM every time it wants a neighbor.<\/p>\n Now imagine the real workload behind this. It is not \u201cfive chunks for one question.\u201d It is multiple specialized micro-agents \u2014 each one running its own RAG hops, each one needing Top-K against the same corpus, each one currently paying its own PCIe bill on every tool call. Part 1<\/a> of this series killed the prefill round-trip. Part 2<\/a> made the GPU shareable across those multiple agents. Part 3 says: now that they are sharing the card fairly, stop making each one of them drive back to the host to look up a neighbor.<\/p>\n 2. Why does Top-K retrieval exist? (a one-minute crash course)<\/p>\n Skip this if you already have a background on this. For everyone else, who are new to this field, here is a short, amateur explanation.<\/p>\n A modern agent does not stuff the whole knowledge base into the prompt. It retrieves. For every reasoning step that needs grounded context, it embeds the query into a fixed-dimensional vector (D floats \u2014 typically 384, 768, or 1024), scores that vector against every row of a corpus of pre-embedded chunks (N rows, also D floats each), and returns the K corpus rows with the highest similarity. That\u2019s it. That\u2019s Top-K vector search. Retrieval-Augmented Generation is just a polite way of saying \u201cTop-K plus a prompt template.\u201d<\/p>\n Two flavors of similarity show up everywhere. Dot product is the cheap one: a single fused multiply-add per dimension, N\u00d7D work in total. Cosine is dot product divided by the product of L2 norms, which becomes a free dot product if you pre-normalize the corpus once at ingest. Most production vector stores do the pre-normalization trick and call it \u201ccosine\u201d while running raw dot-product math at query time. The CUDA-TopK-Retrieval kernel supports both \u2014 it just multiplies by a precomputed per-row norm pointer when cosine mode is on.<\/p>\n Mainstream tools (FAISS, hnswlib, the Python side of cuVS, your favorite SaaS vector DB) all do this scoring + Top-K work. Most of them do it well. The problem is where they do it. Almost every agent framework on the planet calls into the retriever from Python, and the moment Python is on the hot path, the retrieval step is no longer a GPU operation \u2014 it is a PCIe operation with a GPU on one end.<\/p>\n The fix is not \u201ca better algorithm.\u201d It is \u201ca much shorter road trip.\u201d<\/p>\n 3. The \u201cjust keep the corpus on the GPU\u201d lightbulb (and why it\u2019s harder than it sounds)<\/p>\n The pitch is simple<\/p>\n Upload the corpus to VRAM once at ingest.<\/p>\n For every incoming query, cudaMemcpy a tiny D-dimensional float embedding to the device.<\/p>\n Launch a scoring kernel where one CUDA thread per corpus row computes the dot product.<\/p>\n Launch a partial Top-K kernel where each block scans a disjoint row range to emit its own local top candidates.<\/p>\n Finally, launch a merge kernel to walk the per-block heads and emit the global Top-K in best-first order.<\/p>\n You cudaMemcpy exactly 2K numbers back to the host: K indices, and K scores.<\/p>\n This is the \u201ctreat memory retrieval as a hardware primitive, not a software API call\u201d paradigm. The only reason this takes more than a 30-line PyTorch script to achieve is that three tedious edge cases will immediately break the naive approach.<\/p>\n Problem A: Top-K on a GPU is structurally awkward<\/p>\n Scoring the vectors is the easy part. It\u2019s just matrix multiplication, and your GPU was literally born to do that\u2014it is the hardware\u2019s love language. Selection, however, is where the romance dies. Asking a GPU to do a full O(N log N) sort just to grab the top K results is computationally offensive; it\u2019s like alphabetizing your entire recycling bin just to find a single receipt. You could try an O(N) argpartition, but that requires a tree-walk, which shatters GPU memory coalescing into a million unaligned reads. Tournament selection is fast, assuming you want to spend your weekend debugging edge cases. And the moment you cave and reach for a thrust or cub sorting primitive, congratulations: you have just infected your lightweight, standalone C++ pipeline with a massive build dependency.<\/p>\n The architecture picks the boring answer on purpose. It relies on a tiny per-block O(K2) bubble sort over a disjoint row range, driven by a single thread per block, and capped off with a serial multi-way merge. On paper, this sounds terrible. In practice, it works beautifully, for the exact reason stated honestly in the kernel\u2019s comments:<\/p>\n \/\/ Single-threaded per-block scan that materializes a local Top-K list for its row partition. This is the design contract: V1 picks auditability over cleverness. The whole kernel is small enough that a reviewer can read it end-to-end in a coffee break, line it up against the CPU oracle, and convince themselves the GPU output is bitwise correct. The day someone wants a 2\u00d7 kernel, this gets replaced with a warp-specialized tournament selector \u2014 the article\u2019s section 9 promises that explicitly. For K \u2264 32, the single-lane bubble sort is genuinely fine. For K=100, it falls off a cliff. That cliff is documented and benchmarked. (Section 9 again.)<\/p>\n Problem B: GPU and CPU must agree, bit-for-bit, on the tiebreak<\/p>\n Well, like any World Cup Football group league match, ties happen more than often. Two corpus rows have the same score down to fp32 precision. Which one wins?<\/p>\n If the CPU oracle and the GPU kernel disagree on the tiebreak, you can never trust a benchmark, because every \u201cmismatch\u201d alarm is now ambiguous: did the GPU score a row wrong, or did the GPU just break ties differently? You will spend a week in a 3 AM Slack channel.<\/p>\n The fix is to define the comparator in one sentence and implement it in two places \u2014 once on the host, once on the device \u2014 and make those two implementations literally the same expression. <\/p>\n On the host side it looks like:<\/p>\n \/\/ Lexicographic “better” relation for (score, index) pairs under float equality semantics. On the device side it looks like:<\/p>\n \/\/ Device-side replica of the host comparator to avoid cross-TU linkage issues for __device__ code paths. That is the whole tiebreak policy in five lines, twice. Higher score wins; on exact tie, smaller corpus row index wins. The CPU oracle uses it in std::partial_sort, the GPU uses it in the bubble sort and in the multi-way merge, and the benchmark harness will not start timing until the GPU output matches the CPU output exactly \u2014 same indices in the same order, scores within a small fp32 tolerance.<\/p>\n That single comparator is the reason the article can quote a speedup at all. Without it, \u201cthe GPU is 8\u00d7 faster\u201d is just \u201cthe GPU is 8\u00d7 faster at being wrong differently.\u201d<\/p>\n Problem C: VRAM is precious, and the worst place to do malloc is the hot path<\/p>\n Allocating GPU memory per-query is like signing a new car lease every time you need to drive to the grocery store. It is the easiest way to turn a 1-millisecond search into a 50-millisecond traffic jam.<\/p>\n Instead, GpuTopkEngine::initialize buys the car upfront. It runs all the cudaMalloc calls during engine startup, sizing the buffers for the worst possible configuration. Once the engine is actively serving queries, the hot path is completely free of memory management. It is just fast kernel launches and tiny data copies. No fragmentation, no negotiating with the allocator, and cudaMalloc is permanently banned from showing up in your performance traces.<\/p>\n 4. The four-stage pipeline (the actually-cool part)<\/p>\n Step 0: Engine init \u2014 eight cudaMallocs sized for (max_n, max_d, max_k) (GpuTopkEngine::initialize) Let\u2019s walk through each one with the real code. The snippets are even shorter than the tiny source files.<\/p>\n Step 1 \u2014 Upload the corpus once<\/p>\n This is the boring step that makes the rest of the article possible. The corpus goes up exactly one time per ingest, and stays there for the lifetime of the engine:<\/p>\n cudaError_t GpuTopkEngine::upload_corpus_rowmajor(const float32_t* const host_corpus_rowmajor, const index_t N, And that is the entire ingest API. At 1024 dimensions, one million vectors is exactly 4 GB of float32 data, sliding perfectly into the 8 GB VRAM of a vintage GTX 1080. What happens when your corpus hits 10 million vectors? That becomes a distributed systems problem, not a kernel problem. If your data exceeds VRAM, you need a sharding strategy, which we cover in Section 9. But for now, we are here to solve the compute bottleneck, not to invent a new database.<\/p>\n Step 2 \u2014 Score N rows on the device<\/p>\n One CUDA thread per corpus row. 256 threads per block. Each thread accumulates the dot product across D dimensions and writes one float into the dense scores[N] buffer:<\/p>\n \/\/ Row-major dot-product with optional cosine normalization; coalesced reads along D are sacrificed for clarity in v1. One thread per row is the simplest possible mapping. The code comment is honest about it: a follow-up can tile D across warps to raise the arithmetic intensity. For V1, this gives the auditor a one-to-one correspondence with the CPU loop and lets them sleep peacefully at night.<\/p>\n Step 3 \u2014 Each block builds its own local Top-K<\/p>\n Now comes the most awkward part. Picking the top K out of N is conceptually \u201csort and slice,\u201d but a full sort wastes most of the work. We split the row range across P blocks (capped at 128), each block walks its disjoint slice with the tiny bubble sort from Section 3, and writes its own local top-K list out:<\/p>\n const int P = std::min(kMaxPartialBlocks, std::max(1, (static_cast(N) + 4095) \/ 4096)); >>(d_scores_, static_cast(N), static_cast(K), P, d_partial_scores_, One thread per block. Yes, that is wasteful on paper. It is also why a human can audit this kernel in twenty minutes \u2014 the policy if (threadIdx.x != 0 || blockIdx.x >= P) return; collapses the whole intra-block reasoning down to \u201cthis block\u2019s lane 0 owns rows [start, end).\u201d Every block\u2019s s[] and ids[] arrays live in registers \/ local memory, sized by the compile-time kMaxSupportedK = 256 cap.<\/p>\n Step 4 \u2014 Merge the partials into the global Top-K<\/p>\n Finally, one thread on one block walks P cursors over the per-block lists. Each list is already best-first. Pick the best head; emit; advance that one cursor; repeat K times:<\/p>\n for (int out = 0; out < K; ++out) { The merge is ruthlessly efficient: at most P * K reads and exactly K writes, executed by a single thread. To prevent floating-point chaos, the device_is_better comparator enforces strict determinism\u2014if two heads tie on score, the lower corpus row index wins, mirroring the CPU oracle perfectly. Finally, two microscopic cudaMemcpy calls shuttle the K winning indices and scores back to the host. The agent ingests them, and the RAG loop fires again.<\/p>\n That is the entire hot path: one H -> D embedding transfer, three kernel launches, and two tiny D -> H result copies. No Python host loops, no framework overhead, and absolutely no PCIe vacations.<\/p>\n 5. The receipts (i.e., the numbers)<\/p>\n Let\u2019s now evaluate it against the baseline, and see if it was worth the trouble.<\/p>\n A quick note on methodology, before the benchmarking police arrives: every comparison below runs on the same GPU as Part 1<\/a> and Part 2<\/a> (NVIDIA GeForce GTX 1080, Pascal sm_61, 8 GB), driver 535.309.01, CUDA 12.2, host CPU Intel Core i7-8700K, compiler flags -O3 -march=native –expt-relaxed-constexpr. Three trials, one warmup, seed 1, fixed RNG (std::mt19937 with std::normal_distribution), Gaussian embeddings, L2-normalized in dot-product mode. The full sweep is N \u2208 {10k, 50k, 100k, 500k, 1M} \u00d7 D \u2208 {384, 768, 1024} \u00d7 K \u2208 {8, 32, 100} \u2192 45 configurations, all measured via cudaEventElapsedTime with cudaDeviceSynchronize bracketing every interval. Code in src\/host\/bench_main.cpp; raw numbers in examples\/example-run-results\/benchmark_run_results.csv.<\/p>\n The two paths timed:<\/p>\n GPU-resident (treatment). Corpus is already on the device. Each timed iteration: cudaMemcpy query H\u2192D (D floats) + score kernel + per-block partial Top-K + merge + cudaMemcpy K scores D\u2192H + cudaMemcpy K indices D\u2192H. End to end.<\/p>\n CPU round-trip (baseline). Models the default agent flow: cudaMemcpy query D\u2192H + CPU brute-force scoring + std::partial_sort with the same comparator + cudaMemcpy indices H\u2192D + cudaMemcpy scores H\u2192D. End to end.<\/p>\n Both paths run inside the same process, use the same query bytes and the same comparator. The only difference is where the work happens. If you have ever held the position \u201cPCIe is fine, we benchmark the kernels in isolation,\u201d this is what it costs you when you stop pretending the round-trip is free.<\/p>\n Headline (GTX 1080, three trials, mean ms, ratios computed from per-trial means):<\/p>\n Config (N \u00d7 D, K)Baseline mean (ms)GPU mean (ms)Speed-up10,000 \u00d7 1024, K=89.561.357.10\u00d7100,000 \u00d7 768, K=870.6625.702.75\u00d7500,000 \u00d7 1024, K=8483.9069.796.93\u00d71,000,000 \u00d7 1024, K=8977.80114.128.57\u00d71,000,000 \u00d7 1024, K=32973.89125.467.76\u00d710,000 \u00d7 384, K=1003.37155.25GPU is 46\u00d7 slower1,000,000 \u00d7 384, K=100329.49682.38GPU is 2.07\u00d7 slower<\/p>\n Yes, you are reading the numbers right.<\/p>\n The first five rows are the article\u2019s point: at K=8, the GPU-resident path wins on every single configuration in the sweep (all 15 of them), by ratios ranging from a polite 2.43\u00d7 at N=50k, D=384 to a loud 8.57\u00d7 at N=1M, D=1024. At K=32 it wins on 13 of 15 \u2014 the two losses are both at the smallest N (10k), for D=384 and D=768, where the round-trip itself only costs ~3\u20137 ms and the GPU\u2019s three kernel launches barely have room to amortize. By the time you reach realistic agentic-corpus sizes (N \u2265 50k), K=32 also wins comfortably, peaking at 7.76\u00d7. The big speedups are not \u201cmagic kernel\u201d speedups \u2014 they are \u201cwe stopped shipping the corpus back to host RAM for no reason\u201d speedups. The GPU was always going to win this race; the only reason it ever lost was that we kept making it commute unnecessarily.<\/p>\n The last two rows are where this article earns its right to call itself honest. At K=100, the single-lane bubble sort per block becomes O(K\u00b2) = O(10,000) sequential comparisons per block, and the serial merge walks P \u00d7 K head positions. The CPU\u2019s std::partial_sort is heap-based, vectorized by the compiler, and effectively O(N log K) \u2014 much friendlier to K=100. So the GPU loses on 14 of 15 K=100 configs, sometimes by 2\u00d7, sometimes by 46\u00d7. (There is exactly one K=100 config where the GPU still wins \u2014 N=1M, D=1024, 1.44\u00d7 \u2014 because by then there is enough scoring work to dominate the selection ceiling. One row out of fifteen is not a save; it is a curiosity.) That is not a bug; it is the V1 design statement (\u201cauditability over cleverness\u201d) meeting its first concrete consequence. The fix is in Section 9, and it is a warp-specialized tournament selector \u2014 not a frantic refactor.<\/p>\n One more honest caveat in the numbers above: in this committed snapshot, the GPU clocks were not locked. That means absolute milliseconds move slightly with thermals and DVFS; the ratios stay stable. The repo ships scripts\/lock_gpu_clocks.sh for anyone who wants to reproduce the table with locked clocks on a GTX 1080. Needless to mention, the structural finding does not change.<\/p>\n 6. \u201cOK, but how is this different from FAISS \/ cuVS \/ hnswlib?\u201d<\/p>\n A very reasonable question, and worth answering directly, because the vector-search world has a lot of overlapping primitives and an HPC reader will ask this in the first comment.<\/p>\n FAISS (CPU index). The default in most agent frameworks. Excellent library. Lives on the CPU. Every query an agent makes pays the PCIe round-trip this article exists to delete. If you\u2019re already on IndexFlatIP and you\u2019re CPU-bound on retrieval, you are the target audience.<\/p>\n FAISS (GPU index). Solves the GPU residency problem, with a much more mature kernel suite than this repo. The point of CUDA-TopK-Retrieval is not \u201cI out-engineered FAISS-GPU\u201d \u2014 it never was and neither does it try to be. The point is to show, in 343 lines, what the actually-cool retrieval primitive looks like and why agentic pipelines feel slow when it isn\u2019t there. If you need a production index today, use FAISS-GPU. If you want to understand the small hot path that makes the difference \u2014 one tiny H\u2192D copy, three kernel launches, two small D\u2192H copies \u2014 read this kernel.<\/p>\n NVIDIA cuVS \/ RAFT. The serious, production-grade in-GPU vector-search stack. Bigger, faster, more algorithms, more dependencies. Same caveat as FAISS-GPU: this kernel is the pedagogical \/ single-binary version, not a competitor.<\/p>\n hnswlib and friends (approximate nearest neighbor). Different shape of trade-off entirely \u2014 they trade exactness for sublinear query time on huge corpora. CUDA-TopK-Retrieval is exact brute-force scoring + selection; the speedup story is purely about residency, not about skipping work.<\/p>\n The point of this repo is not \u201cbuild your production vector DB on it.\u201d The point is: the agent\u2019s retrieval hop wants to stay on the GPU, and once you accept that, even a tiny hand-written kernel beats a hosted-on-CPU brute force across most realistic K values on a 7-year-old card.<\/p>\n 7. So\u2026 how do I actually try it?<\/p>\n Clone the repo, then build with CMake (the -DGGML_CUDA=ON flag mirrors the llama.cpp build ergonomics from earlier parts of the series):<\/p>\n cmake -S . -B build -DGGML_CUDA=ON -DCMAKE_BUILD_TYPE=Release Then run the demo and the benchmark exactly as the README does:<\/p>\n .\/build\/topk_demo # tiny smoke story (GPU required) topk_demo is the tiny smoke story \u2014 4,096 corpus rows, 128 dims, K=8, prints the neighbor IDs. topk_bench is the harness that emits the TOPK_BENCH_JSON line the Python campaign script consumes. For the full 45-config sweep on canonical hardware:<\/p>\n python3 scripts\/benchmark_campaign.py.example # full sweep (GPU required; writes under examples\/benchmark-campaign-runs\/run-*)<\/p>\n Before publishing-grade runs, lock the GPU clocks (the repo provides the script):<\/p>\n sudo bash scripts\/lock_gpu_clocks.sh<\/p>\n Requirements: Linux, CUDA toolkit, an NVIDIA GPU (Pascal or newer), and the patience to read a CMake file once. Artifacts land under examples\/example-run-results\/ for the quick path or examples\/benchmark-campaign-runs\/run–\/ for the full sweep, and the README is explicit that committing .nsys-rep databases is forbidden \u2014 PNG timelines only.<\/p>\n 8. Plot twist \u2014 this is just 5G beam selection in a CUDA costume<\/p>\n I should probably confess at this point: I am still not a \u201cGPU person\u201d by training. I came up through telecom \u2014 5G NR with a foot creeping firmly into 6G research \u2014 and I keep noticing that every infrastructure problem in agentic AI is a problem which was already solved at the radio layer, maybe around twenty years ago.<\/p>\n For readers without a 3GPP background: in a modern 5G base station, the antenna does not radiate equally in every direction. It forms a codebook of directional beams \u2014 narrow lobes of radio energy \u2014 and at any given moment your phone is being served by the one beam (or the small handful of beams) whose received power is highest on your device. Choosing the right beam, fast, is one of the most-studied retrieval problems in wireless. The UE measures L1-RSRP (a per-beam received-power score) across the candidate beams the gNB (5G base station) has told it to measure, then reports back the best handful via the CSI feedback channel. The gNB uses those reports to decide which beams to schedule (Well, this was as simple as I could make it, the real story is really ugly!).<\/p>\n That is Top-K vector search, in radio costume. The candidate beams are the corpus. The instantaneous channel measurement is the query. The score is received power. K is the number of best beams the report carries back. The UE does the scoring down at the baseband DSP layer \u2014 it does not ship the I\/Q samples back to a central CPU farm and ask a Python script which beam is best, because doing so on a per-millisecond loop would melt the air interface.<\/p>\n Look at this side-by-side and tell me with a straight face these are different problems:<\/p>\n 5G NR beam selection (at the UE \/ baseband)CUDA-TopK-Retrieval (at the GPU)Codebook of candidate beams (fixed at config)Corpus of pre-embedded chunks (uploaded once)Instantaneous channel measurementQuery embedding for this hopL1-RSRP per candidate beamCosine \/ dot-product score per corpus rowTop best beams reported back to gNBTop-K row indices returned to the agentPer-beam score lives in baseband DSP, not in a host CPUPer-row score lives in VRAM, not in host RAMDoing this on a CPU round-trip would melt the air interfaceDoing this on a CPU round-trip melts the agent\u2019s throughput<\/p>\n A quick aside to two very different audiences<\/p>\n To my HPC and CUDA-first friends reading this: I hear you. None of the mathematical primitives here are novel. We all know cuBLAS runs matmuls faster, cuVS handles Top-K at datacenter scale, and a highly tuned tournament selection will crush this per-block bubble sort. But the goal here isn\u2019t to reinvent NVIDIA\u2019s enterprise libraries. The value is the zero-dependency packaging. This is a 343-line architectural proof\u2014complete with a strict CPU oracle and a 45-config benchmark sweep\u2014designed to run entirely on a vintage 8 GB consumer GPU. It\u2019s the kind of end-to-end engineering artifact you build to prove you actually understand hardware memory bottlenecks, rather than just knowing how to call a framework API.<\/p>\n To my telecom friends: if \u201cTop-K vector search\u201d sounded like a foreign language until ten minutes ago, you are not behind \u2014 you are early. For twenty years our world was FPGAs, ASICs, PRBs, and constellation diagrams. We optimized spectrum, not silicon. Then AI-RAN, NWDAF, NVIDIA Aerial, and the 3GPP Rel-20 study items all happened too fast within a few months, and the next decade of telecom careers now demands being bilingual between spectrum-world and GPU-world. The intuition translates cleanly. You have been doing receiver-side Top-K under hard real-time constraints since the first MIMO codebook. Same animal, just in a new zoo.<\/p>\n 9. Honest caveats (because the comments are coming)<\/p>\n If you came here to find what is wrong with this project \u2014 congratulations, you are this article\u2019s first careful reader. Straight from the README\u2019s LIMITATIONS section and the inline code comments:<\/p>\n K=100 is where V1 loses. The partial Top-K path uses single-lane-per-block selection for auditability; it is not yet a warp-specialized tournament-selection kernel. At K=100 the O(K\u00b2) bubble sort dominates, and on the CSV the CPU pulls ahead on 14 of 15 K=100 rows (sometimes by 2\u00d7, sometimes by 46\u00d7). The lone GPU win at K=100 \u2014 N=1M, D=1024, 1.44\u00d7 \u2014 is the scoring work finally being big enough to swamp the selection ceiling, not the selector getting any better. This is documented; the fix is a known follow-up.<\/p>\n GPU clocks are not locked in the committed receipt. The committed environment.json reports gpu_clocks_locked: false. Absolute milliseconds shift with thermals on a consumer card; the ratios in the headline table are durable. The repo ships scripts\/lock_gpu_clocks.sh (persistence mode + lock to 1607 MHz application clocks for a stock GTX 1080) for anyone who needs publication-grade numbers.<\/p>\n Numeric tolerance, not exact float equality. GPU vs CPU score comparisons use a small fp32 tolerance per score; ties are still resolved deterministically by index. This is a real-world necessity \u2014 fp32 reductions associate differently on a GPU than on a CPU \u2014 and the harness will not start timing until indices match exactly.<\/p>\n Synthetic embeddings. The benchmark uses Gaussian-random vectors (std::normal_distribution, seed 1) to isolate the residency-vs-round-trip signal from content effects and keep trials reproducible bit-for-bit. Real embeddings will produce noisier per-trial absolute timings; the structural ratio between PCIe-vacation and on-device math does not move.<\/p>\n One CUDA architecture class. All numbers come from one Pascal-class GTX 1080. On Ada \/ Hopper the absolute milliseconds will shrink for both paths; the structural finding (PCIe round-trip cost dominates a CPU-side retrieval) becomes more important on faster GPUs, not less, because the kernel time shrinks faster than the round-trip does.<\/p>\n RAG slice, not a full vector DB. This is a similarity + Top-K slice. No compression (PQ, OPQ), no filtering, no multi-GPU sharding, no concurrency between queries inside one engine instance. It is the retrieval-hop primitive an agent calls \u2014 not a replacement for FAISS-GPU or cuVS.<\/p>\n Everything on this list is on the roadmap. None of it changes the headline result. The point of putting it in writing is that you should not have to dig for it \u2014 and the moment a benchmark blog post hides its caveats is the moment its numbers stop being trustworthy.<\/p>\n 10. Wrap (and the setup for the final part)<\/p>\n If you build agentic pipelines for a living: please go and look at your retriever. Open whatever profiler you trust. Time one tool call end-to-end. If your GPU utilization drops to zero while a Python host process grinds through a similarity search, you have already won the diagnostic battle. The fix is on GitHub.<\/p>\n If you write CUDA for a living: Yes, the O(K2) bubble sort is intentional. A warp-specialized tournament selector is on the roadmap.<\/p>\n If you build telecom infrastructure for a living: Yes, you caught me. This is the exact same baseband retrieval primitive you have been writing in DSP code for twenty years. The AI industry just changed the vocabulary; the math hasn\u2019t budged. <\/p>\n Coming up next: How to Stop Your Agents from Trauma-Dumping on Each Other<\/p>\n CUDA-TopK-Retrieval proves you can stop bouncing every retrieval hop off the GPU. But if you reread caveat #1 plus the K=100 rows in Section 5, you have already spotted the next ceiling: the per-query work is independent across queries.<\/p>\n Every retrieval hop starts cold. The corpus is on the device, sure. But the agent\u2019s state \u2014 the embeddings of its previous decisions, the keys and values that it would naturally attend back over \u2014 gets dropped and rebuilt on every hand-off. The GPU stays warm; the agent\u2019s memory stays cold.<\/p>\n That is fine for a one-shot RAG step. It falls apart the moment you run the workload this series was built for: multi-hop reasoning across a swarm of specialized agents. At that scale you stop caring about \u201cdid we keep the retrieval on the GPU\u201d and start caring about questions a single-shot kernel cannot answer:<\/p>\n When agent A hands off to agent B, can B resume with A\u2019s accumulated context instead of cold-starting?<\/p>\n How small can the per-hop persistent state be, and still be useful?<\/p>\n What is the latency cost of restoring that state on the next agent\u2019s GPU?<\/p>\n How do we make hand-offs not lose information?<\/p>\n To get the answer of these questions, you will have to do exactly what your CPU does during a cudaMemcpy: sit there patiently and wait for the next part.<\/p>\n See you in Part 4, the final one.<\/p>\n Disclaimer: The illustrations in this article were generated using AI (Claude Opus 4.8). They are illustrative, not photographic, and any labels visible inside the images are stylized rather than authoritative \u2014 refer to the article body and the code itself for precise function names, metric values, and architecture details.<\/p>\n","protected":false},"excerpt":{"rendered":", 343-line tour of CUDA Top-K retrieval. This kernel, CPU oracle, and benchmark suite prove that the standard…\n","protected":false},"author":2,"featured_media":79734,"comment_status":"","ping_status":"","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[6],"tags":[179,7493,25,12533,2459,38873],"class_list":["post-79733","post","type-post","status-publish","format-standard","has-post-thumbnail","category-agentic-ai","tag-agentic-ai","tag-agentic-artificial-intelligence","tag-artificial-intelligence","tag-cuda","tag-deep-learning","tag-editors-picks"],"_links":{"self":[{"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/posts\/79733","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/comments?post=79733"}],"version-history":[{"count":0,"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/posts\/79733\/revisions"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/media\/79734"}],"wp:attachment":[{"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/media?parent=79733"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/categories?post=79733"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.europesays.com\/ai\/wp-json\/wp\/v2\/tags?post=79733"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}![\"Overview](\"https:\/\/www.europesays.com\/ai\/wp-content\/uploads\/2026\/06\/part3img2-1024x683.png\")
Overview of CUDA TopK Retrieval<\/p>\n
\n\/\/ This is not the fastest global selection, but it is easy to reason about and matches the CPU ordering rule exactly.
\n__device__ void bubble_downward(float* const s, int* const ids, const int n) {
\n \/\/ Tiny O(K^2) sort acceptable because K is capped (kMaxSupportedK) and this runs on a single lane per block.
\n for (int i = 0; i < n – 1; ++i) {
\n for (int j = 0; j < n – 1 – i; ++j) {
\n if (device_is_better(s[j + 1], ids[j + 1], s[j], ids[j])) {
\n const float ts = s[j];
\n s[j] = s[j + 1];
\n s[j + 1] = ts;
\n const int ti = ids[j];
\n ids[j] = ids[j + 1];
\n ids[j + 1] = ti;
\n }
\n }
\n }
\n}<\/p>\n
\n\/\/ We use strict weak ordering for std::partial_sort: higher score wins; on exact tie, smaller index wins.
\nbool is_better_score_pair(const float32_t score_lhs, const index_t idx_lhs, const float32_t score_rhs,
\n const index_t idx_rhs) {
\n \/\/ Primary key: similarity score (higher is better for retrieval).
\n if (score_lhs != score_rhs) {
\n return score_lhs > score_rhs;
\n }
\n \/\/ Deterministic tie surface: prefer the smaller corpus row id to mirror stable DB primary keys.
\n return idx_lhs < idx_rhs;
\n}<\/p>\n
\n__device__ bool device_is_better(const float score_lhs, const int idx_lhs, const float score_rhs, const int idx_rhs) {
\n \/\/ Same ordering semantics as topk::is_better_score_pair for bitwise-identical tie surfaces.
\n if (score_lhs != score_rhs) {
\n return score_lhs > score_rhs;
\n }
\n return idx_lhs < idx_rhs;
\n}<\/p>\n
\nStep 1: Upload corpus once into VRAM (upload_corpus_rowmajor)
\nStep 2: Per query \u2014 H\u2192D the embedding (search_resident, first line)
\nStep 3: Score N rows on device (row_dot_scores_kernel)
\nStep 4: Per-block partial Top-K (partial_topk_block_kernel)
\nStep 5: Multi-way merge into global Top-K (merge_partial_topk_kernel)
\nStep 6: D\u2192H the K indices + K scores (search_resident, last lines)<\/p>\n
\n const index_t D) {
\n if (N > max_n_ || D > max_d_) {
\n return cudaErrorInvalidValue;
\n }
\n const std::size_t corpus_bytes = sizeof(float) * static_cast(N) * static_cast(D);
\n const cudaError_t st = cudaMemcpy(d_corpus_, host_corpus_rowmajor, corpus_bytes, cudaMemcpyHostToDevice);
\n if (st != cudaSuccess) {
\n return st;
\n }
\n resident_n_ = N;
\n resident_d_ = D;
\n return cudaSuccess;
\n}<\/p>\n
\n\/\/ Microarchitectural note: one thread per row is simple; a follow-up can tile D across warps to raise arithmetic intensity.
\n__global__ void row_dot_scores_kernel(const float* const corpus, const float* const query, const float* const row_l2,
\n const float query_l2, const int N, const int D, const int cosine_enabled,
\n float* const scores) {
\n \/\/ Map each CUDA thread to exactly one corpus row to keep the reduction logic easy to audit against the CPU reference.
\n const int row = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x);
\n if (row >= N) {
\n return;
\n }
\n float acc = 0.0F;
\n const int base = row * D;
\n for (int col = 0; col < D; ++col) {
\n acc += corpus[static_cast(base + col)] * query[static_cast(col)];
\n }
\n if (cosine_enabled != 0) {
\n const float denom = query_l2 * row_l2_fetch(row_l2, row);
\n scores[static_cast(row)] =
\n denom > 0.0F ? (acc \/ denom) : -std::numeric_limits::infinity();
\n } else {
\n scores[static_cast(row)] = acc;
\n }
\n}<\/p>\n
\n partial_topk_block_kernel<<<\/p>\n
\n d_partial_indices_);<\/p>\n
\n int best_p = -1;
\n float best_s = -std::numeric_limits::infinity();
\n int best_i = std::numeric_limits::max();
\n for (int p = 0; p < P; ++p) {
\n if (heads[p] >= K) {
\n continue;
\n }
\n const float s = partial_scores[static_cast(p * K + heads[p])];
\n const int idx = partial_indices[static_cast(p * K + heads[p])];
\n if (best_p < 0 || device_is_better(s, idx, best_s, best_i)) {
\n best_p = p;
\n best_s = s;
\n best_i = idx;
\n }
\n }
\n out_scores[static_cast(out)] = best_s;
\n out_indices[static_cast(out)] = best_i;
\n heads[best_p] += 1;
\n }<\/p>\n![\"CUDA](\"https:\/\/www.europesays.com\/ai\/wp-content\/uploads\/2026\/06\/part3img3-1024x825.png\")
<\/p>\n
\ncmake –build build -j
\ncd build && ctest –output-on-failure<\/p>\n
\n.\/build\/topk_bench –n 20000 –d 384 –k 32 –trials 3 –warmup 1 –seed 1 –metric 0<\/p>\n![\"Intent](\"https:\/\/www.europesays.com\/ai\/wp-content\/uploads\/2026\/06\/github-social-preview-1024x683.png\")
<\/p>\n