Why WideEP Inference Needs Data-Parallel-Aware Scheduling - Maroon Ayoub & Tyler Michael Smith

This talk covers the development and functionality of LMD, a distributed large model inference framework designed for optimal production environments using technologies like Kubernetes. The speaker discusses the significance of data parallel wire scheduling for efficient inference, highlighting intelligent inference scheduling that combines load awareness and KV cache management. The session emphasizes how typical workload characteristics differ between standard cloud services and AI workloads, which are often stateful. The introduction of PD disaggregation and mixture of experts models is explored, alongside the challenges and solutions related to deploying these technologies effectively in Kubernetes environments. The talk concludes with insights into performance improvements and upcoming features in LMD.

So, I'm Tyler Smith, chief architect for inference engineering at Red Hat. I also work on vLLM. I'm a core maintainer and I lead the large scale serving SIG. Um, and I also work on LMD, mostly in vLLM in components that we use in LMD. Marine? Yeah. Marine Ier from IBM Research. I work on LMD. I'm a core contributor and uh, KV disaggregation SIG. And before we dive

into YDB, why it needs data parallel wire scheduling, we're going to go over uh, LMD very quickly, give a brief overview. And then Tyler will dive from there. So, LMD is a distributed LM inference uh, framework. So, it solves So, with vLLM you're getting uh, best state-of-the-art inference for one inference engine. And from there to take it into production, you need something like LMD to handle your

uh, distributed setup and make sure uh, maximizing throughput, maximizing uh, hardware utilization. So, it's a Kubernetes native solution. It's built on top of uh, Kubernetes. It uh, it contributes a lot, similarly to how we contribute to vLLM, we contribute a lot uh, to the Kubernetes ecosystem. Uh, specifically like the in Kubernetes gateway inference extension, uh, LWS and so on. It's a cross platform vendor neutral uh, solution.

So, uh, you can work with AMD, you can work with Nvidia TPUs and so on. And it's a it's a wide open source and widely open effort where it's founded by IBM, Google, uh, Red Hat, Mistral, CoreWeave, Nvidia and many others. uh, widely backed and with a wide set of partners that are not mentioned here. But the main goal of LMD is cluster scale uh, production state-of-the-art

distributed inference. And one core concept in LMD is valid paths where we prepare a set of uh, guides for common scenarios, common workloads, common deployment patterns in which we test uh, and benchmark pretty much everything. And you get it working out of the box. Like get a a working reference and from there tune it for your uh, workloads or your environments. Uh, we won't be going over

all the valid paths. Simply, we're just highlighting uh, a few here. Starting with intelligent inference scheduling. So, why really Why is this LM LM inference workload different from standard cloud workloads? So, in here you see uh, classic On the left you see classic cloud workloads. You have uniform HTTP requests or whatever for which pretty much all of your microservices or all of your servers uh, can handle

evenly. Like everything looks very uniform. Every everything looks very similar. And that worked well with classic uh, scheduling and classic load balancing solutions. But with LLMs you have uh, a very diverse set of uh, workloads. very different characteristics. Uh, so long pre-fills, long inputs, short inputs, short long outputs and so on. And you have a new kind of uh, state. classic cloud is stateless. AI is stateful

and that's the KV cache. And you need awareness of that and you need intelligent orchestration on top of that to serve such workloads at scale and efficiently. And that's where intelligent inference scheduling comes in. It's the path that highlights AI aware scheduling. And with AI aware scheduling we have uh, two main domains. So, you have uh, load awareness, which is classic load balancing. You're looking at your

load characteristics. So, uh, waiting queues across your instances, uh, GPU or accelerator memory And on the left side you're seeing prefix prefix aware routing. And that's state or AI state aware routing your scheduler knows where these prefix cache or KV cache states reside on in the distributed uh, cluster. And make sure that if you already computed state somewhere, you'll be maximizing its use. You're not going to

be duplicating and so on. And intelligent inference scheduling brings these two uh, together in a well-balanced manner such that you're not having hotspotting that comes out of prefix cache aware routing. And you're not just uh, evenly distributing load like a naive or a classic load balancer. And in a very simple uh, test, we see that uh, so, take these numbers with a grain of salt. So, uh,

the point here is given a workload that on paper you know your hardware can serve easily. So, in this case it's a it's a workload that demands only 73% KV cache available on 16 H100 GPUs on when serving 8 coin 32B models. So, on paper again, you can easily serve or you should be able to serve this workload. But you see that uh, uh, naive and uh,

primitive load balancing just destroys your uh, TTFTs and throughput. While uh, the theoretical uh, possibility. A second valid path uh, or one of the major ones is KV cache management. So, in this one we we look into uh, orchestrating KV cache and not just scheduling on top of it. So, uh, we have two main pillars here. One is north-south KV cache management in we expand the KV

cache pool beyond just the GPU. This work is mostly in vLLM where we contributed CPU or vLLM native CPU offloading. There's also storage offloading in in LMD that enables east-west KV cache management or compute once, reuse anywhere where your nodes can share KV cache and further minimize duplication and recomputation. >> [clears throat] >> All right. Now, I'm going to talk about um, PD disaggregation. So, the third

valid path here. Um, so in PD disagg, um, you know, request comes in. Uh, we do a call to the endpoint picker. The endpoint picker decides, are we going to do disaggregation or not? If it does, then we go back to gateway, uh, which then routes to a sidecar, which sits on the decode pod. Um, then routes the request to the pre-fill pod. We do pre-fill there.

Comes back to the sidecar, gets routed to the decode and then we pull pre-fill to the decode. Um, spending a bit of time here because uh, PD disaggregation is like this is like an a core ingredient of YDB. So, we're working towards multi-node serving of large MOE models. Uh, when we do this, we always uh, disaggregate pre-fill and Oh, no back bucket. No back button. Okay. So,

when we pull the KV cache from the pre-fill worker to the decode worker, uh, we're using a library called Nickel from Nvidia. Um, which uh, uses RDMA. So, typically InfiniBand or Rocky. Uh, and we pull the KV cache directly from the pre worker to the D worker. Uh, Nickel sits on top of a library called UCX. Uh, and then UCX is is really like cross-platform and also

uh, handles a lot of different transports. So, uh, typically we don't want to use uh, TCP. Uh, but when we do, like it works. Uh, it'll just be very slow. Um, uh, we do this uh, in a totally async manner uh, and zero copies. And there's zero memory overhead as well because we pull directly from the KV cache and insert it into the decode worker's KV Okay.

There we go. and we can see a pretty good performance speed up using PD. it is a thing where you only see a speed up if you kind of look at jointly the optimization between uh, uh, like interactivity. So, tokens per second per user. Um, and then also throughput or efficiency in this case in terms of tokens per second per GPU. Um, and we really see good

speed ups in these sort of like middling ranges. if you're really optimizing for one or the other, uh, typically we see aggregated deployments uh, win out. Uh, the other kind of like disclaimer on PD here is you really have to tune not only your like the number of pre-fills and the number of decodes to your workloads, you also have to tune like the size of your pre-fills

and decodes a bit. So, typically in this like really classical case, we see it makes sense to have like one or a few like big decoders, and then you want to stamp out pre-fillers that are really small. So, in this case, we're doing >> uh like one decoder that's TP4, and then we have four pre-fillers that's TP uh TP1. Um and this is basically because uh like

decode just scales better with tensor parallel sizes compared to pre-fill. And you And it gives you more KV cache space, uh which is really important because the requests are much longer lived on decode than pre-fill, where they're they're kind of ephemeral. There's like one request running at a time. So, that is an essential ingredient in wide expert parallelism, uh where we are basically taking a MoE model,

and we are spreading it over mhm say 32 GPUs or more. Um in LMD, we use something called leader worker set, uh where each pod is sits on a node and runs some number of uh vLLM instances. Uh I'm going to spend a bit of time talking about uh mixture of experts Um so, essential components here, you know, so we've got attention uh layers, and then every

other layer is an MLP layer. In MoE models, we have mixture of experts layers instead of dense MLPs. Uh the key thing here is we have each token goes to a router, uh which then gets assigned to K experts. Uh and then each one of those experts is like essentially uh an MLP layer. Uh so, you can think of it as an activation sparse MLP. most uh

big models, uh essentially all of them, uh over the past, you know, couple of years that have come out have been uh MoE models, and I expect this to continue for a few reasons. Uh one of these is they they tend to uh just perform better uh than dense models when keeping the size number of activated parameters uh constant. Uh this is an ablation from DeepSeek, uh

where as they made their expert segmentation more fine-grained, they saw uh better and better performance. And then shared experts are like will help you dump the kind of like common uh parameters or like common information that's important to all tokens. So, that gives you even better uh separation. Um and then as we've sort of gone from like a dense dense MoE, like a Mixture-of-Experts to the sparser

ones like DeepSeek V3, a bunch of things have changed. So, in dense MoEs, you can get away with tensor parallelism. Um you know, TP shards the inner dimension, this intermediate size. These This is like too small to shard on an MoE model. So, you have basically have to do a different strategy. So, instead, you have hundreds of small Uh and like we we don't we can't get

away with like a simple fused MoE uh kernel that's like pretty easy to implement and tune. Instead, we find that because of the size of the models, we need to scale out to multi-node. Um the nice thing is these like MoE like the routing enables us to use these sparse all-to-all ops. So, they just scale better to multi-node than dense models do. So, kind of like pros

and cons there. You can, but you have to uh go multi-node. Okay. Um So, yeah. Key concepts here, we go from TP to data parallel attention and expert parallel MoEs. The data parallel attention is like particularly important for uh DeepSeek models, where MLA means that if you use tensor parallelism, you're replicating the KV cache. Uh then we use sparse all-to-all dispatch and combine for the MoE layers.

So, we hit the router, uh we dispatch based on the top K weights, and then after the MLP layer, we combine them. Um and then we have to add a couple of optimizations to deal with the fact that we have these uh sparse all-to-alls. Okay. One of these is EPLB, where we find that um there are like heavy hitter experts, and so we do online rebalancing uh

basically replicate the heavy hitter experts, so that we get um basically s- like better balance across the expert parallel ranks. Um and then another one that's really important for H200s, but is less important for GB200, is called dual batch overlap, where we break the forward pass into two microbatches, and then run the dense or the the computation ops at the same time as the communication ops, um

so that you can get, you know, roughly a 2x speedup um you know, when the communication and uh computation are like in balance. Uh for GB200, the communication ops are much cheaper, so we don't see that in practice helping out. And then one question is like how wide do we need to go? really like the primary benefit from from from WDP is that it it spreads the

weights out across the ranks, and the size of your KV cache increases super linearly as you increase the size of your deployment. So, in this graph here, we we're like going from 4 to 48 uh DP Um and just the the Y axis here, that's the number of KV cache tokens per DP rank. So, that's like a an efficiency metric. And that's not like just the KV

cache size, that's the KV cache size per especially for decode, as you're increasing the number of ranks, like the num- your KV cache size increases uh immensely, especially going from four uh GPUs, which is one node on a GB200, up to like 16. Um Yeah, and we're using the Flash-Decoding-and-Attention one-sided all-to-all because it uh uses less memory than DPP, which actually has a bit higher throughput, um

as we've seen. >> And the other reason we need to go high is MoEs are like really uh concurrency hungry. So, you need a really high concurrency to uh saturate GPU compute. So, we see more than 2x uh speedup going from 64 to 256 concurrent requests. Um and we really don't top out until we get to like 1024 concurrent requests, and then we're really in like a

throughput offline inference regime, rather than like uh an online inference. Um and there is a trade-off between like concurrency, throughput performance, and uh like long context. So, you you kind of you can kind of get one or the other. Um and then here, like showing that we can scale this out up to 96 GPUs, even on H200s, uh which are connected via InfiniBand. So, much slower than

the uh than the uh NVLink on GB200s. That's you. Okay, I'm going to run through these slides to leave time for questions, but uh so, going back to the scheduling and Kubernetes aspect of these things. So, now we have kind of a new deployment, a new kind of a deployment that we need to support in Kubernetes and in like for inference. And what's different here is uh

so, you have this one deployment spanning multiple uh accelerator nodes, but given DP uh data parallelism, you have multiple entry points. So, this one deployment is still one model, but you have multiple entry points, and you can externally affect the the growth of KV cache for each node. So, you're going to use use intelligent inference scheduling to efficiently manage the distributed KV cache across this one deployment.

So, for instance, if you know you have a hit uh on rank zero on certain GPUs for one request, you want to try and really send it there instead of sending it to a different rank and duplicating this KV cache on a different set of GPUs. So, with this change comes uh a few challenges, especially in Kubernetes, and the standard uh like endpoint or pod model, cuz

now uh and based on like other choices, such as do you have like a one rank on one pod, or do you have one pod per node exposing multiple ports and that's actually what we went with uh in LMD where you have each node is just one pod multiple ports and your scheduler can schedule to each. So that's uh a new kind of endpoint or pod representation

that we needed to uh extend the gateway API inference extension with so uh it's able to differentiate between different endpoints and it's able to use intelligent scheduling on top of these different endpoints. This is kind of very Kubernetesy maybe uh not too relevant for this crowd but uh the point is supporting this kind of deployment in Kubernetes and enabling intelligent inference scheduling in a Kubernetes native manner

uh came with a set of challenges and work that had to be done across the ecosystem especially again the gateway API uh its implementers such as STU and K gateway. And uh today in LMD we're at a point in the latest LMD uh release we're at the point where the scheduler is really effective in load balancing or or externally load balancing for this kind of deployment and

Tyler will show some uh on that. I'm going to skip the technicals to leave some But uh what you should know is the scheduler or the LMD scheduler is used in uh load balancing across these different uh entry points or different ranks of this one deployment. Okay so um to show the benefits of using uh intelligent inference scheduling in a YDP context um I put together a

kind of a gentic-ish uh like multi-turn uh workload. So what I did here is okay so this is uh 4 by EP8 prefill. So here we have enough uh KV cache size that we can uh reuse um uh uh prompts in cache. Um and an EP32 decode so that's kind of like the smallest deployment where we would say yeah that's a wide EP uh deployment. Um so

10K input tokens uh 2K output tokens so think like a reasoning model um and then 500 on like subsequent ISL. So the workload is sort of like you have a long prompt and then there's thinking on every turn and then there's like a little bit of uh input from the user uh maybe some small tool calls on every turn. Um and then and then five turns so

we're reusing in cache. Uh in the blue uh we have KV aware routing the red we have random and on the x-axis we're sweeping across concurrency. So as we get to like 4096 we're kind of filling up the KV cache and you can see uh performance falls off quite a bit once we get to 8192 that's because we're really uh we've kind of blown through uh cache

space. so and and here we're looking at output token throughput. Um so this is kind of like the throughput of the system overall. So the question is why does output token throughput increase so much when we have uh KV aware routing for the prefillers um which you can see here. So here you can see the prefix cache rate on three of the prefill runs or three of

the the KV aware runs compared to three of the random runs. Um so you can see prefix caching is working here. Um and then over here you can see the number of waiting requests and so basically this is a proxy for when we see a lot of waiting requests in the prefillers basically the prefillers can't keep up and so uh the decode is basically just starving for

work which is why the output token throughput is much lower on the random random routing case. I think Yeah you want to just hit next slide. briefly what's coming next uh we're working on YDP for AMD GPUs. Uh this EPLB step I mentioned before can take up to a quarter second so we are pushing that to become asynchronous and uh lower overhead. Um right now there's a

problem in LMD where you have to pick the the decode and the prefill instance at the same time which is why we have this like awkwardness with the sidecar so we'll be working on uh deferred disaggregation in LMD. Um and then we'll be disaggregating everything so we have uh prefill decode I already talked about that. We also have like encoder EPD disaggregation as well to uh disaggregate

the input encoders. and then other things uh we're working on trying to make production ready elastic EP so if you want to scale up and scale down the size of a running decode instance in a YDP setup or if you want to add uh fault tolerance to elastic EP as well. Okay no time for questions but you can join us on Slack drop your questions there and

if anything you're interested you and you you'd like to contribute there is space for contributions all across the stack all across LMD. Thank you. >> [applause]