Section Insights
Introduction to Context Management
What is the focus of the presentation?
The presentation focuses on context management for long context agents, particularly in the realm of large language models (LLMs).
- Context management is crucial for long context agents.
- The talk will cover architectures used in state-of-the-art models.
- Understanding the growth of agent context is essential.
Computational Challenges in Long Context Models
What are the computational implications of long context in language models?
The computational requirements grow quadratically with the length of the context, leading to significant challenges in processing.
- State-of-the-art models aim for a context length of up to a million tokens.
- Quadratic growth leads to trillions of comparisons, making computation intensive.
- Strategies exist to limit the number of comparisons to manage computational load.
Attention Mechanisms in Long Context Models
How do attention mechanisms adapt to long contexts?
Attention mechanisms can be adapted to allow for varying decay rates of context elements, enabling some information to persist longer than others.
- Different decay rates for context elements enhance model expressiveness.
- Sparse attention focuses on a limited number of previous context elements.
- Content-dependent sparse attention can further optimize processing.
Advanced Attention Techniques
What are some advanced techniques for managing attention in long contexts?
Techniques like 'yarn' allow for varying degrees of attention across different context elements, optimizing the processing of long sequences.
- Yarn technique allows for selective attention based on distance.
- Different methods exist for managing attention in long contexts.
- Understanding these methods is crucial for effective model training.
Cache Management in Language Model Serving
How does cache management impact the performance of language models?
Effective cache management is essential for optimizing the performance of language models, influencing cache hit rates and eviction policies.
- Cache-aware routing improves access to relevant data.
- Different providers have varying cache hit rates based on their infrastructure.
- Understanding cache eviction times is important for maintaining active conversations.
Transcript
0:07 Okay. hello everyone. Can you hear me in the back? All good. Okay. So, well, welcome back everyone. Am I echoing? All good. Okay. So, today I'd like to talk about context management for long context agents. And I at first thought about titling this as like context management for long context LLMs because most of the stuff I'm going to be talking about today is kind of more general LLM stuff, but long context is super super important for when we're working on agents and it's not necessarily covered very well in all of the LLM prerequisite classes. So, I'm going to cover it in a fair amount of detail today. I'm also going to try to go through kind of the long context architectures that are used in many of the kind of state-of-the-art open models so you can kind of see the best practices that everybody is using in order to handle a very long context.
1:10 So first I'd like to talk a little bit about how long an agent context can grow. So basically the input is quadratic from the or the number of token accesses that you do is quadratic based on the linear history growth. So if each call you make for an agent takes about you know five tokens and let's say you start out with a an initial system prompt of one token you'll get six tokens in the first call 11 16 21 and 26. you know, this will increase and increase and increase.
1:50 and so by the time you have gotten to the end, you will, if you do things completely naively, you will have had to process ADK tokens for these five calls. And you know, this can continue for a very long time. I I don't know the longest anybody has used a single coding agent session or something like that, but I've certainly done them for like multiple days where I've been working with it for almost the entire day. So you could imagine this could get very very large. the actual computational complexity of this is is actually even larger than this.
2:30 well no sorry. So yeah the the computational complexity is is quadratic here. So why do agent prompts get so long so quickly? I did a little bit of data analysis on the composition of prompts that I used in 1,500 sessions from our coding agent openhands. And each each session was relatively short. These are not like the long multi-day sessions that you do. This is more like solving a single task. But each one used about 78,000 tokens on average.
3:10 23% of this was a system prompt in all of the tool descriptions. this is relatively large for a coding agent. I think our system prompt is somewhere on the order of like 15,000 to 18,000 tokens because we're very opinionated. the open ants agent is very opinionated about how people do work. these tend to vary from like 8,000 to 8,000 to 20,000 tokens depending on you know how many tools you you pull in and stuff like this. separately from this 9% were user messages. So this is the user either saying the initial task or following up with additional information after that. 9% were reasoning traces and you know where the model is thinking through things. 2% were model replies. Sorry, this isn't showing on the slide, but 2% were the model replying directly to the user. So, it's spending about five times as many tokens freezing as it is spending communicating with text to the user. And 20% were tool calls.
4:10 so, you know, you know, writing code, reading code, that sort of thing. And 37% were tool results. So, this is you know, getting the result of, you know, reading a file or something like this. So you know each step can will consist of all of these and that tends to add up very quickly. So when you're handling this very long context there's two big challenges here. the first one is capacity and so whether your model is able to handle this large amount of evidence in the first place and this is a challenge because as I'll talk about in a little bit training models to handle very long context is hard and so in order to manage this you need things like architecture position encoding training data and you also need to be able to evaluate whether you're doing a good job of this.
5:05 And then the second thing is efficiency and whether the system can afford to serve it. And there's a bunch of techniques to do that as well. And I I think some people might be taking this for granted if you're just using, you know, cloud code or codeex or something like this, but they're doing some pretty monumental feats of engineering to take a, you know, multi- trillion parameter model and serve it within seconds to over and over and over again for mil I don't know codeex has 25 million users. So you can imagine 25 million daily users hitting it every few seconds for this huge model. if you want to you can calculate in your head how many flops per per second that might be.
5:51 It's a lot of flops. Cool. so any questions or things? Probably pretty straightforward I guess. Okay. So let's go into inference. And so in order to make this happen the you basically need to be able to calculate prefill and decoding. So calculate your inputs and your outputs at each step. And this you know I'm sure everybody has done in your previous classes for generating from you know a language model. some things you may or may not have thought about in your previous classes. If you took how many people took like language model systems or deep learning systems or something like that actually more than I thought. So if you took the systems course you you probably handled this.
6:47 but other things you need to think about are how do you schedule multiple concurrent queries? and also how do you save the results in the KV cache so you can reuse them and not have to reuse them over and over again. not have to calculate over and over again. So going into a bit more detail here, we have a few phrase phases in inference. The first one is called prefill. you could also call it encoding. but I think pre prefill is kind of the like preferred term of art nowadays. And so you take in all of your input tokens that you don't need to generate and you generate the key and value pairs from each of them. And then when you're decoding you basically take the prompt and key value pairs generate an output and append the key value pair that you get from the output. So you do this kind of incrementally whereas this can be done in a batched sequence.
7:53 And so when we think about this as a serving system, there's a bunch of ways we quantify how good a serving system is. the first one is time to first token. So this is how long you need to start you need to wait between when you start entering a query and when you start seeing the first output. And so like I'm sure everybody you know who uses chat GPT or something like that you put in a long output you notice it takes longer to start getting a response than if you put in sorry put in a long input you see it takes longer to start getting a response than if you put in a short input and that's like your TTFT is higher. then your time per output token is the time like the amount of time it takes between generating each token. And if you go to a website like open router t typically here this will be your time per output token there I am pretty sure they're not counting prefill here and it's just like the time between each of the tokens that might be a little bit small with the black background but you can see that for GLM 5.2 to you have like 63, 48, 24, that sort of thing. And you know, some providers are drastically better than others. Like this provider, not to pick on Mistl, but they're getting, you know, three three tokens per second and then these other ones like Parasel and Coreweave are getting, you know, hundreds of tokens per second.
9:32 So so basically you know this makes a huge difference in your experience using using agents. one other interesting thing is there's different ways people can use agents. And so in some cases you might really hear about this. So like let's say you're working together with a an agent on a interactive task in your favorite coding CLI. you may really really care about, you know, TTFT and Tapot. if you're running an agent in the background to do some background tasks that run overnight, you might not care at all. and there's trade-offs between them, between like throughput and cost and how fast you can get the the tokens. And like I you know also serve coding agents with my company and when I talk to providers they'll often often ask which do you care about more so we can optimize the system for the one you care about more and so there's there's ways you can do that and I'll I'll talk about those too.
10:46 Cool. yeah so sorry just to finish the slide. So throughput is the completed tokens or requests divided by the wall clock time. So this is kind of seen usually on a more holistic level over the whole system. And then the cost is you know how much it costs. So as a refresher I I hopefully everybody should know this already. but you have the query key and value calculation in typical attention and this is so you have this here and then you take the softmax divide it by this. You might have some sort of you know positional bias and then you calculate your output and typical full attention looks like this. you mask all the previous tokens and and or you consider all the current tokens and you mask all the the future tokens.
11:44 so this is of course very expensive because it's quadratic in the sequence length. and so a lot of the stuff I'm going to be talking about today is about how to fix that. so this is on the efficiency side and then on the effectiveness side just because you have a model or just because you have a serving system that can take in you know a million tokens doesn't necessarily mean that you can use that a million tokens effectively and this is a very like old school example of this I guess this is a needle in the haststack evaluation how many people have seen this figure for maybe some but not not that many. Okay, so to explain what this looks like, basically what this is saying is it's a very weird evaluation. Basically, they they took a whole bunch of Paul Graham essays by the like famous entrepreneur and head of Y Combinator Paul Graham and they stuck in something about like the best place in San Francisco to eat a sandwich is something or other and then they ask a question where's the best place in San Francisco to eat a sandwich and then the language model has to answer that correctly. and they made the context longer and longer and longer and then they evaluate models with the longer context and eventually the weaker models start getting worse at these very very long contexts and there's a bunch of examples of things like this. ruler and helmet are kind of the the common benchmarks that people use here. but basically the long the summary is that if you have very long context your models might start getting worse at processing them.
13:28 So has anybody noticed this when you've used a coding agent for a long time like are there common or used any agent for a long time are there common like failure cases that you see? What what kind of things what kind of things have bothered you? Anybody? >> Yeah. >> Forget the previous context. >> Yeah. Forgets the previous context. Did that cause any like problems for you?
13:58 >> Yeah. with me telling it not to do something in the middle of the conversation and then it goes and does it. there there was a really bad example where I think actually Daniel talked about it here where the person had said do not delete this and then it went and deleted it anyway because it had forgotten the context. So yeah, this can be a really big problem in agentic settings because you know the user is giving these incremental instructions. So so I'm going to be talking about several different things. I'm going to be talking about architecture. I'm going to be talking about training inference and compaction. So going first into architecture. this is both about efficiency and accuracy.
14:49 And so as I mentioned standard attention is global. So it makes pairwise comparisons across all of the tokens here. And the problem with this is like this is quadratic obviously. And does anyone know the like standard length of the of models nowadays of these of like state-of-the-art language models? >> 256 to a million.
15:19 >> 256 to a million. Yeah. And I I'd say most people are aiming for a million now because they don't want to be the, you know, the model that has 256 when everyone else has a million. So very quick math in your head, how many pairwise comparisons is that? >> 10 12 >> 10 10^ the 12 or in English. >> A trillion. Yeah, that's a trillion. That's a trillion comparisons every time. like when you're generating a sequence of a million long and each comparison requires you to multiply large vectors by each other over and over and over again. So you're starting to get up into the you know 10 to the I guess it'd probably be 10 to the 20ish which is a lot that that's a lot of compute for a single sequence.
16:12 So there there's ways that people solve this and basically almost all of them fall into the category of taking this quadratic computation and putting some sort of constant limit on the number of comparisons that are made at each time step. And that brings your n squared down to owl where I I wrote w for window but it doesn't necessarily need to be a window. It just needs to be the cap on the number of comparisons that you're making.
16:43 so there's the idea of local models and global models. And so local models are are basically models where this cap on computation happens within a local window. And then global models you have a cap on you might or might not have a cap on computation but that cap on computation is not localized. So you can pass information throughout the entire sequence. And the I'd say almost all of the models that exist nowadays or are widely popular nowadays are some variety of hybrid model where they have many local computations and then one global computation then many local computations one global computation and specifically I'm going to talk about six models that we have here. So the most recent Quen model has local computation of gated delta net.
17:37 I'll explain all of these or most of these in a bit. And then sparse retrieval for the global computation and the ratio is 3:1. GLM 5.3 is also a variety of delta and sparse attention with 3:1. Kim K3 is a variety of Delta and this is a dense dense attention 3:1.
18:08 Neimotron is Mamba which is kind of like recurrent style model and dense attention 4:1. to window sliding window attention dense attention 5 to one and deepseek's a little bit different. It kind of has window windowed attention but then some sparsity sparse stuff with interle branches. So basically the the standard way of doing things is hybrid model with some sort of local thing and either sparse or dense global attention.
18:50 So I'm I'm going to be talking about this. I'm not going to go into like a ton of details for all of them, but it should be enough that you could be able to look them up later. So the first one is sliding window attention. And this is maybe the the simplest conceptually. And what you do is you take chunks or you take chunks of the output and basically you're able to attend to your current chunk and all of the previous chunks and or not all of the previous chunks, a fixed number of previous chunks. So if this is your current chunk, it might be a chunk of like 512 tokens or something like this. You would have an additional 512 and 512 tokens in the previous chunks and you would attend to those. so pretty simple, no changes to attention, but you have something like this. But the nice thing about this is on the first layer you're passing information from the previous one or two chunks, but on the next layer the previous chunks also have information from their previous chunks. And so you kind of get this gradually expanding window.
20:02 so this this is pretty widely used. but the more popular way of doing things nowadays is using some variety of linear attention. And so the way linear attention works is basically we take our standard attention and standard attention has this softmax operation included in it. And this is the reason why it's hard to compute because you're taking the exponent and and dividing. And because you do this there's you basically have no choice but to expand all of the key and query computations like key query multiplication computations in and be quadratic quadratic complexity.
20:57 What linear attention does is you basically start from this kind of standard attention formulation but you take you remove the softmax and the normalization here or I mean the normalization doesn't really matter but you remove the softmax and you get just the key query multiplication and this is definitely less expressive because basically what you're doing is you're moving from something that takes in a bunch of like numbers.
21:52 Anyway, like pretend pretend that this bottom one adds up to one, but you're going from a thing that takes a bunch of real numbers and normalizes them into probability like numbers that are between zero and one and add to one. so this nonlinearity allows you to focus on like only some of the previous like tokens. And this is kind of the good point of attention. That's the selling point of attention that allows you to focus on the relevant previous context. And you're basically taking away that ability, right? You're you're removing the relative comparisons and just getting the the dotproduct between the queries and keys. But it allows you to do something really nice which is that you can use the associative you can use associivity and basically rearrange so that you can convert the sum of all the key and value multiplications before into a single state matrix.
22:56 And so this state matrix is essentially a function that map it's a linear function that maps from a query to a value. So attention is also a linear function that maps from a query to a value. So if I go back to this here you're basically calculating you're taking the query you're calculating the attention over all the keys you're multiplying it by the values and you're getting a value here. and this state matrix is also that. But the nice part about it is this state matrix can be decomposed. So you can update it. You can keep this state matrix around and update it in constant time every time you get a new value. So this essentially becomes like a recurrent neural network.
23:40 So it it allows you to update the state matrix in linear time without having without requiring the quadratic computation. So this basically solves your quadratic problem and it solves it so much that O of N squ becomes O of N. There's no like there's no even W term in this because each time you're you're updating based entirely on the previous state. but there's there's caveats here. So is this clear though so far?
24:17 Okay, cool. so the big caveat is exactly the caveat I mentioned before. Like this is just a less expressive function. and it doesn't allow you to make like relative comparisons between the previous tokens, that sort of thing. So there's a bunch of workarounds around this to try to make a more expressive function that nonetheless is useful in in kind of doing attention-like calculations. And so most of these are if you dig deep into the math in them, they're basically trying to become better predictors of the final value that you would get from attention.
24:59 and you you can dig into the papers to see how that works, but I'm going to explain them mostly mechanistically here. And so the first thing is Delta. And the idea behind deltaNet is you can predict what you can essentially predict what value you would want to be getting from the from the calculation of KT and S. And so you know KT here is the key at time step t and you are comparing it with the matrix at time step t t minus one and you're calculating this v hat t and on the other hand you so basically this this is the prediction that you would get of kt But on the other hand, you know what the output VT should be based on what the query times or based on what the Hang on one second.
26:20 I don't want to I don't want to get this wrong. So you calculate KT with respect to the matrix. And you can view this matrix as a way of calculating the the way of calculating the mapping from a query to a value. And you can also get the the actual value and subtract out what the value you get when multiplying it from the previous together with the previous matrix is. And this is essentially the error that you get here. And then you want to reduce the error. And so in order to reduce this error, what you do is you essentially add in this this term which is KT times the error term. And after you do this, if you multiply in KT to the new matrix, you will no longer have any error. But you don't want to make an update that is too big. So you add this beta term. This is essentially a learning, it's similar to a learning rate. And this learning rate will make the updates smaller. And so if you if you do this, this allows you to basically nudge the matrix in the direction of the in the direction that will cause a query that looks like the current key to output the value that you want essentially. So it's it's reducing the error reducing the error in the prediction of the value that you want.
28:00 So as I said this is going to be a little bit high level. the the papers are are pretty good descriptions of this. so after this this basically what this is doing is this is pushing you in the direction of the of making the the matrix a better predictor of the the value. But one problem with this is you can build up lots and lots of context from the previous spots and linear attention doesn't have any mechanism to prevent very far away cont context from having a bad influence on your next predictions. And so if you have like a very weird key far back in the future, it might have a very big influence on you in the f in the future. And so the way they fix this is basically they have a decay term to make to make the matrix decay a little bit at every time step. And by decaying the matrix a little bit at every time step, they're basically discounting the the far in the past context. And so you multiply the matrix by alpha t and then you add this st tilda into into your calculation here and update the matrix.
29:33 And so the stronger you make this decay basically the the more quickly it will forget the previous context. And so you can make this more local or or less local. And so this is gated delta is used in the Quen models. So it's one of the more the more popular methods now. but both GLM and KI which are kind of two of the state-of-the-art models use a slightly more expressive version of this where they basically have a per element decay term. So each element can decay at a different speed. And because each element decays at a different speed, you can have some portions of the context that you know carry over very long time periods, some portions of the context that carry over very short time periods.
30:28 Cool. So hopefully this was more or less more or less legible. any any questions here? Long long story short, it's basically a way to turn into it's a way to turn the idea of attention to something that you can calculate in context. Mamba is something that's very similar but it's a little bit more mathematically involved and I I decided to skip over it. I covered it more in the language model description course but similar idea but just a bit more complex in the math.
31:14 So the next thing I'd like to talk about is sparse attention. And so sparse attention is basically looking up a not attending to everything in the previous context but only attending to some things in the previous context. And essentially you limit the number of things that you look up in the previous context. There's a few ways you can do it. originally in there was a sparse attention paper by OpenAI quite a while ago like 5 years ago or something like this. and basically what they did was they had sparse attention every you know 512 tokens or something like this. And what this does is it means that you still have a fixed it's still quadratic in the sequence length but it's 512 times less than if you attended to every token and that's like good enough that you can manage to calculate it over the context that you're interested in.
32:14 Another option is to do content dependent sparse attention. And basically what this does is you have a very cheap way of calculating attention. Like for example, you can down project the attention vectors to be very very small and then you do attention over these very very small attention vectors. You pick the top k ones and then you do full attention over the top k. So that's another way that you can do this and this is what's used in deepseek.
32:43 Yeah. >> How do we define the metric for which we do the top similarity compared in this in the shrine space we just do like coine distance. >> Yeah. >> Yeah. I I think it's typically not cosine but it's just product but yeah >> I I I think although maybe maybe they're normalizing beforehand. They actually don't know that detail. Any other questions? So basically these th this is a very very cheap way to achieve global attention because you can just write a kernel that will do this for you. this over here will require you to have two passes, but you know it's potentially more accurate.
33:35 So I I've seen both of these used pretty widely. So a final a final thing is KV compression. How many people saw this in a language modeling course? Maybe some people. Okay. so I will I'll explain this. So basically this is standard multi head attention. And so in standard multi-head attention, you have a bunch of heads and each one has values and keys and all of these are a certain dimension. So they're like dimension, let's say they're dimension D and you have you have like a certain number of heads. Let's say that's H.
34:20 so the two main ways that people solve or help this problem of having very large and expensive KV caches is grouped query attention. And grouped query attention basically just uses a smaller number than h you usually h divided by some power of two number of keys and values. And because it's the KV cache, it's not the KVQ cache.
34:55 if you reduce the number of keys and values that are used, you reduce the amount of memory that you need to hold the KV cache. This is really important because if it gets very large then it limits the number you can fit on one GPU and so this can usually be reduced by like four or by eight or something like that. multiquery attention is just like grouped query attention with a single one. So you can kind of almost skip over that. You already know what it means. multi head latent attention the the is essentially down projecting the size of the keys and values. So it's not reducing the number of heads, it's making each head smaller. And so this is similar to what I talked about with the top K attention. and of course you can combine both of these together as well. You could have group query attention with you know smaller keys of values.
35:57 any any questions here? Okay, cool. so the these are the method the architectural methods that people use to handle very long contexts and with these you can handle up to a million tokens. So like this is you know this makes the efficiency possible. The next thing is training long context models. And this is actually pretty hard. And the reason why is because if we want something to handle 128k tokens or even more, there's actually not that much data on the internet that is 128k coherent tokens.
36:48 it's actually a lot of a lot of data. it's more than most Wikipedia articles, you know, you're starting to get in the in the order of like having a book that you need to have and there's not that many books on the internet. So, typically what you do is you start out with like short pre-training. you do more continued training with longer documents and then you do long adaptation. I wrote packed sequences here. but packed sequences are are useful can be useful for efficiency, but they're not super useful for learning long contexts because every time you like let's say you pack 10 documents into a sequence, you're not learning to attend across the documents and in fact you probably should not be. So you really need coherent long sets of data to do adaptation to very long contexts.
37:41 so I'm going to talk about a few things that we do here. the first thing I'm going to talk about is handling positional encodings. so absolute positional encodings I I hope everybody has seen in their language modeling class. but basically each individual entry in the each individual word embedding gets a position vector and that position vector is typically calculated by some sort of function like using signs and cosiness and other stuff like this. So hopefully you know most people have seen this. So I'll I'll kind of skip over that. Also maybe most people have seen rope. Are people generally familiar with rope? Okay, I see a fair number of people. I see actually less people than I would have expected, but maybe people are just not interested in raising your hand. but but basically rope rope is pretty elegant mathematically in my opinion.
38:48 basically the problem they were trying to solve is they wanted to have a positional embedding. where the positional or positional encoding where the positional encoding is entirely determined by the relative distance between two tokens. And the reason why they wanted to do this is they didn't they didn't want the absolute position to have an effect on they didn't want the absolute position to have an effect on whether two tokens would attend to each other.
39:30 And so this is this has been quite effective and a lot of a lot of models use rope. the the math I'm not going to go through here because it's kind of not the main point, but they do some interesting stuff with rotations and imaginary numbers and stuff like this. but there's also this like super counterintuitive nope. which is don't don't use positional encodings at all. and if you took your class on transformers you might have heard this scary stuff which is like you need positional encodings for transformers because otherwise something in position one in position five will have exactly the same representation. So if you have like this this is a cat but this is not a cat it will attend to both of the cats the same amount or something like this.
40:36 theoretically when the transformers were originally created this might have been true but now actually it's it's not true anymore based on the transformers we use now. anybody have an idea why this is the case? I've kind of given a hint in the slide but you might also need to know a little bit of the history of transformers which is before we were using encoder decoder models where the encoders encoded all of the input and then the decoders generated the output any idea >> because we're doing causal masking exactly so to elaborate on that a little bit if I go back to our transformer equation we have the keys and the queries. And if we look at the queries, even if you have cat at position two and cat at position five, the cat at position two is only attending to the previous token where the cat at position five is attending to the previous four tokens. and so both of them are like actually different. starting at least in the first layer of attention. And so you can tell like even if they are exactly the same word, the embeddings starting in the first layer of attention will will not be different. And there's kind of this counterintuitive paper that demonstrates that you don't need positional encodings at all. And this is especially true this is in auto reggressive transformers. In non-auto reggressive transformers where you're attending all the tokens are attending to each other. You still do need positional encodings.
42:21 this is especially true when we start talking about like gated delta because gated delta has like the decay term and other things like this which causes each token to be different. So once we start having those sort of local biases, we actually don't need positional encodings anymore. and if you if you look we the first two here have rope or at least partial rope. The second two have nope no positional attend encodings. the the fourth one relies entirely on implicit order and mamba and some of them in and inkling has relative positional encodings. So you can see that it's kind of like half split between having relative positional encodings and not at all. But like let's say we did have positional encodings.
43:16 one thing you need to worry about for long context is position interpolation or position extension. And what this means is if you have a pre-training range, you need to essentially extend your model so that it's able to handle longer sequences. And so the the promise of rope was that rope relied only on the relative position between the different tokens in the output. But that that's fine as long as the relative position is only well okay there there's two reasons why this is not fine. So number one if you're training on sequences of length L and then you extend to sequences of length SL then it's never seen two tokens attend to each other over a distance of more than L. So the tokens that are more than L away from each other are going to have problems attending to each other and you're not going to be able to manage that very well. the other problem is that actually relative positional encodings are only one part of the story and like actually the auto reggressive mask has an effect in all these other things. So you can't just take rope and expect it to work really well on longer sequences.
44:41 And so the typical way you fix this is there's a parameter in rope that basically let me see if I have it in my slide. So I I don't I don't really have it in my slide, but there's basically a a parameter in rope theta that tells you how much you rotate at each time step. and you divide this theta by the multiplier by which you extended the sequence. And so basically you're stretching the rope position so that the rope position at position zero would or the rope position at position like five would be 5s would be equivalent to the position 5s in your extended context.
45:30 And there's also a more sophisticated way of doing this called yarn which basically is similar in spirit to like the delta the ki delta attention which basically said that we want some of them to be very long distance and some of them to not be very long distance. In yarn essentially you have some of them be extended more and some of them be extended less and so you don't you don't necessarily interpolate all of them at the same frequency.
46:14 So yeah in the interest of time I'm not going to go into like a lot of detail about this but basically these are these are the various methods that you can use here. So, I'd like to talk a little bit now, taking a step back from the detailed modeling stuff and going into the data that you can use to train this. So, as I mentioned before, packed examples, you can just pack a whole bunch of stuff in there, a whole bunch of documents in into your context length. And this will work at training models to be able to process things near the end of your context window. But what it won't work for is the kind of global coherence. And so there's a number of different ways that people do this. One way is just to find very long documents like books or something like this. There aren't a lot of these but there are books. There are code bases and code bases tend to be coherent. So you can just download a big codebase and that will definitely have more tokens.
47:20 multi-document corpora of some kind. So, you know, like you can get a whole bunch of documents that are on the same topic and then at least you'll have topical coherence between them. I've also seen people train on patents because patents have like there's lots of public data for them and you can kind of link link them together. A big source nowadays is agent trajectories because you know agents can make very very long trajectories if you just interact with them for a little while. So that's another one. And then you can also create synthetic tasks and a lot of the evaluations are synthetic tasks. So you can do that as well.
47:59 and I I have some examples of how the models do this. a lot of the model training companies don't like to share these details with you because this is like if they release their well so number one the closed models like anthropic open AAI you know those places Gemini they will not release even their architectural details the open model companies have to release their architectural details otherwise you can't run their model so what do they keep secret they keep their data mixes and their their training strategy secret so this there's a lot less detailed stuff that we can say. but there's all of them explain all of the ones in here explain a little bit about their data mixtures and they are things like cleaned up sampled documents and video documents, repository files, issue, PRs, long document QA, scientific papers and technical reports.
48:56 And then they also create a bunch of synthetic tasks. And typically they have a scaling strategy of like 8 64 256 1 million that that sort of stuff as well. Oh yes. one other logistical detail that you might need to deal with or might not need to deal with depending on the size of the model you're training. like for example for your project if you want to start dealing with larger models is context parallelism and basically what this means is there's a bunch of ways to do parallelism when training models I wrote context parallelism up here what are two other ways that you can parallelize models during model training or serving >> yes >> data parallelism >> yeah you you said a more specific example of the second one but data parallelism and I would say model parallelism in general and model parallelism can have lots of varieties.
49:59 It can have expert parallelism where you have different experts on different machines tensor more generally tensor parallelism where you split the the parameters and you can also have pipeline parallelism where you have different layers on different different devices. But context parallelism starts becoming import more important for agents because when you train agents you want to train them on very long contexts. And so the way this works is you will compute like one part of the context the keys and values. So you send it to the next device calculate the keys and values for the next for that part send it to the next device and etc etc. So this is very useful. It's really important if you start getting above like 32k 32k context or so. but it's also a pain because you have the you have all the modeling people saying I want my I want my KI delta attention and then you have all the info people saying but wait I need to implement a context parallel kernel for KI delta attention and make sure it actually works and you can back prop through it and all of these other things. So I I personally had an experience where I was trying to train Quen for like a smaller quen model like 135 35b and I I wanted to go do use context parallelism on it and it's like there's no kernel for gated delta net anywhere on the internet and so you need to you need to either go in and implement that kernel or decide to do something else. So th those two can conflict some of the times. So be be aware I guess the the good news is for like standard attention and sliding window attention and stuff like that all of those will have kernels more easily implemented. The hard parts are where you're like not doing a standard attention operation.
52:06 Cool. Okay. So the next part I'd like to talk about is prompt and KV caching. This is very important even if you're not the person who's implementing the language model, but also if you're just a person who's like running agents yourself, but I'm going to talk about the the technical details of how it's implemented. So, you know, at least you you have a good idea about how that works. So basically when we I I simplified this a little bit but when we're generating an agent call we have an input we have an output so this could be like the the messages and tool calls and then we have the tool result and so so as we calculate the input we do prefill and as we calculate the output we do basically decoding Right? And so during prefill and decoding, we have to calculate the KV cache. This is all done in a batch. this is done every time we generate a token. But after we have that, we have the key and value for each of the positions up until the end of the output. And so what we can do is we can throw that in a cache and then we don't need to calculate it anymore.
53:24 The tool result is new because you know we don't calculate over the tool result. So this becomes a new input and we generate an output and we get another tool result. And so in step two we need to calculate the KV cache over the new input and we during prefill and the output but we can reuse the cache and then we do that over and over and over again. And so as you imagine as you can imagine I just showed steps 1 2 3 but by the time we get to step 50 if we haven't done anything tricky with our output this is going to be 50 times larger than the new input in the in the output right so that's pretty significant savings if we can cache this properly and here are some examples from actual papers that you know examine this and basically What we can see is this is like the the token length on different GPUs and this is the amount of of compute that you need with caching and no caching. And so the upper one of the same color is or the lower one with the same color is caching and the the upper one is no caching. So basically as the sequence gets longer the non-cashed one becomes quadratic and the cache one is is mostly linear.
54:50 This also shows up in how much you pay to run an agent. And so I I pulled all of these numbers from the official pricing. And basically you can see here that if you go through all of these all of these numbers including the commercial ones typically cashed tokens are onetenth of the price of uncashed tokens. And so what this means is you need to be very aware of whether your tokens are being cached or not.
55:18 because if they're not being cached you're going to be paying a lot of money to be running your agents. so does anyone have an idea of what a good cash hit rate is if you're running an agent? Yeah. >> 80 to 90%. >> 80 to 90%. >> Yeah. >> yeah. So I I heard 80 to 90% I heard like more than 95%. I my my impression is 90 to 95 is everything is fine. Like I'm I'm good with 90 to 95. If it's less than that, you might be fine, but you might be might be losing out somewhere.
56:03 so, yeah, so basically you can you can cache most of your tokens. one other interesting thing is the output also tends to be five times more than the input. So this is there's a 50x difference between output and cached input which just reflects the amount of the amount of cost savings that you can get from serving it. So why why is caching so much cheaper is maybe kind of obvious. You just need to look up you need to look up the outputs. you just need to look these up instead of like actually performing the calculation over them.
56:49 especially the feed forward network calculation is pretty expensive in transformers and so being able to skip that is is particularly a win. but how do we actually manage that? So the first thing is how do you hold it in memory? And there's two like most popular open-source two most popular open-source LLM serving libraries VLM and SGLANG and both of them are based on an innovation around the how you manage KV caching.
57:26 so VLM's main technical in innovation and like of course they have a lot more stuff in there now. They're like general purpose serving libraries, but like the main technical innovation that led to VLM was something called page detention. And basically the way it works is you have physical GPU memory blocks and you rear you arrange them like kind of memory blocks or or disk blocks on like more typical like CPUbased machines. And so you might have block zero block one block two that you actually need to calculate and they put that you know in like block one and then block seven and then block three or sorry block seven block one block three and then you have an index to each one of these as you're generating.
58:20 And the reason why they wanted to do this is they want to be able to generate in parallel many sequences in parallel and when you finish generating you can like evict that from your memory and so you can evict whole blocks at once and this makes the calculation more more efficient. So this is one technical trick that they used to make this you know make parallel serving work. the technical innovation that kind of led to SG lang in particular was something called RAIX attention and this is specifically focused around like re reusing already cached out inputs. And so what they do is they form a a tree basically it's a radics tree and they use this to look up prefixes that have been used already. And so this when it first came out this was specifically focused around chat bots like chatgpt because you know you can have somebody start a conversation and then continue the conversation and so you'll be caching some portion of the already done conversation. But for now it's like much much more important for agents because agent conversations are much longer and much more frequent. So we need this looking up more frequently.
59:46 so another another thing that you need to think about which is handled in newer versions of sglang for instance is let's say you're doing multi-achine serving. So you're serving on you know different machines you know maybe you have a data center with like 64 nodes or 128 nodes or something like this. every time a request comes in, you want to route it to the machine that has your KV cache that has the best KV cache for you to be accessing and so that you actually have cache aware routing route it to the appropriate place and then have the worker with the best KV cache match process set.
60:32 And so I love getting all my my serving statistics from Open Router because they have so much data and they cover all of the things I want to talk about, but I believe they even have a ranking of providers by cash hit rate. I'm sorry mist the data the data doesn't lie but you can see mistrol has the worst cache hit rate and then you have other one other ones like Alibaba cloud or dicard or silicon flow that have higher cash hit rates this is not entirely reflective of the provider being bad or something like that it could be just the workloads that they're running but there's a number of things that like actually factor into into this cache hit rate and one of them one of them is how good is the cash routing. The other one is how fast is the cash eviction. so the providers will not hold your cash for you forever.
61:38 they will you know evict your cash after a certain amount of time. Does anyone know how long this is for like popular providers? Have you ever taken a look at that >> yet? Close to an hour. >> Close to an hour. I I think the long like I'm not I'm not like super up to date on this, but I know anthropic used to be five minutes. And so it's basically if you have a an active conversation going on with them and you don't pause your conversation, you'll stay in the cache. Otherwise, if you like stop the conversation, you get evicted from the cache. I think you at least have an option to go up to an hour now. but I don't know if you but like it costs a little bit more if you go up to an hour or something. but yeah, ba basically they don't want to keep this around because there this is a huge amount of memory that they're keeping in memory and so they want to use it if it's working and and evict it if it's not working. So that that's another policy that could change your cash hit rate. The bad news for us as users is the pricing doesn't change no matter how good the provider is at routing or or cash eviction. So you you might be paying just because they evict you evict you too quickly which seems a little bit unfair.
62:54 >> yeah >> do they use like cash layering where you can store it on device in memory that's like adjusted to device and then on this >> yeah so the that's a great question. So the question was, do they use a cache layering where you have like a fast cache that's stored in the device and then a slower cache that's stored on disk or something like that? I actually don't know the answer to that question. That'd be a great like a great follow-up for us to do. the open source libraries also are probably not as sophisticated as the closed libraries because if you look at all the people on open router their job is to be really good at serving and so you know they're some of them might be using open source libraries but they'll be building stuff on top of them and that's kind of their you know competitive differentiation. So like I I think the internal solutions that inference providers are likely to be you know a bit better than the ones that are available in open source too with respect to that.
63:54 >> They usually offload to CPU from GPU. >> They usually offload to CPU from GPU. Yeah. >> Yeah. Because your GPU memory is so precious. So yeah. Yeah. >> Cool. so yeah. Oh, sorry. I skipped a important part. So, what does this mean? If you're just designing a harness and you're not deploying models or or anything like this, it means you need to be super super careful to not break your cache unintentionally.
64:27 So, remember that your cache only works if the entire prefix is exactly the same. And so there's a bunch of things that we've thought about doing but we decided are a very bad idea because we don't want to break the cache. So one thing that's completely forbidden you should not do is don't change the system message every time you you take a step because that means like no caching essentially. you don't have a single message that you're able to cache. So one thing one very logical thing that you might think of is okay I want to append the current time and date to the system message so that the agent knows what time it is. Do not do this because the time and date changes every step and so that will break your cache. So if you do need to append it, append it to the final user message or the final tool call result that you get. and like changing your changing your tools dynamically in the middle of the conversation. You can do this. That's okay. another thing is a very logical thing that you might think of doing is routing to a different model every time you every time you take a step. So it's like this step looks like a really good model for you know a really power a really good step for a really powerful model. like like I want to do this step on Fable and then I want to do the ne next three steps on on a cheaper model to save money. But the thing is you're still paying you're still paying the you know 10x cost for Fable the next time you use Fable. So you really need to be very frugal in the amount of switching you're going to be doing if you want to be doing switching to save cost for instance.
66:23 Cool. so the final thing I'd like to talk about is context compaction. how many people know know about this? Use this in your everyday coding agent use. How many people are really excited when it says compacting context? That's nobody good. So the the basic idea behind context compaction is that you have before and this is you know all your observations and actions that you might have had. then after what you tend to do is you tend to take your previous context and you crush it down into just the essence of the previous context.
67:12 and so the reason why people are like not excited about using this, perhaps obviously is if you do a bad job at this, your agent forgets things. It starts doing something completely different and, maybe deletes your file system. So we can view compaction as essentially trying to come up with a previous representation where your future action sequence you take your history, you compact that, you get a working state and you want your future action sequence to be the same or you know similar or possibly better if after compaction than it was before.
68:01 at the same time you also have a constraint which your state needs to be bounded. So it needs to be small enough that you're happy to continue work. so if we look at what survives compaction there's different ways of doing this. I'm going to cover the different ways that people do this, but a common strategy is to take some anchors and these can be the very first like the very first things the user said when they started the session for instance.
68:42 so, one of the answers to the things that I heard agents forget when I asked the question, what do agents forget was sometimes they forget the things at the very beginning of the session? And that actually might be a design bug in the agent harness that you're using. because some good agent harnesses like kind of intentionally avoid deleting the stuff at the very beginning of the session because they think it's more important. but it could also be, you know, impossible to avoid depending on the situation. also you want the stuff in the middle that's encoded and put back into the context.
69:21 And then the recent tail is often kept around. So often people will keep around the the most recent actions. Separately from this, I don't see this a whole lot, but you can also add the evidence to an evidence store. Actually, no, sorry. I I'll correct this. the evidence is typically added to an evidence store. and at the very least, if you're using a CLI coding agent, usually your full history is saved on disk somewhere. And so, if you know where on disk, your coding agent has saved all of this history. Even if you lose stuff through compaction, you can tell the coding agent to go back and read the history, and it can, you know, read the history for you and find the the previous results. but there there are also possibly more structured ways of doing this and so then you know other things can be discarded.
70:15 so there's a few ways to decide compaction policy. the first one is the trigger and there's essentially two ways of doing this. one is a token threshold. So you limit the token threshold. so you can say I'm I'm only going to use up to 200k tokens. Once I go over 200k tokens, I'm I'm going to compact the context. Another thing is like when you get forced to do that by the provider. so that's kind of like a hard limit on the number of tokens you can have.
70:53 finally, you can also have a manual request. And so a lot of kind of coding agent toolkits for instance have a slashcompact like command that you can use to compact the context. then based on that you'll select you might protect the beginning and keep the recent tail and compact the older region. replace it with a summary. that summary can be a structured summary. So sometimes they require a tool call that says this and this and this and this must be in the compaction summary. Sometimes it's just free text.
71:30 there's also I think this is only open AI possibly. But basically they they compact into something encrypted that you cannot read. So you can't see how they're compacting because they don't want you to steal their their precious compaction algorithm. so so there's that. and then yeah, you can you can keep on on working on that.
72:01 so one of the risks here is sometimes you can compact over and over and over again and the first time it works but as you walk through more you you end up losing more and more. So again this is one of the reasons why people don't like compaction because it can cause the agent to go off the rails. so ba basically the way we evaluate this is typically in a holistic state. So you can vary the compaction algorithm and evaluate it against whatever benchmarks you're evaluating against. I would like to warn you about this. I have a a kind of like story from when we were first developing open hands which is we actually I think made one of the first like production grade compaction algorithms. We called it compression then and we evaluated it against bbench which is a very common you know benchmark for for coding agents and we worked pretty hard on it like as a research project for about a month or something like that and we were like we found this great compaction strategy. It doesn't reduce our score at all. it's saving us a lot of tokens. And then we actually tried it in our we tried it in our our system and it made these really silly mistakes like sending a pull request every time it compacted. So it forgot that it had sent a pull request to GitHub. And so you know we would get five pull requests for the same functionality because it had forgotten. or you know it would forget user instructions that were sent in the middle of the conversation because we had only evaluated on single turn conversations that had a single user message and we were compacting away the multiple user messages. So it's not just good enough to benchmark unless you have really really robust benchmarks that represent every use case. You kind of need to test them in the wild as well.
74:03 Cool. so we're we're getting close to time. but yeah, basically what I covered today is you know agent loops grow very quickly. we have a lot of inference challenges in doing so. we can handle efficiency problems through architectures and and caching and compaction is kind of the way we handle this once we expand beyond the level that we can handle. So that's all I have. if there is any questions happy to happy to cover them and next time we're going to be talking about skills and memory. So this is like even longer context over your entire your entire time. are there any questions before we all all leave or >> Yeah.
74:52 >> Yeah, great question. the question was, do we consider sub agents as a way to manage context? And yes, but we're going to talk about sub aents later, so I didn't cover them this time. Yeah.
Summary
- Long context agents can quickly accumulate tokens, leading to quadratic growth in computational complexity.
- Effective context management requires addressing capacity (model limitations) and efficiency (system performance).
- Inference involves prefill and decoding phases, with key metrics including time to first token (TTFT) and time per output token (TPOT).
- Architectural strategies include local and global attention mechanisms, linear attention, and hybrid models that balance efficiency and accuracy.
- Training models for long contexts is challenging due to limited coherent data; strategies include using long documents and agent trajectories.
- Context compaction techniques aim to summarize and retain essential information while managing token limits, but they can risk losing critical context.
- Caching mechanisms significantly reduce computational costs, with a good cache hit rate being crucial for efficient agent performance.
- The presentation concludes with a preview of future discussions on skills and memory in agents.
Questions Answered
What is the focus of the presentation?
The presentation focuses on context management for long context agents, particularly in the realm of large language models (LLMs).
What are the computational implications of long context in language models?
The computational requirements grow quadratically with the length of the context, leading to significant challenges in processing.
How do attention mechanisms adapt to long contexts?
Attention mechanisms can be adapted to allow for varying decay rates of context elements, enabling some information to persist longer than others.
What are some advanced techniques for managing attention in long contexts?
Techniques like 'yarn' allow for varying degrees of attention across different context elements, optimizing the processing of long sequences.
How does cache management impact the performance of language models?
Effective cache management is essential for optimizing the performance of language models, influencing cache hit rates and eviction policies.