Katkat

Design an LLM Inference Serving System

BasetenSoftware Engineer · Technical Screen

Last updated

Design an end-to-end, high-performance LLM serving platform that manages request admission, continuous batching, KV cache memory, token streaming, and GPU fleet autoscaling under tight latency SLAs.

Scenario

You are tasked with designing the core inference serving engine and orchestration layer for a high-performance LLM cloud platform (such as Baseten). The system must serve state-of-the-art open-weights models (e.g., Llama 3 70B, Mixtral 8x22B) to multi-tenant workloads, delivering streaming token generation with strict latency constraints while maximizing GPU utilization.

Requirements

Functional Requirements

  • Streaming API: Support Server-Sent Events (SSE) and bidirectional gRPC streaming for token generation.
  • Batching & Scheduling: Implement continuous (iteration-level) batching and chunked prefill to handle concurrent requests of heterogeneous prompt and output lengths.
  • Memory Management: Manage dynamic GPU memory allocation for Key-Value (KV) caching without fragmentation.
  • Lifecycle & Cancellation: Support prompt cancellation, client disconnect detection, warm/cold model rollouts, and automatic crash recovery.
  • Fairness & Multi-tenancy: Prevent noisy neighbors from exhausting memory or GPU compute.

Non-Functional Requirements & Scale

  • Throughput & Concurrency: Support 5,000 queries per second (QPS) across the fleet, with prompts ranging from 128 to 32,768 tokens.
  • Latency Targets:
    • Time to First Token (TTFT): < 250 ms (p95) for prompts < 2k tokens.
    • Inter-Token Latency (ITL): < 30 ms (p95) per token during decode.
  • Availability & Resilience: 99.95% uptime; instantaneous eviction of failed worker nodes without dropping non-affected inflight requests.

Key Areas to Cover

  1. System Architecture: High-level topology covering the API gateway/router, scheduler, inference engine workers (e.g., Tensor Parallelism across 8x H100 nodes), and object storage for weights.
  2. KV Cache & Memory Subsystem: Allocation strategies (e.g., PagedAttention), handling KV cache exhaustion mid-generation (preemption, swapping, recomputation), and prefix caching.
  3. Scheduler Mechanics: Admission control, interaction between prefill-heavy and decode-heavy iterations, and prioritization policies.
  4. Streaming & Cancellation: Mechanics of propagating client abort signals down to GPU execution batches to immediately reclaim token memory.
  5. Autoscaling & Observability: Metrics used to autoscale (e.g., KV cache saturation vs. queue wait time vs. raw QPS) and key health indicators.

Hints

Hint 1

Differentiate the operational profiles of the prefill phase (compute-bound) versus the decode phase (memory-bandwidth-bound), and consider chunked prefill to avoid decode starvation.

Hint 2

Think of GPU memory allocation like OS virtual memory: using fixed-size blocks (paged memory) avoids external fragmentation and allows sharing of prompt tokens across requests.

Hint 3

For cancellation, ensure the scheduler updates its active sequence masks before the next iteration step so that discarded tokens do not consume kernel launch slots or memory.

Follow-up questions

  • How would you architect prefill/decode disaggregation, and what network overhead is introduced when transmitting KV cache states across nodes?
  • How would you implement multi-tenant prefix caching across concurrent sessions sharing a large system prompt or RAG document?
  • If an unexpected GPU memory spike causes an imminent Out-Of-Memory (OOM) mid-batch, what deterministic preemption policy minimizes wasted compute across the batch?