Code Walk F · layer 06deep dive · retrieval

RAG
retrieval, and where the GPU actually is

Layer 06 of the Frontier Stack, gone deep. Retrieval-augmented generation is how a fixed model answers about things it never trained on: find the relevant documents, then let the model read them. Most explanations stop at the diagram. This one asks the question the company asks — where is the GPU, and what is it bound by at each step. The surprise: the famous part (vector search) is cheap CPU bandwidth work, while the real GPU cost hides in the embedding forwards and, most of all, in the ballooned context that retrieval hands to the generator. Two interactive pieces.

the idearetrieve, then generate
the five stagesembed · search · rerank · augment · generate
where the GPU isembed / rerank (compute) · generate (memory)
the hidden coststuffed context = ballooned KV cache

00 Five stages, and the GPU is not where you think

A model’s weights are frozen at training time, so it can’t know your codebase, today’s news, or your company’s docs. Retrieval-augmented generation fixes that without retraining: keep the knowledge in a searchable index, pull the relevant pieces at query time, and paste them into the prompt so the model can read them. It’s five stages — embed, search, rerank, augment, generate — and the interesting question for an observability company is which of them actually touch a GPU, and what each is bound by. Click through:

The five stages of a RAG query. Each tag says where the work runs (GPU or CPU) and what limits it (memory-, compute-, or I/O-bound). Note the shape: the stage everyone pictures — vector search — is CPU, memory-bound, and cheap; the GPU cost is in the embedding/rerank forwards and, decisively, in the final generate, where retrieved context becomes prefill and KV cache.

The takeaway that reframes RAG: “vector search” is the cheap part. It’s a memory-bandwidth scan, which is why it lives on CPU and why billion-scale search leans on approximate indexes. The expensive, GPU-bound reality is the transformer forwards at the ends — and the long stuffed context thatlayer 05 then has to prefill and cache for every token it generates.

01 What “semantic” search actually is

The heart of stage 02 is one idea: an embedding model maps every document — and the query — to a point in a high-dimensional space, arranged so that things that mean similar things land near each other. Retrieval is then just geometry: return the points closest to the query. Drag the query around this 2-D projection of an embedding space, or pick a preset, and watch which documents light up:

A toy 2-D projection of an embedding space; each dot is a document, coloured by topic, and clusters form by meaning. Drag the query point or pick a preset — the top-k nearest documents light up. That’s all semantic search is: nearest neighbours in vector space. Real embeddings have hundreds to thousands of dimensions, and real corpora have billions of points — which is exactly why you can’t afford to scan them all.

querytop-k

Two things this makes concrete. First, retrieval quality is the embedding model — if it places the query in the wrong neighbourhood, no amount of clever generation recovers. Second, the naive version of stage 02 compares the query to every point: an O(N) scan, pure memory bandwidth. At a handful of documents that’s instant; at a billion it’s the whole cost, which is why real systems replace the full scan with an approximate index (HNSW’s navigable graph, IVF-PQ’s quantized buckets) that visits a tiny fraction of the points.

So what does “visit a tiny fraction” actually look like? The full scan compares the query to every vector —O(N), a straight sweep of memory. HNSW instead lays the points out as a layered graph and walksit: start at a fixed entry on a sparse top layer, greedily hop toward the query, drop a layer when no neighbour is closer, and repeat down to the dense bottom. It touches a handful of nodes — O(log N)-ish— and lands on the approximate nearest neighbour. Watch the walk:

An HNSW graph, three layers deep — sparse on top, dense on the bottom, with a few “elevator” nodes wiring them together. The search greedily hops toward the query and descends layer by layer, greening only the nodes it lands on. The point: a full scan is memory-bandwidth-bound over every vector; HNSW reads a handful of nodes instead, trading a sliver of recall for orders-of-magnitude fewer memory reads. And on a utilization dashboard, neither version looks like anything but a “busy” machine — the difference in real work done is completely invisible.

Green = visited. Watch how few of the drawn nodes the walk actually touches.

02 The bill retrieval hands to the GPU

Here is the part the architecture diagrams leave out. Retrieval’s job is to put more text in front of the model — often thousands of tokens of documents. All of it becomes prefill, and all of it becomes KV cache that every generated token must then read. So the cost of “just retrieve more context” doesn’t land in the cheap CPU search step; it lands in the memory-bound generate step, the most expensive workload in the stack.

  what you tune            where it lands on the GPU
  ─────────────────        ──────────────────────────────────────
  retrieve top-3    →      +1.5k ctx tokens → small prefill, small KV
  retrieve top-20   →      +10k  ctx tokens → big prefill, big KV cache
  “stuff everything” →     +100k ctx tokens → prefill dominates, KV blows
                                              the batch size you can serve
  cheap CPU knob    ─────────────────────▶   expensive memory-bound bill

This is why RAG and the KV-cache tricks of layer 02 are the same conversation. Grouped-query and latent attention exist precisely so that a long retrieved context doesn’t blow the memory budget. Retrieval sets the context length; the attention architecture and the serving stack decide whether that length is affordable.

03 Why an observability company cares

the plasmient thread

RAG is a lesson in looking past the obvious box. The stage with all the attention — vector search — is a memory-bound CPU scan that a roofline explains completely. The stages that actually load your GPUs are the embedding and rerank forwards (which may or may not reach the compute ceiling, depending on batch size) and the ballooned generate (memory-bound, and now more so). Whether your retrieval stack is efficient is entirely a question of which stage is bound by what — and that is invisible on a utilization dashboard that shows one number for the whole pipeline.

That per-stage, per-bound view is Plasmient’s whole pitch. Every forward here sits on the roofline; the generate step is layer 05; the context length it inherits was set by the attention design inlayer 02. Retrieval is where the cheap knob writes the expensive check — and knowing that requires joining the layers, which is the product. Back to theFrontier Stack.

↗ Full references — go to the source

Lewis 2020 names RAG; Karpukhin’s DPR is the embedding recipe; HNSW and the Faiss paper are the search index. Nogueira & Cho is the rerank stage. Start with the RAG and HNSW papers, then Faiss is the code it all runs on.