Lightning Talk: Inside VLLM's KV Offloading Connector: Async Memory Transfers for... Nicolò Lucchesi

This talk covers native key-value (KV) cache offloading in the VLM project, which is recognized as the open-source standard for serving large language models. The speaker, Nicola, discusses moving data in the KV cache from GPU high bandwidth memory (HBM) to lower-cost CPU RAM. The motivation for offloading is to reduce costs and improve the time to first token by skipping pre-fill computations, which is especially beneficial in scenarios with overlapping context. The implementation involves a KV connector interface that allows for flexible management of KV cache and offloading processes. The connector aims for minimal overhead while optimizing throughput as CPU cache hits increase. The talk also mentions ongoing work such as hybrid memory allocator support and multi-tier storage systems.

Yeah, so hi everyone. Nicola, I'm a senior ML engineer with Rada and I'm also a committer with the VLM project. And this this talk is about native KV cache offloading in VLM. And I'm also going to cover the work from Or and his team from IBM research which unfortunately couldn't make it here today. Um So for those of you who have not heard at all about VLM

just like a usual intro. I think it's like very unlikely but still it's the open source de facto standard for serving large language models and it is just the only thing I want to mention is like a very global global effort very distributed more than more than 2,000 contributors from over 50 major companies and just as many in the academia. Um yeah, so let's get right into

it about what is KV cache offloading. I think like very simply all we want to actually do here KV cache think about moving data in the KV cache from GPU HBM to a lower cost storage medium. So we just for the purpose of this talk we're going to mostly talk about CPU offloading. That is moving KV data from HBM to host RAM but it's actually going to

ongoing work to extend that to obviously other higher hierarchy hierarchy where you might want to just offload to like remote storage or SSDs um and and so on. So I guess the main point of this slide is like why would you want to offload in the first place? few points here. So cost of something like very trivial here, but the uh storage media that you want to

offload to is one that's widely more available and has a much lower cost, such as CPU RAM. that's the case for like all the data centers setups that we we see today. what you want to achieve and it's basically you want to skip the pre-fill computation and thus you would like to see a better improvement in time TFT in time to first token. when is this interesting?

Well, we'll see a few use cases about ideally whenever there's overlapping context between requests for one reason or another. One thing you should consider also is that we're going to have a competitive transfer medium which we're going to move the KB data KB bytes around. Like in this case, if you're thinking about CPU offloading in particular, it's going to be PCIE some of you might not be

enthusiastic about, but it's actually we're talking about one order of magnitude difference with respect to NVLink for instance. Actually a bit more than that, more like 20 20X with newest NVLink, but still you also want to consider that as the um transfer medium also gets better with like newer protocols that KV cache is also shrinking. Things like MLA quantization and other compression techniques. So that's just one

way to think about it. And then when say it's just like use cases that everyone probably already knows about what if you have think about multi-turn chats where agent take we have like this very very long context that's just going to be repetitive and most of the time it's going to be shared with things like AI-powered IDs where code base is the same. These are all things

that you can in theory skip and you can in theory offload and and then reload when whenever you have to do pre-fill. Um and let's see it be like just a framing real quick how how one would go about implementing this We have a uh interface which is very much like a contract that VLM provides which is called a KV connector and that this interface allows one

to either implement their own connector which can either be like out of tree or in tree like the one we see here. the idea is to provide a flexible interface for people to tap in into the KV cache management system of VLM. So anything that has to do with the KV either offloading or seeing here or transferring as you might have heard of in PD degradation cases

has to go through this interface. on top of this interface we see here uh this is where the offloading connector resides. So you've got VLM, you've got the connector API so the contract that VLM exposes and the offloading connector which is in this case it's an in-tree connector but it might just also be a out-of-tree connector where you implement your own logic regarding offloading. So the offloading

connector what it does it's it's basically it's a simple offloading API which offer which is based on like a few key points which is going to be like stateless obviously block-based so always going to be thinking in terms of like moving blocks so that the the minimal unit of transfer is going to be like one block. You're not going to move less than one block. why would

you care? It's like interesting if you have cases where you don't have a block which is like fully filled. You only have like couple of tokens in there. So, that's one case where you may not want to like deal with sub sub unit in terms of sub tokens in the block. Um it does design it is specifically designed to be a to have a pluggable back-end system

that is one may want to again not offload to CPU but offload to something like a shared file system. basically this is um the way this functions that is a offloading API which is not something that you'll find like explicitly defined in code. It's more it's more like again a sub contract that offloading connector provides and that is you think about it as a uniform way of

just dealing with a which abstracts a lot of the details that one has to work with when dealing with the internal KV cache management system in VM. And all of this is I mean very much obviously synchronous. So, it's like synchronous base synchronous first. The transfers will be posted and then it will be executed asynchronously with respect obviously to the model execution and the scheduler step. The

way you can use it, it's like it's pretty straightforward in a since a couple of maybe more than that say a couple of months can enable it by default by just adding this KV offloading size argument. So, what what it does behind the the curtains is that it's just going to basically create this connector and the and and and create a VM instance with this connector enabled.

Um some performance numbers maybe slightly different from one may actually want to think about it because what I want to highlight here is actually on the x axis x axis you will see a CPU heat rate. And I already want to and then then a y axis obviously the the throughput so how many tokens per second. But what I actually want to highlight here is already like

at an area which is usually not very interesting but in this case it like 0% so you got like no heat rate at all. Um the CPU offloading connector wants to achieve near zero overhead. That means it wants to be connected that you want to have around but in the worst case like when you have like zero prefix cache hits so it's like totally new request totally

like out of distribution if you want. You don't want your connector to have overhead. So you just want it for the good part but it it don't want it to basically have any effect on your end-to-end throughput. And and this is the case for the offloading connector. And then again the benefits is it's it's quite straightforward because as you as you increase the CPU heat rate you

just going to you just going to be able to basically skip the prefix altogether and just reload from CPU host memory into HBM. And another part of these graphs that I actually want to highlight is actually the opposite tail. So right at the when you have maximum prefix cache hit which means you can load all the blocks in the request because you it's a it's a it's

a full it's a full cache hit. and basically what happens is that you you'll be able saturate the PCIe throughput which is what we want to see here. That's about 60 gigabytes per second per second in uh PCI current generation. And uh closing up, uh what's next? Uh there's an ongoing support for the hybrid memory allocator support that is uh Mamba models. So, uh makes SSMs with

uh full attention models, as well as the uh sliding window models like um the GPTOSS like sliding window for full attention like pretty standard nowadays. And also uh quick mention to the tearing offload. This is another um multi-tier API that will basically enable you to have even easier way to like write and plug in your back end, which might just not be again about CPU memory, but

about a um multi-tier um storage system where you're also offloading to remote. And that's about it for the time I had today. So, uh thanks a lot for your attention.