Beyond the Theory: What Actually Breaks When You Scale Your Disaggregat... Ekin Karabulut & Ron Kahn

This talk discusses the challenges of scaling disaggregated PyTorch models in production environments. The speakers, Akin from Nvidia and Ron, highlight the complexities that arise when handling different phases of model inference like prefill and decode, which require distinct resource management strategies. They present solutions involving open-source frameworks such as Dynamo and Grove, which facilitate efficient scheduling and resource allocation while ensuring that all components function cohesively. The speakers stress the importance of gang scheduling and topology-aware placements to avoid scenarios that can lead to resource deadlocks or latency issues during inference. They conclude by outlining best practices for implementing disaggregated inference architectures in production settings, emphasizing that no single tool can address all challenges, but a combination of tools can effectively manage the intricate dependencies of these systems.

Hello everyone. Today we are going to talk about what actually breaks in production when you scale disaggregated pytorch models. I know that this was a long day for you, but this is going to be a fun one and we are going to break a lot of things. So stay with us. I'm Akin from Nvidia. I'm a developer advocate for everything run AI open source projects. That means

grow which we will talk about Kai scheduler and run AI model streamer which is right now integrated in VLMNST upstream. Maybe you have heard of it. Um and My name is Ron. I'm a senior software developer. I'm one of the maintainers of Grove. So yeah, today we are going to start with a quick look at why disaggregated serving is out there, the frameworks that are enabling it,

and then we will walk through what a day in production actually looks like and what breaks when you get a lot of traffic and we will show what we learned throughout our experiences and well how we are trying to solve these issues with open source projects. So let's start with disaggregated serving. So you probably heard a lot of talks. I already saw some from LLMD project. So

you probably know what prefill and decode is. So basically we have two phases, right? Prefill and decode in auto regressive inference and we they have fundamentally different needs. So we started dividing these two phases. Prefill is compute bound. You're processing the entire prompt in a single forward pass. It's dominated by large matrix multiplication and decode is basically memory bandwidth bound. So we are dividing these, splitting these

so that the hardware is not underutilized. Sometimes we are even using different parallelism techniques, model parallelism techniques on each side. So you want to deploy this these type of models like mixture of expert models using disaggregated serving. Let's say you chose your inference engine like vLLM or SGLang. They also have production stacks that they are right now working on, but in this presentation we will only look

at inference frameworks that are using these engines like LLMD and Nvidia Dynamo. So, feel free to go and check out those stacks as well. So, let's start with the first one. Um this is a high-level overview of Dynamo. Uh it was built from ground up as a coordination framework for disaggregated serving. Um so, I'm not going to go into detail because we want capture a different problem,

but the requests are coming through an API server basically um with which is compatible with OpenAI and Llama Stack API. There is a router that's cache-aware. It knows which prefill workers already hold relevant cache. Um and routes requests there to avoid redundant computations. And the planner on on the left side is basically handling real-time performance. Um and it has integration with all types of inference engines like

vLLM, SGLang, TensorRT, LLM so you can switch in between. Um MixCell is also a part a part of Dynamo. So, it's basically transferring the KV cache from decode from prefill to decode very fast. It's a transfer layer. And LLMD is another great project backed by community, solves the same problem with a different architectural decisions. So, if you want to have a look, feel free to scan the

QR code. They're also very responsive and in Slack channel so you can ask all types of questions. So, let's assume you chose one of those, right? Both Dynamo and LLMD is trying to solve the hard runtime problems in different ways. So, disaggregated execution, smart routing, efficient data transfer, etc. The question we are going to spend the rest of this talk is what happens when you try to

actually run these models at scale in a cluster. So, we are going to hit five different failure modes today. The first one is scaling. So, the unit of scaling is a bit different. It's a coordinated group now. The scheduling is going to be tricky because we are going to see partial allocations, which will lead GPU deadlocks. Placement is trickier because we have multiple components that are connected

to each others, and it affects SLAs. So, the scheduler might not know your components talk to each other. The fourth one is rolling updates. When you retrain the model or roll any type of updates, the version one might not be compatible with version two, and which component is version one, which component is version two, you might not be like the scheduler might not be aware of it.

And startup ordering, so some components need to start before the other ones. So, you need to pay attention to those. We will cover the first three in this presentation, but I want to also highlight the fourth and the five fifth one because this is also something that we see in production. So, here is our system in production, which I oversimplified to make things simple. Let's assume prefill

is using tensor parallelism, and decode is using expert parallelism. So, front end basically routes the requests to prefill TP group. Prefill produces the KV cache, and then we transfer the KV cache to decode using Nichel. But, there are some things that you need to pay attention to. So, you probably did distributed training before with PyTorch torch run. So, what we have is master and workers, right? If

you divide the model using model parallelism, that means if one worker is not working, you don't have a training. I'm assuming that you're not doing elastic training. So, this is the same. Each prefill and decode decode group has a leader and n workers. So, you deploy you need to deploy basically everything together and if one of these fail, then you don't have a functioning prefill group. So,

this is the first thing. And the second thing is this is a whole system, right? So, prefill is not going to be able to generate a token. So, you need also decode to be there. So, this this whole thing is an atomic unit as well. This is a system entire group. You need to schedule this as one thing or nothing at all. Uh and the other thing

is um leader without a worker or worker without a leader means wasted GPUs. So, it's very expensive for you. And another thing is they all have different scaling signals. So, it's they're not scaling with the same SLAs. The front end scales on tokenization throughput, prefill scales on time to first token, and decode scales on inter token latency. So, these are all super tricky things, right? Let's start

with the worst-case scenario. You deployed your system. Uh so, uh you you had uh you had the front end, prefill, and decode, and there aren't enough GPUs on the cluster. So, prefill got the GPUs, but the decode uh is pending right now. That means you cannot serve even a single request because you're not going to be able to generate the the tokens. So, this is classical resource

deadlock. This is why we need gang scheduling across components. This is what we call gang scheduling because this is a whole gang prefill decode and front end. So let's look at a different scenario. This time we deployed everything is fine, but we have some more traffic the inter token latency increase. So what we do is we increase the number of decode replicas. So we increased it, but

one of the workers is not scheduled because we don't have enough GPUs again. So this means that this decode group which is using expert parallelism is not working. So that means you will still have high latency, but you're not going to be able to function still with the second replica. So we also need gang scheduling like all or nothing scheduling within a component as well, right? So

it's multi-layered. Okay, let's say you eventually got everything scheduled. All components are running, but users are reporting terrible latency. So what is happening? The scheduler put your groups in different zip codes, right? This might also happen. So look where things landed actually the the prefill and the front end groups are in zone A or that might also be a rack and the decode group is in zone

B. So you transfer the KB cache from zone A to zone B. That creates a lot of latency. So your scheduler also needs to be topology aware of what components you have and how they relate to each other I will make it even worse. This looks horrifying. If you already did training with with model parallelism, you know that's horrifying. One is they're in different planets, right? The

code leader is on rack A, workers are in in rack B or like zone B. So this is even worse. We need topology awareness when we deploy these models. So, let's take a step back. We need a lot of stuff here, right? So, what do we need as as a list? So, we need gang scheduling across components. So, prefill and decode are treated as one unit and

don't deadlock on each other. We need gang scheduling within components. That means you never get a leader without its workers or the other way around. We need topology aware placements so your KV cache transfers and collectives actually run on the fast interconnects you benchmark for. Um and we need multi-level auto scaling so prefill, decode, and front end can scale independently on their own SLA metrics, right? Um

we will not talk about the last two, but I still want to highlight smart rolling updates we also need. So, we version one and version two components never try to talk to each other or if they're talking it's on purpose and we actually define that. Um and definable startup ordering so your leader doesn't try to connect to workers when they are not ready. So, that's the wish

list. The question is how are we going to put open source components together and glue them together that actually gives us all this. So, I'm going to hand over to Ron uh who's going to talk about that part. Thank you. Yeah. Oh. Can you Yeah. So, uh we saw what breaks, but let's talk about how we can fix it, how we can solve it. This is very

simplified of how production stack for this aggregated inference looks like. So, on the top level we have the application level. Here lives our inference framework is there LLMD or Dynamo uh which is you probably knows. On the second layer we have the Kubernetes layer. We can divide it to two different but important layers. The workload layer, which define what you want to run, and the scheduler layer

that in charge of where is going to be running it. And on the bottom, we see the hour. This is where your GPU, TPU, whatever accelerator you use is sitting and your workload using to run your model. So, let's start on the scheduler. Unfortunately, Kubernetes native Kubernetes does not support gang scheduling or topology aware scheduling right now. They are working on it, but the community has done

a amazing job to support it with a several good option as a scheduler. We could see Volcano, Queue, Kai, and Unicorn. And there are many more. So, in this presentation, we focus about which is a It's open source by RunAI by the acquisition by Nvidia. Um its main feature is optimized GPU allocation at scale. And it supports gang scheduling cross component, gang scheduling within component, topology aware

scheduling, and many many feature like GPU sharing, fairness, and top of it is a CNCF sandbox project, so you can go check it out. I recommend it. Let move to the workload layer. So, same problem. Native Kubernetes does not support deployment of a disaggregated inference. Luckily again, the community came and created two option. The first one is the little worker set, and the second one is Grow.

Let's dive first for the little worker set. So, what is little worker set? It's an upstream Kubernetes project and it's a community baked. Built for production. It's designed as multi-host inference initially. And it's only one leader. Oh, sorry. It always represent one leader and N worker is one group. You can scale then this group, but you can scale it individually. You cannot scale only part of the

worker in a third replica. It's all or nothing. Um and worker automatically discover the leaders. Gang scheduling and topology aware is achievable by supported a a scheduler like Kai, for example, but many others. And it's always representing one component when we are talking about disaggregated inference. Basically, this is a super worker. Okay. So, if we look back, this is the architecture you already seen of how disaggregated

inference looks like in production. You can see how it looks like from a Kubernetes perspective. We have our deployment. You see one of the limitation that leader worker set is always leader and worker. It's not good for all peripheral uh component like the front end or the router. In the prefill section, we see one leader worker set is a prefill and the second one is a decoder

with two replica scale together. Now, let's talk what we solve using and Kai together. So, the first scenario you probably already remember is the one worker pending. Using Kai and leader worker set, this would never happen. If we don't have enough GPUs for an whole replica, it would be basically pending altogether. So, the entire decoder is pending. We don't consume GPU on partial groups. But, as you

assume, this this could lead us to a different where we don't have enough resources for a one functioning decoder. In this scenario, our prefill doesn't really know that decoder does not even exist. Is pending. From a a component perspective, each component done its job. We have all or nothing for decoder, all or nothing for the prefill, all or nothing for our front end, but it's not it's

not parts, it's a whole system. this is a issue we cannot solve right now with them. Let's move to another uh option that we solve using a little works and guy. Each is the spread part in a group. In before we saw that we could end up with a leader and workers lending in a different zone together all together. And now with topology aware scheduling inside the

component, we said say each component is packed together. And no more spread apart inside a group, as you assume, it does not talk about intergroup communication. Well, you definitely can end up in a scenario one prefill is in zone A and the decoder is in zone B and your SLA is dropped because the KB cache is is slow. this is what is the problem. Little It is

again in first, it's not parts, it's an application. You need to treat it as an application. And this is why we created Grove. Grove knows the entire system. uh let's talk about it. Yeah, basically, it's a flexible API to describe any workload, any type from a single instance to a data center scale. It built to describe this again in first. It's as multi-level gang scheduling semantics and

support topology aware semantics in every level. Uh it include many feature, but it works amazing with Kai and we have uh support for different schedulers like Q and volcano, uh blend on in our what mode. Also, it's open source and it's a model of component used by Nvidia Dynamo we also going to support LMD with a well-lit path document that coming later. So, this is maybe the

hardest part for describing the entire system from a graph perspective. What you need to understand now, this entire this illustration run by one resource. One resource see the entire system. No more stitching around a deployment, a little worker set, all the all the different parts. Grove see them as a one unit. So, you can describe your topology component, your requirement in inside one group or the entire

group. Same for gang scheduling. No more prefill without a decoder because we see the entire application. And this is the one of the advantages of API that describe the entire workload and not just a component base. Okay, so before you leave this room, we want you to take some some key takeaways. So, what are those? First of all, disaggregated inference is different than normal deployments. It's not

a traditional service and this applies to agentic pipelines as well. There are multiple components that you need to orchestrate and deploy and they're all dependent on each other. You cannot treat these as traditional model deployments. Um so, first you we need to understand that. Um second, no single tool will solve it. We need to stitch some things together. Uh yes, unfortunate but still. Uh the first part

is the inference framework that you will need to choose, like Dynamo, LLMD. This is the runtime. Within those you will use vLLM, SGLang, TensorRT-LLM, whatever you prefer. Then you will decide on the workload API that's best for your use case. Grow, leader worker set, maybe single deployments, you know, maybe you are not doing disaggregated inference at all. And then there is a scheduling layer that needs to

understand all these requirements that you are writing in a workload API. And basically enforces the placement on your hardware in your cluster. So these three are super crucial to pay attention to. And of course this is very cliche, but choose your tools for your use case. You know, don't overcomplicate things if you do not need anything super fancy, you can go for simpler deployments etc. But if

you're doing something that's multi-components, multi-level auto scaling with multi-level auto scaling needs you will need a bit more specific API and scheduler. So that's it from our side. We would also like to get some questions if you have some. But all these projects that we mentioned are open source. Please feel free to go to their repos and they all have Discord channels and Slack channels. Get involved

with the development if you need any any specific features. We are from Grove, Kai and Dynamo team. If you have any questions we are super happy to help you on Slack and Discord as well. And thank you so much.