AI Systems & ML Infrastructure — System Design

🤖 14 · AI SYSTEMS & ML INFRASTRUCTURE

LLM Serving & Inference Optimization

Serving large language models at scale requires careful management of GPU memory, batching strategies, and latency optimization. Key metrics: tokens/sec, TTFT (time to first token), and cost per 1M tokens.

▸ LLM Serving Pipeline

KV-Cache Management: Stores computed key-value pairs for previous tokens. PagedAttention (vLLM) manages KV-cache like virtual memory — pages allocated on demand, reducing waste from 60-80% to <4%. Enables serving 2-4x more concurrent requests on same GPU.

Continuous Batching: Unlike static batching (wait for all sequences to finish), continuous batching inserts new requests as soon as any sequence completes. Result: 23x throughput improvement over naive batching. Implemented in vLLM, TensorRT-LLM, and TGI.

Speculative Decoding: Use a small "draft" model to generate N candidate tokens, then verify all N in a single forward pass of the large model. Accepts correct prefixes, rejects wrong ones. Achieves 2-3x speedup with zero quality loss. Used by Medusa, SpecInfer.

Technique	Benefit	Tradeoff
Tensor Parallelism	Split layers across GPUs, low latency	High inter-GPU bandwidth needed (NVLink)
Pipeline Parallelism	Split model stages across GPUs	Pipeline bubbles reduce efficiency
INT8 Quantization	2x memory reduction, ~1% quality loss	Slight accuracy degradation
INT4 (GPTQ/AWQ)	4x memory reduction	2-5% quality loss, not for all tasks
Flash Attention	2-4x faster attention, O(1) memory	Custom CUDA kernels required

Key Numbers: GPT-4 ~100 tokens/sec output. Llama-70B needs ~140GB VRAM (FP16). Batch size 32 = ~3x throughput vs batch=1. A100 80GB can serve Llama-13B at ~1500 tokens/sec. TTFT target: <500ms for interactive, <2s acceptable for batch.

Real-World: OpenAI — custom infra with speculative decoding + continuous batching. vLLM — PagedAttention, open-source. TensorRT-LLM — NVIDIA optimized runtime. Anyscale — Ray-based distributed serving.

🎯 Interview Tip: Always discuss three metrics: tokens/sec (throughput), TTFT (time to first token — user-perceived latency), and cost per 1M tokens (economics). Mention KV-cache as the primary memory bottleneck and continuous batching as the key optimization.

RAG (Retrieval-Augmented Generation)

Combines retrieval from external knowledge with LLM generation. Solves hallucination, stale knowledge, and domain specificity without fine-tuning. The dominant pattern for enterprise AI applications.

▸ RAG Pipeline Architecture

Chunking Strategies: Fixed-size (512 tokens, 50 token overlap — simple, fast). Semantic (split at topic boundaries using embeddings). Recursive (split by paragraphs → sentences → words). Sentence-based (preserve complete sentences). Optimal chunk size: 256-1024 tokens depending on use case.

Retrieval Methods: Dense retrieval — embedding cosine similarity, great for semantic matching. Sparse retrieval — BM25/TF-IDF, great for keyword/exact match. Hybrid — combine both with Reciprocal Rank Fusion (RRF). Hybrid outperforms either alone by 5-15% on most benchmarks.

Reranking: Initial retrieval (bi-encoder) is fast but approximate. Cross-encoder reranking scores query-document pairs jointly — 10-20x slower but significantly more accurate. Typical pattern: retrieve top-50 with bi-encoder, rerank to top-5 with cross-encoder. Models: Cohere Rerank, BGE-reranker, ColBERT.

Component	Options	Latency	When to Use
Chunking	Fixed 512 / Semantic / Recursive	Offline	Fixed for speed, semantic for quality
Embedding	ada-002 / E5-large / BGE	5-50ms/chunk	ada-002 for general, E5 for cost
Vector Search	HNSW / IVF / Flat	1-10ms	HNSW for speed, IVF for memory
Reranker	Cross-encoder / ColBERT	50-200ms	Always use for top-K refinement

Real-World: Perplexity — real-time web RAG with citation. ChatGPT Browsing — retrieval from live web. Enterprise KB — internal docs + Confluence + Slack indexed for employee Q&A. GitHub Copilot — code context retrieval from repo.

Pitfalls: Chunk too small → loses context, incoherent answers. Chunk too large → dilutes relevance, wastes context window. Wrong embedding model → poor retrieval quality. No reranking → irrelevant chunks pollute context. Stale index → outdated answers.

🎯 Interview Tip: Discuss the full pipeline end-to-end. Mention hybrid retrieval (dense + sparse) as the production standard. Know the tradeoff: larger chunks = better context but lower precision. Always mention reranking as a critical quality improvement step.

Embedding Pipelines & Vector Search

Convert unstructured data (text, images, code) into dense vector representations for similarity search. The backbone of RAG, recommendations, and semantic search at scale.

Embedding Generation: Batch pipeline — process millions of documents offline, store in vector DB. Real-time pipeline — embed new documents/queries on-the-fly with <50ms latency. Models: OpenAI ada-002 (1536d), Cohere embed-v3 (1024d), E5-large (1024d), BGE-large (1024d).

ANN Algorithms: Approximate Nearest Neighbor search trades perfect recall for speed:
• HNSW (Hierarchical Navigable Small World) — graph-based, O(log N) search, 95-99% recall, memory-heavy (stores full graph)
• IVF (Inverted File Index) — partition-based, less memory, slower search, good for disk-based indices
• ScaNN (Google) — quantization + anisotropic scoring, best throughput at high recall
• DiskANN (Microsoft) — SSD-optimized, billion-scale on commodity hardware

Algorithm	Recall@10	QPS (1M vectors)	Memory	Best For
HNSW	98-99%	5,000-10,000	High (graph + vectors)	Low-latency serving
IVF-PQ	90-95%	10,000-50,000	Low (compressed)	Memory-constrained
ScaNN	95-98%	20,000-100,000	Medium	High throughput
Flat (Brute)	100%	100-500	Vectors only	<100K vectors, exact

Key Numbers: OpenAI ada-002 = 1536 dimensions, ~$0.0001/1K tokens. Cohere embed-v3 = 1024 dims. Search latency: <10ms at 1B vectors with HNSW. Index build time: ~1 hour for 100M vectors. Storage: 1B vectors × 1536d × 4 bytes = ~6TB raw (compressed to ~1.5TB with PQ).

Real-World: Pinecone — managed, serverless vector DB. Weaviate — open-source, hybrid search built-in. Milvus — distributed, billion-scale. pgvector — PostgreSQL extension, simple integration. Qdrant — Rust-based, filtering + vector search.

Tradeoffs: High recall vs latency — more graph connections = better recall but slower. Memory vs accuracy — Product Quantization saves 8-32x memory but loses 2-5% recall. Build time vs search quality — larger HNSW ef_construction = better graph but hours to build.

🎯 Interview Tip: Know HNSW internals (navigable small-world graph with hierarchical layers). Discuss the recall-latency-memory triangle. Mention that most production systems use HNSW with PQ compression for the best balance. Filtering (metadata + vector) is a key production requirement.

AI Gateway & LLM Routing

A centralized proxy layer between applications and LLM providers. Handles model routing, fallback chains, cost optimization, semantic caching, and observability for all AI calls.

▸ AI Gateway Architecture

Routing Strategies: Complexity-based — classify query complexity, route simple→cheap model, complex→powerful model. Cost-based — route to cheapest model that meets quality threshold. Latency-based — route to fastest available model. Capability-based — code tasks→Codex, reasoning→GPT-4, creative→Claude.

Semantic Caching: Hash similar prompts (not exact match — use embedding similarity with threshold ~0.95). Serve cached responses for semantically equivalent queries. Cache hit rates of 20-40% typical for customer support, FAQ-style workloads. Invalidation: TTL-based + manual flush for updated knowledge.

Cost Control: Token budgets per user/team/project. Model tiering: GPT-4o for premium users, GPT-4o-mini for free tier. Prompt compression (LLMLingua) reduces tokens by 2-5x. Response length limits. Batch non-urgent requests for off-peak pricing.

Real-World: LiteLLM — unified API for 100+ LLM providers. Portkey — AI gateway with caching + fallbacks. Helicone — observability + cost tracking. AWS Bedrock — managed multi-model access. Azure AI Gateway — enterprise routing with APIM.

🎯 Interview Tip: Frame the AI Gateway as an "API Gateway for LLMs." Discuss the three pillars: reliability (fallbacks, retries), cost (caching, routing, budgets), and observability (token tracking, latency percentiles, quality scoring).

Agent Systems & Orchestration

Autonomous AI systems that reason, plan, use tools, and iterate to accomplish complex tasks. The ReAct pattern (Reason + Act) is the foundation of modern agent architectures.

▸ Agent System Architecture

Agent Patterns:
• Single Agent — one LLM with tools, simplest pattern, good for focused tasks
• Multi-Agent (Supervisor) — supervisor agent delegates to specialist workers (researcher, coder, reviewer)
• Hierarchical — tree of agents, each level handles different abstraction
• Debate/Consensus — multiple agents propose solutions, vote on best answer

Memory Systems:
• Conversation Buffer — full chat history in context window (limited by token count)
• Summary Memory — LLM summarizes older messages, keeps recent ones verbatim
• Vector Memory — embed all interactions, retrieve relevant past context via similarity search
• Entity Memory — extract and track entities (people, projects, preferences) across sessions

Tool Calling: Modern LLMs support structured tool calling (function calling). Agent describes available tools with JSON schemas. LLM outputs structured tool invocations. Runtime executes tools and feeds results back. Key: tool descriptions are critical for routing accuracy. Limit to 10-20 tools per agent for best performance.

Pitfalls: Infinite loops — agent gets stuck repeating actions (need max iterations). Tool misuse — wrong tool selection wastes tokens and time. Context overflow — long conversations exceed context window. Hallucinated actions — agent invents non-existent tools. Cost explosion — complex tasks can use 100K+ tokens.

Real-World: LangChain/LangGraph — most popular agent framework, graph-based orchestration. CrewAI — role-based multi-agent. AutoGen (Microsoft) — conversational multi-agent. OpenAI Assistants API — managed agent with tools + retrieval. Amazon Bedrock Agents — enterprise agent orchestration.

🎯 Interview Tip: Describe the ReAct loop clearly. Discuss memory strategies for long-running agents. Know when to use single vs multi-agent (single for focused tasks, multi for complex workflows). Mention guardrails: max iterations, token budgets, human-in-the-loop for high-stakes actions.

Feature Stores (Online + Offline)

Centralized platform for managing ML features across training and serving. Solves training-serving skew, enables feature reuse, and ensures point-in-time correctness for model training.

▸ Feature Store Architecture

Core Concepts:
• Offline features — batch computed (daily/hourly), stored in data warehouse, used for training
• Online features — real-time, served from Redis/DynamoDB with p99 <10ms, used for inference
• Point-in-time correctness — training uses only features available at prediction time (no future leakage)
• Feature freshness — how stale features can be before model quality degrades
• Training-serving skew — mismatch between training and serving feature computation

Store Type	Storage	Latency	Freshness	Use Case
Offline	S3 / BigQuery / Hive	Seconds-minutes	Hours-days	Training, batch prediction
Online	Redis / DynamoDB	<10ms p99	Seconds-minutes	Real-time inference
Streaming	Kafka → Online store	Sub-second	Real-time	Fraud detection, recommendations

Real-World: Feast — open-source, Kubernetes-native. Tecton — managed, real-time features. Hopsworks — full ML platform with feature store. Uber Michelangelo — pioneered feature stores at scale (trillions of features). Spotify — custom feature store for music recommendations.

🎯 Interview Tip: Emphasize training-serving skew as the #1 problem feature stores solve. Discuss the dual-write pattern (batch → offline, stream → online). Know that point-in-time joins prevent data leakage in training. Feature stores are critical for any ML system serving real-time predictions.

Distributed Training Infrastructure

Training large models requires distributing computation across hundreds or thousands of GPUs. Key challenges: communication overhead, memory management, fault tolerance, and efficient data loading.

Parallelism Strategies:
• Data Parallelism — replicate model on each GPU, split data batches. All-reduce gradients after each step. Simple, scales to 100s of GPUs. Limited by model size fitting in single GPU.
• Model/Tensor Parallelism — split individual layers across GPUs (e.g., split attention heads). Requires high-bandwidth interconnect (NVLink 600GB/s). Used for layers too large for one GPU.
• Pipeline Parallelism — split model into stages, each stage on different GPU(s). Micro-batching reduces pipeline bubbles. GPipe, PipeDream patterns.
• Expert Parallelism — for Mixture-of-Experts models, route tokens to different GPU groups.

Memory Optimization:
• Gradient Accumulation — accumulate gradients over N micro-batches before update, simulates larger batch size
• Activation Checkpointing — recompute activations during backward pass instead of storing (trades compute for memory, ~30% memory savings)
• ZeRO (DeepSpeed) — partition optimizer states (Stage 1), gradients (Stage 2), and parameters (Stage 3) across GPUs
• Mixed Precision (FP16/BF16) — 2x memory reduction, 2-3x faster compute on tensor cores

Strategy	Communication	Memory Efficiency	Scaling
Data Parallel (DDP)	All-reduce gradients	Model must fit 1 GPU	Near-linear to 1000 GPUs
ZeRO Stage 3	All-gather params on demand	Partitions everything	Trillion-param models
Tensor Parallel	All-reduce per layer	Split layers across GPUs	8-16 GPUs (NVLink)
Pipeline Parallel	Point-to-point activations	Split stages	Across nodes
FSDP (PyTorch)	All-gather + reduce-scatter	Shards like ZeRO-3	Native PyTorch scaling

Fault Tolerance: At 10K+ GPU scale, hardware failures are frequent (mean time between failures ~hours). Solutions: Periodic checkpointing (every 10-30 min to distributed storage). Elastic training — continue with fewer nodes. Redundant computation — duplicate critical stages. Fast recovery — resume from latest checkpoint in <5 min.

Key Numbers: GPT-4 training estimated at ~$100M, ~25,000 A100 GPUs for ~90 days, 13T tokens. Llama-2 70B: 1.7M GPU-hours on A100-80GB. A100 80GB: ~$2/hr cloud, 312 TFLOPS FP16. H100: 2-3x faster than A100. Network: 400Gbps InfiniBand between nodes, 900GB/s NVLink within node.

Real-World: Meta RSC — 16K A100 GPUs, custom InfiniBand fabric. Google TPU pods — v4 pods with 4096 chips, custom interconnect. NVIDIA DGX SuperPOD — turnkey GPU clusters. DeepSpeed (Microsoft) — ZeRO optimizer, most popular training library. Megatron-LM (NVIDIA) — tensor + pipeline parallelism.

🎯 Interview Tip: Know the three types of parallelism and when to combine them (3D parallelism = data + tensor + pipeline). Discuss checkpointing strategy for fault tolerance. Mention that communication overhead is the main bottleneck — all-reduce scales as O(model_size), not O(num_GPUs).

Guardrails & AI Safety in Production

Production AI systems need multiple layers of protection: input validation (block malicious prompts), output validation (catch hallucinations, toxicity), and human escalation for high-stakes decisions.

▸ Guardrails Pipeline

Input Validation:
• Prompt Injection Detection — classify inputs for injection attempts (jailbreaks, role-play attacks, instruction override). Use fine-tuned classifiers + rule-based patterns.
• PII Filtering — detect and redact SSN, credit cards, emails, phone numbers before sending to LLM. Use NER models + regex patterns.
• Topic Restriction — block off-topic queries (e.g., medical advice in a coding assistant). Embedding-based topic classifiers.
• Token Limits — enforce max input length to prevent resource exhaustion attacks.

Output Validation:
• Toxicity Scoring — score output for harmful content (Perspective API, custom classifiers). Block if score > threshold.
• Hallucination Detection — cross-reference claims against retrieved sources. Flag unsupported statements. Use NLI (Natural Language Inference) models.
• Format Enforcement — validate JSON schema, check required fields, ensure structured output matches spec. Retry with feedback if invalid.
• Factual Grounding — verify output is supported by provided context (for RAG systems).

Advanced Patterns:
• Constitutional AI — self-critique: generate → evaluate against principles → revise
• Red-teaming — adversarial testing with automated attack generation
• Canary tokens — embed markers in system prompts to detect extraction attempts
• Output diversity — detect repetitive/degenerate outputs
• Confidence scoring — LLM self-reports confidence, escalate low-confidence responses

Pitfalls: Over-filtering — too aggressive guards block legitimate queries (false positives). Latency overhead — each guard adds 50-200ms. Evolving attacks — new jailbreak techniques emerge weekly. Multilingual gaps — guards trained on English miss attacks in other languages. Context-dependent safety — same content may be safe/unsafe depending on context.

Real-World: Guardrails AI — open-source, programmable guardrails with validators. NeMo Guardrails (NVIDIA) — dialog management + safety rails. Lakera Guard — prompt injection detection API. Azure AI Content Safety — multi-modal content filtering. Custom classifiers fine-tuned on domain-specific attack patterns.

🎯 Interview Tip: Frame guardrails as a pipeline (not a single check). Discuss the tradeoff between safety and user experience (over-filtering frustrates users). Mention that guardrails should be independently deployable and versioned. Know that prompt injection is an unsolved problem — defense in depth is the only approach.