AI Inference Platform

Problem Statement & Requirements

Design a scalable AI model inference platform that hosts machine learning models (including large language models), serves prediction requests in real-time and batch modes, manages model lifecycle, and handles the unique compute challenges of GPU-based workloads — similar to the OpenAI API, AWS SageMaker Endpoints, Google Vertex AI, Replicate, or Together AI.

Functional Requirements

• Real-time inference — synchronous API: send input, receive prediction/generation in real-time (chat, classification, embeddings)
• Streaming inference — server-sent events (SSE) for token-by-token LLM output
• Batch inference — submit large datasets, process asynchronously, retrieve results later
• Model hosting — deploy models of any size: from 100 MB classifiers to 400B+ parameter LLMs
• Model versioning — deploy multiple versions, canary rollouts, instant rollback
• Auto-scaling — scale GPU instances based on request load; scale to zero when idle
• Multi-model serving — multiple models on one GPU (small models) or one model across multiple GPUs (large models)
• API key management — per-customer API keys with rate limits and usage quotas
• Usage metering — track tokens consumed (LLMs) or requests made (vision/classification) for billing
• Model registry — upload, version, and manage model artifacts (weights, configs, tokenizers)

Non-Functional Requirements

• Low latency — real-time inference p50 < 200 ms for small models; first-token latency < 500 ms for LLMs
• High throughput — serve 100K+ inference requests/second across all models
• GPU efficiency — maximize GPU utilization (target > 70%); GPUs cost $2-30/hr each
• Availability — 99.99% for inference endpoints
• Scalability — support 10K+ deployed model endpoints
• Cost efficiency — GPU idle time is burning money; scale-to-zero and bin-packing are critical
• Multi-tenancy — thousands of customers sharing GPU pools with isolation

Out of Scope

• Model training / fine-tuning pipeline
• Data labeling platform
• MLOps experiment tracking (Weights & Biases)
• Edge / on-device inference

Scale Estimations

Traffic

Compute

Storage

Bandwidth

This first layer defines the system boundary. The next chapters add application logic and data dependencies. After examining the request paths and tradeoffs, we will add operational safeguards and bring the complete architecture together.

API Design

Chat Completions (Streaming)

POST /v1/chat/completions
Authorization: Bearer sk-...
Content-Type: application/json

Request:
{
  "model": "llama-3-70b-instruct",
  "messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Explain quantum computing in simple terms."}
  ],
  "max_tokens": 500,
  "temperature": 0.7,
  "stream": true,
  "top_p": 0.9
}

Response (SSE stream):
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","model":"llama-3-70b-instruct","choices":[{"delta":{"role":"assistant","content":"Quantum"},"index":0}]}

data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","choices":[{"delta":{"content":" computing"},"index":0}]}

data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","choices":[{"delta":{"content":" is"},"index":0}]}

... (token by token)

data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","choices":[{"delta":{},"finish_reason":"stop","index":0}],"usage":{"prompt_tokens":28,"completion_tokens":247,"total_tokens":275}}

data: [DONE]

Embeddings

POST /v1/embeddings
Authorization: Bearer sk-...

Request:
{
  "model": "text-embedding-3-large",
  "input": ["Quantum computing is a type of computation...",
            "Machine learning models can be trained..."],
  "encoding_format": "float"
}

Response (200 OK):
{
  "object": "list",
  "data": [
    {"object": "embedding", "index": 0, "embedding": [0.0123, -0.0456, ...]},
    {"object": "embedding", "index": 1, "embedding": [0.0789, 0.0321, ...]}
  ],
  "model": "text-embedding-3-large",
  "usage": {"prompt_tokens": 42, "total_tokens": 42}
}

Batch Inference

POST /v1/batches
Authorization: Bearer sk-...

Request:
{
  "input_file_id": "file-abc123",        // uploaded JSONL file with requests
  "endpoint": "/v1/chat/completions",
  "completion_window": "24h",
  "metadata": {"job_name": "classify-reviews"}
}

Response (200 OK):
{
  "id": "batch-xyz789",
  "status": "validating",
  "input_file_id": "file-abc123",
  "request_counts": {"total": 50000, "completed": 0, "failed": 0},
  "created_at": "2026-04-03T10:00:00Z",
  "expires_at": "2026-04-04T10:00:00Z"
}

// Poll status or receive webhook when complete:
GET /v1/batches/batch-xyz789
→ {"status": "completed", "output_file_id": "file-output-456", ...}

Deploy Custom Model

POST /v1/models/deploy
Authorization: Bearer sk-...

Request:
{
  "model_name": "my-finetuned-llama",
  "model_artifact": "s3://models/my-llama/v3/",
  "framework": "vllm",
  "hardware": {
    "gpu_type": "h100",
    "gpu_count": 4,
    "min_replicas": 1,
    "max_replicas": 8
  },
  "scaling": {
    "metric": "gpu_utilization",
    "target_value": 70,
    "scale_to_zero": true,
    "scale_to_zero_after_idle_seconds": 300
  }
}

Response (202 Accepted):
{
  "endpoint_id": "ep-abc123",
  "status": "provisioning",
  "estimated_ready_seconds": 120,
  "endpoint_url": "https://api.example.com/v1/models/ep-abc123"
}

Data Model

Model Registry

Endpoint (Deployment)

Request Log / Metering

Core Design Decisions

Decision 1: GPU Scheduling — How to Assign Requests to GPUs

Decision 2: LLM Serving Optimization — Continuous Batching

Decision 3: KV Cache Management — PagedAttention

Detailed Flow Diagrams

Real-Time LLM Inference Flow

Model Deployment & Cold Start Flow

Autoscaling Flow

Model Serving Optimizations

Quantization

Prefix Caching

Multi-LoRA Serving

Cost Optimization Strategies

Two prompts with the same token count can have very different generation lengths. Admission must account for growing attention state rather than budgeting only for the prompt. Test what happens when many long-running sequences occupy a worker, and compare scheduling policies using both throughput and waiting time for short requests.

Handling Edge Cases

Cold Start Mitigation

Problem: Model scaled to zero → first request waits 30-120 seconds

Multi-layer mitigation:

  Layer 1: Local SSD Model Cache
    → Keep model weights on GPU node's NVMe SSD after scale-down
    → Re-load from local SSD: 10-15s (vs 30-60s from S3)
    → Only works if same GPU node is re-allocated

  Layer 2: Predictive Pre-warming
    → Analyze historical traffic patterns per endpoint
    → "This model is requested every weekday at 9am"
    → Pre-warm GPU 5 minutes before predicted usage

  Layer 3: Warm Pool of Generic GPUs
    → Keep a small pool of GPUs with popular base models pre-loaded
    → "Always have 3 GPUs with Llama-70B ready"
    → First request: steal from warm pool (near-instant)
    → Background: provision dedicated GPU to replace warm pool

  Layer 4: Serverless Inference Proxy
    → Accept request → return 202 with polling URL
    → Start model loading in background
    → Client polls until result ready
    → OR: hold HTTP connection with progress updates
    → "Model loading... 15s... 10s... 5s... generating..."

GPU Out of Memory (OOM)

Problem: Too many concurrent requests → KV cache exceeds GPU memory

Prevention:
  1. Admission control: max_concurrent_requests per GPU
     Based on: (gpu_memory - model_size) / avg_kv_cache_per_request
     E.g., (80 GB - 35 GB) / 4 GB = max 11 concurrent requests

  2. PagedAttention: dynamic allocation prevents fragmentation
     But: still bounded by physical GPU memory

  3. KV cache eviction: if memory pressure high:
     → Preempt lowest-priority request (pause, swap KV cache to CPU RAM)
     → Resume when memory frees up
     → vLLM supports this natively

Response when at capacity:
  → Queue request (bounded queue, ~30 seconds max wait)
  → If queue full → return 429 Too Many Requests
  → Trigger autoscaler to provision more GPU replicas

Request Timeout / Hung Generation

Problem: LLM generates endlessly or gets stuck in a loop

Safeguards:
  1. max_tokens limit (hard cap: 4096 default, user-configurable)
  2. Server-side timeout: 120s per request (configurable)
  3. Stop sequences: generation stops on matched patterns
  4. Token-level budget: abort if tokens_generated > max_tokens
  5. Streaming heartbeat: if no token generated in 30s → timeout
  6. Circuit breaker: if error rate > 10% → stop routing to GPU

Client-side:
  → SDK implements timeout (default 60s)
  → Streaming: detect no data for 30s → abort + retry

Tradeoffs & Design Decisions Summary

Reliability & Fault Tolerance

Single Points of Failure & Mitigations

Graceful Degradation

Tier 1 (single GPU worker failure):
  → Router detects health check failure (10s)
  → Reroute to other replicas
  → Autoscaler provisions replacement
  → In-flight requests: retried (idempotent for embeddings; not for
    streaming LLM — client reconnects and re-sends)

Tier 2 (GPU cluster partially degraded):
  → Reduce max concurrency (fewer GPUs available)
  → Queue requests (bounded, 30s timeout)
  → Return 429 when queue full
  → Prioritize paid tier customers

Tier 3 (model loading fails):
  → Retry on different GPU node
  → If model artifact corrupted → alert; roll back to last known good version
  → Return 503 with estimated recovery time

Tier 4 (full GPU pool exhausted):
  → Queue all requests
  → Emergency scale-up (burst capacity from cloud provider)
  → Batch requests redirected to spot instances
  → Real-time requests get priority

High-Level Architecture

Component Breakdown

Full System Architecture (Production-Grade)

Comparison with Related Systems

Interview Tips

1. Start with the GPU economics — "GPUs cost $3-30/hr. An idle GPU is burning money. The entire architecture revolves around maximizing GPU utilization: continuous batching (GPU never idle between requests), PagedAttention (fit more concurrent requests), quantization (fewer GPUs per model), and scale-to-zero (don't pay when idle)."

2. Continuous batching is the most impactful optimization — Draw the static vs continuous batching diagram. Static: short requests wait for long ones. Continuous: requests enter/leave independently, GPU always at max batch size. 2-3× throughput. This is the single biggest improvement over naive serving.

3. PagedAttention solves the KV cache problem — "Each LLM request needs 2-8 GB of KV cache. Pre-allocating max context wastes 60%+ of memory. PagedAttention manages KV cache like virtual memory: allocate 4 KB pages on demand, non-contiguous, near-zero fragmentation. 2-4× more concurrent requests per GPU."

4. Quantization is a business decision — INT4 makes a 70B model fit on 1 GPU instead of 4. 4× cost reduction. Quality loss is ~1-2% — acceptable for chat, coding, summarization. Not acceptable for math benchmarks or embeddings. Know when to quantize and when not to.

5. Scale-to-zero + warm pools handle the cold start problem — Don't just say "use a warm pool." Explain the full mitigation stack: local NVMe SSD cache (10-15s reload), predictive pre-warming (load before predicted traffic), warm pool of pre-loaded base models (near-instant steal). Each layer catches what the previous misses.

6. Multi-LoRA is the multi-tenancy answer — "100 customers each fine-tuned Llama-70B. Naive: 100 full copies = 14 TB of GPU memory. With LoRA: 1 base model (70 GB) + 100 adapters (100 MB each = 10 GB). Swap adapter in < 1 ms per request. 100× more capital-efficient."

7. Streaming is not optional for LLMs — "A 70B model generates ~40-80 tokens/sec. A 500-token response takes 6-12 seconds. Without streaming, user stares at a blank screen for 12 seconds. With SSE streaming, first token appears in 500 ms. Perceived latency goes from 12s to 0.5s."

8. Autoscaling on queue depth, not CPU — "CPU utilization is meaningless for GPU workloads. Queue depth is the leading indicator: if requests are queuing, you need more GPUs. GPU utilization is the efficiency indicator: if < 30%, you're wasting money. Latency is the quality indicator: if degrading, something is wrong."

9. Batch inference uses spot instances — "Batch jobs aren't latency-sensitive. Run on spot GPUs (60-70% cheaper). If spot is reclaimed: checkpoint progress, resume on new instance. At $15M/month GPU spend, spot saves $5-10M."

10. End with cost numbers — "Without optimization: 70B model on 4× H100 FP16 = $72/hr. With INT4 quantization: 1× H100 = $18/hr. With continuous batching: 2.5× throughput = $7.20/hr effective per unit of throughput. With scale-to-zero: only pay during active usage. Combined: 20-50× cost reduction vs naive deployment."