Lightning Talk: KV-Cache Centric Inference: Building a State-Aware... Maroon Ayoub & Martin Hickey

This talk covers the critical concepts of service level objectives (SLOs) in AI inferencing, focusing on latency and cost. The speaker explains the importance of caching, as well as the pre-fill and decode phases in enhancing AI model performance. They highlight the role of attention mechanisms in transformer models and how optimizing these phases can lead to significant improvements in computation efficiency. The session introduces LLMD, an open-source distributed inference platform that emphasizes KV cache orchestration and intelligent inference scheduling. By utilizing these strategies, developers can achieve better resource utilization and throughput, especially in environments utilizing multiple GPUs. Finally, the speakers encourage contributions to the LLMD project, promoting collaboration in optimizing AI inferencing.

Welcome everybody. Bonjour. I hope you're enjoying Paris or if you're from Paris, you're probably just saying it's everyday of the week. So, what do I want to talk about today? Myself Marwan would like to talk about First of all, can anyone tell me what the three major service level objectives are or SLOs of using AI inferencing are? Anyone want to give me an idea? Shout it out.

Yes? Anyone else? Sorry? Yeah. Yes. With the three ones I'm looking at is latency, cost, and error. Because you always want your cost to be down. You want your latency to be fast and you want to have minimal if no errors because as we know with AI models, you're never going to have no errors. But the two ones we're going to talk about today is around the

areas of latency and cost. Okay? >> [snorts] >> And the last number of years, a lot of work has been done around flops. So, how fast your system can perform calculations. We're going to look at today around caching and also look at the pre-fill and decode phases of inferencing. And this is going to be important going forward cuz these performance gains are really helping especially around with

agent workloads and stuff going forward. So, my name is Martin Hickey and I work over at IBM Research and my areas are around AI inferencing and cloud native computing. Marwan Ayub, also from IBM Research. I work on LMD. I'm a core maintainer. And also he dresses better than me. I've two quotes here. So, if we look at the green side at the start, really what we're talking

around here is how important this quote from Manos is and what it can do to your API pricing. And they're saying here that you're looking at a 10x uh savings. The other one is if you look down on the bottom right-hand side is the cost here essentially is what you're talking here is that if you have a huge token prompt, how we can stall everything and basically

under your underuse your resources. And there are the two things we're going to look at. Looking at that costing saving with KV caching and also uh looking at prefilling the code and using your resources better. Now, just bear in mind some of this still it depends on the situation with the prefilling the code as well. So, I'm just going to take a quick step back. Does everyone

know what attention is? Any hands? So, I'll just give it a little brief overview here. So, attention is a mechanism that's used in a lot of models, mostly transformer models, that decides which part of the input matter most when producing an output. So, in other words, when you're generating your token. And at the inference time, most transformer models operate in two distinct phases. Your prefill, which is

your processing of the entire input and getting your context, and your decode, which is around generating the token, one token at a time. And how prefill caching is so important is because if you look at the top diagram, you can see here every time you're trying to generate the token, you're also generating these key token values. Whereas in the bottom diagram here, the bottom half of the

diagram, you can see here that they're being cached. So, there's a huge savings when you think about it is that every time you're not generating those key value tokens. And you can really see here if you look at the graph on the far right-hand side the throughput that you can get once you have these tokens cached. And if you look at here the example of your multi-turn,

so if you're looking at a conversation with a chatbot, etc., that it's keeping this context every time. So, you might start off with saying what's AI, you get a result. But, every question after that or every interaction after that is going to be keeping that level of of context. So, what happens in that situation is you really need cash around that. So, prefill and decode has been

another area in the last few years where researchers so has shown that in certain situations, if you can separate both the prefill and decode, it has a great um uh optimization. And what happens in this situation is that the prefill phase is more compute-intensive and the decode phase is more memory-intensive. that the decode can sometimes uh dominate the CPU uh wait time and you lose a lot

of efficiency you have in uh virtual caching, etc. So, how does this pull in with LLMD and vLLM, etc.? So, the key point here is that hands up anyone that has heard about LLMD in the last few days? Okay, a few hands. So, just briefly, it's an open-source distributed scaling inference platform. It was announced as a CNCF uh sandbox project about 2 weeks ago, I think it

was in in KubeCon. And it's the key point here is it's built on open-source uh tools uh and the stack, etc. So, you're looking here at uh PyTorch projects like vLLM and LLM cache, but also SIO Envoy and it's deployed on Kubernetes. So, there's been a lot of optimization done here around prefix-aware routing, so routing to where your different caches are and the best cache to get

to, and then also around the prefill and decode. So, I'm not going to talk any further on this, but these optimizations have really made helped with throughput, with latency, resource utilization, etc. But, I'm going to hand over to the expert, Marwan, to walk through it with you. Yeah, I'm going to walk over a few of the KV cache-centric valid paths LMD provides. So, LMD itself is a

uh large-scale Kubernetes native serving platform that is super KV cache-centric. And I guess the more the deeper you get into infer- the into inference, the more you understand that everything really is or most of things I'm talking about are around KV cache orchestration and efficient KV cache reuse. And uh I'll go over some of the LMD valid paths or guides. A key concept in LMD to provide

an out-of-the-box working environment to be used almost as is in production uh for common scenarios and common deployment patterns. Starting with Can you flip me? PD disaggregation. So, Martin mentioned the two main phases of inference and how each is uh dominated by different things and each requires different specialization. That's one major uh guide or valid path in LMD that provides the disaggregation where you can deploy specialized

number of prefill instances, specialized number of decode instances, and LMD would orchestrate and orchestrate remote prefilling uh disaggregated PD and so on when needed. Uh the second a second important valid path wide expert parallelism. So, for large uh MOE models, how do you deploy these efficiently over multiple nodes? How do you still uh manage their distributed KV cache? the bunch of optimizations packed within uh the guide.

Intelligent inference scheduling is a no-brainer and a like something you have to take in when serving inference at scale. It provides KV cache aware routing capabilities, so it's aware of the distributed KV cache pool, which otherwise uh if if you're blind to that, pretty much burning your resources. So, vLLM does extremely well on each instance and with prefix cache aware you're unifying this disaggregated pool uh on

scheduling. Plus, for classic load and uh other kind of uh like classic awareness, GPU utilization, uh queuing, and so on. And with intelligent inference scheduling, you get the you get balanced well-balanced inference load balancing. And and you see that in a very simple experiment in which uh we prove that uh so, it's a workload that demands only 70% of uh of the distributed KV cache in this

scenario over uh 16 H100 Nvidia GPUs. And it shows that with intelligent only you're actually able to uh achieve the the hypothetical uh optimal performance. Uh the last well, it path we'll cover, but we have a We have a lot more that you can see on on our site providing uh guides for common workloads. So, KV cache management, uh expanding the KV cache pool, connecting the uh

disaggregated KV cache pool. So, going beyond GPU, you're off you can offload to CPU memory. You have a bigger pool for KV cache uh storing and reuse. And you can offload to uh remote storage or local storage for an even bigger KV cache pool or sharing across instances. So, let's pull it together. So, for the key for everyone, whatever we're in life is what's it going to

cost. And can we get uh quick responses back? And that's what we talked about initially. And so many key areas where those performances are now coming in around having efficient uh pre-fill uh KV caching and looking at your uh pre-fill and decode phases and see if there's a performance issue with that. Finally, the last one is is looking at the LLMD platform, which is an open-source inference

scaling platform that's doing lots of work around this area. Uh as I say, it's open-source. We're uh the link is down here. So, go into the website, check it out, and please come and contribute because as we say, it's open and we're looking for everyone to uh make the contributions uh count. So, thank you very much. Thank >> [applause]