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

, 343-line tour of CUDA Prime-Okay retrieval. This kernel, CPU oracle, and benchmark suite show that the usual Agentic RAG round-trip—bouncing queries throughout the PCIe bus—is the silent killer of your pipeline. By conserving similarity search resident on gadget reminiscence, this structure achieves an 8.6x speedup over optimized CPU baselines even on a 7-year-old GTX 1080.

That is Half 3 of the “Manufacturing-Grade Agentic Inference” sequence. Every half removes one type of redundant work from an agentic LLM pipeline. Part 1 killed redundant prefill. Part 2 killed redundant ready — how a number of micro-agents share one GPU by way of time-slicing. Half 3 (this publish) retains RAG retrieval on the GPU with a customized CUDA Prime-Okay kernel. Half 4 persists agent state throughout hand-offs so the subsequent agent by no means has the cold-start downside.

Key Takeaways

The issue: in agentic RAG, each instrument name that wants context fires a similarity search. A default pipeline ships the question embedding from the GPU to Python, lets the CPU rating N corpus rows and choose the perfect Okay, then ships the reply again. That round-trip is the silent tax. The compute is okay; the journey is the invoice. Everyone knows, journey isn’t low cost, irrespective of the place you wish to go (pun supposed!)

The straightforward repair: add the corpus to VRAM as soon as, then maintain the similarity scoring, the Prime-Okay choice, and the merge step on the gadget. Solely the tiny per-query embedding (D floats) and the Okay outcomes journey throughout PCIe.

The receipts: on the identical 7-year-old GTX 1080 utilized in Components 1 and a pair of, the GPU-resident path runs the retrieval hop as much as 8.57× quicker than a CPU brute-force baseline. At Okay=8 it wins on all 15 sweep configurations (N ∈ {10k, 50k, 100k, 500k, 1M}, D ∈ {384, 768, 1024}) with speedups from 2.43× to eight.57×. At Okay=32 it wins on 13 of 15 configs, peaking at 7.76×. At Okay=100 — the place the V1 selector deliberately stays easy — the CPU wins on 14 of 15 configs. That final sentence is the sincere half (Properly, even when I had lied, you may have simply caught it).

The kicker: the wins aren’t “magic kernel” wins. They’re “we stopped delivery the corpus again to host RAM for no purpose” wins. It’s also precisely the type of “measure many candidates, report solely the perfect Okay again to the buyer” determination a 5G base station and your cellphone have been making each few milliseconds since CSI suggestions turned a factor.

TL;DR: Default agentic RAG treats the GPU as a serving field and the retrieval as a Python concern. Each instrument name ships the question embedding D→H, lets the CPU compute N dot merchandise, type the candidates, choose the highest Okay, and ship indices and scores H→D. For an agent that calls a vector retailer ten instances per reasoning step, that round-trip is the dominant value — not the mannequin, not the embedding, it’s the journey. CUDA-TopK-Retrieval retains the corpus resident on the gadget, runs scoring + per-block partial Prime-Okay + a multi-way merge completely on the GPU, and exposes a tiny C++ orchestrator API (upload_corpus_rowmajor as soon as, search_resident per question). The host-touching bytes per question collapse to 1 D-length embedding up and 2K outcomes down. On a GTX 1080, throughout a 45-config sweep, the GPU-resident path beats the CPU-round-trip baseline on all 15 Okay=8 configs (2.43× to eight.57×, peaking at N=1M, D=1024) and on 13 of 15 Okay=32 configs (the 2 losses are each on the smallest N=10k for D=384 and D=768, the place the round-trip itself is already low cost; big-N Okay=32 wins climb to 7.76×). At Okay=100 the V1 kernel intentionally stays easy — single-lane-per-block bubble type with a serial merge — and the CPU wins on 14 of 15 configs; that ceiling is the article’s sincere punchline and a clear setup for Half 4.

(Fast confession earlier than we begin: I got here at this from a 5G/6G RAN engineering background. Beam choice at a base station seems to be shockingly near RAG Prime-Okay — the UE scores a codebook of candidate beams by acquired energy and stories the perfect handful again over the air. There’s a complete part on that under — part 8 — however it’s also why this kernel exists within the form it does.)

Structure psychological mannequin — maintain this open whilst you learn.

agent.embed(question) → cudaMemcpy H→D (D floats) → row_dot_scores_kernel → partial_topk_block_kernel (P blocks) → merge_partial_topk_kernel → cudaMemcpy D→H (Okay indices + Okay scores)

All the things under is simply commentary on one a part of that line.

1. A confession: each RAG step in your agent is a tiny PCIe highway journey

In Part 2 of this sequence, we efficiently remoted our LLM agent’s inference loop, conserving token technology operating sizzling and quick on the gadget. We designed a system which avoids stalling. However the second we give that agent a instrument to look an exterior data base—the core of any multi-hop Retrieval-Augmented Technology (RAG) pipeline—we silently destroy all that hard-won efficiency and we hit the wall. When you’ve got ever wired an “agentic” pipeline to a vector retailer by way of a Python retriever, here’s what actually occurs on every instrument name (with a little bit intentional dramatization):

You: “Agent, discover me the 5 chunks most related to ‘how do I declare deduction beneath part 80C?’“

Agent: “Certain. Embedding the question on the GPU. ✅”

Agent: “Now bouncing the question embedding again to host.”

(cudaMemcpy D→H, ~1,024 floats) Python retriever: “Received it. NumPy loop. Dot product N instances. argpartition. Prime-5.”

(CPU scores half one million corpus rows, one row at a time, whereas a 9 TFLOP GPU watches) Python retriever: “Achieved. Listed below are the indices and scores.”

Agent: “Cool. Bouncing them again to the GPU now.”

(cudaMemcpy H→D, 10 numbers) Agent: “Prepared. What was the query once more?”

The agent has a wonderfully good GPU. The corpus is sitting in 4 GB of the VRAM. The question embedding was already on the GPU — we simply generated it there. After which, on each single retrieval hop, we ship the question again to the host, do brute-force similarity in NumPy / FAISS-on-CPU / a hand-rolled loop, and ship the reply again.

Your GPU’s utility meter: spends many of the retrieval step idle. Your PCIe bus: will get a exercise it didn’t join. Your agent’s tool-call latency: dominated by one thing that’s neither the mannequin nor the embedding. That’s the joke.

That can also be the soiled secret of each agentic RAG demo that scales previous the toy “ten chunks in reminiscence” stage. The retrieval hop bounces off the GPU and again, each time, and the larger the corpus, the more severe the tax. At one million rows of 1024-dim embeddings, the round-trip alone — not even the scoring, sure, simply the round-trip — eats many of the funds of the retrieval step itself.

CUDA-TopK-Retrieval is what occurs once you resolve the round-trip is non-obligatory and you’ll fairly write 343 strains of CUDA than let the agent trip by way of host RAM each time it needs a neighbor.

Now think about the actual workload behind this. It isn’t “5 chunks for one query.” It’s a number of specialised micro-agents — each operating its personal RAG hops, each needing Prime-Okay in opposition to the identical corpus, each at the moment paying its personal PCIe invoice on each instrument name. Part 1 of this sequence killed the prefill round-trip. Part 2 made the GPU shareable throughout these a number of brokers. Half 3 says: now that they’re sharing the cardboard pretty, cease making every one in every of them drive again to the host to search for a neighbor.

2. Why does Prime-Okay retrieval exist? (a one-minute crash course)

Skip this if you have already got a background on this. For everybody else, who’re new to this area, here’s a brief, novice clarification.

A contemporary agent doesn’t stuff the entire data base into the immediate. It retrieves. For each reasoning step that wants grounded context, it embeds the question right into a fixed-dimensional vector (D floats — usually 384, 768, or 1024), scores that vector in opposition to each row of a corpus of pre-embedded chunks (N rows, additionally D floats every), and returns the Okay corpus rows with the very best similarity. That’s it. That’s Prime-Okay vector search. Retrieval-Augmented Technology is only a well mannered means of claiming “Prime-Okay plus a immediate template.”

Two flavors of similarity present up in all places. Dot product is a budget one: a single fused multiply-add per dimension, N×D work in complete. Cosine is dot product divided by the product of L2 norms, which turns into a free dot product when you pre-normalize the corpus as soon as at ingest. Most manufacturing vector shops do the pre-normalization trick and name it “cosine” whereas operating uncooked dot-product math at question time. The CUDA-TopK-Retrieval kernel helps each — it simply multiplies by a precomputed per-row norm pointer when cosine mode is on.

Mainstream instruments (FAISS, hnswlib, the Python aspect of cuVS, your favourite SaaS vector DB) all do that scoring + Prime-Okay work. Most of them do it properly. The issue is the place they do it. Nearly each agent framework on the planet calls into the retriever from Python, and the second Python is on the new path, the retrieval step is now not a GPU operation — it’s a PCIe operation with a GPU on one finish.

The repair shouldn’t be “a greater algorithm.” It’s “a a lot shorter highway journey.”

3. The “simply maintain the corpus on the GPU” lightbulb (and why it’s more durable than it sounds)

The pitch is easy

• Add the corpus to VRAM as soon as at ingest.
• For each incoming question, cudaMemcpy a tiny D-dimensional float embedding to the gadget.
• Launch a scoring kernel the place one CUDA thread per corpus row computes the dot product.
• Launch a partial Prime-Okay kernel the place every block scans a disjoint row vary to emit its personal native high candidates.
• Lastly, launch a merge kernel to stroll the per-block heads and emit the worldwide Prime-Okay in best-first order.

You cudaMemcpy precisely 2K numbers again to the host: Okay indices, and Okay scores.

That is the “deal with reminiscence retrieval as a {hardware} primitive, not a software program API name” paradigm. The one purpose this takes greater than a 30-line PyTorch script to attain is that three tedious edge circumstances will instantly break the naive strategy.

Drawback A: Prime-Okay on a GPU is structurally awkward

Scoring the vectors is the simple half. It’s simply matrix multiplication, and your GPU was actually born to try this—it’s the {hardware}’s love language. Choice, nevertheless, is the place the romance dies. Asking a GPU to do a full O(N log N) type simply to seize the highest Okay outcomes is computationally offensive; it’s like alphabetizing your complete recycling bin simply to discover a single receipt. You might attempt an O(N) argpartition, however that requires a tree-walk, which shatters GPU reminiscence coalescing into one million unaligned reads. Event choice is quick, assuming you wish to spend your weekend debugging edge circumstances. And the second you cave and attain for a thrust or cub sorting primitive, congratulations: you’ve simply contaminated your light-weight, standalone C++ pipeline with a large construct dependency.

The structure picks the boring reply on objective. It depends on a tiny per-block O(Okay²) bubble type over a disjoint row vary, pushed by a single thread per block, and capped off with a serial multi-way merge. On paper, this sounds horrible. In observe, it really works fantastically, for the precise purpose said truthfully within the kernel’s feedback:

// Single-threaded per-block scan that materializes an area Prime-Okay record for its row partition.
// This isn't the quickest international choice, however it's simple to purpose about and matches the CPU ordering rule precisely.
__device__ void bubble_downward(float* const s, int* const ids, const int n) {
  // Tiny O(Okay^2) type acceptable as a result of Okay is capped (kMaxSupportedK) and this runs on a single lane per block.
  for (int i = 0; i < n - 1; ++i) {
    for (int j = 0; j < n - 1 - i; ++j) {
      if (device_is_better(s[j + 1], ids[j + 1], s[j], ids[j])) {
        const float ts = s[j];
        s[j] = s[j + 1];
        s[j + 1] = ts;
        const int ti = ids[j];
        ids[j] = ids[j + 1];
        ids[j + 1] = ti;
      }
    }
  }
}

That is the design contract: V1 picks auditability over cleverness. The entire kernel is sufficiently small {that a} reviewer can learn it end-to-end in a espresso break, line it up in opposition to the CPU oracle, and persuade themselves the GPU output is bitwise right. The day somebody needs a 2× kernel, this will get changed with a warp-specialized event selector — the article’s part 9 guarantees that explicitly. For Okay ≤ 32, the single-lane bubble type is genuinely high-quality. For Okay=100, it falls off a cliff. That cliff is documented and benchmarked. (Part 9 once more.)

Drawback B: GPU and CPU should agree, bit-for-bit, on the tiebreak

Properly, like all World Cup Soccer group league match, ties occur greater than usually. Two corpus rows have the identical rating right down to fp32 precision. Which one wins?

If the CPU oracle and the GPU kernel disagree on the tiebreak, you’ll be able to by no means belief a benchmark, as a result of each “mismatch” alarm is now ambiguous: did the GPU rating a row mistaken, or did the GPU simply break ties in another way? You’ll spend every week in a 3 AM Slack channel.

The repair is to outline the comparator in a single sentence and implement it in two locations — as soon as on the host, as soon as on the gadget — and make these two implementations actually the identical expression.

On the host aspect it seems to be like:

// Lexicographic "higher" relation for (rating, index) pairs beneath float equality semantics.
// We use strict weak ordering for std::partial_sort: greater rating wins; on precise tie, smaller index wins.
bool is_better_score_pair(const float32_t score_lhs, const index_t idx_lhs, const float32_t score_rhs,
                          const index_t idx_rhs) {
  // Main key: similarity rating (greater is best for retrieval).
  if (score_lhs != score_rhs) {
    return score_lhs > score_rhs;
  }
  // Deterministic tie floor: desire the smaller corpus row id to reflect steady DB major keys.
  return idx_lhs < idx_rhs;
}

On the gadget aspect it seems to be like:

// Machine-side reproduction of the host comparator to keep away from cross-TU linkage points for __device__ code paths.
__device__ bool device_is_better(const float score_lhs, const int idx_lhs, const float score_rhs, const int idx_rhs) {
  // Similar ordering semantics as topk::is_better_score_pair for bitwise-identical tie surfaces.
  if (score_lhs != score_rhs) {
    return score_lhs > score_rhs;
  }
  return idx_lhs < idx_rhs;
}

That’s the complete tiebreak coverage in 5 strains, twice. Larger rating wins; on precise tie, smaller corpus row index wins. The CPU oracle makes use of it in std::partial_sort, the GPU makes use of it within the bubble type and within the multi-way merge, and the benchmark harness won’t begin timing till the GPU output matches the CPU output precisely — similar indices in the identical order, scores inside a small fp32 tolerance.

That single comparator is the explanation the article can quote a speedup in any respect. With out it, “the GPU is 8× quicker” is simply “the GPU is 8× quicker at being mistaken in another way.”

Drawback C: VRAM is valuable, and the worst place to do malloc is the new path

Allocating GPU reminiscence per-query is like signing a brand new automobile lease each time it’s worthwhile to drive to the grocery retailer. It’s the best approach to flip a 1-millisecond search right into a 50-millisecond site visitors jam.

As a substitute, GpuTopkEngine::initialize buys the automobile upfront. It runs all of the cudaMalloc calls throughout engine startup, sizing the buffers for the worst potential configuration. As soon as the engine is actively serving queries, the new path is totally freed from reminiscence administration. It’s simply quick kernel launches and tiny information copies. No fragmentation, no negotiating with the allocator, and cudaMalloc is completely banned from exhibiting up in your efficiency traces.

4. The four-stage pipeline (the actually-cool half)

Step 0:  Engine init — eight cudaMallocs sized for (max_n, max_d, max_k)   (GpuTopkEngine::initialize)
Step 1:  Add corpus as soon as into VRAM                                      (upload_corpus_rowmajor)
Step 2:  Per question — H→D the embedding                                     (search_resident, first line)
Step 3:  Rating N rows on gadget                                             (row_dot_scores_kernel)
Step 4:  Per-block partial Prime-Okay                                            (partial_topk_block_kernel)
Step 5:  Multi-way merge into international Prime-Okay                                  (merge_partial_topk_kernel)
Step 6:  D→H the Okay indices + Okay scores                                       (search_resident, final strains)

Let’s stroll by way of each with the actual code. The snippets are even shorter than the tiny supply information.

Step 1 — Add the corpus as soon as

That is the boring step that makes the remainder of the article potential. The corpus goes up precisely one time per ingest, and stays there for the lifetime of the engine:

cudaError_t GpuTopkEngine::upload_corpus_rowmajor(const float32_t* const host_corpus_rowmajor, const index_t N,
                                                  const index_t D) {
  if (N > max_n_ || D > max_d_) {
    return cudaErrorInvalidValue;
  }
  const std::size_t corpus_bytes = sizeof(float) * static_cast<:size_t>(N) * static_cast<:size_t>(D);
  const cudaError_t st = cudaMemcpy(d_corpus_, host_corpus_rowmajor, corpus_bytes, cudaMemcpyHostToDevice);
  if (st != cudaSuccess) {
    return st;
  }
  resident_n_ = N;
  resident_d_ = D;
  return cudaSuccess;
}

And that’s the complete ingest API. At 1024 dimensions, a million vectors is strictly 4 GB of float32 information, sliding completely into the 8 GB VRAM of a classic GTX 1080. What occurs when your corpus hits 10 million vectors? That turns into a distributed programs downside, not a kernel downside. In case your information exceeds VRAM, you want a sharding technique, which we cowl in Part 9. However for now, we’re right here to resolve the compute bottleneck, to not invent a brand new database.

Step 2 — Rating N rows on the gadget

One CUDA thread per corpus row. 256 threads per block. Every thread accumulates the dot product throughout D dimensions and writes one float into the dense scores[N] buffer:

// Row-major dot-product with non-obligatory cosine normalization; coalesced reads alongside D are sacrificed for readability in v1.
// Microarchitectural be aware: one thread per row is easy; a follow-up can tile D throughout warps to boost arithmetic depth.
__global__ void row_dot_scores_kernel(const float* const corpus, const float* const question, const float* const row_l2,
                                      const float query_l2, const int N, const int D, const int cosine_enabled,
                                      float* const scores) {
  // Map every CUDA thread to precisely one corpus row to maintain the discount logic simple to audit in opposition to the CPU reference.
  const int row = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x);
  if (row >= N) {
    return;
  }
  float acc = 0.0F;
  const int base = row * D;
  for (int col = 0; col < D; ++col) {
    acc += corpus[static_cast<:size_t>(base + col)] * question[static_cast<:size_t>(col)];
  }
  if (cosine_enabled != 0) {
    const float denom = query_l2 * row_l2_fetch(row_l2, row);
    scores[static_cast<:size_t>(row)] =
        denom > 0.0F ? (acc / denom) : -std::numeric_limits::infinity();
  } else {
    scores[static_cast<:size_t>(row)] = acc;
  }
}

One thread per row is the only potential mapping. The code remark is sincere about it: a follow-up can tile D throughout warps to boost the arithmetic depth. For V1, this provides the auditor a one-to-one correspondence with the CPU loop and lets them sleep peacefully at evening.

Step 3 — Every block builds its personal native Prime-Okay

Now comes probably the most awkward half. Selecting the highest Okay out of N is conceptually “type and slice,” however a full type wastes many of the work. We break up the row vary throughout P blocks (capped at 128), every block walks its disjoint slice with the tiny bubble type from Part 3, and writes its personal native top-Okay record out:

  const int P = std::min(kMaxPartialBlocks, std::max(1, (static_cast(N) + 4095) / 4096));
  partial_topk_block_kernel<<
>>(d_scores_, static_cast(N), static_cast(Okay), P, d_partial_scores_,
                                      d_partial_indices_);

One thread per block. Sure, that’s wasteful on paper. It’s also why a human can audit this kernel in twenty minutes — the coverage if (threadIdx.x != 0 || blockIdx.x >= P) return; collapses the entire intra-block reasoning right down to “this block’s lane 0 owns rows [start, end).” Every block’s s[] and ids[] arrays stay in registers / native reminiscence, sized by the compile-time kMaxSupportedK = 256 cap.

Step 4 — Merge the partials into the worldwide Prime-Okay

Lastly, one thread on one block walks P cursors over the per-block lists. Every record is already best-first. Decide the perfect head; emit; advance that one cursor; repeat Okay instances:

  for (int out = 0; out < Okay; ++out) {
    int best_p = -1;
    float best_s = -std::numeric_limits::infinity();
    int best_i = std::numeric_limits::max();
    for (int p = 0; p < P; ++p) {
      if (heads[p] >= Okay) {
        proceed;
      }
      const float s = partial_scores[static_cast<:size_t>(p * K + heads[p])];
      const int idx = partial_indices[static_cast<:size_t>(p * K + heads[p])];
      if (best_p < 0 || device_is_better(s, idx, best_s, best_i)) {
        best_p = p;
        best_s = s;
        best_i = idx;
      }
    }
    out_scores[static_cast<:size_t>(out)] = best_s;
    out_indices[static_cast<:size_t>(out)] = best_i;
    heads[best_p] += 1;
  }

The merge is ruthlessly environment friendly: at most P * Okay reads and precisely Okay writes, executed by a single thread. To forestall floating-point chaos, the device_is_better comparator enforces strict determinism—if two heads tie on rating, the decrease corpus row index wins, mirroring the CPU oracle completely. Lastly, two microscopic cudaMemcpy calls shuttle the Okay profitable indices and scores again to the host. The agent ingests them, and the RAG loop fires once more.

That’s the complete sizzling path: one H -> D embedding switch, three kernel launches, and two tiny D -> H end result copies. No Python host loops, no framework overhead, and completely no PCIe holidays.

5. The receipts (i.e., the numbers)

Let’s now consider it in opposition to the baseline, and see if it was well worth the hassle.

A fast be aware on methodology, earlier than the benchmarking police arrives: each comparability under runs on the similar GPU as Part 1 and Part 2 (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, mounted RNG (std::mt19937 with std::normal_distribution), Gaussian embeddings, L2-normalized in dot-product mode. The complete sweep is N ∈ {10k, 50k, 100k, 500k, 1M} × D ∈ {384, 768, 1024} × Okay ∈ {8, 32, 100} → 45 configurations, all measured by way of cudaEventElapsedTime with cudaDeviceSynchronize bracketing each interval. Code in src/host/bench_main.cpp; uncooked numbers in examples/example-run-results/benchmark_run_results.csv.

The 2 paths timed:

• GPU-resident (therapy). Corpus is already on the gadget. Every timed iteration: cudaMemcpy question H→D (D floats) + rating kernel + per-block partial Prime-Okay + merge + cudaMemcpy Okay scores D→H + cudaMemcpy Okay indices D→H. Finish to finish.
• CPU round-trip (baseline). Fashions the default agent movement: cudaMemcpy question D→H + CPU brute-force scoring + std::partial_sort with the identical comparator + cudaMemcpy indices H→D + cudaMemcpy scores H→D. Finish to finish.

Each paths run inside the identical course of, use the identical question bytes and the identical comparator. The one distinction is the place the work occurs. When you’ve got ever held the place “PCIe is okay, we benchmark the kernels in isolation,” that is what it prices you once you cease pretending the round-trip is free.

Headline (GTX 1080, three trials, imply ms, ratios computed from per-trial means):

Sure, you’re studying the numbers proper.

The primary 5 rows are the article’s level: at Okay=8, the GPU-resident path wins on each single configuration within the sweep (all 15 of them), by ratios starting from a well mannered 2.43× at N=50k, D=384 to a loud 8.57× at N=1M, D=1024. At Okay=32 it wins on 13 of 15 — the 2 losses are each on the smallest N (10k), for D=384 and D=768, the place the round-trip itself solely prices ~3–7 ms and the GPU’s three kernel launches barely have room to amortize. By the point you attain lifelike agentic-corpus sizes (N ≥ 50k), Okay=32 additionally wins comfortably, peaking at 7.76×. The massive speedups aren’t “magic kernel” speedups — they’re “we stopped delivery the corpus again to host RAM for no purpose” speedups. The GPU was all the time going to win this race; the one purpose it ever misplaced was that we stored making it commute unnecessarily.

The final two rows are the place this text earns its proper to name itself sincere. At Okay=100, the single-lane bubble type per block turns into O(K²) = O(10,000) sequential comparisons per block, and the serial merge walks P × Okay head positions. The CPU’s std::partial_sort is heap-based, vectorized by the compiler, and successfully O(N log Okay) — a lot friendlier to Okay=100. So the GPU loses on 14 of 15 Okay=100 configs, generally by 2×, generally by 46×. (There may be precisely one Okay=100 config the place the GPU nonetheless wins — N=1M, D=1024, 1.44× — as a result of by then there may be sufficient scoring work to dominate the choice ceiling. One row out of fifteen shouldn’t be a save; it’s a curiosity.) That’s not a bug; it’s the V1 design assertion (“auditability over cleverness”) assembly its first concrete consequence. The repair is in Part 9, and it’s a warp-specialized event selector — not a frantic refactor.

Yet one more sincere caveat within the numbers above: on this dedicated snapshot, the GPU clocks had been not locked. Which means absolute milliseconds transfer barely with thermals and DVFS; the ratios keep steady. The repo ships scripts/lock_gpu_clocks.sh for anybody who needs to breed the desk with locked clocks on a GTX 1080. Pointless to say, the structural discovering doesn’t change.

6. “OK, however how is that this completely different from FAISS / cuVS / hnswlib?”

A really cheap query, and value answering immediately, as a result of the vector-search world has lots of overlapping primitives and an HPC reader will ask this within the first remark.

• FAISS (CPU index). The default in most agent frameworks. Glorious library. Lives on the CPU. Each question an agent makes pays the PCIe round-trip this text exists to delete. When you’re already on IndexFlatIP and also you’re CPU-bound on retrieval, you’re the audience.
• FAISS (GPU index). Solves the GPU residency downside, with a way more mature kernel suite than this repo. The purpose of CUDA-TopK-Retrieval shouldn’t be “I out-engineered FAISS-GPU” — it by no means was and neither does it attempt to be. The purpose is to indicate, in 343 strains, what the actually-cool retrieval primitive seems to be like and why agentic pipelines really feel sluggish when it isn’t there. When you want a manufacturing index at this time, use FAISS-GPU. If you wish to perceive the small sizzling path that makes the distinction — one tiny H→D copy, three kernel launches, two small D→H copies — learn this kernel.
• NVIDIA cuVS / RAFT. The intense, production-grade in-GPU vector-search stack. Larger, quicker, extra algorithms, extra dependencies. Similar caveat as FAISS-GPU: this kernel is the pedagogical / single-binary model, not a competitor.
• hnswlib and associates (approximate nearest neighbor). Totally different form of trade-off completely — they commerce exactness for sublinear question time on enormous corpora. CUDA-TopK-Retrieval is precise brute-force scoring + choice; the speedup story is only about residency, not about skipping work.

The purpose of this repo shouldn’t be “construct your manufacturing vector DB on it.” The purpose is: the agent’s retrieval hop needs to remain on the GPU, and when you settle for that, even a tiny hand-written kernel beats a hosted-on-CPU brute drive throughout most lifelike Okay values on a 7-year-old card.

7. So… how do I truly attempt it?

Clone the repo, then construct with CMake (the -DGGML_CUDA=ON flag mirrors the llama.cpp construct ergonomics from earlier components of the sequence):

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

Then run the demo and the benchmark precisely because the README does:

./construct/topk_demo                 # tiny smoke story (GPU required)
./construct/topk_bench --n 20000 --d 384 --k 32 --trials 3 --warmup 1 --seed 1 --metric 0

topk_demo is the tiny smoke story — 4,096 corpus rows, 128 dims, Okay=8, prints the neighbor IDs. topk_bench is the harness that emits the TOPK_BENCH_JSON line the Python marketing campaign script consumes. For the complete 45-config sweep on canonical {hardware}:

python3 scripts/benchmark_campaign.py.instance          # full sweep (GPU required; writes beneath examples/benchmark-campaign-runs/run-*)

Earlier than publishing-grade runs, lock the GPU clocks (the repo gives the script):

sudo bash scripts/lock_gpu_clocks.sh

Necessities: Linux, CUDA toolkit, an NVIDIA GPU (Pascal or newer), and the persistence to learn a CMake file as soon as. Artifacts land beneath examples/example-run-results/ for the fast path or examples/benchmark-campaign-runs/run--/ for the complete sweep, and the README is specific that committing .nsys-rep databases is forbidden — PNG timelines solely.

8. Plot twist — that is simply 5G beam choice in a CUDA costume

I ought to most likely confess at this level: I’m nonetheless not a “GPU particular person” by coaching. I got here up by way of telecom — 5G NR with a foot creeping firmly into 6G analysis — and I maintain noticing that each infrastructure downside in agentic AI is an issue which was already solved on the radio layer, possibly round twenty years in the past.

For readers with no 3GPP background: in a contemporary 5G base station, the antenna doesn’t radiate equally in each path. It types a codebook of directional beams — slim lobes of radio power — and at any given second your cellphone is being served by the one beam (or the small handful of beams) whose acquired energy is highest in your gadget. Selecting the best beam, quick, is likely one of the most-studied retrieval issues in wi-fi. The UE measures L1-RSRP (a per-beam received-power rating) throughout the candidate beams the gNB (5G base station) has instructed it to measure, then stories again the perfect handful by way of the CSI suggestions channel. The gNB makes use of these stories to resolve which beams to schedule (Properly, this was so simple as I might make it, the actual story is admittedly ugly!).

That’s Prime-Okay vector search, in radio costume. The candidate beams are the corpus. The instantaneous channel measurement is the question. The rating is acquired energy. Okay is the variety of greatest beams the report carries again. The UE does the scoring down on the baseband DSP layer — it doesn’t ship the I/Q samples again to a central CPU farm and ask a Python script which beam is greatest, as a result of doing so on a per-millisecond loop would soften the air interface.

Take a look at this side-by-side and inform me with a straight face these are completely different issues:

A fast apart to 2 very completely different audiences

To my HPC and CUDA-first associates studying this: I hear you. Not one of the mathematical primitives listed here are novel. Everyone knows cuBLAS runs matmuls quicker, cuVS handles Prime-Okay at datacenter scale, and a extremely tuned event choice will crush this per-block bubble type. However the purpose right here isn’t to reinvent NVIDIA’s enterprise libraries. The worth is the zero-dependency packaging. It is a 343-line architectural proof—full with a strict CPU oracle and a 45-config benchmark sweep—designed to run completely on a classic 8 GB client GPU. It’s the type of end-to-end engineering artifact you construct to show you truly perceive {hardware} reminiscence bottlenecks, fairly than simply understanding easy methods to name a framework API.

To my telecom associates: if “Prime-Okay vector search” gave the impression of a international language till ten minutes in the past, you aren’t behind — you’re 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 research gadgets all occurred too quick inside a number of months, and the subsequent decade of telecom careers now calls for being bilingual between spectrum-world and GPU-world. The instinct interprets cleanly. You may have been doing receiver-side Prime-Okay beneath arduous real-time constraints for the reason that first MIMO codebook. Similar animal, simply in a brand new zoo.

9. Trustworthy caveats (as a result of the feedback are coming)

When you got here right here to seek out what’s mistaken with this undertaking — congratulations, you’re this text’s first cautious reader. Straight from the README’s LIMITATIONS part and the inline code feedback:

1. Okay=100 is the place V1 loses. The partial Prime-Okay path makes use of single-lane-per-block choice for auditability; it isn’t but a warp-specialized tournament-selection kernel. At Okay=100 the O(K²) bubble type dominates, and on the CSV the CPU pulls forward on 14 of 15 Okay=100 rows (generally by 2×, generally by 46×). The lone GPU win at Okay=100 — N=1M, D=1024, 1.44× — is the scoring work lastly being large enough to swamp the choice ceiling, not the selector getting any higher. That is documented; the repair is a identified follow-up.
2. GPU clocks aren’t locked within the dedicated receipt. The dedicated atmosphere.json stories gpu_clocks_locked: false. Absolute milliseconds shift with thermals on a client card; the ratios within the headline desk are sturdy. The repo ships scripts/lock_gpu_clocks.sh (persistence mode + lock to 1607 MHz utility clocks for a inventory GTX 1080) for anybody who wants publication-grade numbers.
3. Numeric tolerance, not precise float equality. GPU vs CPU rating comparisons use a small fp32 tolerance per rating; ties are nonetheless resolved deterministically by index. It is a real-world necessity — fp32 reductions affiliate in another way on a GPU than on a CPU — and the harness won’t begin timing till indices match precisely.
4. Artificial embeddings. The benchmark makes use of Gaussian-random vectors (std::normal_distribution, seed 1) to isolate the residency-vs-round-trip sign from content material results and maintain trials reproducible bit-for-bit. Actual embeddings will produce noisier per-trial absolute timings; the structural ratio between PCIe-vacation and on-device math doesn’t transfer.
5. One CUDA structure class. All numbers come from one Pascal-class GTX 1080. On Ada / Hopper absolutely the milliseconds will shrink for each paths; the structural discovering (PCIe round-trip value dominates a CPU-side retrieval) turns into extra essential on quicker GPUs, not much less, as a result of the kernel time shrinks quicker than the round-trip does.
6. RAG slice, not a full vector DB. It is a similarity + Prime-Okay slice. No compression (PQ, OPQ), no filtering, no multi-GPU sharding, no concurrency between queries inside one engine occasion. It’s the retrieval-hop primitive an agent calls — not a substitute for FAISS-GPU or cuVS.

All the things on this record is on the roadmap. None of it adjustments the headline end result. The purpose of placing it in writing is that you shouldn’t should dig for it — and the second a benchmark weblog publish hides its caveats is the second its numbers cease being reliable.

10. Wrap (and the setup for the ultimate half)

When you construct agentic pipelines for a dwelling: please go and have a look at your retriever. Open no matter profiler you belief. Time one instrument name end-to-end. In case your GPU utilization drops to zero whereas a Python host course of grinds by way of a similarity search, you’ve already received the diagnostic battle. The repair is on GitHub.

When you write CUDA for a dwelling: Sure, the O(Okay²) bubble type is intentional. A warp-specialized event selector is on the roadmap.

When you construct telecom infrastructure for a dwelling: Sure, you caught me. That is the very same baseband retrieval primitive you’ve been writing in DSP code for twenty years. The AI business simply modified the vocabulary; the mathematics hasn’t budged.

Developing subsequent: The right way to Cease Your Brokers from Trauma-Dumping on Every Different

CUDA-TopK-Retrieval proves you’ll be able to cease bouncing each retrieval hop off the GPU. However when you reread caveat #1 plus the Okay=100 rows in Part 5, you’ve already noticed the subsequent ceiling: the per-query work is impartial throughout queries.

Each retrieval hop begins chilly. The corpus is on the gadget, positive. However the agent’s state — the embeddings of its earlier choices, the keys and values that it might naturally attend again over — will get dropped and rebuilt on each hand-off. The GPU stays heat; the agent’s reminiscence stays chilly.

That’s high-quality for a one-shot RAG step. It falls aside the second you run the workload this sequence was constructed for: multi-hop reasoning throughout a swarm of specialised brokers. At that scale you cease caring about “did we maintain the retrieval on the GPU” and begin caring about questions a single-shot kernel can’t reply:

• When agent A arms off to agent B, can B resume with A’s accrued context as a substitute of cold-starting?
• How small can the per-hop persistent state be, and nonetheless be helpful?
• What’s the latency value of restoring that state on the subsequent agent’s GPU?
• How will we make hand-offs not lose info?

To get the reply of those questions, you’ll have to do precisely what your CPU does throughout a cudaMemcpy: sit there patiently and look ahead to the subsequent half.

See you in Half 4, the ultimate one.

Disclaimer: The illustrations on this article had been generated utilizing AI (Claude Opus 4.8). They’re illustrative, not photographic, and any labels seen inside the photographs are stylized fairly than authoritative — check with the article physique and the code itself for exact operate names, metric values, and structure particulars.