Routing Stateful AI Workloads in Kubernetes - Maroon Ayoub, IBM & Michey Mehta, Red Hat¶

Playlists: KubeCon + CloudNativeCon North America 2025

Explains why standard Kubernetes load balancing fails for stateful LLM inference and how llm-d's cache-aware routing fixes it, with benchmark data.

Categories: Networking, Storage & data, AI & ML

Topics: llm-d, kv cache, inference routing, stateful ai, load balancing, kubernetes

Transcript¶

▶ 0:00 Uh hello everyone. Thank you for coming to our talk. We'll be presenting route routing stateful AI workloads in Kubernetes. Uh what's stateful in EI where Kubernetes fails and show some data. I'm Aruna Yub. I work at IBM research. I lead the KV disagregation SIG in LMD a project you'll hear about in this talk. >> And hi, I'm Mishi Ma. I work on performance at Red Hat. >> Okay. So what you'll get today in this talk, you'll learn the single most important metric in inference in AI production. You'll understand why standard Kubernetes load balancing fails to deliver in in for AI uh inference.

▶ 0:45 You'll understand what LMD is and how LMD solves this problem in a Kubernetes native manner. and you'll see real world data showing uh how LMD unlocks uh the potential of your hardware. We'll start with a joke. So, two LMs walk into a bar. The first orders a beer. The bartender goes, "Uh, that will be three bucks." The second say, "Same for me, but I'm paying cash." This joke captures part of the essence of what we we'll be talking about. We'll dive into what paying cash means and why it's super important for inference. So we're starting with the 10x cost problem. Uh what we're presenting really affects your bottom line inference

▶ 1:28 production. Without these optimizations, you'll be pretty much paying 10x than you have to and you'll be under utilizing your resources. So I'm starting with a with a quote from a Manus blog. Manus is an an agent uh startup that popped a while ago. And in one of their blogs they wrote the KV cache hit it is the single most important metric for production stage AI. It directly affects your latency and cost. And you you really see this theme around in uh mainstream API mainstream AI such as open AIS, Antropic uh Google's if you check the pricing pages you can see for GPT5

▶ 2:11 uh you pay uh back end a quarter for per 1 million uncashed tokens and 10x less for 1 million cash tokens. The same can be seen in cloet and pretty much all uh mainstream AI pricing. And this talk will will show you how these how Kubernetes or standard Kubernetes destroys these savings and why it doesn't work out of the box for AI inference. So some background to get the idea the most important factor here is the KV cache. So what is the KV cache really? So to start there I'll go about I'll go over self attention in high level. Self attention is how transformers or transformer models understand context.

▶ 2:56 Attention captures relations between different topics in different possessions in the input like different paragraphs, relations between awards and subwords. And it's a score that is computed between every two tokens in in your input. So it's a quadratic operation. For 10,000 tokens, there will be 100 million of these operations. And this step is called the prefill step. It's uh it's what dominates time to first token the time till you see the uh first token show up on your screen. And the solution for this problem is the KV cache. So instead of calculating KVS, keys and values tensors at every

▶ 3:39 step you're reusing them when you're decoding new output tokens. So it's called KV cache, but it's neither KV key value or a cache. This this is a good animation that illustrates uh this attention process and what KVS are. It's an animation from hugging face. On the top you see uh attention the attention process for four tokens without KV cache. You see this growing matrix which is required at every step and every row or every uh tensors or data for uh tokens you already saw are exactly the same. In the bottom you see the effect of caching where you're only calculating the new

▶ 4:21 data. You're reusing uh the tensors or the KVs from previous tokens and you're only calculating the new ones. So with KV cache understood in high level we'll move to uh VM's magic or uh how they efficiently manage the KV cache and uh what differentiates them pretty much and the standard every other inference engine followed or developed in parallel. So two key techniques VLM has paid the attention the manner in which they uh maintain the physical GPU memory. It's pretty much memory virtualization for GPU and automatic prefix caching which is uh responsible for the management and

▶ 5:06 reuse of the KV cache blocks utilizing uh some sort of prefix hashbased LRU map to maintain this information. But the bottom line is it's how VM um achieves this KV cache reuse and uh avoids recomputation where possible and that's what really drives the TTF gains uh you see across the inference stacks. A simple example showing uh inference over 10,000 tokens on a 132B uh model shows the first time you're inferencing you're paying for the full prefill. you're paying you're waiting 4.3 seconds. The second time when the KV cache is effective, it only takes 6

▶ 5:51 seconds which is about seven times faster. So we saw uh KV cache for attention and the forward pass but uh in the wild it's this KV cache or prefix cache reuse is pretty standard in pretty much most of the standard AI workloads. The first being conversational AI, multi-turn AI. So at every turn, you're adding a new question and your context or your prefix cache is the entire conversation. If you don't that if you don't have that ready in the KV cache, you're pretty much recomputing everything in every turn and the waiting time will grow quadratically.

▶ 6:33 The second being agentic workflows uh standard agent. So the context or prefix cache use potential there is also the history or the steps history. So at every step you have an accumulation of uh actions uh tool descriptions observations and such and every new step appends to that context and reusing it is pretty it's it's super important for practical agents pretty much. So we now briefly we now understand Kevy cache slightly and its its applications in in a single instance and we'll talk about where Kubernetes or standard Kubernetes load balancing

▶ 7:16 fails. So we talked about single instance and how VM wins there or the inference engines win there. But in production you don't have a single instance. You have multiple instances. You have replicas. You have shards or ranks. And this once unified KV cache is now disagregated. And if your load balancer is naive to that fact is blind to the KV cache disagregation, you're destroying these optimizations pretty much. A simple example which is uh pretty much how all your requests will be served in a in a naive load balancing scenario is where you have you've served a request on pod A. You have the KV cache warm there. You get uh to the next

▶ 8:01 stage the next turn in the chat or the next step in the agent and now your uh load balancer chose pod B to distribute the load evenly. You've lost the KV cache. You you missed the warm KV cache. You duplicated the work. So now your pod, your second pod is redoing the exact same operations. Your user experienced uh increased latency for no reason and you wasted your GPU resources. Whenever you're uh duplicating your prefill work, you're not doing the code. And this becomes the standard in uh in your cluster with naive load balancing. So we established that the scheduler really needs uh needs vision into this

▶ 8:46 uh disagregated KV cache uh pool or states and that's exactly what LMD provides. So you get a real-time global view of KV cache across all your instances uh in a Kubernetes native implementation that builds on top of uh the Kubernetes gateway API inference extension in what we call the LMD inferenceuler which implements precise prefix cache aware scheduling as well as pre-filled decode disagregation which is out of scope. a high level diagram on how the scheduler operates. So on the right you see your VLM fleet. Each pod is emitting KV events. A KV event is fired every

▶ 9:30 time a cache block is created or evicted on an instance. And this KV event stream is ingested or consumed by the uh efficient KV cache manager library LMD provides which consists of two main components. the KV events pool which has the logic for consuming these events and index them or feeding them into an efficient index a KV block index and the KV cache indexer which exposes the API for querying this index efficiently such that so that the inference or the inferenceuler is able to query it or use this API to serve requests and the standard uh request serving flow looks

▶ 10:13 like the following. So you query the KV cache index, you get a prefix cache affinity score which is pretty much the percentage of the prefix already cached on your pods which represents the computational work you're saving. We combine that with load aware scoring one that looks at VM queuing KV KV cache utilization and other load metrics and achieve a balanced scheduling strategy. It's very important to have the both because if you're only using prefix cache affinity, you're you'll be overloading your pods. So with this scheduling strategy, you get uh uh uh balanced balanced scheduling that

▶ 10:57 ensures you're reusing prefix cache across the pods while also ensuring that you're you're using all your hardware. And with that, uh Misha will continue with some data. One more. Okay. Okay, thanks Maroon. So we're going to go over two different benchmark scenarios that we ran. Uh this is the first one. This is called we call it the shared system prefix benchmarking data. And so here imagine there's 150 different companies within a company. You let's say always send a standard large prefix in. That's part of your company's protocol. You're going to send 6,000 tokens in and then each employee within the company here we've just

▶ 11:40 modeled five employees per company of 150 different companies. Uh so each of these employees will send 6,000 tokens in plus 1,000 will be the question and another 10,000 will come back. So that's sort of roughly the footprint of this particular workload. Uh we ran this on 16 Nvidia H100 GPUs. Uh the model there is the quen 32B. It was a TP2 deployment. So really we had eight replicas of this workload running and essentially if you do the math and look at all of the tokens that we've just talked about. It fills up let's say three quarters of your available KB

▶ 12:25 cache memory. But the important part is it's three if this is provided you put this all in the right place right that you don't duplicate things. So theoretically this workload fits. Um now let's look at these charts. I don't know how well they're showing up on the screen there but yeah that looks pretty good. So the one at the left is TTFT which is time to first token. It's a really important metric right you ask question when do you get that first response back? That's the one on the left. The one in the middle is what we call time per output token, right? That's each incremental token that dribbles in. How much do you have to wait for that token? And so the combination of these two kind of

▶ 13:10 dictates what you see on the rightmost panel there, which is what is your overall throughput, right? How many tokens are we cranking out per second? Now if you look at those uh those lines the bottom lines remember the the the left two panels the best thing you can do is get lower is better obviously right so uh the bottommost line is the blue line which is the um that's the precise prefixuler and so that's doing really well here you can see it's almost a flat line at the bottom maybe it's about 500 milliseconds for that line then the orange line which is the approximate scheduler is a little worse than that, but it's still, you know, relatively

▶ 13:54 decent. It's towards the bottom. And then you can see the other two schedulers are kind of the top. So, let me once again explain these four lines. The precise scheduleuler is the one where not only do you know where everything went, but if it let's say a particular question went to a particular replica and then it got kicked out of there for some reason, right? It got evicted because it was running low. The precise one knows that, right? it won't make the mistake of sending it there if it knows that it shouldn't be sending it there. The approximate one kind of just keeps track of where the last one went. Uh which means that it could potentially get it wrong because it may think something is still in a replica even though it's been evicted a while back. The other two are random, right? Which

▶ 14:37 is where you literally do it randomly. And then the load aware scheduleuler, those are the two lines at the top that you see. Those are more where it's looking at load metrics like okay how busy is there is the is each replica in terms of is there queuing there or um how how much is in the KV cache right looks at a couple of metrics to do that but you can see the clear differentiation there in the TTFD line uh and then the for this particular workload the t the teapot the middle panel is they roughly track equivalently across these four workloads and then for the one on the right uh which is just the throughput. Once again you see what we would hope to see which is that the

▶ 15:18 orange line is doing very well and the blue line is doing even better and the other two are kind of more towards the bottom. So just two more views of this particular data the for the same workload that we talked about before. Um here we showing you every time you ask you know a prompt comes in ideally you have landed on a replica where you already have that particular uh prompt sitting. So we're looking at the prompt the effective cache throughput which is literally how much did you get for free right how much did you get out of that replica without having to do a lot of work uh because as Maroon shared earlier it's like seven times cheaper to cr to pull it out of your cache rather than

▶ 16:03 doing all of that expensive computation all over again. So once again we're seeing trends here which are what we would like to see. Precise gets the most number of tokens directly out of the cache. um the approximate score uh score gets a little less out of the the cache and then you can see the random one really doesn't manage to get a whole lot out of the cache because remember we've got eight replicas here. So chances are you may have sent it to the wrong place where the question that you previously asked is no longer sitting there. Uh similar pattern here. This is uh what a really important metric which is there's always queuing in a system, right? I mean you're going to have a little bit of queuing and here this

▶ 16:46 shows the magnitude of the queueing and you can see that it's the right direction precisely queuing looks a lot better than all the way on the right hand side where that's the randomuler where there's quite a lot of queuing. Okay, this is now what we talked about before was like a single large system prompt that is shared among multiple users in an organization. This one is more a conventional chat kind of workload, right? So you're asking questions, you're getting answers, you're asking questions, getting answers, and the context keeps growing over time, right? So for this one, we're using essentially a context. So the initial prompt is 10,000 and then we go

▶ 17:31 back and forth five times sending a th00and in thousand come back then everything goes back out and you get so this is just any kind of a chat workload that that you've you've you've all used before. Um and you can see those uh if you look at the the the the four the quadrants there the leftmost one are the metrics for the load awareuler the right one are the metrics for the prefix uh awareuler and I just wanted to call out one thing in all of this data here which is if in each quadrant if you focus on the top left that's really how many prompt tokens are you processing per second versus how many of

▶ 18:16 the ones that you process came out of the cache. So, similar to what I showed you for the previous workload here, we've got them both superimposed. So, you can you notice for the load aware, you can clearly distinguish the yellow and the green lines, which means that there was a gap between what it asked for versus what it got. And that gap represents real GPU work being done. on the right hand quadrant. If you look at that equivalent uh chart, you can't even see the yellow line, it overlays the green line so closely because virtually all the tokens that it asking for are coming out of the cache. So, it did a really good job putting stuff in the right place and then pulling it out of the cache. Okay, so we've we've gone over lots of

▶ 19:05 different uh different ways of doing the scheduling. And let's just quickly recap what the four methods were, right? So first is you all know this like random you just pick some of these eight replicas or whatever number of replicas you have and and you're done, right? You just send it there. Load aware takes you one degree better, right? You're saying, okay, is there queuing? If there's queuing, I won't send it there. I'll send it to the pod that looks least busy. That's one metric. The other one in loadware scheduling is you would look at, okay, is his KV cache kind of is it 60% full? Is it 70% full? Because you may not want to send it to a KV cache that's really using 80 or 90% of the of the place. You may send it to some

▶ 19:49 replica that has more KV cache available. Next we move to approximate previous caching which is where you say okay I will send it to the last place that I know I sent this prompt. It may not be completely right but it's a pretty good guess and in fact does better than the load aware uh strategy that I just talked about. And the final one, the holy grail is really precise one, right? Because you know exactly if it's there, where is it? Let me send that uh question to the precise replica where we know all the KV cache for this particular question is sitting.

▶ 20:31 Okay. So, you know, I've gone over lots of uh you know, data here, but really the most important data is your own data. It's your benchmark the ones the workloads that you are going to run. So here for LLMD uh that link is for what we call the well-lit paths. These are the paths where we carefully enumerate for you step by step how to run these workloads. And so if you go to this place you will find one well-lit path for u approximate uh it's just called inferenceuler and that's the approximate one. There's another well-lit path for the precise prefixuler. There are others there, but these are the two that we're

▶ 21:15 focusing on for this particular talk. Okay, so let's let's wrap up with the key takeaways. First of all, as Maroon talked about right at the beginning, AI workloads are stateful, right? When you run something, it's leaving some debris behind in the replica where that question went. So, you really want to send prompts to the current location. So it's really stateful in that sense. Uh traditional routing that we do would not be aware of this issue of where is the state that you really care about. That's why we have all the stuff that we just talked about. And so LLMD is something that makes Kubernetes AI aware. So thank you all for coming to this talk and we

▶ 22:00 are now open for questions. questions. >> You're all convinced Kash is super. >> Yes. I think the question was do we persist the cache data or is it is it only sitting in memory uh in in in the work that we're talking about here we do not but in theory you

▶ 22:43 can take the KV cache and push it out to main memory if you don't have you know if you want to just keep switching back and forth between main memory and GPU memory and then there's also connectors available to actually store it onto persistent uh storage but that's not part of the project right now. >> I'll I'll elaborate on that. So in LMD we do have uh KV cache offloading. What you're referring to could be called KV cache offloading. We have offloading to CPU memory. So that extends your instances memory. So you know the HBM or the GPU memory is kind of very limited. So you have a you can tar the space or increase it with CPU memory. And we also have storage offloading. So you can offload to persistent storage or local

▶ 23:27 storage or shared storage and there are trade-offs and use cases where you switch these on and off. CPU offloading is a no-brainer. You get it you get a lot of CPU uh on all your nodes. So you can make use of that in terms of persisting KV caches. So these are you know you're extending the memory but uh whether you want to keep these KV caches in or not is uh is it depends on the workload. We have some data on that in the LMD website. Is there any advantage to the randomiz that affect the answers that the LM is

▶ 24:12 producing like if you learned that produce a similar set of answers? >> Yes. So this is this is a very good question. what you're referring to is semantic caching where the model or your system remembers answers and may reuse them for different or similar questions. But the again the KV cache the name is very confusing. It was chosen by ML people. So it's neither KV in the standard sense or cache. It's really the state you need in your GPU memory to generate the next token. And for the same inputs, the same exact inputs or

▶ 24:55 prefixes, it's exactly the same KVs. And that's why you cach them, but you really need the state in order to generate more. So what really affects the output is maybe the temperature and randomness in the generation, but the KVs really need to be there as is every time for every same prefix. >> Okay, thank you. Is there any logic around invalidating or sort of when you know to drop it or sort of you know uh user logs off um it's not worth keeping it around anymore. What's the thinking behind some of that? >> Yeah. So today it's uh so we we refer to VLM and I think it's it's uh standard

▶ 25:40 pretty much everywhere. The KV cache memory is managed in an LRU manner. So KV blocks come in if there's no space you affect the least recently used. But we are uh we do we are looking in LMD at making that more efficient. So maybe cost or context aware LRU management. So that so you send a prompt today all the blocks or chunks are treated the same. But if you know that this part is a system prompt and this part is a unique maybe something that will never repeat uh portion of the prompt then you you would want to treat them differently. And we are looking into exposing APIs for uh granular and explicit KV cache

▶ 26:24 management so that you as a user can define whether this is important to cache or your prompt is important to cache or it's a one-time thing but today it's an LRU uh cache management pretty much across all the tiers. >> Hey, thank you for the talk. Um, I had a question about routing because so you're doing this when the request comes in and I was wondering if there's a possibility that we end up routing a lot of decode heavy requests to the same instance because we don't know how many tokens will be generated. So if we are doing prefix cache matching, we route requests to instance A, but all of those requests generate

▶ 27:07 2,000 tokens. So we don't we end up out of KB cache. And so how we can handle that? >> So that's also a great question. There is the prefilled decode disagregation optimization you can also that we also provide. It was out of scope for this talk but it would help you sort of manage these kinds of workloads. So your decode instances would be optimized and would you'd want them to do just decode. So prefix cach hits only help you there. and pre-fill instances would be optimized for prefilling and you know you'd want they still do that best but we are looking into trying

▶ 27:51 to predict >> the decode length or estimate like how much you'd get and that would definitely help with better scheduling >> also uh does LMD also do autoscaler >> autoscaling uh there is support for autoscaling it's in very early stages but it's something we're actively working on. >> Okay. Thank you. >> When you're using a a distributed cache like LM cache for distributed inference, would you then kind of focus instead of doing the preciseuler? Would you look more at like your Q depth and other metrics to make that load balancing decision?

▶ 28:36 >> Yes. So the dis the disagregation or the distributed cache here comes naturally from spanning up multiple instances. We don't actively try to disagregate it or split it. You would have intentional disagregation in maybe tiered caching where you have remote storage say for example across your node. But the disagregation or distribution here is really that you naturally have multiple pools of KV cache since you have multiple instances and you're trying to be aware of this disagregation so that you don't duplicate uh KVS across these otherwise isolated pools.

▶ 29:20 >> Okay, thanks. >> I think we're on time. Uh thank you everyone. Thanks everyone for coming. Bye-bye.