The Token Slice: Implementing Preemptive Scheduling Via Chunked Decod... Maroon Ayoub & Kellen Swain

This talk discusses the 'token slice' and its role in improving AI workload scheduling through the LLMDA (Large Language Model Deployment Architecture). The speakers, Kellen from Google and Maroon from IBM, elaborate on the inefficiencies in decoding requests, particularly the challenge of managing long vs. short requests in multi-tenant environments. They introduce the concept of LMD, which serves as an orchestration tool for state-of-the-art inference using Kubernetes. Key features discussed include intelligent inference scheduling, KV cache management, and the innovative approach of slicing decode outputs into fixed-size chunks. This enables more efficient workload distribution and scheduling mid-stream, enhancing predictability and performance in AI systems. Future enhancements and the importance of community contributions to LMD are also emphasized.

Uh, thanks for coming to our talk today. We're going to be talking about the token slice. Uh, if that doesn't mean anything to you right now, well, you're in the right place. So, I am Kellen. I work at Google, uh, and this is Maroon. Yeah, Maroon here from, uh, IBM, LMD core contributor, and I'm filling in for Itai Levran who was supposed to give this talk. So,

we'll we'll talk about him in a bit. Agreed. Yeah, Itai did a lot of the work here. Uh, shout out to him. Unfortunately, he couldn't make it, uh, but all of us work on LLMD. Uh, so just to give a quick agenda, we'll talk about the problem statement of what the token slice is trying to solve. What is LLMD? Because that's a key component to solving this

problem. Uh, we'll go over the token slice design, and then some related and future work. So, what problem does the token slice solve? Why does it exist? As many of us know, decode length is an unknown problem. We don't know when we send a request, we don't know how long it's going to be residing in the decoder. If you were in the last talk, we touched on

that previously. That can be challenging because you can have some long-running requests that will live on your decoder as fast ones move. If you accumulate enough long-running requests, they can kind of hog the decoder resources, and then you can have some shorter requests that could potentially have shorter decode, uh, sitting and waiting in a queue. That can be really painful in multi-tenant environments where you want to

have multiple users, and maybe you have one user that, uh, has significantly long decode work, uh, and some users may have short ones. So, that can result in a mismatch of the the resource consumption for a given time slice. So, this does sound like an inference orchestration problem, which sounds like a job for the LLM dude. Uh, this is our non-official LLMD mascot. I made this in

Nano Banana last night. Please note that it used location and put an Eiffel Tower in there, which I thought was really fun. But, anyways, so what is LLMD? LOMD is a coordination across multiple organizations with the idea being that it is a quote unquote the open source place to manage state-of-the-art inference. So, currently we run on Kubernetes. Um I'll touch on it at the end, but we're

pretty interested in Ray or Slurm other environments as well. Um but the way it works this laser pointer doesn't work, so I'm going to use my finger. But we have this inference gateway which is an L7 load balancer and we use the an XProc for the inference scheduler here. So, basically when a request comes in, we send the body to this inference scheduler. We take a look

at this and then we tell the L7 where to route to these pods which we call an inference pool, which is essentially a group of model servers that can serve the same requests. and one thing I should mention about LOMD is it is completely composable. If you want certain points like the inference pool which includes the scheduler and inference gateway, you can take that. If you want

to change your accelerators, please do. Same thing with model servers and then auto scalers. The whole point is to mix and match, take what works for you. Uh and we we show that through the well-lit paths and Rune's going to go over that now. Yeah, so one key concept and I apologize if you've been here you heard these, but we're going to go through them. So, one

key concept in LMD is the well-lit paths or guides in which we uh provide a field tested benchmarked set of uh common scenarios, common workloads, and so on that answer or provide uh uh an inference solution for uh pretty much common common scenarios. So, a part of these are what you see here, but we have a much longer list and it's constantly being refactored and improved to

sit better and like help immediately help whoever is using them. So, a few of them are PD disaggregation, wide expert parallelism, which we just covered in the previous talk. As well as intelligent inference scheduling and KV cache management. So, on prefill decode disaggregation, the point here is disaggregating the two main stages of inference, prefill in which you're calculating context or you're computing the KV cache pretty or

the KVs for attention, and decode which is the autoregressive part in which you're generating tokens. Each have different characteristics. One is compute heavy, like you're doing a lot of matrix multiplication. You want to parallelize and specialize your instances in a certain way. And then you have decode, which is memory bound since you you want this state to available in in KV cache or in in in your

memory as long as this like Kelen said, the possibly very long decode phase is going on. So, with PD disaggregation, we disaggregate and these two phases. on here, you see the scheduling path for PD disaggregation in which the scheduler chooses decides whether PD is needed. So, like if your prompt is short enough or if you know you have a KV cache hit somewhere, then you don't need

to disaggregate. You just send a request to a decode pod. But if you do need to disaggregate, then you make the selection of both instances based on certain metrics and KV cache state. You send it to the sidecar to execute the orchestration. The sidecar pretty much orchestrates the prefilling in the prefill instance, uses leverages VLMs Nixle connector to transfer the KVs to the decode pod, and continue

from there. Then there is intelligent inference scheduling which can which will be maybe very relevant here. Which covers the essentials of scheduling for AI workloads. And the AI workloads are different because there is this new state that KV cache or the state you need to orchestrate well in order to achieve efficient inference and cost-effective inference. So in this one we balance between two methods or two approaches.

One being load aware. So you're looking at the cluster state. You're looking at KV cache utilization or GPU memory utilization. Uh queuing in certain pods. And you're using these load metrics to attempt to evenly distribute load across your instances. Then you have prefix cache aware routing which is super essential for inference. You can't or well depending on your workload you need one form one form of prefix

cache routing or another and we provide uh several granularities. In this one you're actually looking at the state or you're approximating the state even uh of the KV cache across your distributed fleet of model servers. And you're using that information to route requests to where the state or where most of the state is. So this one achieves uh prefix aware routing achieves sticky routing. It tries to

send as many requests as possible to where the state resides. Load aware routing achieves distributive routing. Uh and intelligent inference scheduling finds the balance between the two so that you're uh optimizing on both uh simultaneously and getting the maximum out of your hardware. And in a simple scenario in in benchmark of uh tenants uh like this is a chatbot assistant scenario pretty uh old benchmark of which

is only point is to show that in a scenario where you have uh your workload demands less of what you actually have. So, you have 16 H100 GPUs uh serving eight Qwen 32B uh models and you know that your workload demands 70% of the distributed KV cache. So, on paper you should be very well able to serve this workload with the hardware you have and we see

that uh the red and green lines, which are naive load balancing, just load like load aware or round robin which ignore the state uh can't achieve this uh hypothetical uh efficiency because they they destroy KV cache localities, they destroy a lot of the AI specific optimizations then with intelligent inference scheduling, that two lines here uh you do achieve that uh potential optimal performance and this shows so,

take the numbers with a grain of salt. The main point here is you need KV cache or AI awareness to serve AI workloads. Otherwise, even this you know, workload that doesn't require all your hardware cannot be served. One other valid path is KV cache management in which we uh handle all that's there on KV cache offloading and One part is north-south KV cache management or offloading in

which we expand the KV cache pool beyond GPU memory. So, in this one you see uh CPU offloading which was contributed to vLLM. So, with upstream vLLM you can uh switch a flag and get and use your CPU memory uh for KV uh utilization. So, it it extends the KV cache pool. And that's the vertical path. Then you have the horizontal path, which is which uses storage

a bigger tier or uses peer-to-peer, which is not yet there, but is uh planned or like there's a PD flavor of of peer-to-peer KV cache transfer, in which you try to implement compute once, reuse everywhere. So, you computed If you computed something at one node, you're able to transfer it instead of prefilling or recalculating uh KV cache. And now we get to the uh topic at hand,

and again, this is uh this is mostly Itay Ilan's work. So, if you're interested in this, do reach out in the uh LMD Slack, which we'll reference uh near the end. And uh we'll try to answer questions, but he's the expert on the matter. But the core concept here is to pretty much continue to bring well-established ideas from other worlds like CPU scheduling to AI. So, instead

of letting requests run eternally if if needed till uh the model sends out an end-of-sentence token or like or you reach uh uh max tokens or whatever, we slice the decode outputs into fixed-size chunks or slices, defining a new unit of of or a predictable atomic unit of work. So, now instead of working with open decode, we're working with fixed chunks or slices of decoding. And this

lets the scheduler uh reevaluate request placement or request uh uh or workload distribution midstream instead of, you know, waiting or Kelenier will cover different approaches that are implemented in LMD. So, this really gives the scheduler more a bigger hand in midstream's rescheduling and interleaving execution of different requests instead of like instead of allowing some to block others. It It really makes scheduling of the near term more

predictable. So, it's It was implemented in an experimental design building on the LMD sidecar pattern where the sidecar is handles prioritization across requests given some signals from above. It handles queue management, so when you're slicing when you're slicing outputs, you still need to build up the outcome and only send it once it's ready, so the sidecar manages that. And it also manages preemption. It's in the LMD

inference scheduler as an experimental feature. Once it's validated and we know that at the very least it doesn't damage uh inference in in other scenarios like in when you're not saturated and we do see a clear value with it, it will be moved to the or should be core functionality of the LMD And one challenge there is really So, now we have context switching and if not

managed properly, this can be This can be very costly. It can even It can significant potentially significantly harm inference because So, the this thing comes at a clear cost of Now, you have this operational complexity. You need to manage these uh chunks and implement queuing and so on in scheduling. And you also need to be sure that uh you're not actually paying too much for re prefilling.

So, if you send a request in traditional uh decoding, when you send a request, KV cache is locked on the instance and it keeps building up. When you're slicing it, you're risking uh that you need to recompute KV cache if if it gets evicted between slices. So, there are these things to and other context switching overheads to take care of before this can be uh production ready,

let's say. And uh that's pretty much on token slicing. Again, I hope I did it uh justice, but if you questions, follow-ups, and so on, do reach out to Itai on the LMD slack. Kevin Kevin will be covering uh an alternative we do have in in LMD. Thanks, Marin. I think you did a great job. I hope Itai will like it, too. Uh so, a lot of

what he discussed uh is dealing with multi-tenancy. Uh you probably wouldn't want to preempt if it's all just the same user sending the requests. Uh so, a lot of that is overlapping with flow control. That also manages uh priority or different priority requests, different tenants, uh and fairly allocating those. Uh the reason we went with the sidecar, uh again, if you were at the last talk, we

have some uh design decisions we made that we are going to make some changes to that required the sidecar because sending a request back to the X proc currently isn't supported. Uh but once that is, we will be able to support it in flow control. If your question is, "What is flow control?" that is fair. There we go. Uh flow control is a mechanism within our scheduler

that allows for throttling of requests. And why would we want to do that? In the case of YDP, again, if you were at the last talk, you can actually overload your decode workers potentially, which can cause a throughput collapse. So, you want to do some head-of-line blocking beyond the model servers, and flow control does that. Additionally, flow control will let you take different priorities, and I tried

to show through the thickness of the lines here that we're sending more requests from priority four. So, higher number, higher priority. And these negative requests can get shed. So, maybe this is like an offline batch that can take like 3 days, and you just want to consume consume capacity when your production workload has low traffic. And the way we do that is through saturation detection. Saturation detection

is kind of the magic behind this idea. Essentially, we turn the in in in a PD case we turn like KV cache utilization of the D workers into a single float number, which we call fullness. And then you can use that for interesting things like auto scaling, as well as deciding your dispatch rate. So, flow control, we also have some other work going that may work with

the token slice, or is is a different imagining of the token slice. So, we're actually working on eviction in in flow control, similar to what I just talked about. Let's say you have low production traffic, but you still have some batch workloads you want to send. Well, you can send those batch workloads, and because we don't know the future, we don't know when we're going to get

a spike of production requests. When those product or the high priority requests come in, we can actually evict or preempt the the batch workloads off to the KV cache, or maybe we just evict them entirely, which lets you fully utilize your pool, but also provide premium service to your high high priority workloads. questions are: Is that sufficient? Is this eviction, the simple eviction or preemption, is that

currently sufficient compared to token slice? That's to be determined. Token slice does take classic CPU scheduling uh and also takes the downside of it, which is that context switching of the of the threads. In this case, it's the decode request, that is something to consider. It is valuable in providing consistency, but the trade-off is you may be you may be taking some overhead. So, it really depends

And so, if you would use like use this, please let us know. If you think you have an improvement, also please let us know. A lot of LLMD is taking research questions and turning it into production uh utility. So, if you would like to help, please do. And this is our shill. Uh I think we have time for Q&A. but I'll leave this up, so if people

want to take out their phones, um but if we have time for Q&A uh questions. And questions can be on anything around LLMD or what you saw. Go ahead. So, short question on the flow control. Are you maintaining an ingress cache of the requests coming in for your prioritization of your back pressure backlog and the ability to potentially lower the ingress window size for those requests to

be able to equally handle sessions before they starve and reissue? Okay, so so so the question you you can either drop the sessions when they come in and wait for them to come in again, or you've got or you can maintain a session cache. And with the session cache add intelligence for prioritization, back pressure, and potential throttling to maintain the session in the socket rather than dumping

it and having it retry and over load your system. Exactly. That's exactly right. We keep an in-memory cache. Uh so, with that our deployment model, we are we are currently active passive. So, should the scheduler go down, you will lose those requests, but right, we keep the connection open so that you don't have to retry. Yeah. Of course. Any other questions? Sure. Um this could be applied

to other domains um and frameworks. I just wanted to Are there any books you'd recommend on the sort of general theory at a high level? General theory of scheduling? Oh, man. Specific books? No. Um I'd say Go ahead. scheduling for AI. Like, I'm sure there are already books and there are many books on generative AI, but uh it's a novel field. I think you'd find many papers

on uh scheduling algorithms for AI workloads and so on. But, I think the best place to really like feel the state of the art and actually jump in is places like LMD where you can easily get something working, something even local working, and just join the inference ecosystem and contribute to projects like LMD. And specifically to that generalization that this doesn't necessarily only apply to inference, it's

a fair point. Um I haven't read any book specifically. Uh I try to focus on inference, but we do run into common problems. Thundering herd is one. Uh so, you know, we fix that with power of two, things like that. Uh so, a lot of this is learn as you go. Um inference is kind of drinking from a fire hose, so stopping to maybe learn read one

topic I I found that that is a maybe misallocation of my time just cuz there's so much going on, but yeah. Anything else? Oh. I do have one more thing. I I do want to mention um so this is our inference scheduler that we commonly use. So this is for primarily for Kubernetes. Uh it is expected to work in Kubernetes use the inference pool. This pie inference

scheduler uh I vibe coded it and basically told Claude code to replicate this in Python and it worked pretty well. Go Claude code. But the idea of this was for RL. Uh a lot of RL is done in Ray and so we wanted something that was Ray native. I'm probably going to rewrite this in Rust I don't like the GIL and the GIL doesn't like me. So

we're going to fix that. But anyways, if you're interested in a non-Kubernetes scheduler or interested in RL, uh please join us. This is This is rather nascent, but happy to to work with people here. So. I think that's all we got. Thanks all. Thank you.