Deep dive on LLM Inference at Scale — Harshul Jain, Audible & Tanmay Sah, Independent AI Researcher

Uh, so, good afternoon, everyone. My name is Harshul Jain, and he is Tanmay Shah. And we would like to welcome you all in this 2-hour workshop on the LLM inference. So the goal of this workshop is to understand this domain from the first principles, dive deeper into it, and, like, understand what's going on throughout the industry.

A bit of background about us: so I'm a senior software engineer at Audible, have been building MLAI data platforms for the past 5 years, and on the sides I've been writing this open-source handbook on LLM inference. And Tanmay, he is the senior quantitative modeler at Zion's Cup Bank Corporation.

He recently completed his PhD, and he has been actively doing research in the agent verifiers and the world models.

Uh, so a quick show of hands here. Who here is, like, brand new to the LLM inference? Okay, great. And who here has, like, deployed these models in production? They have been tuning it, they have been serving the production traffic.

Okay, great. So this workshop is targeted towards the beginner and the intermediate level. And all of the slides and exercises, they are in the repo. I will share that soon.

Here is the quick agenda for the workshop. We will start with the problem statement. We will try to understand a few of the pain points around LLM inference. Then we understand what causes those pain points and build our foundations from there.

Then we will dive into, like, two kind of the optimizations that we do, like the model optimizations and the serving optimizations. And then we start learning about different serving engines that are available to deploy our LLM inference solutions in production.

And we will showcase some benchmarks and the decision chart on, like, which engine to use.

Cool. So to understand the pain points, first we need to know what is, like, LLM inference. So I'm probably a lot of us already know this. But yeah, anything that you ask your AI to do, like whether it be generate a video, audio, analyze any text, analyze your medical reports, or like your tax bills, all of that is like an LLM inference.

And this market is like approximately $23 billion today.

SemiAnalysis recently shared that if you want to model, like, a Google search queries with LLMs, you need, like, a profit drain of like $36 billion. And query cost has to be less than 0.5 cents to keep your search business profitable.

On the other hand, the Business Insider mentioned, like, your AI has to be put on diet. And everyone has to start auditing and budgeting their token usage. And all of this is happening, why? Because your hardware is limited, compute is expensive, your inference is expensive.

And with the growing need of, like, more and more AI usage, this inference cost is rising more and more. So this stat, it's an old stat from the OpenAI, but it's still, it's still true. So if you look at the, like, training cost of the GPT-3, it was like around $4.6 million.

It was a one-time cost. But if you see the inference cost, that has been like, it's a recurring cost because it's an operating cost that scales with every user that comes in, every token that comes in, every session that is being initiated on the, like, AI.

And there are only two ways to basically counter this. One way is you reduce your token usage. Alternative is you should try to optimize your inference solutions as an inference service provider for your customers and for yourself. And so this, we have been seeing, like, a lot and lot of, like, new solutions coming out every then and now.

And so the idea would be like, okay, we will try to build those foundations that will help us understand and evaluate, like, whatever ships next.

So yeah, to get started, like, we will do a quick demo. Like, it's a short demo of, like, what are the different pain points around inference. And so this is the repo. I mean, you can pull it or you can also open it on the GitHub.

It's called LLM Inference at Scale. A bit of background here, like, four months back when I didn't know anything on the LLM inference, I started learning it. I saw, like, a lot of resources were scattered. So we started putting it, like, all together in one place so that it could benefit people.

Yeah, so let me actually get out of this slideshow mode and probably go into

I will go to this extended mode.

Okay, great.

Yeah, so in this repository, if you see a README file, there is, like, a link to the slides. So this, it will be like this folder where you have, like, a PPTX and there is, like, a benchmark report in there.

You can always, like, download it. And then for the demo purposes, we have a couple of Jupyter notebooks. We have, like, collaborated with Molab who are the, like, Google Colab alternative. And what they basically provide you is, like, a free RTX 6000 GPU.

It's a 100 GB VRAM GPU. So Molab, and we have, like, already set up these notebooks so that it becomes easy to, like, experiment with. And, like, all of the assets and everything are preset for you.

So we will start with, like, a simple demo A. Probably,

let me just see.

Okay.

Yeah, so when it comes to the inference, you need to do an inference on a certain model,right? For the workshop purposes, we are using a simple Mistral 7B model. It's a small model of around 15 GB in size.

So we are going to, like, load that into the GPU. So, and we would look like some of the GPU stats as well. So we see, like, okay, we are working on this 6000 Blackwell. And you might be thinking I'm not running the cells because I don't trust the Wi-Fi at conferences.

So yeah, so I would probably be just going over the results that we kind of ran previously.

Yeah, so we have, like, a GPU which is like 102 GB. Now, the first thing that comes to my mind is, like, what's my memory consumption looks like when I do the LLM inference. So I load this model and I see, like, okay, I have, like, a 15 GB here.

So I have roughly, like, 87.5 GB. And now when I do the, like, inference here, what I notice is, like, the more the number of inputs I pass, more is the memory that I need. And it's increasing slowly, but it's still increasing.

So imagine, like, if you have a context length of, like, around 4,000 or 16,000 or 32,000 tokens. So this memory could, like, really grow big and it, you could actually get, like, all of those out-of-memory issues. So definitely this is, like, your problem one, like, your memory increasing with the increase in tokens.

So in form of, like, a simple visualization, it looks like this.

The second problem that you would see is, like, the time to your first token, it's very, very slow. We measure it by a metric called TTFT. It's a short form of it.

And when you try to, like, measure the TTFT with the, like, input size, you would see, like, longer the context, you would see, like, this TTFT being slow. So now there are two problems. Your memory increases with the token size, your TTFT increases with the token size.

Sorry, not the token size, the context size.

And then the third is the, like, throughput. The throughput is, like, how many tokens can you serve per second and then how many users can you serve per second. So if you take a very, very vanilla implementation on your local system, it would be, like, very sequential.

So if you send, like, five requests, all those five requests would be catered, like, sequentially rather than parallelly.

And so, like, your request basically takes more time to complete if you have, like, multiple users. So these are the, like, three problems. There is a fourth one I haven't described it here. Probably we will build that intuition as we move forward.

But let's remember, like, these are the three problems: the memory, TTFT, and the throughput.

Cool. I will go back to the slides.

Okay, perfect.

What is the URL for this? The one that you are showing? Can you show the URL?

So it should be this.

Is it visible?

Yeah, there's no directory card.

Oh, you need the, oh, okay, workshops. Okay.

Yeah, so within that repository, if you see a workshop folder, you see that README and then the README has all the links, the slides, and the demos.

Does that work?

Okay, perfect.

Okay, so let's start working through the foundations. Like, let's start understanding what are the reasons behind those pain points. And for that, like, we have to look at this inference pipeline.

So we get, like, an input text. That text could have, like, any number of words. You convert those into the tokens. So for simplicity, you can assume one word equal to one token. Then you kind of convert them into, like, the embeddings.

And then you send it to the, like, transformers. Like, there are 32 layers of transformers, but that's specific to the Mistral 7B. Different models have different kind number of layers. And then you generate a new token. And that token basically goes back to the input.

Then you generate another token. And that keeps on going. Now, in this entire pipeline, you would see, like, 95% of your compute is, like, taken by these transformer layers. So it's worth looking at, like, what goes within this transformer layer.

Within this transformer layer, you would have, like, more layers. You have, like, a normalization layer, you have an attention layer, you have a feedforward layer and all. And attention layer is the one, I think, that has been very, very famous.

Attention is all you need paper. I think that's very well known. So attention is the most compute-intensive layer. And we need to understand what goes within that attention layer. So what does attention do? Attention, so if you have an input text, it needs to find the attention scores of every token with respect to all of the previous tokens.

And to do that, what it needs to do is, like, it needs to project every token into, like, a key query and the value space. So in, like, in a simpler terms, just understand this, like, if you have 10 tokens, then it needs, like, the 10 different query key and the value vectors.

If there are 100 tokens, you would need 100 key and the value vectors. If there are 1,000 tokens, you would need 1,000 key value vectors. And so, like, your number of the key and the value vectors, they increase as you increase the input size.

And if you calculate the, like, KV size per token, for Mistral 7B, it comes out to be 131 KV. This is because, like, you have two vectors, K and V. You have to multiply the size. One vector is, like, 128 dimensions.

You have to multiply it by 32 transformer layers. And then you have to multiply it by the KV heads. For Mistral 7B, it's 8 KV heads. It's not, like, 32 because it uses a different kind of an attention mechanism, which we will talk about for sure.

But yeah, so the KV size per token is, like, your 131 KV. Now, imagine if you have 4K context. So that size becomes, like, half a GB. If you do, like, 16K context, that size becomes 2.1 GB. Now, multiply by it by the users.

Like, assume you can serve multiple users together. At the same time, within that GPU, you could have, like, 42 GB with a 4K context and 80 users. And if your GPU is only, like, let's say, 24 GB, you are already running out of the memory.

So you cannot serve that many users with that many context.

To visualize this, look at a GPU memory. So the GPU memory has, like, a model weights which are pretty fixed. These are pre-trained weights. There is, like, an overhead that is also fixed. That also, like, that changes, but it does not change that much.

Overall, you can assume it's fixed. And then there is, like, a leftover memory. So this leftover memory is what being used by your KV, like, key and the value vectors. So

assume, like, you have one user. You can only serve that many key and the value vectors or that many tokens which can, like, fit in this entire 80 GB, like, memory that is left.

So we can show this with a simple demo too.

Okay, let me,

okay, great.

Okay, great.

Let me see if I can actually run this.

Probably, why the heck is this happening?

Okay, great.

Yeah, so you would see, like, the GPU is attached.

So here we are just trying to confirm the, like, memory based on the maths and based on the intuition that we have built. So the model memory is, like, let's say if you have 7 billion parameters, you are doing a 16-bit precision.

Your total memory comes out to be 14.6 GB. You can basically verify that with the maths. So if you do all that maths, that comes out to be the 14.6 GB. Now comes the KV, K and the KVV size.

So this KV size is, like, your 131 KV per token.

And if you do that maths and you try to, like, visualize this.

Oh, sure.

Okay, and then let's just visualize this. Okay, great.

Yeah, so this is the, like, memory chart. So if you see, like, as your context increases, your memory keeps increasing. Then another thing to realize is, like, as your users increase, then also your memory increases. So if you want to serve, like, 160 users on a GPU, you can support, like, you can only support, like, a lesser context length.

So there is always a trade-off between what context length you can serve versus how much cost you can save by, like, putting your multiple users or the concurrent users into, like, a single GPU. So you have to always take that trade-off.

And we will go through that, like, in a couple of more slides.

Do you create different pools for different context lengths and serve greater on that?

Can you repeat, please?

Do you create different pools for different context lengths or just memory optimization for the high payload?

I'm sorry, I cannot hear you.

Do you create different pools with different context lengths so you can serve with theright memory optimization?

Yeah.

Cool.

Okay, great. So let me pull back. So that was, like, memory. We need to understand why we had, like, a slower time to first token when we increase the context length. So for that, like, we need to understand the two phases of inference.

And those phases are, like, the prefill and the decode phase. I think you would all seen, like, a lot of articles, but we just wanted to explain it. So when you send, like, a lot of, like, when you send these input tokens, what you want to do is you want to build those key and the value vectors that I mentioned for all the tokens.

Then you want to compute the attention scores of every token with respect to the previous token. All this operation that you do, it's a very, very metrics heavy. It's a very, very compute heavy operation. And we all know, like, the GPUs, they are, like, very well suited for a heavy compute workload.

So we call, like, a prefill to be, like, a compute bound. And it does take some time to complete. So whatever time that this phase takes to complete, that's your time to the first token. So if you have, like, more input tokens, you have to generate more key value vectors.

You have to do a lot more attention math. And because of that, your TTFT becomes more, more slower. Whereas if once you generate one token, you need to keep doing this to generate another tokens sequentially, one after another.

But in that process, every time you have to build the key and the value vectors of all the previous tokens, which is same as prefill. Like, you were building key value vectors there also, here also. But in decode phase, you are only computing the attention math for the new token.

And that is why it's a very less, it's lesser compute oriented. And it's also called as memory bound. We will see it shortly why it's called as memory bound.

So in a classic timeline, you would see prefill and decode phase like this. So time taken by prefill, that's your time to first token. And then your time taken by every decode step, that's your, basically, your inter-token latency.

So that's, like, the fourth metric

that you need to worry about, like, what's the time being taken by your decode step.

Okay, cool. Now, why does the, like, decode step or why does decode takes time and why it's being called as, like, a memory bound operation? Let's try to understand that. To understand that, we need to look at how the metrics math basically works on the GPU on a high level.

So GPU has two kind of memories. You have a high bandwidth memory, you have a shared memory. So the high bandwidth memory is a larger size, but a lower, like, lower bandwidth. By lower bandwidth, I mean, like, you can transfer data out of it at a lower rate.

Compared to the shared memory, so the shared memory is smaller in size, but it has a very, very high bandwidth. That means you can transfer data in and out of it with a very fast thing. So when you have to do a metrics math, so you have to pick the data in chunks from the high bandwidth memory, you have to put it into the shared memory, do that math, write back the result into the high bandwidth memory.

For the prefill phase, when you have to do this, you have to do this metrics math only once. But for the decode phase, you have to do this metrics math, like, again and again because you are generating each and every token sequentially.

And so, like, you, it doesn't matter, like, how fast is your decode because now you can transfer your data out of the high bandwidth memory into the shared memory at a certain speed because you are limited by the high bandwidth memory bandwidth speed.

And so that governs your, like, token ceiling, like, at what rate can you actually generate tokens out of the decode step.

If we look at this in the roofline plot,

so there is a left section which is called to be a memory bound. Mathematically, it's governed by the arithmetic intensity. Arithmetic intensity is the number of flip-flop operations that you perform per byte of data being transferred. So for the decode step, decode step, since you are transferring a lot of data, like the key and the value vectors of all the previous tokens, the model weights, but you are doing the, like, less computation because you're computing attention math for only one token.

So its arithmetic intensity is very low. But for a prefill phase, you are transferring the data once, but then, like, you are doing this heavy computation. And so its arithmetic intensity is very high. So now you know, like, in terms of mathematics, like, why the compute, like, why the arithmetic intensity of prefill is very high compared to your decode.

Okay.

So this is, like, another small, small demo.

Let me time I have to, okay.

Okay, great. I hope this is already running. So yeah, again, we are loading the model. Now, this is the, like, the prefill cost. So what we are basically doing is we are getting the, like, the input text, and then we are trying to generate this, the prefill step, the amount of time it takes.

We see, like, as we increase the, like, size of the input tokens, this prefill is increasing. So your, and this is the reason why your TTFT increases. And then, like, your decode time. So the decode time is, like, on average, it says about the same.

And so if it is, like, assuming, like, you ignore the, like, cold start, your decode time is, like, approximately around the average line. It is still impacted by, like, the input size. It's not, like, it's a constant time.

And it is because it still needs to pull the key and the value vectors from the memory for all the previous tokens. So there is still, like, that basically small increase in time that you would see with the decode step.

And then this is the, like, classic roofline plot.

Okay.

Presentation.

Okay.

Five. Okay, great. Okay, so now, now let's try to understand, like, the throughput dimension. You want to understand how many users you can actually serve. And I think we saw, like, a diagram of the GPU memory where we saw, okay, there is some memory that is free for the key and the value vectors to grow.

So assume, like, you have just a single user.

What's the total KV size that you have you can basically support? It's defined by your context limit. The max users that you can support is, like, whatever is your GPU availability, like, whatever is the memory that is available in the GPU, you divide it by the key and the value size per user.

And when you do that, like, it comes out to be, like, your the concurrent users. Now, assume, like, your GPU is fixed, your model is fixed,

so your KV size per token is fixed. There are only two dimensions that are left here, which is context and your concurrent users. If you want to serve more concurrent users, you have to reduce the context length. If you reduce the context length, you could impact your quality.

So these are the two dimensionsright now that we are trading off. Then if we, but can you actually serve the, like, max number of concurrent users? In an ideal world, probably not because every business has, like, a latency SLO that we have to meet.

So if you remember, like, in the decode step, I said the time for the decode still increases if you have more inputs. It also increases if you have more users. So ultimately,

your inter-token latency also gets impacted if you have, like, a higher batch size. And your TTFT also gets impacted. So now there is a third dimension you have to worry about, which is, like, your latency. So the three dimensions that you have is, like, a quality, latency, and the throughput.

So it comes out to be like this trade-off triangle where you have to choose between the two. So for a premium chat application, you would want to prioritize definitely the quality, and you want to prioritize the, like, the latency.

You would not want your users to wait infinitely for the, like, like, not infinitely, but probably for the larger latency. You can always sacrifice the number of users you can support on the GPU and probably take that costed being more customer obsessed.

In form of, like, and, like, if you consider, like, an agentic, sorry, the async agent workload, you would want to, like, prioritize definitely quality and the throughput

because these are the long-running tasks. And you would want to, like, serve as many as concurrent tasks as possible, but with a very, very higher quality.

And often, like, we think, like, okay, if the GPU is, like, a very expensive GPU, that might not be a good fit for us. But it turns out that could actually serve you the lowest cost per million tokens.

But you really have to trust your kind of calculations on the max users that you want, and, like, you really have to make those estimations correctly.

So we do have, like,

let me just, where is this? Okay, great. So for the capacity calculator, there is, like, a link to the Colab because I was facing certain issues with Molab. I had to migrate out the whole widget library, and I didn't have time.

So being lazy, I just picked Colab there. Apologies to Molab.

So my VRAM is connected.

Wait. Okay.

Wi-Fi, probably.

Okay, great.

So what we have done over here is we have, like, shaded some, like, the GPUs with their VRAMs, with bandwidths, the flip-flops, and the cost per hours.

Then we kind of, like, built this simple, like, capacity calculator. This is just a KV visualizer where you kind of, like, when you increase the number of tokens, you see, like, your KV size, it increases. And when you increase the number of users, your size is, like, increasing at a much faster rate.

And then

in this capacity calculator,

let it run. So we have, like, a model which is, like, a 7 billion parameter model that we selected.

We set the, like, precision to be FP16. Now, we decide, the way we basically go by the GPU decision is you have to decide what's your, like, you have to fix one dimension first, which you care about the most.

For premium chat, I mentioned, like, latency is definitely the one. And then, like, for the async workloads, the batch, the minimum batch size that you want to serve for, from, like, a single GPU, that is the second dimension.

So you want to fix these first. So I will go about, like, in a premium chat application.

So I can go ahead with, like, 10 milliseconds latency. Minimum batch size, I don't care. Like, I can serve, I'm okay with, like, probably two.

Okay, so probably with the seven concurrent users on a single GPU. And then, like, my context limit is very important to me because I want to focus on the quality as well.

And so, like, I do see, like, some of the GPUs. So the H108GB, it's like a $8 per hour, but, like, am I 300x? Is it? Yeah. So it's like around $10 per hour. But if you do all that throughput math that we shared in the mathematics before, you could find, like, your cost per million dollar tokens.

That could be very, very, that could be, like, lesser. So you need to do such calculations by fixing those dimensions, and you need to decide your GPU to, like, reduce your kind of inference cost. This is at least the first step that you can take towards optimizing the inference.

Okay, cool.

So the next slide. So let me, okay, great. And so, like, now the next thing is about the model optimization. So we are now basically have built that foundation where we understood some of the pain points, reason behind those pain points, why those were happening, how we could, like, address that GPU capacity thing.

We need to understand what can we do, like, what can we further do about it. So it is about, like, model optimization, and I think I would like to invite Tanmay. He can talk more about these model optimizations, provided he has worked on this, like, during his research times.

Okay.

I can control.

Yeah.

Here. Okay.

Hi, everyone. Mic check. Am I audible at last? Yeah. Okay, so hi, I'm Tanmay Saha. I work as a senior quant modeler and also I'm an AI researcher. My work focuses on agent verification andright now building world models.

So for this one, model optimization, before we start model optimization, so I created a research template so that it will be easy for us to understand all these complex things. So our template is simple. First, we will identify the problem.

Second step, we will solve the problem using two algorithms. These are just fake algorithms. So first algorithm is called ostrich algorithm. Whenever we see, just like ostrich, whenever we see a problem, ostrich put their head into the sand.

So same thing we will do. Whenever we face a problem, we will just ignore it. So this is an important algorithm we should follow. Second one is created, it is called world cup algorithm. For example, we don't know who will win this FIFA World Cup.

So what organizers did, they break the 48 teams into 12 groups, then round 32. So round 32right now is currently going on, then round 16, then quarterfinals, then semifinals, and finals. So what they are doing is that they are breaking it into smaller problems, and the useful results are moving forward.

So same analogy or same algorithm we will use to understand this model optimization, all those things. So yeah, let's start. So I have one H100 GPU. I have to use this open source model, what is called GPT-OSS 120 billion parameter model.

Soright now, I think it's, so they have trained it on BF float 16, and weight is 240 gigabyte. What should I do? This is the problem we have. So first thing, what we have to do is that 240 gigabyte and 80 gigabyte H100.

So, and I have to fit only in one GPU, not in multiple GPU. So what can we do? I think simple step is that just compress it. But how should we compress it? That's the another challenge. So if we compress BF float 16 to FP8, then it will be around 120 gigabyte, but our GPU H100 is still 80 gigabyte.

So what I think they did is that they compressed it into further MX FP4, and I think size is around 65 gigabyte. So this is something we can do, compress, but question, so, and we will use over this ostrich algorithm.

We are assuming that there is no loss in compressing a bigger model into a smaller size. Second thing,

in this one, okay, yeah. So in this one, in this slide, we have used this Mistral 7B, so 7 billion parameters. So it's a small model, 7 billion parameters. So if you multiply it by two bytes, so weight of its around is 14.5 gigabyte, which can easily fit into H100 or even A40.

So next, what we can do is that, like, Mistral 7B, instead of compressing it of floating point 16, we can apply different techniques like int8 or int4 or NF4. So basically, we have to just use ostrich algorithm and just believe that there is no quality loss kind of things.

But somehow we also have to mathematically prove that by doing some kind of testing on some external benchmark that whether it is working or not.

And this comes under post-training quantization kind of thing. One can also do this one during fine-tuning. One can also do this kind of quantization. This comes under a quant-aware training kind of thing. So let's move to our next problem.

So

we have this huge matrices. Just imagine 1000 by 1000 dimension matrix A and another matrix 1000 by 1000. So if we multiply by this two matrices, so number of operations will be 1000 raised to the power cube. And this is kind of a problem in

terms of computing. So we wondered our matrix multiplications should be fast and it should save memory. So what should we do? We have a giant matrix. Okay, let's take this one, Mistral 4096 by 4096. What should we do to solve our problem of speeding up the things and saving the memory?

4096 by 4096. So first thing is that we will use just our world cup algorithm. We can decide a random number, just break the block vertically. It does not matter what you are choosing it. So you have, so let's say we have 4096 columns.

We will break it into a group of 128 column each. So 128, 128, 128, 128 vertically. So we will get a 30, we will get this 32 blocks if we divide this 4096. Then what will happen by doing this thing?

So if we just divide this one vertically, then we can use a multiple GPU to speed up the process. So this kind of thing is called multi head attention. So what else can we do? We have a big matrix, like,

as I have mentioned, that ostrich algorithm. So our main problem is sizing. So what we can do is that instead of having all those 32 vertical blocks, we will throw away 31 blocks and we will assume that one block is sufficient enough that all the queries can handle those blocks.

Our loss will be almost negligible. And we come up with this algorithm, which, and this algorithm is called a multi query attention. So as we can see,right now we are at two spectrum. One is multi head attention where we split it into 32 blocks and use

different GPUs or do some parallel processing. And at the same time, we are just throwing 31 blocks and we are calling this as a multi query attention.

So at both extreme, we should be come up with a middle ground, like something we can say that instead of throwing all the 31, maybe we can group, we can group, we can group some of the blocks together so that, and we can assume that similar blocks will attend to similar kind of queries.

So this kind of technique comes under grouped query attention, which is very popularright now. Even in Mistral or in other models, this grouped query attention works. Soright now we have understand that we have a big matrix. We can divide it the way we want and doing some mathematical calculation prove that loss is almost negligible kind of thing.

So what else we can do?

So after that, after this grouped query attention,

see,

we have a big matrix. One is key and one is value. Let's compress that matrix into a latent vector and then come up with some algorithm to reconstruct from latent vector to our original matrix. So this kind of strategy comes under this one, multi head latent attention.

But again, it has some problems with rope because rope is position dependent and it is position independent kind of thing. So yeah, one needs to also include some index for keys also so that one can map it. But again, main problem is that

why we are multiplying all those big matrices? So because that's how this attention mechanism works that each token will pay attention to every token. So how about, let's don't pay attention to all the previous token, only pay attention to the important tokens, which is important for us.

So this is kind of, this kind of field is evolving. So this comes under sparse, DeepSeek sparse attention. So

yeah, and yeah, yeah. So okay, next, yeah, so next one is flash attention. So

in flash attention, so main problem is that,

so currently, so currently, not currently, soright now almost everyone uses flash attention. But in 2022 or 2023, so that's how it works. That's how it works is that, so this QK, query and key matrices, they were in HBM.

It loads, it first, it loads into this one, Tensor core, and it do some calculation and then it will write it back to HBM. And then this process goes on multiple times. So in flash attention, what they did is that, is that instead of multiplying the whole matrices, so they just divided it into, like our world cup algorithm, divided the bigger matrices into a small tile and only put those small tiles into SBM so that it can process multiplication fast and just keeping track of this some three variables so that they can calculate this online softmax.

Yeah, next one. So yeah, so this is just mathematics. So if we have a multi head attention, if it is 524 KB, then it depends upon how much

grouping we want. And so if instead of 32 KB head, we only want to use 8 KB head. So we can get a compression of 4X times. And this multi head latent attention, this formula depends on the model to model, how many layers your model have.

So in the original DeepSeek paper, I think they have some 128 dimension, 128, I don't remember the exact dimension, but according to that, they have used this one latent vector in which they have used 512 as a dimension and some 64 for rope index.

So, and then they show that it is 56X more compressed than multi head attention.

Okay,

yeah, so yeah, so this is the trade off, trade off diagram. So here I think we have not talked about this linear attention or Mamba. So main problem is just all this matrix multiplication. Right now, everyone is using attention.

Suppose in future, if we don't want to use attention, rather than generating tokens sequentially, just use maybe diffusion models where we can generate everything simultaneously. So all these algorithms will change also. But here I think they have two more.

One is linear attention and one is Mamba. So according to this slide, so if we are not compressing anything, so MHA is just, we are parallelizing the process. So there is no quality loss, so it's a good. And then this grouped query attention, which is I think almost every model is using just GQA and DSA kind of thing.

Yeah, I think same thing we are providing in the attention mechanism scorecard. So I think this one, MHA quality is good. Throughput is okay. And for grouped query attention, it depends upon your use case also.

Though quality is almost similar to multi head attention, but use case also matters a lot. Yeah, multi query attention is just one extreme. We are, I don't know why, but we are just assuming that we only need one block and all the queries will attend to that smaller block.

So quality is not that great for MQA. And this multi head latent attention. So yeah, if you have tried some this DeepSeek model, so I think they are doing great job, yeah, in quality wise. Besides that, sliding window.

So all these are some techniques which, yeah, all these are some techniques like just slide the windows, all those things. And instead of, yeah, instead of multiplying everything, so linear attention is just saying that summarize everything first and then look up into it.

And then Mamba, this is just a state space model. Yeah.

I can cover that.

Okay, good. Thank you, Tanmay. So for the model like optimizations, we also have like two notebooks here.

So there will be, I have to go to this.

Okay, so for the quantization, like the demo,

this is, is this already done? No, let me just run this.

Okay, so we are loading the model, which is like Mistral 7B.

So this one is like with the FP16 baseline.

Wait, did it run? Okay, so it's two milliseconds just run. Did this run? Okay, so yeah, this time it's fetching that model with the FP16 precision.

The Wi-Fi.

It's gonna take time.

Okay.

Yeah, because it's downloading the weights from the Hugging Face.

Huh?

This is on that Molab.

Yeah, so Molab is like running online.

Yeah, so why is that?

Because it needs to make the network call through the Hugging Face and like it. Fetching.

I don't know, like, but it's taking time to download probably.

Okay,

throughput. Okay, so here we see like the memory size is like 15 GB around approximately with the FP16 precision. We are trying to do the 2X compression as Tanmay talked about with the Int8. Let's download. Okay, so we do see like your memory size is now like 0.7.5 GB.

What that means is now you have more memory for your KB to basically grow. That means you can either serve higher context limit or you can serve the higher concurrent users there.

If you do the like in for your basic, you are doing the 4X compression. So that with the 4X compression, it would be more lower. It would be, I think, around

3 to 4 GB. Yeah, 4.5 GB.

And yeah, so this is, wait.

So this is just a basic plot of like, so these are the like theoretical numbers. We are not doing the like any throughput test here, but usually you would see like your memory increases. So probably you would also have like a bit of higher throughput.

From some of the benchmarks that we studied, we saw like the Int8 compression, it does have like a lower throughput.

Okay, and then there is like a demo on the like the attention mechanisms.

So for the attention, okay, I have to run this.

Okay, so it has run. Oh, wait, why does it say no GPU detected?

It should say the GPU should be detected.

Oh, okay.

Wait.

This is surprising.

Yeah, I guess it's not like able to detect the GPU for some reason.

We do have like a GPU here.

Okay, never mind. Yeah, so but the like basic idea here was more like

as you try to move towards like compressing the computation, like by using different attention mechanisms, like moving from the multi head to the grouped query attention and then to the MLA, you would start seeing some optimizations. I think yesterday night we were doing some benchmarking.

I wanted to correct this part. So it wasn't like 56X, it was 14X. Basically, the demo had a mistake of like a computation where it did not multiply the number of layers.

Yeah, so apologies for that. So this MLA is like a 14X savings work in comparison to like your multi head attention.

So now that we have understanding of the pain points, the foundations, the one side of the optimizations, which is the model optimizations, we want to talk about what can you do on the like the serving side. So

the first thing is we saw like when you perform like a simple decode step, you are pulling it, you are basically pulling the model weights and then you are recomputing the key and the value vectors for all the previous tokens, even though you already computed those vectors for the tokens.

So there is definitely like a lot of compute wastage.

And if you kind of analyze the time complexity of it, it would come out to be O of n squared. And the way to resolve that is like a classic trade off against the memory. You can maintain a memory of those vectors against the tokens and you can reference that memory.

So that memory was called as like KV cache. And the like the flow looks something like this.

And then based on this KV cache, there were like four optimizations that were really possible.

The first one is about the paged attention. So what's the different, what's the problem today? So when you send like multiple requests as the input to the GPU, these requests are in a batch.

Every request is allocated like a continuous memory storage, let's say of, I'm just taking an example, like let's say 2 KB. However, like your request needed only, let's say 1 KB. So there is like

a 50% of that memory fragmentation.

And this fragmentation basically leads to the memory wastage. That means there was a space in the memory where you could have served more requests, but you could not because you were looking for that contiguous block of the memory.

So an inspiration too was being taken from like how the OS works. Like you maintain a logical memory and you basically have a physical memory. So in the logical memory, it would still feel like

that the KV vector for the like every token is like a contiguous, but it will be mapping to a different physical address.

So that really helped like saving a lot of memory. And it was only possible because they considered like memory as a set of blocks and you would be dynamically allocating those blocks as the request need, as the like new tokens comes in and they need that kind of memory.

Another lever is like when you are sending multiple requests in the batch, GPU is like taking those requests, but it does not accept the new batch unless all the requests in that batch gets completed. So the diagram looks more like a paged attention, but here it is more about like when is GPU available to take the next batch.

So there is a time period where GPU is like sitting really idle and you want to like resolve for that. And for that, like the idea was like, okay, let's do that continuous batching.

So the continuous batching also really helped with like throughput because now you can ship more requests pretty quickly, keep making sure like GPU always get, it's always like occupied and it's not like sitting idle. So you are saving on that compute.

The third is the like prefix caching. So you remember like the KV cache held you, save the computation for a single request across the tokens. But what if like you have the same tokens across multiple requests? How do you basically save against that?

So the prefix caching, which was introduced by the VLLM, exactly counters that.

And then the third is like we talked about the fourth actually. So we talked about quantizing the model, but you could also, you can also like quantize the KV weights. So that means now you need like a lesser space for your key and the value vectors.

That means you can serve more key and the value vectors in the memory. And that means like you can serve more tokens. That means you can serve more context limit. And that means like you can serve more model quality.

And all of this is like

already present in the VLLM. You don't really need to reinvent that wheel.

And you can like deploy this VLLM in production and you could see that basically growth. So next we have like a benchmark that we did. So this benchmark was, let me see if I have that

here.

The demos.

So doing this benchmark takes like around one hour because you have to continuously stop and like restart the VLLM servers and you have to load the models and all. So it does take a lot of time in doing the testing, but I can like really tell you here what we are doing.

So we have kept the model as same, like the Mistral 7B. And then we have like the set of input questions that we are sending. Consider them as the prompts. Then we have a couple of helper functions here, like checking the server is up or not.

The server is the VLLM server. Then there are helper functions to get the VLLM metrics. And I will talk about like what those metrics are. Then there are like a lot of the benchmarks and all. And then you have to measure the KV usage and all.

So these are the like helper functions. So the baseline is very simple. Like we have a Hugging Face baseline. This is the raw, like sending the text to the LLM, getting back the response. We see some results here.

We saw like Hugging Face has a throughput of like around 51 tokens per second. Time to first token was like 54 and then the inter-token latency was 19. This was all run on the H100.

And then we start like a very default VLLM server. So by default, VLLM provides you the paged attention, continuous batching, and the KV caching. So three things are present by default.

And when you try to compare those benchmarks, you see your throughput is like almost 15X. You are able to serve more tokens per second. Then your time to the first token,

that also rises.

And then the inter-token latency, it kind of goes down. And then your KV versus users and the versus context it increases for sure.

Now when you apply the prefix caching to it, so with the prefix caching, you see like your throughput increases more. Your TTFT decreases. Your inter-token latency is approximately the same. And then your KV cache usage versus the users, it's kind of going down.

The versus the context, it's not going down. It's approximately the same. I think this is also approximately the same. It's like not that big of a deal. When you apply the like KV quantization on top of it, so it becomes like, so you see like your throughput is like almost similar.

Your time to first token is similar. Your token latency is similar. But then your KV usage actually goes down. This is because like you have quantized your key value space. And then there is a concept of speculative decoding that Tanmay will talk about.

So when you try to benchmark those, so you also see like there is like a bit of like the less KV usage there. Although like the results are approximately the same.

So yeah, I mean overall like these are the like the metrics across probably I should

zoom out. Okay, zoom out is not working. Great. So yeah, this is the like VLLM benchmarks. It's your production default by the way. We will also share the decision tree when we try to talk about like the other engines.

So

yeah, so we should talk about like what are some of the other inference optimizations we can do on top of it and what were some of the other solutions that came out. So I would like to again invite Tanmay.

He is going to talk about like some of these optimizations.

Oh, sorry. I'm so sorry. I didn't enable the slides.

What was the, okay, great. Perfect.

Which one?

The speculative decoding.

Yeah. Thank you, Harshul. Yeah. So all these are like speculative decoding. All these are the, so what we say, different flavors of same kind of soda. So this technique comes under decoding accelerator. So first one, so we are only talking about this speculative decoding, but there are other variants like self-speculative, eagle, medusa.

I only like, I think this one eagle algorithm. Personally, I don't think speculative decoding works because the main problem is alignment. Okay, so let's start with what is speculative decoding. Main problem is that in transformer architecture, all these tokens are generated sequentially one by one by one.

How about just use a smaller model and let a smaller model to generate maybe let's say four or five tokens and this teacher model, or we can say according to our world cup algorithm, we can say referee. So referee will decide how many tokens it accept.

And this loop keeps on going on. And our assumption is that there are certain domain where this kind of things will work, like maybe in decode, maybe in coding where almost there is no creativity. Each code or syntax is almost similar.

So maybe it can help it. But based on personal testing, I didn't find this speculative decoding useful at all. But other techniques like self-speculative decoding where teacher model also have one head, auxiliary head, and it will do same similar kind of things what this base model or small model is doing it.

But then this eagle came, eagle one, two, three, I don't know how many versions are. But it is just saying that instead of generating tokens, let's train a small model, train a small model and just take features from one of its main models layer so that instead of generating token, it will generate this feature.

So eagle is

better compared to this other kind of technologies. And then another one is medusa, which is just saying that just generate all these tokens parallelly. Okay, so here in this slide.

Yeah.

Okay.

Okay, yeah. Okay, now we come to, now we will come to this one prefix caching. So I don't know whether people are using this one static prefix caching or not. But thing is that main problem with prefix caching is that sometimes we type and make a small kind of mistake.

And this standard static prefix caching is basically it takes a prompt, does some hashing, and then next time when user ask similar kind of question, it will try to match the hash. So

if hash is equal, then it will instead of recomputing all those key and weight, will just take it from the storage. But you know that sometimes we make a mistake or maybe we can just change a word or letter, something like that, then we have a very higher cache

miss hit rate. So that's why this one regex tree. So regex tree is becoming very popular and also because of agent. So I think almost everyone is doing agent and most of the computation is going during test time, inference kind of thing where we keep on asking same kind of questions and prompt.

For example, you are an expert software engineer multiply by 200 times. This kind of loop keeps on going inside this agentic kind of things where it is all necessary to keep or store similar kind of things in a regex tree.

So regex tree is

just, so regex tree is just an advanced version of this prefix tree where

we will just collapse a node if it does not have any branch. And for this kind of work where keep on repeating same thing, this regex tree helps a lot and SGLang use this kind of algorithm for prefix caching.

Okay, yeah, then there is another thing. One is Tensor RTLLM. This is very confusing. When I first started, I was just confused what is Tensor RTLLM. So yeah, so Tensor RT

is just a standard SDK kind of thing. Tensor RTLLM is just an inference engine just like VLLM, SGLang. But problem is that it is related to NVIDIA. They optimized each and every layer and every problem. As I mentioned in our world cup algorithm, they just break everything and optimized everything at hardware level also.

So yeah, okay, next. Yeah, so for this workshop, we also did some benchmarking like which is best. So our setup was something similar. So we did two kind of testing. First one is without agentic testing where we just, so we use the shared GPT, this one data set and just ask those questions using VLLM and SGLang.

Okay.

Yeah, okay.

Let me just zoom it up.

Okay, great. Okay, yeah, so yeah, for this workshop, we used H100 and our first testing was that we just asked, we take questions from shared GPT and put it into VLLM, SGLang and we found that actually there's no statistical difference between which one is better.

So both have almost similar kind, so both are fulfilling similar kind of request per second, TTFT and latency. But

only difference we have seen during agentic branching. So what we did was that we asked that similar kind of question that you are the best, this one software engineer in the world, just solve the problem of traffic congestion in this city kind of thing.

Then we put this into LLM, LLM generates some output. Then we did another round two also. So once this LLM generates this output, then in round two, we have especially mentioned that provide review the proposal and give ratings from one to 10.

So this two times we did and this loop keeps on repeating it. What we found is that for this kind of

workflow where everything is standard, all those prompts and context engineering comes into the picture. If we do proper this agentic branching, then I think this SGLang is three to four times better. But again, this depends upon the different setup.

Maybe if you do it, you may get different results. Okay. Yeah, so I think, did we uploaded it on GitHub?

Yeah.

Okay.

Yeah, so the PDF is like also in the drive. It's the same link as the slides. So a quick summary here. So on a standard API workload throughput, you would see like a VLLM and the SGLang would be the same.

So if you don't have, if you have like a standard workload, definitely go with VLLM. It's the production default anyways. But what Tanmay was also saying is when you try to like make it like agentic workloads, that is where like your SGLang really shines and it kind of like provides you all the benefits.

So yeah, keep like VLLM as a default, but if you have agentic workloads, probably try to move towards the SGLang if you're not happy with the VLLM. But

okay, let me

wait.

Okay.

And then like there is like a comparison that is done at the 120 billion, like for the GPT OS is 120 billion. This is a benchmark that was prepared by Clarify. So there is like a blog link here.

Oh, nice. Okay. Yeah, so they did the similar benchmark and they included like a Tensor RTLLM in it.

Definitely, you can always go through these benchmarks and try to understand which basically suits your use case. As we mentioned like Tensor RT, they tried to optimize the hardware side as well, having the peak hardware performance.

Wait, this is

okay.

Yeah. And then like in terms of when you want to depict like your engines, once you figure out like between VLLM, SGLang, Tensor RT, so there are some new engines that are popping up, NVIDIA Dynamo for sure. So they are also for the agentic session routing.

Hugging Face is always there. It's a simple no server. Then there is like an MStar engine that was recently proposed by Stanford. They are for like the multi model.

So definitely you could explore those. And when you try to basically just to like give a quick summary, we start with like a baseline. We try to find what model could fit our use cases.

So you could pick like DeepSeek, you could pick like, don't pick like a Mistral 7B. I mean, it's not good. But yeah, so you pick your model and you want to like have a smaller memory and you want to try to fit that bigger model into smaller memory so that you could save cost on the GPU cost.

So you can do like all those quantization. Then you can apply all those serving optimizations by using theright serving engine under the hood. So that can really provide you that throughput that you really want.

And

now something that you can do after going back home because we cannot like actually go over all the material here is definitely reading about some of the source informations like different attention mechanisms, different like these engines, like try to just read the different benchmarks which are present online as well.

And then there are a lot of like in-depth guides or the next phases of it, which is like learning about some KV eviction strategies. So world is moving towards having a separate KV cache engineering domain. So you want to understand what's going on in there.

So KV eviction, cache compressions, hybrid memories. So there are like a lot of solutions that are happening around there. So always try to stick to those foundations or like the fundamentals or the first principles and try to see like which solution basically solves what problem and whether you actually need that problem to be solved for your use case.

And then there is like distributed LLM inference, which is like a different pain point altogether. You would probably need like a two hour workshop there as well to like go over like all the internals, do all the hands on.

Yeah, and this is something we are trying to propose for the AI engine in New York session, which is to like dive deeper into the advanced sections of the LLM inference. So this workshop was more for the like beginner and the intermediate level.

So in this form, we do have like a feedback as well, plus also the interest. If you think like we need certain improvements on certain sections, definitely give that feedback as well. And if you want to see this workshop in like New York

fair, I mean, definitely feel free to enroll your interest.

How is it possible?

Let me just check.

URL works,right? Not the QR code? Okay, probably I forgot to link those two together.

Okay, cool. Yeah, so if you can give that feedback, let me just, okay,

that will be fine. And yeah, I think we would like to wrap this workshop then. And I'm sure like a lot of you would be having a lot of questions, so we can take all those like offline. We can meet and we can like talk about those questions.

Yeah, sure.

Thank you everyone. Thanks for joining. I think it was really meaningful and all of you like came here. Thanks a lot.

Yeah, thanks.