Embedding similarity search pipeline

Every RAG system, every semantic search bar, and every recommendation feed depends on the same two-step pipeline: turn things into vectors, then find the closest vectors fast. I have been surprised by how much inference engineering goes into making this pipeline work at scale, so today I am building it end to end.

Step 1: embedding model inference

Embedding models like E5, GTE, BGE, and Cohere Embed are encoder-only transformers (usually based on BERT or its descendants). Unlike autoregressive LLMs, they process the entire input in a single forward pass and produce a fixed-size vector. This has important implications for serving:

No decode phase. Every request is a single prefill. There is no autoregressive loop, no KV cache that grows over time. The GPU does one forward pass and is done.
Uniform compute. Requests with similar token counts have nearly identical latency. This makes batching much simpler than LLM serving.
High throughput potential. Because there is no sequential decode, you can push very large batch sizes and saturate GPU compute.

A typical embedding model (e.g., GTE-large with 326M parameters) processes a batch of 256 sequences of 512 tokens on an A100 in about 80 ms. That is over 3,000 embeddings per second on a single GPU. Compare that to an LLM generating 50 tokens per second per request, and you see why embedding workloads need a different serving architecture.

The pooling step

The transformer outputs one vector per token. To get a single embedding for the whole input, you need a pooling strategy:

def mean_pooling(token_embeddings, attention_mask):
    """Mean pooling: average all non-padding token embeddings."""
    mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size())
    sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
    sum_mask = mask_expanded.sum(dim=1).clamp(min=1e-9)
    return sum_embeddings / sum_mask

def cls_pooling(token_embeddings):
    """CLS pooling: take the [CLS] token embedding."""
    return token_embeddings[:, 0, :]

Most modern embedding models use mean pooling because it captures information from all tokens. Some older models use CLS pooling. Always check the model card. Using the wrong pooling strategy will silently produce bad embeddings with no error message.

Normalization matters

After pooling, most embedding models L2-normalize the vectors so that cosine similarity equals dot product. If you skip normalization, your similarity scores will be wrong and your retrieval quality will degrade. Always normalize unless the model explicitly says not to.

Step 2: vector indexing

Once you have embeddings, you need to store and search them efficiently. Brute-force search (compute dot product with every vector in the database) is exact but scales linearly. At 10 million vectors of dimension 1024 in FP32, that is 40 GB of data and roughly 10 billion multiply-add operations per query. Possible on a GPU, but not practical for sub-10ms latency at scale.

Approximate nearest neighbor (ANN) algorithms trade a small amount of recall for massive speedups. The three dominant approaches:

IVF (Inverted File Index). Cluster the vectors into, say, 4,096 centroids using k-means. At query time, find the nearest centroids and search only those clusters. Reduces search space by 100x or more.
HNSW (Hierarchical Navigable Small World). Build a multi-layer graph where each node connects to nearby vectors. Query traverses the graph from top layer to bottom, narrowing the search at each level. Very fast queries (sub-millisecond) but high memory overhead because you store the graph structure.
Product Quantization (PQ). Compress each vector from 1024 x 4 bytes = 4 KB down to, say, 64 bytes by splitting it into subvectors and quantizing each subvector to a codebook index. Massively reduces memory and enables fast distance computation via lookup tables.

In practice, you combine them. FAISS (Facebook AI Similarity Search) supports composite index types like IVF4096,PQ64, which means: cluster into 4,096 cells, then PQ-compress vectors within each cell to 64 bytes.

Building the pipeline with FAISS

import faiss
import numpy as np

# Assume we have 10M embeddings of dimension 1024
d = 1024
n = 10_000_000

# Step 1: Train the index on a sample
nlist = 4096  # number of IVF clusters
quantizer = faiss.IndexFlatIP(d)  # inner product for normalized vectors
index = faiss.IndexIVFPQ(quantizer, d, nlist, 64, 8)
# 64 subquantizers, 8 bits each -> 64 bytes per vector

# Train on a representative sample (100K to 1M vectors)
training_vectors = load_sample(100_000)  # np.float32, shape (100000, 1024)
index.train(training_vectors)

# Step 2: Add all vectors
for batch in load_vectors_in_batches(n, batch_size=100_000):
    index.add(batch)

# Step 3: Search
index.nprobe = 64  # search 64 of 4096 clusters
queries = encode_queries(["what is context parallelism?"])  # shape (1, 1024)
distances, indices = index.search(queries, k=10)  # top-10 results

The nprobe parameter controls the recall/speed tradeoff. At nprobe=1, you search only the nearest cluster (fast but may miss results). At nprobe=64, you search 64 clusters (slower but recall above 95%). Tuning nprobe is the single most impactful knob.

GPU-accelerated search

FAISS supports GPU-accelerated search, which makes a dramatic difference at high query throughput. Moving the index to GPU:

gpu_resources = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(gpu_resources, 0, index)
# Same search API, but 10-50x faster for batch queries

On an A100, GPU FAISS can process 10,000+ queries per second against an index of 100 million vectors. The bottleneck shifts from search to embedding inference, which is why you want to batch the embedding model aggressively.

End-to-end latency breakdown

For a single semantic search query hitting a system with 10M documents:

Tokenization: 0.1 ms
Embedding inference (single query, GTE-large, A100): 2 to 5 ms
L2 normalization: negligible
ANN search (IVF4096,PQ64, nprobe=64, GPU): 0.5 to 1 ms
Document fetch (top-10 from storage): 1 to 5 ms
Total: 4 to 11 ms

That is fast enough for real-time search. The embedding inference dominates. If you are serving high QPS, batch queries together and amortize the GPU kernel launch overhead. A batch of 32 queries takes only slightly longer than a single query.

The embedding pipeline is the one place in AI inference where you can genuinely achieve single-digit millisecond latency at scale. No autoregressive decode, no KV cache management, just encode and search.

Tomorrow we look at the opposite extreme: speech-to-speech latency, where every millisecond of delay is perceptible to the user and the pipeline spans multiple models.