Your RAG pipeline works beautifully in the notebook. Queries return accurate, grounded responses. The retrieval quality is solid. Then you deploy it, real users start hitting it, and the first complaint arrives: "Why does it take four seconds to respond?" You look at your traces and realize the LLM inference is only 800ms — the retrieval step is eating the other 3.2 seconds. You've built a retrieval-augmented generation system that's bottlenecked by the retrieval.
This is the production reality that most RAG tutorials skip entirely. Getting retrieval correct is a different problem from getting retrieval fast. The two objectives sometimes pull against each other, and navigating that tension — knowing exactly which knobs to turn, how far to turn them, and what you're sacrificing — is what separates a proof-of-concept from a system that can serve thousands of concurrent users under SLA. Sub-100ms retrieval is achievable for most production workloads, but it requires understanding what's actually happening inside your vector store, your index structures, and your application's request lifecycle.
By the end of this lesson, you'll be able to design and tune a retrieval layer that consistently delivers results in under 100ms for collections in the tens-of-millions range, implement multi-tier caching that eliminates redundant compute without sacrificing freshness, and make principled trade-offs between recall and latency when your use case demands it.
What you'll learn:
Before tuning anything, you need an honest model of where time actually goes. Most engineers assume retrieval latency is dominated by the vector similarity search itself. It usually isn't. Here's a more realistic breakdown for a typical self-hosted Qdrant deployment serving a 5M-vector collection:
Network round-trip (client → vector store): 8ms
Embedding inference (query text → vector): 45ms
Vector similarity search: 12ms
Payload fetching + serialization: 9ms
Network return (vector store → client): 8ms
─────────────────────────────────────────────────
Total: 82ms
The vector search itself is 12ms. The embedding inference is 45ms. If you spend all your energy tuning HNSW parameters, you're optimizing the second-smallest component of your latency budget. This is why a holistic view matters — you need to attack every component, not just the one that sounds most technical.
The numbers above shift dramatically based on your infrastructure choices:
Your first step in any optimization effort is to instrument your pipeline and decompose actual latency, not estimate it. We'll cover how to do that in the benchmarking section.
Approximate nearest neighbor (ANN) search works by trading guaranteed recall for dramatically reduced search time. Two index families dominate production usage: Hierarchical Navigable Small World graphs (HNSW) and Inverted File Indexes (IVF). Understanding their internal mechanics is prerequisite to tuning them intelligently.
HNSW builds a layered graph structure. Imagine a city: at the top layer, you have a sparse graph connecting only major hubs — like international airports. At the bottom layer, you have a dense graph connecting every point to its nearest neighbors. When you query, you enter at the top and greedily descend: find the best neighbor in the sparse layer, jump down, find the best neighbor in the next layer, and so on until you reach the bottom, where you do a local exhaustive search.
The two key parameters are:
M (number of bidirectional connections per node): During index construction, each new vector is connected to M existing vectors. Higher M increases graph connectivity, which improves recall but increases memory usage and construction time. For 768-dim embeddings, values between 16 and 64 are typical. Lower than 16 and you'll see recall drop significantly; higher than 64 and you're paying memory costs that rarely justify themselves.
ef_construction (search depth during index building): When inserting a vector, the algorithm searches for ef_construction candidates to connect to. Higher values produce better-quality graphs (better recall at search time) at the cost of slower index building. Typical range: 100–400.
ef (search depth at query time): This is the most important runtime parameter. It controls how many candidates are explored during search. This is the primary lever for the recall/latency trade-off. Higher ef → more recall, more latency.
import hnswlib
import numpy as np
# Building an HNSW index for a document corpus
# Assume we have 1M vectors of dim 768 (e.g., from all-mpnet-base-v2)
dim = 768
num_vectors = 1_000_000
index = hnswlib.Index(space='cosine', dim=dim)
# ef_construction=200, M=32: good balance for production
# Building: ~15 minutes on a 16-core machine
# Memory footprint: ~4.5GB
index.init_index(max_elements=num_vectors, ef_construction=200, M=32)
# Index construction — do this once, persist to disk
vectors = np.load('corpus_embeddings.npy') # shape: (1M, 768)
ids = np.arange(num_vectors)
index.add_items(vectors, ids, num_threads=16)
index.save_index('production_index.bin')
# At query time, load and set ef
index.load_index('production_index.bin', max_elements=num_vectors)
# ef=50: ~92% recall, ~4ms search on 1M vectors (CPU)
# ef=200: ~98% recall, ~14ms search on 1M vectors (CPU)
# ef=500: ~99.5% recall, ~35ms search on 1M vectors (CPU)
index.set_ef(50) # Set based on your recall requirements
query_vector = np.random.randn(1, 768).astype(np.float32)
labels, distances = index.knn_query(query_vector, k=10)
The critical insight here is that ef is a runtime setting you can change without rebuilding the index. This means you can A/B test recall/latency trade-offs in production without downtime — an important operational advantage.
Inverted File Index works by clustering your vector space into nlist Voronoi cells during index construction. At query time, you compute the distance from your query to each cluster centroid, then only search within the nprobe closest clusters.
nlist (number of clusters): The square root of your corpus size is a reasonable starting point. For 1M vectors: nlist=1000. More clusters → finer partitioning → faster search (you search fewer vectors per probe) but worse recall if vectors are near cell boundaries.
nprobe (number of clusters to search): This is the IVF equivalent of HNSW's ef — your primary runtime latency/recall lever. nprobe=1 is fastest but gives poor recall. nprobe=nlist is exact search. For most production use cases, nprobe between 10 and 100 gets you to 90-99% recall.
import faiss
import numpy as np
dim = 768
num_vectors = 1_000_000
nlist = 1024 # ~sqrt(1M)
# IVF with product quantization for memory efficiency
# IVF1024,PQ32 is a common production recipe
quantizer = faiss.IndexFlatIP(dim) # Inner product (equivalent to cosine on normalized vecs)
index = faiss.IndexIVFPQ(quantizer, dim, nlist, 32, 8)
# 32 = number of sub-quantizers, 8 = bits per sub-quantizer
# Training is required for IVF — needs representative sample
# Rule of thumb: 39 * nlist training vectors minimum
training_sample = np.load('training_vectors.npy') # ~50k vectors
index.train(training_sample)
# Add all vectors
vectors = np.load('corpus_embeddings.npy')
index.add(vectors)
faiss.write_index(index, 'production_ivfpq.index')
# At query time
index = faiss.read_index('production_ivfpq.index')
index.nprobe = 32 # Search 32 of 1024 clusters
query_vector = np.random.randn(1, 768).astype(np.float32)
faiss.normalize_L2(query_vector) # Normalize for cosine similarity
distances, labels = index.search(query_vector, k=10)
HNSW vs. IVF: When to choose which?
| Dimension | HNSW | IVF |
|---|---|---|
| Memory per vector | Higher (graph edges + vectors) | Lower (especially IVF+PQ) |
| Search latency (low concurrency) | Better | Slightly worse |
| Search latency (high concurrency) | Similar | Better (more parallelizable) |
| Recall at fixed latency budget | Better | Slightly lower |
| Index build time | Slower | Faster |
| Dynamic updates | Good (add without rebuild) | Poor (requires rebuild or IDMap) |
| Corpus size > 100M vectors | Challenging (memory) | Better with PQ compression |
For most RAG applications where the corpus fits in RAM and updates are incremental, HNSW is the better default. IVF+PQ becomes compelling when you're dealing with hundreds of millions of vectors or when memory cost is the binding constraint.
The most common mistake in latency optimization is benchmarking with a single thread, measuring mean latency, and declaring victory. Production systems don't work this way. You need to measure:
Here's a benchmarking harness you can actually use:
import time
import concurrent.futures
import numpy as np
from collections import defaultdict
import hnswlib
def benchmark_retrieval(index_path: str, query_vectors: np.ndarray,
k: int = 10, ef_values: list = None,
concurrency: int = 20, num_requests: int = 1000):
"""
Benchmark HNSW retrieval across different ef values under concurrent load.
Simulates a production environment where multiple requests arrive simultaneously.
"""
if ef_values is None:
ef_values = [32, 64, 128, 256]
dim = query_vectors.shape[1]
index = hnswlib.Index(space='cosine', dim=dim)
index.load_index(index_path, max_elements=5_000_000)
results = {}
for ef in ef_values:
index.set_ef(ef)
latencies = []
def single_query(query_vec):
start = time.perf_counter()
labels, distances = index.knn_query(
query_vec.reshape(1, -1), k=k
)
end = time.perf_counter()
return (end - start) * 1000 # ms
# Sample query vectors randomly to simulate realistic traffic
query_sample = query_vectors[
np.random.choice(len(query_vectors), num_requests, replace=True)
]
with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
futures = [executor.submit(single_query, q) for q in query_sample]
latencies = [f.result() for f in concurrent.futures.as_completed(futures)]
latencies.sort()
results[ef] = {
'p50': np.percentile(latencies, 50),
'p95': np.percentile(latencies, 95),
'p99': np.percentile(latencies, 99),
'max': max(latencies),
'qps': num_requests / (sum(latencies) / 1000 / concurrency)
}
print(f"ef={ef:4d} | p50={results[ef]['p50']:6.1f}ms | "
f"p95={results[ef]['p95']:6.1f}ms | p99={results[ef]['p99']:6.1f}ms | "
f"QPS={results[ef]['qps']:6.0f}")
return results
# Typical output for 1M vectors, 768 dims, 20 concurrent threads:
# ef= 32 | p50= 3.1ms | p95= 5.8ms | p99= 8.2ms | QPS= 4821
# ef= 64 | p50= 5.9ms | p95= 10.1ms | p99= 14.3ms | QPS= 2890
# ef= 128 | p50= 11.2ms | p95= 18.7ms | p99= 24.1ms | QPS= 1654
# ef= 256 | p50= 21.8ms | p95= 34.2ms | p99= 42.6ms | QPS= 891
Notice that p99 latency is often 2-3x the p50. This is the number you need to design against for SLAs. An application that "usually" responds in 15ms but has a p99 of 40ms will still frustrate users on 1 in 100 requests.
Also notice the non-linear relationship: doubling ef doesn't double latency because of how the graph traversal works at different depths. This is why you need to benchmark your specific index rather than extrapolating from published benchmarks.
Warning: Never benchmark retrieval by sending requests sequentially from a single thread. HNSW and most vector stores show dramatically different behavior under concurrent load due to cache effects, NUMA topology, and thread contention. Your single-threaded benchmark might show 3ms p99; your 20-thread benchmark might show 25ms p99. Test at production-realistic concurrency from day one.
If your embedding inference is taking 40-80ms on CPU, fixing that alone gets you within budget. Here's how:
Converting your embedding model to ONNX and quantizing it to INT8 typically reduces inference latency by 3-5x on CPU with less than 1% quality degradation for most sentence transformer models.
from sentence_transformers import SentenceTransformer
from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer
import numpy as np
import time
# Export to ONNX with INT8 quantization
# Do this once, save the artifacts
model_name = "sentence-transformers/all-mpnet-base-v2"
# Using Hugging Face Optimum
from optimum.onnxruntime import ORTOptimizer
from optimum.onnxruntime.configuration import AutoQuantizationConfig
# Export base ONNX model
ort_model = ORTModelForFeatureExtraction.from_pretrained(
model_name,
export=True
)
ort_model.save_pretrained("./embedding_model_onnx")
# Quantize to INT8
optimizer = ORTOptimizer.from_pretrained(ort_model)
qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)
optimizer.quantize(
quantization_config=qconfig,
save_dir="./embedding_model_int8"
)
# Benchmark: original vs quantized
tokenizer = AutoTokenizer.from_pretrained(model_name)
texts = ["What are the key features of transformer architecture?"] * 32 # batch of 32
# Original model
original_model = SentenceTransformer(model_name)
start = time.perf_counter()
for _ in range(100):
embeddings = original_model.encode(texts, batch_size=32, normalize_embeddings=True)
original_latency = (time.perf_counter() - start) / 100 * 1000
print(f"Original model: {original_latency:.1f}ms per batch")
# ONNX INT8 model
ort_quantized = ORTModelForFeatureExtraction.from_pretrained("./embedding_model_int8")
tokenized = tokenizer(texts, padding=True, truncation=True,
max_length=512, return_tensors="pt")
start = time.perf_counter()
for _ in range(100):
outputs = ort_quantized(**tokenized)
# Mean pooling
embeddings = outputs.last_hidden_state.mean(dim=1).detach().numpy()
onnx_latency = (time.perf_counter() - start) / 100 * 1000
print(f"ONNX INT8 model: {onnx_latency:.1f}ms per batch")
# Typical results:
# Original model: 187.3ms per batch
# ONNX INT8 model: 41.2ms per batch
For many RAG applications, query diversity is lower than you think. User questions about a customer support knowledge base cluster around a few hundred common phrasings. Caching the embedding computation for recent/frequent queries is often the highest-ROI optimization available.
import hashlib
import json
import time
from functools import lru_cache
import redis
import numpy as np
class EmbeddingCache:
"""
Two-tier embedding cache: in-process LRU + Redis for distributed deployments.
The in-process cache handles hot queries with zero network overhead.
Redis handles cross-process sharing and survives pod restarts.
"""
def __init__(self, redis_client: redis.Redis, ttl_seconds: int = 3600,
local_cache_size: int = 1024):
self.redis = redis_client
self.ttl = ttl_seconds
self._local_cache = {}
self._local_cache_size = local_cache_size
self._local_hits = 0
self._redis_hits = 0
self._misses = 0
def _cache_key(self, text: str, model_name: str) -> str:
"""Deterministic key that captures both text and model version."""
content = f"{model_name}:{text}"
return f"emb:{hashlib.sha256(content.encode()).hexdigest()}"
def get(self, text: str, model_name: str) -> np.ndarray | None:
key = self._cache_key(text, model_name)
# L1: in-process cache
if key in self._local_cache:
self._local_hits += 1
return self._local_cache[key]
# L2: Redis
cached = self.redis.get(key)
if cached is not None:
self._redis_hits += 1
vector = np.frombuffer(cached, dtype=np.float32)
# Promote to L1
self._set_local(key, vector)
return vector
self._misses += 1
return None
def set(self, text: str, model_name: str, embedding: np.ndarray):
key = self._cache_key(text, model_name)
# Store as raw bytes — more efficient than JSON
self.redis.setex(key, self.ttl, embedding.astype(np.float32).tobytes())
self._set_local(key, embedding)
def _set_local(self, key: str, embedding: np.ndarray):
# Simple LRU eviction: evict oldest when at capacity
if len(self._local_cache) >= self._local_cache_size:
oldest_key = next(iter(self._local_cache))
del self._local_cache[oldest_key]
self._local_cache[key] = embedding
@property
def hit_rate(self) -> dict:
total = self._local_hits + self._redis_hits + self._misses
if total == 0:
return {'l1': 0, 'l2': 0, 'miss': 0}
return {
'l1': self._local_hits / total,
'l2': self._redis_hits / total,
'miss': self._misses / total,
'total_requests': total
}
class CachedEmbeddingService:
def __init__(self, model, cache: EmbeddingCache, model_name: str):
self.model = model
self.cache = cache
self.model_name = model_name
def embed_query(self, text: str) -> np.ndarray:
# Normalize the text before caching — whitespace differences shouldn't miss
normalized = ' '.join(text.lower().strip().split())
cached = self.cache.get(normalized, self.model_name)
if cached is not None:
return cached
embedding = self.model.encode([normalized], normalize_embeddings=True)[0]
self.cache.set(normalized, self.model_name, embedding)
return embedding
Tip: When caching embeddings, normalize your input text before generating the cache key. Trim whitespace, lowercase, and collapse multiple spaces. A user typing "What is RAG?" and "what is rag?" should hit the same cache entry. For many domains, this alone doubles your cache hit rate.
Embedding caching is just one layer. A complete RAG caching strategy has three tiers:
Tier 1 — Embedding Cache: Cache query_text → embedding_vector. TTL can be long (hours to days) since the same embedding model produces the same output for the same input forever.
Tier 2 — Retrieval Result Cache: Cache embedding_vector → [doc_ids, scores]. TTL must account for corpus updates — if you're adding new documents hourly, your retrieval results can be stale. This is the most dangerous cache to get wrong.
Tier 3 — Full Response Cache: Cache query_text → full_LLM_response. Very effective for FAQ-style queries with stable answers; terrible for queries that need current information.
import hashlib
import time
import redis
import numpy as np
from dataclasses import dataclass
from typing import Optional
@dataclass
class RetrievalResult:
doc_ids: list[int]
scores: list[float]
retrieved_at: float
cache_hit: bool = False
class RetrievalCache:
"""
Caches retrieval results keyed by a quantized embedding fingerprint.
Quantization is critical: even tiny floating-point differences in the
query vector shouldn't cause cache misses.
"""
def __init__(self, redis_client: redis.Redis, ttl_seconds: int = 300,
similarity_threshold: float = 0.99):
self.redis = redis_client
self.ttl = ttl_seconds
self.threshold = similarity_threshold
def _vector_fingerprint(self, embedding: np.ndarray, precision: int = 3) -> str:
"""
Create a cache key from an embedding by rounding to `precision` decimal places.
This allows semantically identical queries with minor float differences to hit cache.
Precision=3 gives ~0.001 per-dimension variation tolerance.
"""
rounded = np.round(embedding, decimals=precision)
# Hash the rounded vector bytes for a compact key
return hashlib.md5(rounded.astype(np.float32).tobytes()).hexdigest()
def get(self, embedding: np.ndarray, k: int) -> Optional[RetrievalResult]:
key = f"retrieval:{self._vector_fingerprint(embedding)}:k{k}"
cached = self.redis.get(key)
if cached is None:
return None
data = json.loads(cached)
return RetrievalResult(
doc_ids=data['doc_ids'],
scores=data['scores'],
retrieved_at=data['retrieved_at'],
cache_hit=True
)
def set(self, embedding: np.ndarray, k: int, result: RetrievalResult):
key = f"retrieval:{self._vector_fingerprint(embedding)}:k{k}"
data = json.dumps({
'doc_ids': result.doc_ids,
'scores': result.scores,
'retrieved_at': result.retrieved_at
})
self.redis.setex(key, self.ttl, data)
def invalidate_on_corpus_update(self, updated_doc_ids: list[int]):
"""
When new documents are indexed, invalidate retrieval caches.
For bulk updates, nuclear option: flush entire retrieval cache namespace.
For surgical updates: scan and delete keys where result contains updated IDs.
"""
# Nuclear option for large batch updates
pattern = "retrieval:*"
cursor = 0
while True:
cursor, keys = self.redis.scan(cursor, match=pattern, count=500)
if keys:
self.redis.delete(*keys)
if cursor == 0:
break
The retrieval cache TTL deserves careful thought. Here's a framework:
Warning: Don't cache retrieval results if your corpus is updated frequently and users expect to retrieve new content. A user who just uploaded a document expects to be able to query it immediately. Retrieval caching with a 1-hour TTL will make them think the system is broken.
If you're using a managed or self-hosted Qdrant deployment (a common choice for RAG), several configuration decisions significantly impact latency that the documentation buries.
# qdrant/config/production.yaml
storage:
# CRITICAL: Use mmap for collections that don't fit fully in RAM
# Allows OS to handle paging efficiently while keeping hot data in memory
on_disk_payload: false # Keep payload in RAM if your RAM budget allows
# Performance mode: speed (RAM) vs memmap (disk-backed)
# Use "always_ram" for collections < available_ram * 0.7
vectors:
on_disk: false
service:
# Enable gRPC — significantly lower overhead than HTTP for high-QPS
grpc_port: 6334
# Thread pool: rule of thumb is 2x CPU cores for I/O-bound workloads
# For embedding-heavy workloads: match to CPU core count
max_workers: 32
# HNSW index parameters per collection (set at collection creation)
# These cannot be changed without rebuilding the index
hnsw_config:
m: 32
ef_construct: 200
full_scan_threshold: 10000 # Fall back to exact search for small collections
max_indexing_threads: 0 # 0 = use all available cores
on_disk: false # Keep HNSW graph in RAM
from qdrant_client import QdrantClient
from qdrant_client.models import (
VectorParams, Distance, HnswConfigDiff,
OptimizersConfigDiff, QuantizationConfig, ScalarQuantization,
ScalarType
)
client = QdrantClient(
host="localhost",
port=6333,
grpc_port=6334,
prefer_grpc=True, # Use gRPC for lower latency
timeout=5.0
)
# Create collection with production-optimized settings
client.create_collection(
collection_name="knowledge_base",
vectors_config=VectorParams(
size=768,
distance=Distance.COSINE,
),
hnsw_config=HnswConfigDiff(
m=32,
ef_construct=200,
on_disk=False, # Keep graph in RAM
),
# Scalar quantization: reduces memory by 4x with ~1-2% recall loss
# INT8 quantization is almost always worth it at scale
quantization_config=QuantizationConfig(
scalar=ScalarQuantization(
type=ScalarType.INT8,
quantile=0.99, # Clip outliers at 99th percentile
always_ram=True, # Keep quantized vectors in RAM
)
),
optimizers_config=OptimizersConfigDiff(
indexing_threshold=20000, # Start indexing after 20k vectors
memmap_threshold=200000, # Switch to mmap if segment exceeds 200k vectors
)
)
# Query with tuned parameters
def retrieve_documents(query_vector: list[float], k: int = 10,
ef: int = 128) -> list:
results = client.search(
collection_name="knowledge_base",
query_vector=query_vector,
limit=k,
search_params={
"hnsw_ef": ef, # Runtime ef override
"exact": False, # Always use ANN, never fall back to exact
},
with_payload=True,
with_vectors=False, # Don't return vectors unless needed — saves bandwidth
)
return results
Tip: Qdrant's scalar quantization (INT8) is one of the most impactful settings you can enable. It reduces memory usage by 4x, which means more of your index stays in RAM rather than being paged from disk. On collections of 10M+ vectors, this difference can mean the difference between 8ms and 80ms retrieval times.
Vector search is fundamentally a memory-bandwidth-bound workload. Understanding the hardware your index lives on is not optional for sub-100ms at scale.
Modern multi-socket servers have Non-Uniform Memory Access (NUMA) topology: each CPU socket has "local" memory it accesses quickly and "remote" memory on another socket that costs 1.5-3x more latency. If your HNSW index is loaded into memory on NUMA node 1 but your retrieval threads are pinned to NUMA node 0, you're paying this penalty on every cache miss.
# Check NUMA topology
numactl --hardware
# Bind your retrieval service to a specific NUMA node
# This ensures both the memory allocation and the execution threads are local
numactl --cpunodebind=0 --membind=0 python retrieval_service.py
# Or use taskset to pin threads explicitly (useful when running in containers)
taskset -c 0-15 python retrieval_service.py # Pin to first 16 cores
For very large HNSW indexes that you're memory-mapping from disk, Linux's transparent huge pages (THP) can help — but it's context-dependent. THP reduces TLB pressure for large sequential access patterns. For the random-access nature of graph traversal in HNSW, explicit huge pages often help more than THP.
# Check current huge page status
cat /proc/meminfo | grep -i huge
# Pre-allocate 2MB huge pages for a 40GB index
# Number of pages = index_size_bytes / 2097152
echo 20480 > /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages
# Mount hugetlbfs
mkdir /mnt/hugepages
mount -t hugetlbfs nodev /mnt/hugepages
# In your application, use mmap with MAP_HUGETLB
# FAISS supports this via the resource manager
import faiss
faiss.omp_set_num_threads(16) # Set FAISS thread count
# Read index into huge-page-backed memory
# (This requires mmap support in your loading code)
index = faiss.read_index('large_index.bin', faiss.IO_FLAG_MMAP)
Both FAISS and hnswlib use OpenMP for parallelism. The right number of threads depends on your concurrency model:
import os
import faiss
import hnswlib
# For high-concurrency serving, limit parallelism per query
# and rely on the server's concurrency instead
os.environ["OMP_NUM_THREADS"] = "2" # 2 threads per query
faiss.omp_set_num_threads(2)
# Then serve with a thread pool that matches your CPU count
# FastAPI + uvicorn with multiple workers handles this naturally
Putting the pieces together, here's a retrieval service architecture that achieves sub-100ms end-to-end:
import asyncio
import time
import numpy as np
import redis
from fastapi import FastAPI
from pydantic import BaseModel
from contextlib import asynccontextmanager
from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer
import hnswlib
import logging
logger = logging.getLogger(__name__)
# ─── Configuration ────────────────────────────────────────────────────────────
class Config:
MODEL_PATH = "./embedding_model_int8"
INDEX_PATH = "./production_index.bin"
INDEX_MAX_ELEMENTS = 5_000_000
EMBEDDING_DIM = 768
DEFAULT_K = 10
HNSW_EF = 64 # p95 < 15ms, recall ~95%
REDIS_URL = "redis://localhost:6379"
EMBEDDING_CACHE_TTL = 7200 # 2 hours
RETRIEVAL_CACHE_TTL = 300 # 5 minutes
LOCAL_CACHE_SIZE = 2048
# ─── Global state (loaded at startup) ────────────────────────────────────────
app_state = {}
@asynccontextmanager
async def lifespan(app: FastAPI):
# Load everything into memory at startup — not on first request
logger.info("Loading embedding model...")
app_state['tokenizer'] = AutoTokenizer.from_pretrained(Config.MODEL_PATH)
app_state['embedding_model'] = ORTModelForFeatureExtraction.from_pretrained(
Config.MODEL_PATH
)
logger.info("Loading HNSW index...")
index = hnswlib.Index(space='cosine', dim=Config.EMBEDDING_DIM)
index.load_index(Config.INDEX_PATH, max_elements=Config.INDEX_MAX_ELEMENTS)
index.set_ef(Config.HNSW_EF)
app_state['index'] = index
logger.info("Connecting to Redis...")
app_state['redis'] = redis.from_url(Config.REDIS_URL, decode_responses=False)
app_state['embedding_cache'] = EmbeddingCache(
app_state['redis'],
ttl_seconds=Config.EMBEDDING_CACHE_TTL,
local_cache_size=Config.LOCAL_CACHE_SIZE
)
logger.info("Retrieval service ready.")
yield
# Cleanup
app_state['redis'].close()
app = FastAPI(lifespan=lifespan)
# ─── Request/Response models ──────────────────────────────────────────────────
class RetrievalRequest(BaseModel):
query: str
k: int = Config.DEFAULT_K
ef: int | None = None # Optional per-request ef override
class RetrievalResponse(BaseModel):
doc_ids: list[int]
scores: list[float]
latency_ms: float
cache_hit: str # "embedding", "retrieval", "none"
# ─── Core retrieval logic ──────────────────────────────────────────────────────
def embed_query(text: str) -> tuple[np.ndarray, bool]:
"""Returns (embedding, cache_hit)"""
cache = app_state['embedding_cache']
model_name = "all-mpnet-base-v2-int8"
normalized = ' '.join(text.lower().strip().split())
cached = cache.get(normalized, model_name)
if cached is not None:
return cached, True
tokenizer = app_state['tokenizer']
model = app_state['embedding_model']
inputs = tokenizer(
[normalized], padding=True, truncation=True,
max_length=256, return_tensors="pt"
)
outputs = model(**inputs)
embedding = outputs.last_hidden_state.mean(dim=1).detach().numpy()[0]
# L2 normalize for cosine similarity
embedding = embedding / np.linalg.norm(embedding)
cache.set(normalized, model_name, embedding)
return embedding, False
@app.post("/retrieve", response_model=RetrievalResponse)
async def retrieve(request: RetrievalRequest):
start = time.perf_counter()
# Step 1: Embed (with caching)
embedding, emb_cached = embed_query(request.query)
# Step 2: Search (with caching)
ef = request.ef or Config.HNSW_EF
retrieval_cached = False
# Check retrieval cache
cache_key = f"ret:{hashlib.md5(embedding.tobytes()).hexdigest()}:k{request.k}:ef{ef}"
cached_result = app_state['redis'].get(cache_key)
if cached_result:
data = json.loads(cached_result)
doc_ids, scores = data['doc_ids'], data['scores']
retrieval_cached = True
else:
# Actual ANN search
index = app_state['index']
if ef != Config.HNSW_EF:
# Per-request ef override — note: this isn't thread-safe with hnswlib
# Use Qdrant or a separate index instance per thread for this
index.set_ef(ef)
labels, distances = index.knn_query(embedding.reshape(1, -1), k=request.k)
doc_ids = labels[0].tolist()
scores = (1 - distances[0]).tolist() # Convert distance to similarity
# Cache the result
app_state['redis'].setex(
cache_key, Config.RETRIEVAL_CACHE_TTL,
json.dumps({'doc_ids': doc_ids, 'scores': scores})
)
elapsed_ms = (time.perf_counter() - start) * 1000
cache_hit = ("retrieval" if retrieval_cached
else "embedding" if emb_cached
else "none")
return RetrievalResponse(
doc_ids=doc_ids,
scores=scores,
latency_ms=round(elapsed_ms, 2),
cache_hit=cache_hit
)
Warning: The
index.set_ef()call in hnswlib is not thread-safe. If you need per-requestefoverrides in a multi-threaded server, use Qdrant (which acceptshnsw_efper-query) or maintain a pool of index instances. Using hnswlib directly in a FastAPI app requires careful thread management — or simply fixefat load time and expose it only as an admin setting.
You'll optimize a retrieval pipeline from ~400ms end-to-end latency to under 100ms. Here's your starting point and the optimization path.
Setup: Clone the exercise repository (or set up your own with FAISS and a sentence transformer model). Use the MS MARCO Passage Retrieval dataset — it's a realistic RAG corpus with 8.8M passages and provided query sets.
Step 1: Establish your baseline
# baseline_retrieval.py — intentionally unoptimized
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import time
# Load a large corpus index (pre-built for you in the exercise)
index = faiss.read_index('msmarco_flat.index') # Exact search — no ANN
model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
queries = [
"What is the capital of France?",
"How does photosynthesis work?",
"What causes inflation in an economy?",
# ... 50 more representative queries
]
latencies = []
for query in queries:
start = time.perf_counter()
embedding = model.encode([query], normalize_embeddings=True)
distances, labels = index.search(embedding.astype(np.float32), k=10)
latencies.append((time.perf_counter() - start) * 1000)
print(f"Baseline p50: {np.percentile(latencies, 50):.1f}ms")
print(f"Baseline p95: {np.percentile(latencies, 95):.1f}ms")
Step 2: Convert the flat index to HNSW with M=32, ef_construction=200. Re-run the benchmark. You should see search time drop by 90%+ — but embedding inference still dominates.
Step 3: Export the embedding model to ONNX INT8. Re-run. Total latency should drop significantly.
Step 4: Add the embedding cache. Run the same 50 queries twice (simulating return visitors). Measure second-run latency. With a high cache hit rate, you should see the entire retrieval pipeline complete in 5-15ms.
Step 5: Run the concurrent benchmark (concurrent.futures.ThreadPoolExecutor with 20 workers, 500 total requests). Ensure your p99 stays under 100ms. Tune ef down if needed — 32 is often sufficient for many RAG applications where you're re-ranking anyway.
Expected results after optimization:
Mistake 1: Rebuilding indexes without downtime planning
HNSW index parameters (M, ef_construction) cannot be changed without rebuilding. Rebuilding 5M vectors takes 20-40 minutes. The pattern: maintain a "hot" index serving traffic and build the new index offline, then do an atomic swap (rename the index file and restart the service, or use Qdrant's collection aliases).
Mistake 2: Ignoring payload retrieval cost
You optimized the vector search to 5ms. But then fetching the actual text chunks from your database takes 80ms because you're doing 10 individual row lookups. Fix: batch the payload fetch in a single query, or denormalize the payload into the vector store itself (most vector DBs support storing payload alongside vectors).
Mistake 3: Setting ef too high "just to be safe"
Engineers who haven't measured recall carefully often set ef=500 because they fear missing relevant documents. In practice, most RAG pipelines use a re-ranking step (a cross-encoder) downstream, and getting 20 candidates with 92% recall is identical in terms of final quality to getting 20 candidates with 99% recall — because the same documents end up in the top-5 after re-ranking. Measure your end-task quality, not just ANN recall, and you'll often find you can drop ef dramatically.
Mistake 4: Not accounting for GC pauses in JVM-based vector stores
Elasticsearch and OpenSearch implement their own ANN search (HNSW via Lucene) in the JVM. Garbage collection pauses can add 50-200ms spikes to otherwise fast queries. If you're seeing bimodal latency distributions — mostly fast with occasional 200ms+ outliers — GC is the likely culprit. Tune JVM heap size, use G1GC or ZGC, and monitor GC logs.
Mistake 5: Caching with stale model versions
You upgraded your embedding model. All your cached embeddings are now from the old model. You'll get cache hits but return wrong vectors. Always include the model version identifier in your cache key. When upgrading models, either flush the cache entirely or use a new key namespace.
Mistake 6: Testing on localhost and declaring victory
Retrieval on localhost hits a loopback interface with essentially zero network latency. In production, your embedding service, vector store, and application server might be in different VPCs or AZs. Measure in an environment that matches production network topology before finalizing your latency budget.
Getting retrieval under 100ms isn't a single optimization — it's an architectural discipline. The key insights from this lesson:
Decompose before you optimize. Measure where time actually goes: embedding inference, network, index search, payload fetch. The biggest wins are almost always in embedding inference and caching, not in ANN parameter tuning.
HNSW and IVF serve different needs. HNSW is better for dynamic corpora under 50M vectors that fit in RAM. IVF+PQ is better for enormous corpora where memory is the binding constraint.
ef is your most important runtime lever. You can tune it without rebuilding. Tune it by measuring recall on your actual query distribution, not just ANN recall metrics. And tune it at your p99 under production-realistic concurrency, not single-threaded.
Caching has multiple layers. Embedding cache (long TTL, high value), retrieval result cache (short TTL, high value for hot queries, dangerous for frequently-updated corpora), and response cache (only for truly stable content). Each layer attacks a different part of the latency budget.
Hardware decisions multiply your software optimizations. ONNX + INT8 quantization is a 3-5x win on embedding inference. NUMA awareness can be a 2x win on large indexes. Huge pages help for memory-mapped indexes.
Next steps to deepen your expertise:
Learn about re-ranking integration: Sub-100ms retrieval is often paired with a cross-encoder re-ranker. Understanding how retrieval quality requirements change when you have re-ranking downstream will help you make better ef decisions.
Explore multi-vector retrieval: ColBERT-style late interaction and multi-vector representations require different indexing strategies. They're more expensive but achieve significantly better retrieval quality.
Study corpus-aware sharding: When your corpus exceeds what comfortably fits in a single machine's RAM, you need horizontal sharding. Understanding how Qdrant's distributed mode and Pinecone's pod architecture handle this is essential for 100M+ vector deployments.
Implement observability: Integrate OpenTelemetry spans around each component of your retrieval pipeline. Production optimization requires continuous measurement, not one-time tuning.
Explore GPU-accelerated ANN: FAISS with CUDA support and purpose-built GPU vector stores like Milvus can deliver sub-millisecond ANN search at extreme QPS, if your query volume justifies the infrastructure cost.
The retrieval layer is the foundation of any RAG system's user experience. Invest the time to do it right.