Designing an AI Chat Service

What you'll learn

• Problem Statement

  Design a conversational AI service where users send text prompts and receive streamed responses generated by a large language model (LLM). The system must support 200 million daily active users, maintain low time-to-first-token (TTFT) la…

• Requirements

• Functional Requirements

  1. Send a text prompt and receive a streamed response from an LLM 2. Create, list, and resume conversations (chat sessions) 3. Persist all messages (user prompts and assistant responses) within a conversation 4. Resume a conversation wit…

• Non-Functional Requirements

  1. Low TTFT — time-to-first-token must be under 1 second for paid tier users 2. GPU efficiency — maximize throughput per GPU dollar (continuous batching, KV-cache reuse) 3. 200M DAU scale — handle millions of concurrent generation reques…

• Extended Requirements

  1. Multi-model routing (different model sizes for different use cases) 2. Content moderation and safety filtering 3. Plugin/tool use within conversations 4. Multi-modal inputs (images, files) ---

• Core Entities

  | Entity | Key Fields | Notes | |--------|-----------|-------| | User | userId, tier (free/paid), dailyTokensUsed, dailyTokenBudget | Tier determines priority and limits | | Chat / Conversation | chatId, userId, title, createdAt, modelId…

• API Design

• Create a New Chat

• List User's Chats

• Get Messages in a Chat

• Send a Message (Streaming Response)

  Why Server-Sent Events (SSE) over HTTP/2? - Unidirectional streaming (server to client) is all we need — the client sends one request, then just listens - Built-in reconnection semantics (Last-Event-ID header) - Works through proxies, lo…

• Cancel a Generation

  ---

• Token Streaming Architecture

• How the Stream Flows

  1. Client sends POST with the user's message 2. API Gateway authenticates, applies rate limits, forwards to Chat Service 3. Chat Service persists the user message, assembles the context (conversation history), and submits a Generation Re…

• Minimizing Latency in the Stream

  | Principle | Implementation | |-----------|---------------| | No buffering | Flush each token immediately at every proxy layer | | Heartbeat during prefill | Send event: heartbeat every 500ms while the model processes input (keeps conne…

• Why TTFT Matters

  Time-to-first-token (TTFT) is the elapsed time from when the user submits a prompt to when the first output token appears on screen. It is the single most important latency metric for a chat service.

• The Two Phases of LLM Inference

  TTFT = Queue wait time + Prefill time + First decode step

• Optimizations to Reduce TTFT

  | Technique | How It Helps | TTFT Reduction | |-----------|-------------|----------------| | Prefix caching / KV-cache reuse | If the conversation history (system prompt + prior turns) was computed before, reuse the cached key-value tens…

• Why KV-Cache Reuse Is Critical for Chat

  In a multi-turn conversation, the model must "see" all prior messages to generate a contextual response. Without caching: Without caching, TTFT grows linearly with conversation length. With KV-cache persistence between turns, only the ne…

• GPU Orchestration and Scheduling

• The Scheduling Problem

  GPU workers are the most expensive and scarce resource. A single A100 GPU costs ~$2/hour. At 200M DAU, the system may need tens of thousands of GPUs. Every millisecond of idle GPU time is wasted money.

• Request Queue with Priority

• Continuous Batching

  Traditional batching waits for a full batch of requests, processes them together, and waits for ALL to finish before accepting new ones. This wastes GPU cycles when some requests finish early. Continuous batching (also called iteration-l…

• Affinity Routing

  When a user sends a follow-up message in the same conversation, the GPU worker that handled the previous turn likely still has the KV-cache in GPU memory. Routing the new request to the same worker avoids re-prefilling the entire convers…

• Worker Health Monitoring

  | Signal | Action | |--------|--------| | No heartbeat for 10s | Mark worker unhealthy; stop routing new requests | | GPU memory > 95% | Stop accepting new requests; let current ones drain | | Inference latency spike | Reduce batch size…

• Rate Limiting and Tiered Access

• Multi-Level Rate Limiting

  | Layer | What It Limits | Why | |-------|---------------|-----| | API Gateway | Requests per minute per user (token bucket algorithm) | Prevents API abuse, DDoS protection | | Scheduler | One active generation per user at a time | Preve…

• Tier Configuration

  | Parameter | Free Tier | Paid Tier | |-----------|-----------|-----------| | Queue priority | P2 (low) | P0/P1 (high) | | Max context window | 4K tokens | 128K tokens | | Daily token budget | 50K tokens | 2M tokens | | Requests per minu…

• Weighted Fair Queuing

  To prevent heavy users from starving others even within the same tier: This ensures that even during peak load, no single user (or small group of heavy users) can monopolize GPU capacity at the expense of others in the same tier. ---

• Context Management as Conversations Grow

  As conversations extend to dozens or hundreds of turns, the total context can exceed the model's context window. The system needs strategies to manage this.

• Strategy Comparison

  | Strategy | How It Works | Pros | Cons | |----------|-------------|------|------| | Truncation | Drop oldest messages beyond context window | Simple, predictable | Loses early context; user notices "amnesia" | | Summarization | Periodic…

• Recommended Hybrid Approach

  When to trigger summarization: - When total conversation tokens exceed 75% of context window - Summarize the oldest batch of turns - Store the summary as a special "system" message in the conversation ---

• Data Model

• PostgreSQL Schema

• Why PostgreSQL?

  - Messages are relational (belong to chats, chats belong to users) - Strong consistency needed for message ordering (sequence numbers) - ACID transactions ensure a user message and its corresponding assistant response are persisted atomi…

• Partitioning Strategy

  Partition the Messages table by chatid hash. This ensures all messages for a single conversation reside on the same shard, enabling efficient retrieval without cross-shard queries. ---

• Cancellation

  When a user clicks "Stop generating," the system must halt inference, persist whatever was generated so far, and release resources.

• Cancellation Flow

• Why Check Between Decode Steps?

  Each decode step (generating one token) takes ~10-30ms. Checking a cancellation flag between steps adds negligible overhead but ensures the system responds to cancellation within one token generation cycle. This is far better than waitin…

• Design Evolution & Trade-offs

  How this system grew from an MVP to what a production version looks like — and the choices you'd defend in an interview.

• Stage 1 — MVP: One box, one database

  The simplest thing: single API server calling an LLM provider, storing conversation history in Postgres. What breaks first — Latency (LLM calls are 2-30s), streaming needs, DB write contention on chat history table. ---

• Stage 2 — Scale-out: Multiple pods behind a load balancer

  - Stateless API layer, N replicas behind an L7 LB. - WebSocket / SSE for streaming tokens back to clients. - DB stays single-node. What breaks next — LLM provider rate limits, cost explosion (every call costs money), cold-start latency f…

• Stage 3 — Bottleneck: Conversation context & memory

  The FIRST hard problem for AI chat: how to give the model conversation history without blowing token limits or latency. | Option | How it works | Latency | Cost | Complexity | Best for | Fails at | |--------|-------------|---------|-----…

• Stage 4 — Second bottleneck: LLM rate limits, cost, and failover

  Options: | Option | How it works | Latency impact | Cost impact | Complexity | Best for | Fails at | |--------|-------------|---------------|-------------|------------|----------|----------| | A. Multi-provider routing | Route to OpenAI,…

• Stage 5 — Mature architecture

  Full diagram: CDN → LB → API pods → (Redis session cache, vector DB for RAG, Postgres for durable history, Kafka for events, LLM router with fallback, moderation service). New components: | Component | Purpose | Why now? | |-----------|-…

• Migration Path (MVP → Enterprise)

  Step-by-step migration in dependency order: 1. Add streaming (SSE) — Users see tokens as they generate (perceived latency drops 10x). No backend changes yet. 2. Horizontal scale API layer — Multiple pods behind LB. DB connection pooling…

• High-Level Architecture

  ---

• Scaling

• Autoscaling GPU Workers

• Separate Pools per Model Size

  | Pool | Model | Use Case | Instance Type | |------|-------|----------|---------------| | Pool A | Large (70B+ params) | Complex reasoning, paid tier | A100 80GB / H100 | | Pool B | Small (7B-13B params) | Simple tasks, free tier, high t…

• Cost Optimization

  | Strategy | How It Saves | |----------|-------------| | Spot/preemptible for free tier | 60-70% cost reduction; free tier can tolerate occasional retries | | Warm pool instead of cold scaling | Avoids 5-10 minute GPU initialization on s…

• Multi-Region Deployment

  ---

• Fault Tolerance

• GPU Worker Crash

• Client Connection Drop

• Database Failure

  | Failure | Impact | Recovery | |---------|--------|----------| | PostgreSQL primary down | Writes blocked (new messages can't be persisted) | Promote replica to primary; queue writes during failover | | PostgreSQL replica down | Reduced…

• Summary: Key Design Decisions

  | Decision | What It Solves | Trade-off | |----------|---------------|-----------| | SSE over HTTP/2 for streaming | Unidirectional token delivery with reconnection support | Requires proxy configuration to avoid buffering | | Continuous…