System Design Interview: Design an AI Inference Stack | Part 3

Hi Inner Circle!

This is part three of the system design questions I think you should be ready for this year.

Design an AI inference stack.

Inference is the part after training — a finished model takes new input and produces an answer. You already know how to call one. You send a prompt, and the response streams back a word at a time.

The interview isn't testing whether you can call a model. It's testing whether you know what's happening on the GPU while those words are appearing.

There are five pieces to be ready with.

The stack, at a glance

• Gateway ~ auth, rate limiting, and a token budget

• Router ~ picks a model server using live load and cache signals

• Inference engine ~ vLLM, SGLang, or TensorRT-LLM

• KV cache ~ living in GPU memory, next to the model

• Continuous batching ~ many conversations sharing one card

• Kubernetes ~ scaling and self-healing, with warm capacity

Why the gateway needs a token budget

Auth and rate limiting are standard. AI adds one more.

You're not just limiting how many requests someone sends.

You're limiting how much text they can generate.

A single request asking for a thousand words costs far more GPU time than a hundred requests asking for a sentence.

That's a token budget, and it's what stops one user from consuming your entire GPU capacity with a handful of well-crafted prompts! :)

Why routing isn't round-robin

Once you're past the gateway, something has to decide which model server handles this request.

Round-robin is the wrong answer here, and saying so is worth points.

The right signals are real-time metrics:
→ How deep is the server’s request queue?
→ How heavily loaded is it right now?
→ How much KV cache is already being used, and what cache capacity is available for reuse?

Cache-aware routing matters more than people expect.

If a server already holds the KV cache for a long, shared system prompt - routing the next request with the same prompt to that server lets you reuse the cached entries instead of recomputing them.

The router picks the server. Now ~ what actually happens on it?

The two phases

This is the core of the answer, and it's the part most candidates never reach.

Generating a response happens in two phases, and they behave like completely different workloads.

Prefill. Your prompt gets split into tokens, and the model processes all of them at once, in parallel. It computes the query, key, and value projections for every token together. Heavy matrix multiplication, GPU fully occupied — this phase is compute bound. Out of it comes the KV cache.

Decode. Now the model generates the answer, one token at a time. Each token depends on the one before it, so there's nothing to parallelize — it's sequential by nature.

And here everything flips. Decode does almost no math per step, but it has to read the entire KV cache from memory on every single step. The GPU is waiting on memory, not on math. This phase is memory bandwidth bound.

That's why the response streams instead of arriving all at once. You're watching decode happen, one token at a time.

And decode is where almost all of your time and money goes.

The KV cache

When the model reads a token, it turns it into two vectors ~ a key and a value.

To generate the next token, the model needs the K and V of every token that came before it. And those never change. Token five's key and value are identical whether you're on token six or token six hundred.

So you store them instead of recomputing them. That's the KV cache, and it's what makes decode cheap.

It's also what fills your GPU.

The cache grows with every token generated, and it grows per conversation. Your GPU is holding the model weights plus a cache for every conversation running right now.

That's the number that decides how many users fit on one card. Not FLOPs. Memory.

Why continuous batching is the highest-leverage thing here

A GPU serving one person is mostly idle. So you run many conversations at once, in a batch.

But they don't finish together. One person wants two lines, another wants three pages.

If you wait for the whole batch to finish before starting the next one, the short requests sit there holding a slot and doing nothing while the long one runs.

So instead, the moment one request finishes, a new one takes its place. The batch refills continuously, at every generation step.

Same GPU, same model, several times the throughput. If you only get to make one optimization argument in the interview, make this one.

Scaling, and the thing that breaks autoscaling

Traffic goes up, you need more servers. One dies, you need it replaced. Kubernetes handles both without anyone doing it by hand.

But GPU servers are slow to start. Pulling a large model and loading it into memory takes minutes, not seconds. You cannot scale up fast enough to catch a spike — so you keep spare capacity running ahead of demand and plan for it, rather than reacting to it.

And you don't scale on CPU utilization here.

You scale on queue depth and remaining KV cache.

Two numbers to watch:

Time to first token. That's what users feel — the wait before anything appears. It's dominated by prefill and queue time.

Cost per request. The number that quietly kills you.

Let's follow one request

1. The request hits the gateway. Auth, rate limit, and a check against the user's token budget.

2. The router looks at queue depth, current load, and free KV cache across the model servers — and whether any of them already holds the prefix this request needs. It picks one.

3. The inference engine tokenizes the prompt and hands it to the scheduler. It joins the running batch at the next step, rather than waiting for a slot to free up.

4. Prefill runs. Every prompt token processed in parallel, keys and values written into the cache. The first token comes out — that's your TTFT.

5. Decode begins. One token per step, each step reading the full KV cache back out of memory, interleaved with every other conversation in the batch.

6. The request finishes. Its KV blocks are freed, and the next queued request takes the slot immediately.

One request. Two phases. Hundreds of memory reads per word.

Interview takeaways

→ Be specific about what you’re optimizing. TTFT often involves prefill and queueing, while throughput often depends on decode and batching.

These aren’t hard rules; the bottlenecks vary by workload, hardware, and optimization techniques already in place.

→ Prefill and decode have different hardware demands. Disaggregating them across GPU pools can optimize each phase independently, but adds communication and orchestration overhead.

For smaller workloads, weigh the trade-offs before adopting it.

→ KV cache limits concurrency. Context length and the number of active requests determine KV cache demand, affecting how many concurrent requests a GPU can support. Cache fragmentation can further reduce memory efficiency.

→ Cold starts need capacity planning. If loading a model takes minutes, warm pools and spare capacity matter more than simply scaling on CPU utilization.

Final thoughts

This design features vLLM, SGLang, and TensorRT-LLM as the inference engine options.

The engine changes what you tune, but the underlying principles remain the same: every engine deals with the prefill and decode phases, along with KV cache constraints.

The same applies to other optimization techniques I haven’t covered, such as quantization, speculative decoding, chunked prefill, and prefix caching. Each addresses different bottlenecks, and understanding which phase you’re optimizing helps you reason about when and why to use them.

That wraps up the three-part series.

Parts 1 and Part 2 are available here if you missed them.

See you in the next edition!

Stop typing AI prompts. Start talking.

You think 4x faster than you type. So why are you typing prompts? Wispr Flow turns your voice into ready-to-paste text inside any AI tool. Speak naturally, tangents and all, and Flow cleans it up. Available on Mac, Windows, iPhone, and Android.

Try Wispr Flow Free