Inference

Why a lecture on inference

00:04 So this is lecture 10. We're going to take a brief respite from scaling laws. And we're going to talk about inference. So the question of inference is a very simple one. Given a fixed model that we've trained, generate responses given prompts. First, we're going to start by understanding what the implications of inference are. And the workload that it entails. And then we're going to talk about ways of making inference faster. And throughout this lecture, you're going to see that there's a lot of-- inference is a very deep topic. Actually, we didn't do inference last year in lecture. So this is the first year we're doing it. But there's actually many, many topics that could span multiple lectures which I'll try to condense into one. So inference shows up in multiple different places. The most obvious place is if you actually want to use a model, you want to use it to chat. You're using Cursor or something to do code completion. If you're running batch data processing job using your language model, all of these cases demand inference because you need to generate tokens from your actual model. But it also shows up in other contexts. If you want to even evaluate your model, let's say on instruction following need to do inference. There is a lot of interest in test-time compute, which means thinking more before you actually output the final answer. And that's also more inference because thinking is basically generate tokens. And then, finally, even training itself.

01:44 If you're using reinforcement learning, you need to sample responses and then evaluate them based on some reward and that also requires inference. So inference isn't just I want to put up a chatbot demo. Inference actually is going to underlie many of the basic functions of a language model. And even though it's one lecture, I want to stress how important it is for many things. And we'll probably come back to this when we talk about an alignment later in the class. So now inference is important.

The three metrics, and why generation is sequential

02:18 So the theme of this class is efficiency. And efficiency clearly matters. Training is a one-time cost, but inference you repeat multiple times. So here's some anecdotal stats on why inference is a big deal. So Sam says OpenAI generates 100 billion words a day, which is quite a lot. And even Cursor, which is not that new of a product, is allegedly generating a billion lines of accepted code each day. So it's just giving you an idea of how much inference is accounting for and it costs of inference compared to training are definitely increasing. So how do you measure what good inference looks like/ So there's Time-To-First-Token, TTFT. So this is how long an individual user needs to wait before any generation happens at all. And this matters clearly for interactive applications. If you have a big prompt and then you have to wait there for 10 seconds, that may not be a good user experience. Latency is how fast tokens are arriving after maybe the first token. This also matters for interactive applications. Throughput is something a bit different. Throughput is how many tokens in general are generated per not for overall users. So this is particularly useful in batch processing applications. So you can think about the throughput is high throughput doesn't mean low latency,

03:56 because some requests might just take a very long time and you still have high throughput. Latency is like the worst case of an individual user. So what do you need to think about when you think about the efficiency of inference? So in training, the key idea is that you get to see all the tokens, at least a supervised training, which means that you can parallelize over the sequence. This is exploited heavily in the transformer. So you've done the transformer training. You know that you basically construct these tensors over the entire sequence. And it's just like tensor matmuls and then you get your output. But the key defining feature of inference, at least for transformers, is that you have to generate sequentially. You can't parallelize because the generation of a token depends on all of the past. So this is going to be the key thing that's going to make inference a lot harder. And in particular, it's going to be harder to utilize all the compute that's available. And it's going to be memory-limited, as we'll see in detail later. So a lot of people are doing inference. Anyone who's actually has a product and platform quickly realizes that these costs in doing large models is going to go up. So they spend a lot of time and engineering effort trying to reduce that time. So both providers serving close models and providers serving open weight models pay a lot of attention to inference. More so than I think the average academic, because we're not

Notation, and the arithmetic-intensity yardstick

05:36 actually serving any models. We're just training and getting a score and putting in the paper. But people who are actually serving models pay a lot of attention to inference. So there's also a bunch of open source packages which are interesting to look at as well. So I want to understand the inference workload in detail. So I'm going to review briefly this transformer math that you did in assignment one and we talked a little bit about it during the first week of class. So this is from the JAX-ML scaling book, which is something you guys should really take a look at. I think it does an excellent job of outlining many of the key concepts here. And they have this really nice diagram that shows essentially the computation graph taken in input and having it go through the attention and the MLP layers. In particular, we're going to use this notation so just to review this quickly. So B is the number of sequences in your batch. L is the number of layers. T is the sequence length. You can think about it as the number of tokens you're going to generate or query using. S is also the sequence length, but how many you're kind of conditioning on in your prompt. V is the vocabulary. D is the dimensionality of the model. F is the MLP hidden dimension, which is usually four times D. H is the attention head dimension. N is the number of query heads. So generally N times H equals D, and then

07:19 in GQA, grouped-query attention, you have a different number key value heads as query heads. Usually K is smaller than N And G is the number of groups. So K times G equals N. And this diagram shows that you take your X. You feed through the QKV matrices and you do a bunch of things. So remember that the FLOPs required for a feedforward pass is 6 times the number of tokens, which is B times T times the number of parameters. Plus for the attention there's another order of T. So T times T is T squared dependence. So let's also review arithmetic intensity, which is going to help us characterize when something is compute-limited versus memory-limited. So just to start with the basic matmul. So let's take a matrix X which is B by D and a matrix W, D by F. And just to give some color to this computation, B is the batch size, D is the hidden dimension, and F is the projection matrix in the MLP. So now let's do count the number of FLOPs and memory read and writes for just doing X times W. So we're going to start with initialize to 0. And what one has to do for this is we're going to read X from HBM. So that means it incurs a memory cost of 2 times

09:10 B times D assuming everything is in bf16. You also read W. So that's 2 times D times F. Then you do the matmul and that incurs 2 times B times D times F FLOPs. So remember, this is from the first lecture. So hopefully this is review. And then you have to write it back out which is you have to pay another transfer. So the total number of FLOPs is just the matmul. And the number of bytes transferred is essentially the size of all the matrices that are read and written. And arithmetic intensity is basically the ratio. So the ratio is this expression. And in general, just to simplify things a bit, generally the batch size is much less than D and F. B may be 100 and a D and F might be thousands or tens of thousands. So I'm using SymPy here just to keep myself from making silly mistakes. So basically, I'm letting C go to infinity and D scales as C times B and F scales as C times B. And that gets you a simplified equation of B. So the arithmetic intensity is B for this particular matrix multiplication. So the way to interpret this is how many FLOPs are done per byte that was transferred?

10:52 So now the second part is you look at the accelerator, which for H100 flops per second is 989 teraflops, memory bandwidth 3.35 terabytes per second. And you divide, and that gives you what is called the accelerator intensity. And if you look at the computation intensity, which is B, If it's greater than accelerator intensity, that means you're compute-limited. That means you're able to use all the GPUs or TPUs. And if you're less than that, then you're memory-limited, which is bad. And so you're compute-limited in this matrix multiplication case, if B is greater than 295 for a H100. And all this is a bit idealized. The actual details. This is giving you a first-order approximation. So in an extreme case, so generally if you that means if you use batches of size let's say 300, then you'll be going to be able to saturate the GPU. But what happens if your batch is really small? So in particular B equals 1, which essentially corresponds to a matrix vector product. Then the arithmetic intensity is basically 1. And that is really, really bad. That means you're going to be memory-limited which kind of makes sense because basically you're reading and writing this D times-- or actually you're just reading this D times F matrix and you're performing essentially the same number of FLOPs. So the ratio between the FLOPs and the reads

12:32 is the same, which gives you 1, and 1 is bad. You want a lot of FLOPs to be done for any memory read because memory reads are slow. But this is in essence, what happens with generation. Because you're proceeding token by token, we'll see that basically your arithmetic intensity is going to be 1. And that's why generation is going to be memory-limited and not compute-limited. So this is a very simple example that I think gets at the core of why generation is going to be slow. So maybe I'll pause and take any questions on this just to make sure everyone's clear. Yeah. When doing inference, if you have a batch size of larger than 1? I think I heard the question, why don't we have a batch size of more than 1? So I'll get to why you can. But there's batch size is going to mean batch size times sequence length later. So in summary, matrix multiplications are the core computation. So we just studied a matrix multiplication and counted the number of FLOPs it requires over the number read and writes. And we show that ratio, which is the accelerator intensity depends on one of the dimensions in case in this case, the batch dimension. And that's why big matrices are good because that can saturate your compute. Whereas if you have even a thin matrix, B equals 1, that's really bad because you're spending a lot of time reading from memory and not doing that much compute.

Naive decoding, the KV cache, prefill vs generation

14:19 So now let's talk about the arithmetic intensity of inference. So let's just get more into the weeds of what inference looks like. So the naive thing you can imagine doing and all these nice pictures are taken from this book is that you have a transformer. You give a prompt in. It gives you logits over the vocabulary of the next token. And you just sample from that. And then once you get that attach it to the prompt. And then you feed it through the transformer, and you look at the logit sample again and you repeat. So that's the most naive thing to do. And the complexity here is pretty bad because each token you generate is like N squared or T squared, computation through the transformer. So that's no good. But if you look at this closely, you'll notice that you're doing a lot of redundant work. All the work in basically encoding the prefix basically stays the same. So this for a bidirectional transformer would be different. But at least for autoregressive causal transformer, it is the case that you should be able to share a lot between prefixes. And so the solution is you cache and you cache in the high band

16:05 HBM because that's where you have enough space to store stuff. So this is what it looks if you have a KV cache schematically. So you take your prompt. The prefill step is you feed it through the transformer and you compute this KV cache, and then you generate the logits over the next token. And then you put that into-- take that generate a token and the cache, and then you can feed it through the transformer. But you've already computed these, so you don't have to do that again. You just need to compute this new KV vector for this token. And now that allows you to more quickly generate the next token, and so on. So basically, you're filling up this KV cache which corresponds to the tokens that you've either prefill with or that you've generated so far. So instead of T squared per token, it's going to be more like T. So concretely the KV cache is for every sequence in your batch, for every token in your sequence, for every layer of the transformer, for every head, you're going to store a H-dimensional vector. So you might think that this is going to take a lot of memory and you wouldn't be wrong. So there's two stages of inference. So prefill is you're given your prompt encoded in a vector. So this is just like what you do in training. It's parallelizable. It's fast. You're compute-limited and life is good. And then you're doing generation, which is you're generating response tokens one by one sequentially. And this is the part that's going

MLP layers: intensity is B·T

17:46 to give us a lot of trouble in terms of efficiency. So now let's compute the FLOPs and memory I/O both for the transformer. So we're going to break it down into MLP layers and the attention layers. And just for notationwise, we're going to do this computation with S being the number of tokens we're conditioning on. Think about the length of the prompt. And T is the number of tokens we're generating or querying using. And in prefill, T is going to be S because we're not generating T tokens, but we're querying using each of these tokens. And a generation where T is just 1. So hopefully the matrix multiplication is so fresh in your head because this is going to be essentially that but a little bit more complicated because it's a transformer. So we're going to count the FLOPs and bytes generated. So first, we're going to take X which is a T by D matrix. I think maybe these T should be S's. But anyway, that involves doing a bunch of transfers, Basically the size of that matrix times 2 because bf16. Then there's the three-way matrices-- the up projection, the gate, and the down projection. They're all the same num size up to transposition. So you need to transfer those. Then you do the up projection. That's some number of flops.

19:35 So B times D times T times F. So whenever you multiply two tensors basically the contracting dimension only gets counted once whereas the other dimensions you just kind gather together. You need to write it out. You also have the gate which is this same thing. You write it out. You compute your non-linearity. You multiply some stuff and you down project. And that's a B times T times D times F, which is basically the same number of FLOPs, and you write out the result. So if you look at the counting, I guess maybe you can check the results. Actually, you don't need to check it because this is SymPy and it's guaranteed to be correct. But again, we're going to assume that B times D is much smaller than D and F. And we get that the intensity is B times T. So this is analogous to the matrix multiplication case where the arithmetic intensity, which we want to be high depends on how large your batch is and how many tokens you're essentially generating. So now if you look at the two stages, prefill, life is good remember, because we can just make BT large enough. You use a batch size. Even a batch size of 1 actually is maybe OK if you have long enough sequence. So that's not a problem. Now generation, this is where it becomes

21:15 a little bit harder because you're generating one token at a time. So T is 1. So if T is 1, that means for BT to be large you need B to be large. And B is essentially the number of concurrent requests. So this is of interesting because your efficiency depends on having large batch sizes. Because I mean, intuitively, it makes sense. If you can take a lot of requests, batch them together, then you can get better efficiency, at least throughput. But this also depends on what B is.

Attention layers: intensity is ST/(S+T), and B cancels

21:53 Because if you're only getting a few requests at a time, then you're not going to be able to use your hardware very efficiently. And this speaks to the very dynamic aspect of inference, which we'll come back to later in the lecture. So now what about attention? It turns out attention is even worse for reasons I'll try to get into. So let's do the counting FLOPs bytes transferred. So I'm going to read the QKV matrices from HBM. I'm going to compute the attention, which is a matrix, which is Q times K and the number of FLOPs is B times S times T times D. So remember S and T are the same during prefill. So that's your sequence length squared times B times D. And then I'm only looking at the matrix multiplications because the FLOPs from the other steps don't really matter. And then you project out to-- sorry, you take a combination of this and V. So actually, this is mathematically incorrect because there's some softmaxes there. But the essence of the math models are the same. So that's the same number of FLOPS. And then you write to HBM. And here I'm assuming there would be more bytes transferred if you didn't use FlashAttention. FlashAttention means that you don't have to keep on writing back to HBM into intermediate steps. But the order is actually not really affected. So qualitatively, it doesn't really

23:43 matter whether you use FlashAttention or not. But the math here depends on the-- the constants matter. But let's look at the flops and the bytes transferred. And if you divide and simplify, you get this rather nice expression. I mean nice in that it's simple, not nice in that it's good efficiency, which is S times T divided by S plus T. So let's try to interpret this a bit. So in prefill, T equals S. So that means you're prefill intensity is order S. So that's good because as long as you have long enough sequences, then you're good to go. And generally, the sequences you assume are kind of long enough. During generation, however, you'll see that the intensity is essentially 1. S over S plus 1. But that's basically one. And remember one is really bad. But notice like what there's no dependence on B at all. So unlike in MLPs, remember in MLPs, the generation the prefill was BT, which is great. And then arithmetic intensity was B, which was not great because it depends on the whims of your users and workloads but still could be larger than 1. Whereas for attention, it's actually just always less than 1, no matter how long your sequences are how many users there are, it's always 1.

25:30 So why is this intuitively that there's no dependence on B, the batch dimension? So the reason is that in the MLP layers, intuitively, every sequence hits the same MLP weights whereas in the attention layer each sequence has its own KV cache. Because the KV cache is sequence-specific, which means that you can't really, use-- In MLP case, you can read all the weights and then you process a batch intuitively, whereas in an attention case every sequence kind of requires additional memory. You don't get any kind of savings if you batch them up. Mathematically, I guess you can look at it through here where the number of FLOPs, there's a B here, which is expected, but the number of bytes transferred is B times-- there's a scaling and B, so when you divide that B cancels. Whereas over here, there is a B here. But we're assuming that DF dominates. So when you divide, basically, there's no B essentially left in the denominator. So you can look at it mathematically or you can just kind of reason about it intuitively as for the attention, the KV cache is every sequence its own unique snowflake. So the summary is prefill is compute-limited whereas generation is memory-limited.

Napkin math: Llama 2 13B on one H100

27:17 The MLP arithmetic intensity is B, which to make good enough, you need a bunch of concurrent requests. But attention intensity is 1, and it's also impossible improve that. OK, I'll pause a bit for any questions. So let's move on. So now we know that inference, thanks to generation, is memory-limited. Let's try to study the throughput and latency at least in theory. So let's focus on let's see. So we're going to make some assumptions. So all of this is napkin math is a little bit stylized, but it gives you roughly the right kind of scaling and the right way to think about things. So we're going to assume that communication and compute can be perfectly overlapped, which is obviously false but it's good enough for making these qualitative estimates. So what we're going to do is we're going to instantiate the latency and throughput for a Llama 2 13B on H100. So for a 13B, here are the values. Let's just put the sequence length to be 1,000, hidden dimension to be model dimension to be 5,000. 4 times-- Actually, I don't know if that's not 4 times. But anyway, F is some multiple of that number of heads, number key value, number of query heads number of key value heads, which for Llama 2 is the same.

29:08 We'll get to that point later, and so on. And for the memory bandwidth of H100, that's the number. So that's the config. And we're going to compute the memory latency and throughput. So first, let's just quickly get the number of parameters. You guys did this in assignment one, so I won't belabor this but it's some expression that depends on all the different variables, and to store the parameters. We're going to use bf16 because inference is generally going to be 16-bit not 32-bit. So we're going to multiply by 2. So that's the memory that the parameters take. We don't need gradients. We don't need optimizer states because we're not training. But we do have to store the KV cache, which are the sum of the activations. Not all the activations, but some of them for every sequence of length S. And how much do we have to per store per sequence it's basically the sequence length times the number key value heads times the dimension of that head times the number of layers times basically 2 for basically both the key and the value, and 2 for bf16. So that's how much the cache size takes. And so the total memory is the batch size times the cache per sequence, plus the memories, plus the parameter size.

30:55 So now latency is going to be determined by memory I/O. Remember it's memory-limited, so we're just going to compute how much memory it needs to be transferred into the GPU to do this computation. And simply memory over the memory bandwidth. And throughput is essentially the inverse of latency, but scaled up by B because we're looking at generating B tokens in parallel. So now if we substitute our Llama to config, we'll see that the number of parameters checks out. It's 13 billion roughly. The memory latency and throughput have these expressions. So memory grows obviously. This is the parameter size. This is the key value cache size times B. Latency also goes up as a function of B. Throughput increases, but you'll see that it increases up to a point. The B shows up in both the numerator and the denominator. So there's limits to how much you can stretch throughput even if you could fit everything in memory. So those are the expressions for latency, throughput and memory for this particular model. So now let's instantiate with different batch sizes. So if B equals 1, then the latency is about 8 milliseconds. So every 8 milliseconds you generate a token and the throughput is 124 tokens per second OK. So that's 13B on a H100 if you're using a batch size of 1.

32:44 So now what happens if you use a batch size of 64? So you'll see that the memory usage increases because you need to store the KV cache for all 64 sequences now. The latency goes up because you have to instead of just processing one, you have to wait for everything to finish. But the throughput also goes up actually quite a lot. So you're seeing this immediate trade-off between latency and throughput if you want low latency you just use B equals 1. But if you want high throughput, you want larger B in general. What happens if you use a batch size of even larger, so 256? You'll see that the latency goes up, the throughput goes up, but you see that the throughput isn't going up that much because you get diminishing returns after a while. But the most-- you can actually do this on a H100 because if you look at the memory, it's 240 gigs, so it doesn't even fit. So batch size you can only increase to a certain point because of memory. So just to recap, there's a trade off between latency and throughput. Smaller batch sizes, yield better latency. Larger batch sizes yield better throughput. And finally, last week, we talked about parallelism for training. And it was kind of a complicated, annoying. At least one type of parallelism for inference

Replicas, sharding, and the TTFT knob

34:29 is really, really nice and simple. You just launch M copies of the model. No communication because the model you don't need to update the models. The latency is the same and the throughput increases by M. So that's pretty good. So always remember that don't forget easy things. Now there are cases where if you have a large enough model, then maybe it doesn't even fit on a single GPU and you need to shard the model. And in this case, you also want to start sharding the KV cache in some cases to get better efficiency. For more details, check out this book chapter. So the time-to-first-token, which is a metric I mentioned earlier, is essentially a function of the prefill. It's basically, how long does it take to encode the prompt. And usually, this is compute-limited, so you're basically going as fast as you can. And there's not much you can do about it given a fixed architecture. Sorry, you can improve it if you reduce the batch size still. But if you want to improve the throughput, you have to increase the batch size. So any questions about that? So this was on computing the throughput and latency. And because of the memory limited argument

36:12 that I gave in the previous part, I just focus on memory and compute how many bytes need to be sent. And that gives me a rough bound on the latency. In practice, there are some regimes where compute does matter, but I'm ignoring that just to keep things simple. OK, question. Is this assuming a single GPU? Yes, this is assuming a single GPU. While creating a batch from multiple users, each of these prompts have to be completed within the [INAUDIBLE] number of tokens generated. So there is empty slots in these values. Yeah, so the question is if you have multiple users and you're batching them together, they might arrive at different times. They're going to finish at different times. So we're going to get to that. That's going to be a special issue that we're going to have to deal with. Any other questions? So now we have a good handle on what the inference workload

Shrinking the KV cache: grouped-query attention

37:21 looks like. We looked at the arithmetic intensity. We looked at the transformer inference with respect to arithmetic intensity. We saw that it was memory-limited thanks to the attention where the KV cache has to be special for every sequence. And then using that, we can compute throughput and latency, which are the main inference metrics that we care about. Now how do we make things better? So there are some things that you can do on that are lossless. You can write better kernels. You can improve your systems. But I would say that the kind of there's a lot you can do if you're willing to take shortcuts. And these are kind of really interesting because technically this lecture is on inference, but secretly it's on model architectures, because what you'll see is that a lot of the changes in model architecture are going to have direct impact on inference, and we're actually inspired by needing to do inference quickly. So the big bottleneck here is the KV cache. because remember memory-limited, which means that the less memory stuff takes then the faster you go. Not just because of FLOPs even though that's part of it. But mostly due to memory, because it's mostly about memory transfers. If that's one thing you take away from this lecture, it's like all about the memory for speed. So the problem is that if you just start walking away the KV cache, you might lose accuracy.

39:06 So how can you make sure you don't lose too much accuracy but still maintain your KV cache small. So there's a bunch of ideas I'm going to go through that all essentially try to change the architecture to reduce your KV cache. Some of these ideas I think you've seen, but I'll go through them kind of in this more systematic way. So there's this idea called grouped-query attention. So multi-headed attention, which is the vanilla transformer, keeps around basically a number of heads. And for each of that number, you have same number of keys, values and queries. There was one time a multi-query attention, which you only have one key and one value. Basically, one key value had. Turned out that was not very expressive. So there was a intermediate point where you have a reduced number of keys and values, and then you have more queries. So why are we doing this? Well, remember we want to reduce the KV cache size. So the fewer keys and values there are, the better. So the batch size and the sequence length doesn't get changed, but it's in the dimensionality of these vectors don't change, but it's the number key value heads that we're reducing. So that's basically the idea. And this paper shows that you do get latency and throughput improvements, so times per sample.

40:47 And as you increase the number of groups up to eight or so, basically there's a negligible-- It's really fast compared to the full attention. And as you increase the number of groups, obviously you end up at the original. So that's latency and throughput improvements. And just to actually do this more rigorously, we have our Llama 2 13B model. And if we compute the statistics, this is using a batch size of 64. Remember, this is what we got. I guess I should have printed out latency here as well. And then if you run it with GQA, you see that the memory is reduced and the throughput goes way up. So this is actually great. So this is what happens if I take the Llama 2 13B architecture and I just reduce. So for every key value head, I have five query heads. That's 1-to-5 ratio means. This also means we can use a larger batch size. Because remember last time we tried to do 256, it didn't even fit in an H100's memory. So now we can actually comfortably fit into the H100 memory and then

42:24 we can further improve the throughput by using a larger batch size. So you can see a lot of different effects here. By reducing the number key value pairs, the memory of the KV cache reduces. That means the throughput and latency go up automatically because fewer memory transfers. And furthermore, as a secondary effect, I can increase the batch size within the GPU and that further improves the throughput. So that's wonderful. We have to also make sure the accuracy doesn't drop. So this is original paper that shows that this is full attention. This is GQA. The time is much less, but the accuracy is basically the same. Now what actually happened? So Llama 2 did not use this ratio. But Llama 3 actually picked up GQA and probably motivated by the kind of inference costs.

MLA, cross-layer attention, and local/hybrid layers

43:33 Actually Llama 2, I think, the 70B — the large model did have GQA but not the smaller ones. So that's GQA. There is another way to reduce the key value cache. And this comes from DeepSeek So this is actually from the DeepSeek V2 paper and it's called multi-head latent attention, which Tatsu lectured about previously, but I'll try to talk about it in the context of inference and its implications. So the basic idea is here's full attention. And GQA says, I'm going to use fewer keys and values. MLA says I'm not going to change the number key and values. I'm going to project these into a lower dimensional space. So it's another way of shrinking the KV size, but just, I guess, in a different dimension. So instead of using N times H dimensions for the KV cache of each token, I'm going to project down to C dimensions. And this is what DeepSeek did. It's actually quite aggressive reduction from 16,000 to 512. Only wrinkle is that this is not compatible with rope. So they need to add a few more dimensions to input rope back in, but overall, this is actually quite promising from a KV reduction perspective. I'm not going to do the math, but you can just trust me that you can see how the KV cache would be reduced a lot. And you get to the same kind of latency and throughput advantages.

45:13 And in terms of accuracy, they actually showed that compared to GQA-- actually, maybe I'm showing the wrong thing here. Image. I meant to show that the MLA actually improves, but this table does not show that. So I'll have to dig that up later. But anyway it MLA does preserve the accuracy as well. So there's another idea which says, well, GQA basically you can think about it as a sharing key value vectors, within a token and within a sequence. But we can also look at something called cross-layer attention, which there's a paper on this, but I think many people have been kind of thinking about this and doing this. So I don't know if this is actually the first paper. But basically, if you look at the transformer diagram, you have the key value projection of one layer, and then you have the next layer and these key value vectors are separate, usually. But the idea here with CLA is that we're just going to use the same key value projection across layers. That's why it's called Cross Layer Attention. So just as GQA shares across heads, CLA shares across layers. So here they show that they empirically improve the Pareto frontier of accuracy and the KV cache size. So KV cache size, which relates to throughput and latency,

47:09 you want it to be small. And you want perplexity also to be small. So they're able to improve that. So notice that, for example, 64 heads, the cache size it gets reduced but the validation perplexity does go up a little bit. But overall the there's advantage in making that trade-off. OK, so there's yet another way to do things. So local attention, which has been explored actually quite a bit since even kind of the there's a Longformer, there's an OpenAI paper, and then Mistral and I think many others use this as well. It's a very, I guess, a natural idea. Instead of if you look at a full attention diagram, it's dense N squared, and that's where a lot of your complexity comes from. And basically, the idea is you're going to just attend to only the past K tokens, which means that in a KV cache, as you're generating the sequence, you don't have to remember everything. As soon as the token kind of falls outside your window that you have attention, you can just throw it away. So local attention is very could say that the KV cache size remains constant as opposed to growing with a sequence length.

48:49 So this is really good because that means that for even long sequences you can have quite a small cache. But the problem is that this still hurts accuracy, because if you just think about it like, why are we doing attention instead of RNNs, is that we needed to have long range model and long range dependencies. And this is in some sense even to call it attention is a little bit kind of overselling this. This is only looking at the local context, which is not very expressive. So what you do here is you can interleave local attention with full global attention hybrid layers. So for example character.ai used, for every six layers, they had one global attention global layer and five local layers. So it looks something in addition to cross layer attention. So it looks something like this where full attention, every layer you have to store the KV cache. And what they did is that for every six layers you have the full attention, but in between you have this local attention. And on top of that, they have KV cache sharing locally, both for the local attention and the global attention. So this is like all the tricks kind of-- not all the tricks, but many of the tricks combine together. So in summary, these are a few ways

50:34 to reduce the KV cache size. Because remember inference is memory-limited. So you want to reduce the cache size, but you don't want to hurt accuracy too much. And there are many ways to do it. You can lower the dimensionality of the KV cache. You can have few KV cache vectors. You can reduce the dimensionality of a keV vector. You can share the KV cache across layers, and also you can use local attention on some of the layers. Any questions about this set of tricks for reducing the KV cache? Yeah. [INAUDIBLE] question about the quality. And I feel like the weights are all that you shared across the layers. Do you just have one set of weights or just like KV And that's shared across the [INAUDIBLE]. The question is, are the weights shared. So the KV cache is shared but the weights are shared. So what happens is the weights for doing the projection need to be shared. So there's some consistency. There was another question. The context size is too large and then it increases the KV cache as well context size when we prompt the [INAUDIBLE] the context that is given to the translation model. it can become too long and it increases the size. So we try reducing the size summarizing the context. So the question is if you have really long contexts,

52:18 let's say your prompt is huge, that's going to intrinsically take a lot of KV cache. So all these tricks can try to reduce that. You can do more aggressive things that there's ideas like gist tokens or ways to summarize the prompt, which we're not going to talk about in this class, but there's ways to have that address the long prompt situation as well. So now I'm going to talk about even more radical ways of making inference go faster by changing the transformer.

Beyond the Transformer: SSMs, linear attention, diffusion

52:57 So the KV cache, these are basically variants of the transformer. But maybe you can actually go outside the transformer and do better because the transformer wasn't really designed with heavy inference workloads in mind. They were just trying to train a good model that efficiently. It was mostly about training efficiency. And the autoregression, as we pointed out, is really causing this bottleneck here with the autoregression plus the full attention. So we're going to talk about two directions-- state space models and diffusion models. This is going to be fairly quick. So the idea of state space models is actually drawing ideas from signal processing and control theory. Initially, the motivation was trying to model long context sequences without suffering the N squared blowup. So it wasn't necessarily about inference speed. But it turns out if you solve that problem, you get faster inference, too. So there's a kind of early paper on S4, which uses classical state-space models, which are basically these linear dynamical systems, which are you've been used to model long contexts and shoehorning them into the kind of a modern neural setup. This work is nice in that it has this RNN kind of interpretation due to the linearity structure, and also has a convolution. Interpretation as well.

54:43 So they published this paper and show that I think worked really well on these long context synthetic tasks. But what they found is that, well, I guess what was discovered is that they don't really work well for language modeling. And that's obviously kind of a disappointment, because a lot of the value of transformers is being able to do language well. So in a series of papers, they identified a set of synthetic tasks that captured the essence of why these models weren't working well. And that's basically these associative recall tasks. So here's a synthetic task where you're given basically a sequence key value pairs. And the goal is to predict basically lookup the key and output the value. So in some sense, it's kind of a logically a trivial task. But it's long sequence because I can have a lot key value pairs. And I'm going to have to look far back. It can be arbitrarily long dependence. And you can see that local attention is not going to work very well, because it's just going to remember the last few sequences. And well, the problem with state-space models is that they're good for these kind of signal processing tasks. But really, you need to go isolate a particular key value pair. And pull out the answer for those type of tasks it didn't really work. So there's a bunch of work I'm not citing. There's Hyena, H3, and then Mamba, which basically tweaked or changed the SSM to basically handle these associative recall

56:29 tasks. And eventually, it worked better. Up to 1B scale was matching transformers. And the idea of Mamba has been popular in scaled up even to a 52B MoE by the AI21 folks (Jamba). Notice that in this case they still had to use a transformer. So a transformer, but only I guess eight layers. They had a transformer. The rest of them were Mamba layers. And so that still led to a fairly big savings and speed up. But, more recently, there's this kind of revival of this older idea called linear attention, where instead of-- let's see if I can make this bigger. It's actually a very simple idea. So you know local attention or sliding window attention is. Linear attention is this idea that you essentially-- So basically in the attention computation, there's a key and a query and you dot product and you take the exp of that, which is basically giving you a exp kernel. So you can basically take a Taylor expansion of that and write that computation as basically dot products of some nonlinear map. So then what you essentially have is you can think about for every key value position, you are basically applying some of nonlinearity, blowing it up into some space and then doing

58:13 some linear computation over it. And because it's linear attention, it actually kind of behaves like an RNN and it's linear in the sequence length rather than quadratic. I know that was a little bit fast, but I just want to give you the taste of it. And so this idea has been actually scaled up quite successfully. So there's this organization called MiniMax that's training pretty legitimate models, up to 456 billion parameter MoEs. And they use this basically linear attention idea. Now they have to use full attention still once in a while it seems. I don't think people have been able to get around having some full attention, but at least it seems like people have been able to get rid of most of the full attention in, most of the layers are not full attention anymore, they're just either linear layers or local attention layers which are much, much more efficient. OK, so the linear plus local attention now are actually yielding serious state of the art models. And it's probably safe to say that, well, I don't know what exactly the close model providers are doing, but I would suspect that there will be at least as kind of advanced in terms of as efficient as this and leveraging sparsity. So it's kind of an interesting question when people ask like, well, attention all

59:57 you need and it's transformers? Well, yes and no. I mean, I guess in some sense there is still the sense it's square there. Maybe it will be able to get rid of it. But most of the transformer has been like pretty radically changed by having other much lighter weight components. And you're still able to get much of the same kind of accuracies. And all of this is really helpful for inference because on these non full attention layers, you're basically replacing the ordered T KV cache, which grows as a sequence length with something that's constant. And there's papers that follow up-- I think they on the BASED paper. Where did that go? Either in this paper or in follow up work analyzing basically the trade off between the KV size and the ability to do various types of recall tasks, which makes sense because if you don't store very much won't be able to solve certain tasks. But there is this trade-off curve that you can try to play with. So that's all I'll say about the state-space models. So now let's talk about a completely different style of generation models. Diffusion models. So diffusion models have been very popular in image generation, but they turn out to be fairly tricky to get working in text, although there recently have been some advances here. So the idea of diffusion is that instead of generating autoregressively, you just generate every token

61:38 in parallel. So obviously you only do that, via some simple layer. It's not going to be very good. You can't generate all the words in parallel and expect it to be coherent. But what you do is you iterate and you keep on refining this generation until it gets to your final generation that you output. And the idea behind generating in parallel, you're no longer autoregressively bound, and that generating all tokens in parallel well can be done in parallel. So you get to saturate your GPUs relatively easy as long as your context length is large enough. So recently, there's this Inception Labs has produced some pretty interesting models. There's not much written about them, but you can see a demo of the generation and process. It's just kind of generates code instantaneously, but it's obviously kind of broken code. And then it's kind of refines over time. And this is one of their benchmarks that show that, at least on coding, I'm not sure about other tasks, that if you look at the tokens per second, these models are way out here in terms of speed compared to anything that's transformer. Even Jamba, remember, it was like a hybrid Mamba transformer architecture is quite slow compared to these diffusion models. So now whether diffusion models will be kind of general purpose and powerful enough in all of these, that remains to be seen.

63:22 But I think you have such a lead on the token speed here that even if you I think you can put more compute and recover some of the accuracy losses if you need to. So the summary here is that I think this whole kind of novel architecture thing is actually really exciting for inference because they allow you to sidestep kind of fundamental obstacles. So if you're dealing with attention, you just have this fundamental KV cache obstacle that you can quantize, you can optimize, but it's still there. And so by making a kind of state-space model, you're shrinking that to a constant size. And as long as you can keep up the accuracy, which is big if, then you win big time. Same with the diffusion models. Autoregressive generation is a key bottleneck. Now if you just generalize it, generate things in parallel. Now all of a sudden, you change the game completely. So there's much more work to be done here

Quantization, pruning, distillation

64:31 in improving inference. So as you can see now, inference is the inference game is much broader than it seems at first sight. It's not really about necessarily the system's optimizations to make it fast, although you obviously need those. But I think the real gains are coming from real radical changes in architecture. OK, so I have about 10 minutes left. I'll go through these quickly. Quantization and model pruning. So quantization the key idea is just reduce the precision of the numbers. Very easy to do. And the thought is that less memory means a less bytes transferred higher latency-- Sorry, this should be lower latency, higher throughput. And you do have to worry about accuracy. Of course, that's the trade-off. If you look at the different types of formats, FP32 used for training not used for inference. Really, BF16 is the default for inference. You can go down to FP8 or INT8, which now is less accurate, but much cheaper than even FP8. So people do do a bunch of inference in INT8, which if you look at the range, I mean, it's an integer between 127 negative 128, which is not that-- it's pretty low precision. And people even go down to INT4, which is they're not. INT4 is pretty low. There's also other ways you can do. So once you just decide that you want to quantize,

66:22 I guess you could do several things. You can train with the quantization, but obviously that means you need to retrain models and more, I guess, commonly you do post-training quantization, where you take an existing model and you try to quantize it and try not to screw things up too much. So there's a paper called LLM.int8(), which I'll talk through briefly. So in quantization, basically what happens is that you take your vector, which is let's say FP16, and then you need to figure out the dynamic range. If you want to pack it into INT8, you need to figure out what the largest value is. And once you figure that out, you can divide by that and multiply by 128. And then you get your integers. And then if you need to dequantize, then you go the other way. So basically quantization means that remember memory is a bandwidth right. So a bottleneck. So all your transfers are happening in int8. But when you actually do, I guess you sometimes have to upcast to a floating point to actually do the arithmetic. So the problem with INT8 is that not everything fits nicely. And you have these outliers which appear in larger networks that screw things up. So what this paper did is that you take this matrix, you identified a really large outlier values. And then you handle them separately using in full 16-bit precision and then do the most the vast majority INT8. So this works well, but it's actually a bit slower.

68:11 So the motivation here wasn't the inference speed, but more the even being able to fit your model into memory. There's another paper called Activation-aware quantization. And here the idea is that you're kind of quantizing the weights, but you're going to figure out which weights to quantize based on the activations. Really quickly, you're going down to actually INT3, and this obviously reduces memory by quite a bit and leads to a 3x speed up. So the general idea here is that you get a trained model. And it just happens that some of the weights or activations are going to be abnormally large. So for those you handle separately and then everything else you can work in low precision. I'll talk about model pruning ideas. Very light quantization. The basic idea is very simple. You just rip out parts of an expensive model to make it cheaper, and then you fix it up. So in this NVIDIA paper, what they do is they first identify important either layers or heads or hidden dimensions using a small calibration size. They use some simple scores to compute that. And then you just remove the unimportant layers or hidden units or heads. And then now if you just take that model, it's going to be clearly worse. But so then the last step is you distill the original model

69:51 into the pruned model. So you repair the model from the initialization, which is your prune. So you're not starting from scratch. You're starting from something that's worse but hopefully not worse, and hopefully retains a lot of the same structural properties of the original model. It's just maybe not kind of calibrated in some sense. And the results are pretty good on that. So they have these 15 billion parameter models that they're able to reduce to 8B with hardly any drop in this, I guess, at least according to MMLU, and then down to 4B with some drop, but you're also going down quite a bit to 4B model. OK, so maybe to summarize this taking shortcuts idea. You can reduce inference complexity without hurting accuracy. You can do it from scratch where you just define a fresh architecture that's by construction fast and just train it. Or you can distill you define that architecture. You can take a slow model and you figure out some of scheme to initialize the new model with the old model.

Speculative decoding: exactness for free

71:02 And then you basically do distillation. So now all of these are a little bit unsatisfying because they're lossy. So you get massive speed ups, but you always wonder, well, maybe this model isn't as actually as good as original. So speculative decoding or speculative sampling allows you to basically have your cake and eat it, too. So recall there's two stages of inference. You prefill which you're given a sequence, you encode all the tokens in parallel. There's a compute-limited, which is great. Notice that this also gives you log probabilities for each of the tokens. And then there's generation, which is one token at a time. It's memory-limited. It's slow. So in other words, checking is faster than generation. So intuitively this makes sense. But hopefully now you also appreciate the math behind why this is true. And the speculative sampling idea is actually really, really elegant. It was proposed in parallel by these two independent teams from Google. And the idea is to use a cheap draft model p to just run ahead and generate some tokens. And then you're going to evaluate those tokens with a target model. And because evaluation of given tokens is just prefill, so you can do that in parallel, which is fast. And then you accept it. If it looks good. So this is what it looks like in real life. So if you're using a big model generating one token at a time, that's slow.

72:45 But in speculative decoding, you have the draft model that's racing ahead, generating a lot of tokens and using the big model to essentially verify. And sometimes it will reject and sometimes it will accept. And the acceptance rate basically determines how fast of a speed up you have. So here is the more formal algorithm. So you're going to have a look ahead of K. So you're going to use your draft model and generate K tokens autoregressively. So this is hopefully fast because your draft model is small. And then you're given these K tokens that you generated. And I'm going to score them based on I'm going to compute the probability under the target model Q. Now I'm going to decide whether I want to accept this or not or not. So I go through each token and I'm going to essentially accept it with probability Q over P. And the 1 just makes sure this probabilities are between 0 and 1. This kind of looks like if people are familiar with Metropolis–Hastings it's kind of where this kind of comes from. So intuitively, you're sampling with P. So you need to divide that out because you don't want P. You want Q. So this is an importance weight. So if you accept it, then great. You move on and you look at the next draft token and so on. And if you don't accept it, then you're going to sample from the target model, the slow model, but you do this correction where you've already tried to sample using p.

74:26 So you don't need to do that anymore. You subtract it out and you sample from Q. So this is basically kind of a rejection sampling with a proposal P and a target Q. The only difference is that you are sampling — in rejection sampling, if you reject and you reject and you just try again and you try again. And here, we don't want to keep on looping forever because if you reject we're just going to say, OK, fine, we'll bite the bullet and just sample from the more expensive model. So the cool thing here is that you're guaranteed to get an exact sample from the target model. So those of you familiar with sampling, this is what shouldn't be too surprising. You're able to use prior information to speed up sampling. But in The language modeling context, this is kind of nice. I'm going to skip this. This is not really a proof. This is just some derivation to show that for a case of vocab too, why this these formulas give you the right unbiased sampling procedure. And it works pretty well. So the accuracy should be actually the same since it's the same model. But maybe there's some randomness there. But the speedup is you're getting a factor of 2 speed up, essentially. So in practice what you do is you take you have something like a 70B model your draft model is much, much smaller. And if your target model is 8B, then your draft model might be

76:08 1B. And you generally want to make the draft model as closest to the target as possible. And so if you're doing some distillation that could make it even better. There's a bunch of-- this is a pretty hot area of research and inference. There's a lot of ways to improve this process. You can use Medusa, which is this way, to have the draft model instead of generating autoregressively sample multiple tokens in parallel, or EAGLE, where you're actually taking high level features of the target model and pumping them into the draft model to generate. So the draft model doesn't actually have to stand alone. It can be kind of glommed on to the target model to help it generate. So summary exact sampling from the target model thanks to math. And this exploits the symmetry between checking and generation. Or prefill and generation. And there's actually a lot of room for innovation on the draft model, which can everything that we've talked about before, where

Live traffic: continuous batching and PagedAttention

77:17 you can have different radical architectures, different ways of quantizing, all of those apply. The only thing is that you get to basically guarantee that you're getting exact sample. So now I'll go-- I'm out of time, but quickly go through the question that came up earlier, which is that in practice when you're serving, there's live traffic. Requests come at different times. They finish at different times. Some of them have shared prefixes. Some of them don't. They have different lengths. So it's very heterogeneous in comparison to training where you get basically a dense block of tokens, and you're basically going to push it through your GPU at full speed. So what do you do in this case? So there's a series of papers that explore this. And the basic idea is if this is so the last part is more of systems-level contribution. So the idea is that you don't wait for batches to-- the train leaves. There's no the train doesn't wait for you. So when a new batch comes you're just going to put it in. Which means that the worker that's generating tokens needs to hand control back to the scheduler every step. So you generate a token, come back to the scheduler and say if there's new requests then, they get stuck in. And then it kind of continues. So you're kind of not wasting any s time waiting around

78:56 for requests. Now, there's a kind of a problem with batching, I think, which is behind the question. Batching works when everything is at the same dimensionality, but every quest might be a different length. So there's this idea of selective batching where you basically break up your computation for attention. Everything has to be handled separately. But for MLPs, remember, which are the bulk of the computation, you can actually take tensors of different sizes and you just flatten them. And because they don't interact, they can just be kind of along for the ride in the batch dimension. OK, I know that was fast, but I'll just quickly go over PagedAttention now. This is the paper behind vLLM, which some of you probably have used. And this addresses the memory usage problem. So if you have a KV cache and prompts are coming in and finishing, then your cache is going to get fragmented. So you're going to have-- you're going to allocate a bunch of space for a request, but you don't know how many tokens are going to generate it. So there's going to be internal fragmentation. And then there's also going to be external fragmentation where there's padding between the request and responses. So that's no good. So the PagedAttention basically says remember operating systems? We have and how virtual memory works. We divide the KV cache into a sequence of contiguous blocks. And then we just put them wherever we find white space. So if you have two requests coming in, then they might just the first request might be here, here, and here and the second request might be here and here. So the blocks are the things that you're going to keep contiguous.

80:37 And that's going to give allow you time to coalesce your memory. So you can also play these tricks where if you have a sharing of prefixes then there's another idea from our operating systems, which is copy on write. So you basically maintain reference counters for how many basically sequences are using this particular block. And then if you need to diverge and have blocks go in different directions, then you copy and you reduce the reference count. There's a bunch of other vLLM optimizations, which I won't go through, but basically the summary is remember your operating systems classes. You can apply them to inference as well. So quick summary. Inference is really, really important. and the characteristics are distinct from training. You're memory-limited and it's also dynamic, which leads to a bunch of new challenges. We saw a whole host of different techniques around new architectures quantization, pruning, distillation, speculative decoding. There's ideas from systems which can allow you to better use your memory overlap, communication, and compute and things like that. But I would say that there's probably even more opportunity in the modeling and architecture, because if you think about it all, inference narrowly is inference in a particular model. How do I run this particular model?

82:15 But who cares about that particular model? You care about delivering good accuracy given your resource budget. So a lot of these ideas that are trying to reduce the KV cache, changing the transformer are basically ways to sidestep the problem and say, well, I have something that's more efficient. And I can train it in a way that gets me better accuracy then I win. OK, so that's all I have and I will see you next time. Then we're back to scaling laws.