raging inference
What even is inference? Everybody's talking about it. I try to find out in this post. (Please don't come at me if some of this is outright wrong, i am learning _>)
Inference engineering is all about making AI models work well in production. Fast, cheap, reliable and at scale. Suppose you did it. You have a 1 trillion parameters model. The holy grail of LLM. Nothing beats it and it's always right. It can write an OS from scratch (although i am pretty sure you trained it to write it), write one word email replies on threads and everything in between. But what good is such a powerful model if nobody is using it. So you release it for masses. And now your model (and very very expensive GPUs) has to serve millions of users making requests asking it to migrate their entire codebase from Python to Rust because it's trending or they want to get promoted. Simplest way is to do:
Request -> Model -> Output
Suppose your model = weights + architecture. It will have billions of parameters, transformer layers, attention, feed forward networks etc. Training has already happened at this point.
Now the request comes in, "Explain HTTP". Model needs to calculate `P(next token | "Explain HTTP") and then repeatedly
P(next token | "Explain HTTP <token1>")
P(next token | "Explain HTTP <token1> <token2>")
...
Inference is the process of performing those calculations using the trained model to produce an output.
So:
Training:
Data -> Learning -> Model weights
Inference:
Request + Model weights -> Computation -> Output
So why is just Request -> Model -> Response not enough? because model is not a function. You cannot give it an input and expect a deterministic output. To produce an output, computer has to perform billions/trillions of mathematic operations, use hundreds of GB of memory and GPUs and take seconds to generate output while generating response token-by-token.
Model does not generate response all at once. It is not:
Training:
Data → Learning → Model weights
Inference:
Request + Model weights → Computation → Output
Instead it does something closer to:
|-> "TCP"
│
"Explain TCP" ---|
|___> probability distribution
|
V
select token
|
|--------V--------|
| |
"TCP is" probability
|
select token
|
"a"
↓
...
So, it generates one token at a time. So a request isn't one computation.
Coming back to our problem statement, we have to serve millions of users and potentially thousands of them at the same time. We can simply do:
Request 1 → Model → Response 1
Request 2 → Model → Response 2
Request 3 → Model → Response 3
...
But GPUs have great ability to do parallel processing. So, we are wasting a lot of capability. That makes a case for inference engineering. Instead of sending each request to one GPU, we can do:
Request A -|
Request B -|
Request C -|--> GPU
Request D -|
...
GPU processes this requests together. So the actual system becomes:
Requests
|
v
Queue
|
v
Scheduler
|
v
Batch
|
v
GPU
|
v
Responses
But LLM generation makes baatching harder. Suppose we have:
A - needs 10 tokens
B - needs 500 tokens
C - needs 20 tokens
D - needs 100 tokens
After we batch them and process them, after 10 tokens:
A - DONE
B - still running
C - DONE
D - still running
Now our batch is:
[ B D]
We are wasting capacity. A good sophisticated inference server can instead do:
[B D E F]
This is continuous/in-flight batching. But now we need a scheduler constantly making decisions about which request should be running on GPU at this moment.
Now that we have decided which request to process, in what order, and in which batch, we have another problem to solve. Memory which my sources tell me is very expensive today.
A well trained model can have more than 70B parameters. That's a huge amount of memory. On top of this, each request needs it's own temporary state. For this, during generation, model maintains a KV cache for each conversation. At high level:
Model weights
+
Request A KV cache
+
Request B KV cache
+
Request C KV cache
+
...
=
GPU memory
For thousands of concurrent requests, we need to decide few things:
- Where do I store the state associated with every request?
- When can I free it?
- What happens when GPU memory is full?
- Which request gets admitted?
- Can I reuse cached computation?
To top it off, model might not even fit on one GPU. Suppose a model requires 140 GB memory but our GPU has only 80 GB. We need to use two GPUs to serve one request and both of these GPUs need to communicate with each other.
Request
|
v
Scheduler
|
v
GPU 1 ------- GPU 2
| |
\---- GPU 3 --/
|
v
Response
All in all, we have now introduced networking, synchronization and memory management in our inference. But there's one more thing. Latency When we ask "What' 2+2?" we expect an answer immediately. But an LLM does:
Request
|
v
[500ms]
|
v
first token
|
v
[50ms]
|
v
next token
|
v
[50ms]
|
v
next token
|
v
...
THere are different things to otpimize here. For example:
Request --------> "The"
<- 400ms ->
Token generation speed:
The -> answer -> is -> ...
^
|
tokens/sec
Total latency:
Request -> complete response
In a way, this whole thing is similar to Database. In database we are provided abstraction such that we do:
Query -> Database -> Result
But underneath the surface we have multiple components like Parse, query planner (something i am planning to implement in Rust, do follow me on Github), index selection etc. Similarly for model we have a logical abstraction like:
Request -> Model -> Output
But as mentioned in this post, there is a lot underneath the surface.
At this point, i have gone way too deep down the rabbit hole and might as well go further down. So, maybe you can come back and see what is KV cache, Prefill, batching, scheduling and all the other stuff. And i am going to implement this. In rust probably. After i figure out how to make LLM calls for free coz me broke!