Route, Serve, Adapt, Repeat: Adaptive Routing for AI Inference Workl... Nir Rozenbaum & Kellen Swain

This talk focuses on adaptive routing for AI inference, specifically through the use of the inference gateway. The speaker, Kellen, who is a software engineer at Google, elaborates on the integration of existing L7 proxies, particularly Envoy, to optimize routing for AI models based on request content. The session also highlights the challenges posed by inference routing, such as the thundering herd problem, and presents a pluggable framework that allows for flexible routing decisions. Additionally, the talk introduces two core components for adaptive routing: async detectors, which monitor endpoint metrics, and adaptive configurators, which dynamically adjust routing scores. The results indicate promising performance benefits and resilience against high load scenarios, while opening avenues for further exploration in enhancing user experience and system efficiency.

All right. Hello, hello. Thanks for coming to my talk today. Today we'll be talking about adaptive routing for AI inference, specifically in inference gateway or the gateway inference extension, as it's also known. For an introduction, I'm Kellen. I'm a software engineer at Google. I specifically work on the inference team on GKE. My co-speaker Near was unavailable to attend, but we'll hear from him in a recording in

just a minute. So, first things first, I hope that you all figure out how inference gateway works, what makes it great and useful today, some of the shortcomings we found, why what we're doing today, the adaptive routing, can help, some other ways we're also exploring this problem, and then if you want to join, want to help, please do, and I'll give you links there. what is this?

Why do you care? This this talk is intended to be to be targeted to anyone at KubeCon. So, my quick elevator pitch is that inference routing requires looking at the body, specifically for things like prompt content or the model name if you're doing PEFT routing like LoRA. So, to do that, we wanted to leverage existing L7 proxies, specifically Envoy in this case due to its capabilities with

xProc, and we wanted that because we wanted a global view of the model server state so that we can do this efficient routing. So, how did we integrate? Well, the yellow if if you've seen inference gateway, you've seen this picture before. The yellow is what is existing today or existed before inference gateway came into being. So, we have our L7 proxy, wrong button. Sorry. We have our

L7 proxy here. The request comes in, assuming the green doesn't exist, the request comes in to the model server, and then it's sent back all the all with the L7 managing the connection. Typically, this is going to be through Gateway API. What we add is using the extension protocol provided by Envoy supported by by other gateways. We treat the endpoint picker service as an advisory role. So,

it can take a look at the the contents of the body like we mentioned earlier, and it can provide routing hints that the L7 will use and route to the model server that the endpoint picker suggests. So, how does the endpoint picker select? uh I will mention and people probably here know that uh inference and specifically inference routing is not a solved problem. So, to combat an

unknown future, we built a pluggable framework allows us to make decisions, create routing algorithms, and uh be wrong. So, other people can come up with other ideas. So, at a high level, we'll we'll assume we have a single profile. There's multiple profiles here if you're doing something more sophisticated like PD. For our case, we'll assume we have one scheduling profile. So, to go through that, you would

have filters that would filter off nodes. Uh maybe you this specific request is a high priority customer, so you want to filter to your B200s or your your latest hardware. You would have your scores, which is unique individual plugins that analyze the model servers and score based off a specific state. Uh things like Q depth on the model server, KV cache utilization, and prefix match of the

prompt with the given model server. And then after that, we pick. Uh the most naive way to pick is just take the max score and pick off that. Uh that is something we have default today. You can also do weighted random if you want some distribution. Uh and then there's a a variant of power of two also could be considered epsilon greedy where you will often pick

the max score, but then sometimes filter it off so that you can break the cache a little bit. And we'll get to that in a minute. inference routing crash course, what is the biggest win that we've seen? Uh it is easily prefix aware routing. Uh so, this is a very, very simple prompt, but assume you have these two prompts, what is the weather like in Amsterdam? What's

the weather like in Phoenix? We can see that there is a common prefix there apart from the cities themselves. Uh again, this is a very short prompt, but in production, you're probably going to have longer matches. Well, we can treat this like a dynamic programming problem because of the way that the KV cache is generated and how attention works. So, if we know where these prefixes are,

we can route to those model servers saving prefill time, saving your HBM for for a decode, etc. creating that routing algorithm, we got some pretty decent perf improvements. Now, I will mention this is a benchmark intentional to showcase how prefix routing can help. So, take these numbers with a grain of salt. Uh K8s is meant to be a baseline, so basically your L4 load balancer or a

Kubernetes service that you would just deploy naively. Uh when we use the inference gateway with a very heavy prefill workload, you can see the benefits are clear. Uh you get quite a bit delta on your output tokens per second because you're spending less time prefilling. Your TTFT is much lower because you don't need to prefill. Uh again, this is this is a benchmark to showcase the utility,

but your production workloads, should they be prefill heavy, should see the same benefit. it's not all it's not all sunshine and roses over at inference gateway. There are some issues with the way we currently do things. our scores, like I mentioned, we currently weight them so that you can normalize the score and then build a composite of these different scores. The weighting is manually tuned, and with

if it's not tuned correctly with your prefix scoring, you can have the thundering herd problem as as people who are load balancing experts would know. And that can often cause your model server, a single model server to be overloaded if you have something like a long system prompt. That would cause your one model server that has that system prompt to repeatedly get hit while your others remain

idle, which is not ideal given how costly these accelerators are. Additionally, we don't want end users to have to be concerned with our pluggable framework. The framework is for contributors and power users, people who say, "I don't think you're doing this quite right. I want to tweak that." We give them tools to do so. So, we don't want this config based labyrinth. We don't want a user

to go, "Okay, well, you know, I do have a lot of prefill, but it's not so much. How should I tune the weights?" That's not a That's not a problem we want the end user to to have to solve. So, that is also a current issue with our with our framework. Uh we do have the well-lit path algorithms in LLMD, but they may not cover the the

nuances of your traffic patterns. how can we help? Why am I standing here today? Uh I want to mention the just to summarize this problem that again, for the user, we don't want them to have to worry about this. We want this to reliably work out of the box and adding new novel scores, should someone decide to contribute, should be easy and the weight management should be

easy as well. So, that leads us to adaptive routing. So, some folks over at LLMD, if Maroon's in the audience, him specifically, made made the observation that we can categorize our scores into two groups. So, we have the affinity based scores, for for a metaphor, magnetize the request towards certain model servers. Your prefix match, if there's a Laura adapter already loaded in HBM. You also have your

distribution. So, is this model server overloaded? Is it too full? That's your KV cache utilization, your Q depth, running requests. Uh we we're looking at token load, things like that. If we categorize these two, then we can we can come with a heuristic that if we take a look at the distributive scores of all the pools and see that there's a significant imbalance, perhaps we shouldn't worry

so much about the affinity, and we should weight distributive scoring higher. So, rather than having these manual weights, we we use a function to you can seesaw between the two of if you're significantly significantly variant, then you want to be more distributive to balance the pool again. But, if you're equally equally distributed, it's okay to use your affinity and be affinity based. So, with that, I'll let

Near take it away. I'd like to talk about adaptive routing concept. Starting with the design principles we try to follow. We don't want to do any change to the request path. Meaning, there should be no per request overhead. We also want adaptive routing to be deterministic and explainable. So, if we look at the dashboard and see behavioral changes, we want to know why. Another important point, we

want to keep it completely asynchronous and decoupled from the other layers of inference And last, we want to follow the same principles as the other layers of inference gateway and make sure the mechanism can be extended and customized. As Kellen explained in previous slides, today, when a request comes in, it is being sent from the user to the gateway and then to the scheduler, which is selecting

the endpoint based on metrics and the scheduler config. This is the data plane. We want to keep the data plane as is without adding any complexity during the request serving. In this diagram, we can see at the very high level the flow of adaptive routing We introduced two core components. First, async detector, which conceptually consumes the endpoint metrics and calculate the signal representing some state of the

pool. For example, it could be imbalance of the pool, burst, pressure detection, and more. The second new component is the adaptive configurator, which consumes the signal that are produced by the async detectors and can dynamically adjust the scheduler This logic is part of the control plane and doesn't affect the data plane. The only required change to data plane was to snapshot the scheduler config at the beginning

of request handling to make sure the config is consistent within the scope of a single request. So, as I mentioned in the high-level diagram, the two core abstractions that we added are async detector and adaptive scheduler configurator. Let's dive into them. Starting with async detector. Here are some key properties of this component. First, it should be asynchronous. It periodically calculates a pool state like imbalance or burst

based on the metrics that are collected from the endpoints of the pool. It is also rate-limited. There is no point running the detector faster than the rate that we collect metrics. So, we tied it to the metrics scraping cadence. And it is stable. It produces a signal between zero and one with smoothing mechanism and asymmetric thresholds for entering and exiting a state. So, think about an example,

a burst We want to make sure we use asymmetric thresholds for entering burst mode and exiting burst mode to avoid flip-flop. Following the explanation of what is an async detector conceptually, I want to use a specific example to walk you through. We implemented in our code imbalance detector that conceptually reads per end point load metrics like heavy cache utilization, the waiting request, and the active request. Then

it calculates the pool imbalance using some statistic methods, and eventually it emits a signal. When the signal is zero, that means the pool is perfectly balanced. All the end points in the pool have the exact same number of requests they are handling and the exact same heavy cache utilization. And the bigger the signal is, the higher imbalance we have in the pool. The second core component that

we added is adaptive configurator. It periodically consumes one or more async detectors. It is then doing some smoothing and monotonic translation from the signal to the new scorer weights. And at the end eventually, it dynamically reallocates scorer weights budgets based on the scorer category. To prove the concept of adaptive routing, we implemented it in inference gateway and ran some benchmarks. The initial benchmarks were done with imbalance

detector only. This imbalance detector uses normalized coefficients of variation to calculate imbalance of the pool. Means this is reactive imbalance detection. The adaptive configurator consumes it and uses generalized sigmoid function to ensure it is monotonic and smooth. So, generally speaking, when we have low imbalance, we want to favor affinity scores. And when we have high imbalance in the pool, we want to favor low distribution to make

sure that balance is restored. Great. Thank you, Nir. Hey. So, let's take a look at those we we wanted to look at a few things like I mentioned earlier. We want to make sure that we can deal with this stampeding herd problem, so an end user doesn't have to concern themselves with We want to make sure that we don't regress, definitely. And then we want to the

scores to hopefully reduce the amount of waiting config that you need to do. So, first Nir took a look at the thundering herd problem. The way he did this is he had a low concurrency of a of requests that had a very very high prefix match. So, to simulate that that single hot node. He then spiked it with these different concurrencies here. Don't mind the gaps. We

didn't Those That's not missing data. We just used different concurrencies. So, as you can see, the control group, which is our out-of-the-box inference gateway configuration, uh at these these medium concurrencies here of 125, it performs worse because of that hot node. But as we move over to the very high concurrency here, the distributed scores kick in with their with their weights, and they overpower the uh prefix

match, and then they act basically similar. We have a little bit better TTFT here, but our our uh throughput, but our TTFT is roughly the same. we we can assume based off this evidence that we are now be able to solve this hot node problem without causing any different configuration for the user and it can support more traffic patterns, which is great. I also wanted to make

sure that we met our baseline. LMD has some posted benchmarks with Inference Gateway and LMD have a very very close partnership. So, expect a few more LMD shells here. We wanted to run this baseline. So, I did that. We we have the setup and again, we used a control group and we see that their performance is basically the same. I will mention that the control group had

a few few more dropped requests than the adaptive did. I didn't want to tune those because I just ran three three runs of each and I didn't want to augment the data. So, I left that. So, but you can see where that failed here. Those drops in throughput or where those drops in failures are. But we can see that our time to first token is basically the

same with some error margin and same with our ITL. So, we can assume that we also have not not regressed with our implementation here, which is great. Additionally, when using this imbalance detector, we treated each distributed score as equal and we found that we were able to not regress and get that performance. So, hopefully reducing some of that config So, just to summarize, we have comparable perf.

It is fully async. It doesn't cause any performance issues on the system as it is today and it has explainable adaptation because we are just simply using the the distribution and the variance. So, we have promising results, but there's still a lot to explore. And so, what's next? So, with adaptive, we could do more robust benchmarking. We want to make sure that not only are we we

performing against the benchmarks that we use to showcase our tool, but also against production workloads. We want to take a look at adding additional scores. Is our adding additional scores is that easy to do or is there a config that needs to be managed? Those are questions we want to ask. Additionally, the interactions with our other systems, uh, such as flow control based filtration, which I'll get

to in just a moment. Uh, and I want to mention this isn't the only way to solve this problem. Uh, we I have some folks on my team uh, and some some others in LMD now are working on this the predictive latency. They take a different approach where they use uh, machine learning, you know, very inception, but it's good idea. They use machine learning to predict the

latency and then use that as a stack rank. So, that's another option which requires low configuration. Uh, they trade a little bit operational complexity and compute, but this is very promising. Uh, and if you're interested in this, please check it out. They just had a blog come today. Uh, and the blog will let them talk about it. Also, we have flow control. This is flow control something

that I've been really interested in this year. Uh, with saturation detection, the idea being that flow control happens before. I didn't explain it, but flow control happens before scheduling. So, for flow control to throttle requests at the inference pool, it requires a good knowledge of the saturation of that pool to know it's it's dispatch rate. So, we can use that same idea of the saturation detection to

know when we should filter off pods. Because if we can know when a pool is full, then we should know when each endpoint is full. So, does that solve the problem? Uh, additionally, with prefill and decode, uh, I expect that to become more and more common as more and more uh, frontier models are MOE. So, does decomposing the problem make this easier? That's uh, that's another thing

that we want to explore and we are actively exploring. Again, promising results, but there is a lot a lot a lot to look into. So, please join us. Here's my LMD show. Uh, LMD is joining CNCF. So, we are we are siblings now. Uh this is the QR code for the L and MD scheduler where we do some of the PD work. Uh if you're interested in

RL, uh we we've put together a pie scheduler that is based on the root of the inference gateway. Uh if the the core first commit was vibe coded and told to look at inference gateway. It does work, don't worry. It's not It's not horrible code, but it was generated. Uh and then obviously we have inference gateway here. Uh if you want to join any of these efforts,

please do. Uh we would we would love to have you. Uh and I tried to leave quite a bit of time for questions because I went over some content, but that's all I got.