Katkat
Back

Design a Dynamically Batched Inference API

AnthropicMachine Learning Engineer · Onsite
HardMachine Learning EngineerOnsiteSystem Design

Last updated

Design a scalable, multi-tenant LLM inference serving platform that balances Time-to-First-Token (TTFT) and throughput using dynamic iteration-level batching, KV cache memory management, and strict admission control.

Scenario

You are tasked with designing the core inference serving layer for a high-volume LLM API (such as the Anthropic Claude API). Clients submit varied requests ranging from short conversational turns (few hundred tokens) to massive document analysis (100k+ tokens). The service must maximize token throughput across an expensive cluster of GPU nodes while maintaining interactive responsiveness, strict fairness across tenants, and avoiding out-of-memory (OOM) failures.

Requirements

Functional Requirements

  • Inference Streaming: Stream generated tokens to clients in real-time via HTTP SSE or gRPC.
  • Dynamic / Continuous Batching: Form iteration-level batches dynamically rather than waiting for fixed-size static batches.
  • Early Cancellation: Instantly halt computation and free resources when a client terminates the connection.
  • Multi-Tenant Fairness: Ensure bursty enterprise clients do not starve interactive free or standard tier users.

Non-Functional Requirements

  • Latency: P95 Time-to-First-Token (TTFT) under 400ms for short prompts; P95 Inter-Token Latency (ITL) under 35ms.
  • Reliability & Safety: Zero GPU worker crashes due to memory exhaustion (OOM), even under sudden traffic spikes.
  • Efficiency: Maintain high GPU compute utilization (>70% sustained MFU) across heterogeneous prompt/generation lengths.

Scale & Assumptions

  • Peak traffic: 25,000 requests per second (RPS) globally.
  • Input prompt lengths: 50% short (<1k tokens), 40% medium (1k–16k tokens), 10% long-context (16k–128k tokens).
  • Output lengths: Uniform distribution between 50 and 2,000 tokens.
  • Infrastructure: Fleet of 2,000 modern 8-GPU nodes interconnected with high-bandwidth fabrics.

Key Areas to Address

  1. API & Request Lifecycle: Define client-facing and internal worker contracts, including token-by-token streaming and cancellation signals.
  2. Batch Scheduling Engine: Detail iteration-level (continuous) batching, splitting the prefill (prompt computation) and decode (token generation) phases.
  3. Memory Management: Explain how GPU KV cache memory is allocated, fragmented, and reserved (e.g., paged/virtual memory abstractions).
  4. Admission Control & Preemption: Describe load-shedding mechanisms, priority queues, and strategies for when decode memory exceeds physical capacity (e.g., swap to host vs. recomputation).

Hints

Hint 1

Differentiate between the compute-bound prefill phase and the memory-bandwidth-bound decode phase when scheduling iteration steps.

Hint 2

Model the KV cache allocation using non-contiguous memory blocks (similar to paging in OS kernels) to avoid external fragmentation and accurately predict memory limits.

Hint 3

Design a client disconnect propagation mechanism that informs the GPU scheduling loop at the start of the next decode step to reclaim active KV cache pages immediately.

Follow-up questions

  • How would you integrate Prompt/Prefix Caching into your scheduler to avoid redundant prefill passes for shared system prompts across different tenants?
  • If the KV cache pool fills up entirely mid-generation across all active requests in a batch, what specific preemption strategy would you execute and why?
  • How does disaggregating the prefill cluster from the decode cluster impact latency, throughput, and networking overhead compared to collocated workers?