LLM Inference API with Dynamic Batching

0) Problem Restatement

Design a high-concurrency inference API for a large language model that runs on a limited pool of GPUs. Many independent requests arrive all the time. A GPU is far more efficient when it processes many requests together (a batch). So the system should group compatible requests into shared GPU calls, increasing throughput, while keeping each request's extra waiting time small, and still stream each answer back to its own caller. Anthropic asked this twice.

1) Requirements

• An API: POST /v1/generate { model, prompt, max_tokens, temperature, stream }.
• Batch compatible requests (same model; compatible settings) onto GPUs.
• Stream tokens per request as they are generated.
• Bound added latency: e.g., a request waits at most ~10–20 ms to join a batch.
• Timeouts, cancellation, priorities (paid vs free), and a clear "overloaded" response.

1.1 Why batching matters (simple math)

Generating one token for one request loads all the model weights from GPU memory, and that's the slow part. Generating one token for 32 requests at once loads the weights once and does 32× the useful work. So throughput can grow almost linearly with batch size, until the GPU's compute or memory limits.

2) Architecture

flowchart LR
    C["Clients"] --> GW["API Gateway - auth, rate limits"]
    GW --> RT["Router - picks model replica"]
    RT --> Q1["Replica 1 queue"]
    RT --> Q2["Replica 2 queue"]
    Q1 --> B1["Batcher / scheduler"]
    B1 --> G1["GPU worker - model loaded"]
    G1 -->|"tokens per request"| B1
    B1 -->|"stream"| GW
    Q2 --> B2["Batcher / scheduler"]
    B2 --> G2["GPU worker"]

• Router: sends each request to a replica of the right model with free capacity (least queue depth).
• Batcher (per replica): collects requests from its queue and forms batches for the GPU.
• GPU worker: runs the model on the batch, and returns the new tokens for each request every step.
• Streaming back: the batcher routes each generated token to the right caller's connection (SSE).

3) Batching Policies

3.1 Static (simple) batching

Wait until B requests are waiting or T ms have passed since the first one, whichever comes first. Then run the batch until all of them finish.

• Problem: a short answer (10 tokens) must wait for the longest answer (1,000 tokens) in its batch, and new requests can't join until the batch ends. The GPU sits partly idle as requests finish.

3.2 Continuous (in-flight) batching (our choice)

The model generates one token per step for every active request. So at every step:

• Remove requests that finished (hit an end token or max_tokens) and free their slots.
• Add waiting requests into the free slots (first process their prompt, the "prefill", then join the generation steps).
• The batch always stays full, short requests leave early, and new ones start almost immediately. This is how vLLM and TGI work.

3.3 What limits batch size

Not just the request count: each active request needs KV-cache memory on the GPU (its stored attention state), which grows with prompt + output length. The scheduler admits a new request only if there's enough KV memory for it. Paged KV memory (as in vLLM) reduces waste.

4) Key Flows

1. A request arrives → auth and rate limit → the router picks a replica → it waits in that replica's queue (with a deadline).
2. At the next step with space, the batcher admits it, running the prefill (maybe in chunks, so a very long prompt doesn't stall everyone).
3. Each step, new tokens are sent to each request's stream.
4. On finish, cancel or timeout: remove it from the batch and free its memory. If the client disconnects, cancel right away, since otherwise GPU time is wasted.

5) Fairness, Priorities and Overload

• Priority queues: paid or interactive traffic is admitted first, and batch/offline traffic fills leftover capacity.
• Per-tenant limits: tokens per minute and concurrent requests, so one customer can't fill every batch.
• Queue deadlines: if a request can't start within, say, 5 seconds, return 429/503 "overloaded" quickly instead of timing out later. Clients retry with backoff.
• Bucketing by length: keeping very long and very short prompts in separate pools or replicas can improve latency.

6) Scaling the Pool

• Each replica = one model copy on 1+ GPUs. Throughput per replica comes from load tests (e.g., 2,000 tokens/sec at p95 time-to-first-token of 500 ms).
• Replicas needed = peak tokens/sec ÷ tokens/sec per replica, plus headroom.
• Autoscaling is slow (loading weights takes minutes), so scale on forecasts and keep a buffer.
• Track separate SLOs: time to first token (users feel it) and tokens per second per request.

7) Trade-offs & Alternatives

8) Wrap-Up

Batching lets a GPU serve many requests for about the cost of one step each, so put a batcher in front of every model replica. Use continuous batching: at every generation step, drop finished requests and admit waiting ones, limited by available KV-cache memory, and stream each request's tokens back to its caller. Add priorities, per-tenant limits, queue deadlines with fast "overloaded" errors, cancellation on disconnect, and capacity planning based on measured tokens per second per replica.