Retrieval-Augmented Generation as a Distributed System

"The model asked me one question. I asked nine billion vectors. We agreed to split the difference inside nine hundred milliseconds, and somehow nobody noticed how many of us it took."

A Vector Shard That Answered in Time

A language model knows only what was in its training data, frozen at a point in time, and it cannot cite a source it never saw. Retrieval-augmented generation, or RAG, fixes both problems by giving the model fresh, specific text at query time: before the model answers, a retrieval system finds the most relevant passages from a corpus and places them in the prompt. The idea is simple to state and is often demonstrated on a laptop with a few thousand documents. That demonstration hides the engineering. A production RAG system serves a corpus that does not fit on one machine, answers more queries per second than one server can handle, and must do both inside a strict latency budget. The moment any of those three pressures binds, and in a real deployment all three bind at once, RAG stops being a prompt-engineering trick and becomes a distributed system. This chapter is about that system; the quality of the retrieved text matters, but the distribution is the point.

Those three pressures are the same ones that opened this book in Section 1.1: data too big for one machine, model too big for one device, throughput beyond one server. Retrieval makes all three concrete. The corpus and its embeddings are the data; a billion passages at 768 dimensions in 32-bit floats is roughly three terabytes of vectors alone, before any index structure, which is why the index must be sharded across machines. The language model is the model, already distributed across a fleet by Chapter 24. And every user-facing query competes for the same index and the same model, so every stage must also be replicated for throughput. RAG is therefore not one axis of distribution but several, stacked, which is exactly why it earns its own chapter rather than a paragraph.

1. The Four Stages, Each Its Own Subsystem Beginner

Trace a single query through the pipeline and four distinct distributed systems light up in turn. Figure 25.1.1 lays them out left to right; we walk them in the same order. The first stage is embedding. The raw query text is meaningless to a vector index, so a transformer encoder maps it to a dense vector in the same space as the corpus embeddings. That encoder is replicated across machines because query volume exceeds one accelerator's throughput, and the corpus side of the same embedding job, encoding billions of passages offline, is a large distributed batch pipeline in its own right, the subject of Section 25.2.

The second stage is search. The query vector is compared against every corpus vector, in effect, to find the nearest neighbors, except that exact comparison against billions of vectors is far too slow, so the system uses an approximate nearest-neighbor (ANN) index that trades a little recall for orders of magnitude less work. That index does not fit on one machine, so it is partitioned into shards, each holding a slice of the corpus, and a query is scattered to all shards and the partial results gathered back, a pattern Sections 25.4 and 25.5 develop in full. The third stage is reranking. ANN search is fast but coarse; a heavier cross-encoder rescores the top candidates with the query and passage read together, which is more accurate but too expensive to run over the whole corpus, so it runs only over the few dozen candidates that search returned, on its own replicated service (Section 25.7). The fourth stage is generation: the reranked passages are concatenated into the prompt and sent to the distributed language model of Chapter 24, which produces the answer.

The crucial observation in Figure 25.1.1 is that the four boxes do not share an algorithm; they share a deadline. The embedding service knows nothing about graph-based indexes; the ANN shards know nothing about transformer decoding. What couples them is that a user is waiting, and the four stages run in series, so their latencies add. That sum, and not any single stage's cleverness, is what the architect must hold under control. The next section makes the sum precise.

2. The End-to-End Latency Budget Intermediate

Because the stages run in series, the user-perceived latency of one query is the sum of the per-stage latencies plus the network hops between them. Writing $t_{\text{embed}}$, $t_{\text{search}}$, $t_{\text{rerank}}$, and $t_{\text{gen}}$ for the four stage latencies, the end-to-end time must fit a service-level objective (SLO):

This is a budget in the literal sense: a fixed quantity of milliseconds to be allocated across stages, where spending more on one leaves less for the others. The subtlety is that the relevant number is not the average latency of each stage but its tail. A query waits for the slowest of the $K$ ANN shards, not the average shard, so $t_{\text{search}}$ is governed by the tail of the shard-latency distribution, and the more shards a query fans out to, the worse the slowest-of-$K$ tends to be. The same logic applies wherever a stage fans out. The honest budget is therefore written in tail latency, typically the 95th or 99th percentile:

and the gap between the expected maximum and the maximum expectation is exactly the tail-amplification effect that Section 5.3 analyzes as the central hazard of fan-out systems. Budgeting RAG in average latency and then being surprised by the p99 is the most common way these pipelines miss their SLO. We allocate the budget at the tail and measure it at the tail.

To see where the budget actually goes, and how it shifts as the corpus grows, we model the four stages with simple closed-form latencies and sum them across corpus sizes from one million to ten billion vectors. The model is deliberately coarse; its job is to show the shape of the budget, namely which stage dominates and how that stays stable as data scales, not to predict any specific deployment to the millisecond.

import math

# Per-stage latency models (milliseconds), all simple closed forms so the
# numbers are reproducible and the dominant term is easy to read off.

def embed_ms():
    # One query embedded on a replicated encoder: fixed forward pass + network
    # hop to the embedding service. Independent of corpus size.
    forward = 6.0          # transformer encoder forward pass for one short query
    network_hop = 1.5      # client -> embedding replica round trip
    return forward + network_hop

def search_ms(n_vectors, shards):
    # Sharded ANN index. Each shard searches n_vectors/shards points; graph-based
    # ANN visits ~ c * log(points_per_shard) nodes, then a scatter-gather merges
    # the per-shard top-k. Fan-out adds one network round trip plus tail-straggler
    # slack that grows slowly with shard count.
    per_shard = n_vectors / shards
    visit = 0.45 * math.log2(max(per_shard, 2.0))   # ANN node visits per shard
    fanout = 1.5                                     # scatter to shards round trip
    straggler = 0.25 * math.log2(max(shards, 2.0))   # slowest-of-S tail slack
    merge = 0.4                                       # top-k merge of shard results
    return visit + fanout + straggler + merge

def rerank_ms(candidates):
    # Cross-encoder reranker scores the candidate passages in a batch on a
    # dedicated replica; cost is roughly linear in the candidate count.
    per_candidate = 0.35
    network_hop = 1.5
    return per_candidate * candidates + network_hop

def generate_ms(prompt_tokens, output_tokens):
    # Distributed LLM serving (Chapter 24): prefill over the retrieved context,
    # then autoregressive decode. Decode dominates and is corpus-independent.
    prefill = 0.02 * prompt_tokens     # context (query + retrieved passages)
    decode = 11.0 * output_tokens      # 11 ms per output token, served distributed
    return prefill + decode

SLO_MS = 900.0          # end-to-end p95 latency budget for the answer
SHARDS_PER_BILLION = 100  # we add index shards as the corpus grows
CANDIDATES = 50          # top-k passages handed to the reranker
PROMPT_TOKENS = 1200     # query + retrieved passages fed to the LLM
OUTPUT_TOKENS = 64       # generated answer length

corpus_sizes = [1_000_000, 10_000_000, 100_000_000, 1_000_000_000, 10_000_000_000]

print(f"{'corpus':>14} {'shards':>7} {'embed':>7} {'search':>7} "
      f"{'rerank':>7} {'gen':>7} {'total':>8} {'dominant':>9} {'SLO':>5}")
print("-" * 86)
for n in corpus_sizes:
    shards = max(1, round(n / 1_000_000_000 * SHARDS_PER_BILLION))
    e = embed_ms()
    s = search_ms(n, shards)
    r = rerank_ms(CANDIDATES)
    g = generate_ms(PROMPT_TOKENS, OUTPUT_TOKENS)
    total = e + s + r + g
    stages = {"embed": e, "search": s, "rerank": r, "generate": g}
    dominant = max(stages, key=stages.get)
    ok = "ok" if total <= SLO_MS else "OVER"
    print(f"{n:>14,} {shards:>7} {e:>7.1f} {s:>7.1f} "
          f"{r:>7.1f} {g:>7.1f} {total:>8.1f} {dominant:>9} {ok:>5}")

        corpus  shards   embed  search  rerank     gen    total  dominant   SLO
--------------------------------------------------------------------------------------
     1,000,000       1     7.5    11.1    19.0   728.0    765.6  generate    ok
    10,000,000       1     7.5    12.6    19.0   728.0    767.1  generate    ok
   100,000,000      10     7.5    13.2    19.0   728.0    767.7  generate    ok
 1,000,000,000     100     7.5    14.0    19.0   728.0    768.5  generate    ok
10,000,000,000    1000     7.5    14.9    19.0   728.0    769.4  generate    ok

Output 25.1.1 carries the chapter's central practical lesson. As the corpus grows ten-thousand-fold, from one million to ten billion vectors, the total latency barely moves, because the search stage grows only logarithmically once the shard count tracks the corpus size, exactly the property that Section 25.5 engineers on purpose. The stage that dominates the budget is generation, the distributed language model from Chapter 24, and it is corpus-independent. This is liberating: it means the retrieval half of RAG, the part this chapter is about, can scale to web-sized corpora without eating the latency budget, provided the index is sharded well. The danger is not that search gets slow as data grows; it is that an unsharded index makes search grow linearly, at which point a few hundred million vectors blow the budget on their own. The whole point of distributing the index is to convert that linear growth into the flat line you see above.

3. Why Retrieval Must Scale Out at All Beginner

It is worth pausing on why retrieval cannot simply live on one large machine, because the answer determines the shape of the whole chapter. Two ceilings force it out, and they are the same data and throughput ceilings from Section 1.1, now wearing retrieval's clothes. The first ceiling is the corpus and its embeddings. A web-scale text corpus runs to billions of passages, and each passage becomes a dense vector. At 768 dimensions and four bytes per number, a billion vectors is about three terabytes before the index adds its own overhead, and a graph or quantized index can multiply that. No single machine holds that in memory at search speed, so the index is partitioned: shard $j$ holds a disjoint slice of the corpus, and a query is searched against all shards in parallel and their results merged. This is the same partitioning-and-sharding arc the book has followed from Chapter 11's sharded embedding tables onward, now applied to a search index.

The second ceiling is query throughput. A single index replica can sustain only so many searches per second; a service answering thousands of user queries per second must therefore replicate the index, running several copies behind a load balancer so that capacity scales with traffic. Sharding and replication are orthogonal: sharding splits one logical index across machines so it fits and so each query does less work, while replication duplicates the sharded index so more queries can run at once. A production index is laid out as a grid, sharded one way and replicated the other, which is precisely the design space Section 25.5 opens up. The classical foundations for the search inside each shard, locality-sensitive hashing and approximate nearest neighbors, were built in Section 12.7, and the distributed top-$k$ gather that merges shard results is the same reduce pattern as the MapReduce aggregation of Section 6.7; this chapter scales those ideas to the index and binds them to a latency budget.

4. The Map of the Chapter Beginner

Each stage of the pipeline in Figure 25.1.1 expands into its own section, and a few cross-cutting concerns get sections of their own. The embedding stage, including the offline pipeline that encodes the whole corpus into vectors, is Section 25.2. The store that holds those vectors and answers searches, the vector database, is Section 25.3. The approximate-nearest-neighbor algorithms that make search fast inside one shard, HNSW graphs, inverted files (IVF), and product quantization (PQ), are Section 25.4, and the way those indexes are sharded and replicated across machines is Section 25.5, the structural heart of the chapter.

Beyond pure vector search, Section 25.6 covers distributed hybrid search, combining dense vectors with classical keyword retrieval so the two signals reinforce each other. The reranking stage and its distributed, multi-stage refinement is Section 25.7. Because the same queries and the same hot passages recur, Section 25.8 treats distributed caching for retrieval, the layer that lets a system skip work it has already done. Finally, Section 25.9 closes the loop with evaluation, how to measure recall, latency, and freshness across a sharded, replicated retrieval fleet, building directly on the evaluation methodology of Chapter 5. Once the pieces are in place, the web-scale and agentic case studies of Chapter 36 and Chapter 40 assemble them into complete systems serving real traffic.

# Each component is a handle to a distributed service, not a local object.
from llama_index.core import VectorStoreIndex
from llama_index.core.postprocessor import SentenceTransformerRerank

index = VectorStoreIndex.from_vector_store(sharded_vector_store)  # 25.3-25.5
reranker = SentenceTransformerRerank(top_n=40, model="BAAI/bge-reranker-base")  # 25.7

query_engine = index.as_query_engine(
    similarity_top_k=200,          # ANN search breadth, fanned out to shards
    node_postprocessors=[reranker],# rerank stage trims 200 -> 40 candidates
    llm=distributed_llm,           # Chapter 24's served model
)
answer = query_engine.query("How do index shards bound search latency?")

5. A Frontier in Motion: Retrieve, or Just Read Everything? Advanced

The premise of RAG, that you retrieve a small relevant slice rather than feed the model the whole corpus, is under active pressure from two directions at once, and the resolution is shaping how these systems are built right now. It is worth knowing where the argument stands before committing a design to it.

Whichever way these debates settle, the distributed machinery is the durable part. A longer context window still needs a fast way to choose what goes in it; an agent that retrieves ten times per task needs each retrieval to cost a tenth of the budget. The sections that follow build that machinery stage by stage, starting with the pipeline that turns a corpus into the vectors everything else searches. We have the pipeline, the budget that governs it, and the map; Section 25.2 opens the first stage, the distributed embedding pipeline.