Improving LLM Throughput via Data Center-Scale Inference Optimizations

Key Takeaways

Hello everyone. Uh now let's talk about inference. My name is Maxim. Um I'm part of um Nvidia effort called Dynamo. And what we are doing we are focusing on data center scale inference optimizations. Production scale inference is challenging. So today we are serving models that can contain trillions of parameters. They generate enormous amount of tokens especially in thinking mode. They consume a lot of tokens when you send them huge prompts and contexts for production grade system. You want a lot of uh things you want them to be responsive, right? You don't want like a model just wait for a long time until it generates the first token, right? Um you want to make sure that you can serve multiple users at the same time. So it provides high throughput. You want it to be intelligent in the sense that there is no performance degradations or accuracy degradations if compared to the train model checkpoint due to quantization uh or other different um reasons you want to be you want it to be cost effective and energy efficient. So there is so many different things that you need to take into consideration when you're running production grade entrance workloads. That's exactly what we are trying to focus on uh in Dynamo how we can serve LLMs at data center scale. So let's zoom into uh what is inference and um essentially inference contains consists of two different workloads. preview when you take the your prompt and you process it and then decode when you start generating token by token. And an interesting observation here is that these two different phases they have different bottlenecks. Prefill is typically compute bound while decode is memory bandwidth bound. And the main reason being is that for every new token that you generate, you need to pull full model weights from HBM into your GPU cores. So the way we've been running inference so far was collocating these two phases on the same uh GPU on the same uh uh host and node. But what happens in practice especially when you're running multiple users and you want to make sure you have continuous batching which means like you're in the process of generating tokens when a new request comes in you need to add it to the inflight batch which means you need to run this prefill phase and during that time generation is paused for all the other inflight requests. So people have been trying to solve to address this problem applying different techniques such as chunk prefill but a better solution could be just apply disagregation. So disagregate these two phases and make them run independently and the way you can do that is you assign you allocate certain amount of GPU resources to just run prefill. So you process your prompt, you compute KV cache, then you transfer that KV cache to the decode nodes that only run decode. So by applying this trick, this type of optimization, you can achieve sometimes significant throughput uh increase. So this disagregation is just one of the uh optimizations you might think of when you're talking about data center scale inference. There is other things that you should probably consider uh scheduling, how you can route requests between different worker replicas, memory management, how you deal with KV cache uh management, avoiding it computation, how you transfer data between multiple nodes in your cluster, etc. So in Nvidia Dynamo we are applying we're trying to apply all possible known techniques and the first one that I already briefly described is u disagregated serving when you run prefill on one set of GPUs and decode on another set of GPUs and that way you can achieve significant throughput increase. It's very important to take into consider consideration uh your target GPU architecture. Just an example when we were uh optimizing deepcar1 model uh for black wall and compared that to hopper applying different optimiz optimization techniques we were able to achieve 15x throughput uh increase in performance. So here on this chart you can see xaxis is interactivity which means like how fast you can generate tokens. It's like it's measured in tokens per second per user. On the y-axis you have throughput normalized by GPU. So how many tokens you can generate per second per GPU and there is a TR this is like a typical trade-off. You you can be super interactive like generating token super fast but you need to spend more GPU resources and usually you want to find a sweet spot on this curve and decide what is your configuration that you want to deploy. Another interesting observation in uh LLM inference workloads is that a lot of requests they share part of the prompt right just an example you might have seen some leaked system prompts from uh open GPT models or entropic models they contain instructions what model is supposed to say on do and what it's not supposed to do and those systems problems they are constant they are there for every single request and it doesn't make sense to recmp compute KV cache for this chunk of the data instead you want to reuse it another great example could be you're trying to reason about certain data set right you're sending some documents to large language models asking different questions so the context remains the same it's just like the set of documents but questions are So one interesting optimization you can apply here is that reuse computed KV cache. For example when a new request comes in a component called router contains information about all the prefixes that have been cached on all the workers and it can decide which worker contains this specific prefix so that it makes sense to send it there. and save a lot of like avoid a lot of recommutation. It's called KV aware router. So router takes into consideration which prefixes are cached where in addition to overall um load on any specific worker replica if it's KV cache super huge it's running out of memory so it's balancing between these two kind of criteria. So when you're running things in production your uh traffic patterns might change and in Dynamo we have a component called SLO based planner. Essentially the main idea is that you want your workload to guarantee certain service level objectives. You want your time to first token to be reasonably low. You want your inter token latency to be also reasonably low. And you want to maintain these characteristics over time. And what as low best planner does it subscribes to different metrics emitted by workers and makes autoscaling decisions based on those metrics. For example, initially it can decide to um assign n number of GPUs to decode and m number of GPUs to prefill. But as time goes on, we start seeing that uh prefill becomes the bottleneck. it cannot keep up with the number of requests and prefill like takes consumes a lot of time while decode workers they just like uh waste their resources waiting for prefill to complete. So in that case planner can rebalance GPUs across these two categories of workers and assign more GPUs to prefill while in a while it might happen that now the code becomes the bottleneck again based on the matrix planner can make can make another rebalancing decisions and assign more uh GPUs to the code worker and less to prefuel. So the main idea is to observe real time metrics and make autoscaling decision accordingly. We have talked we've talked about uh the idea of reusing KV cache. There is one interesting challenge there. U number of like amount of GPU memory is limited. You cannot store all KV caches computed so far in GPU memory. you need to free it from time to time. An interest an interesting idea could be let's just take these KV caches premputed and offload them to CPU memory and when it fills up to local disk or even we can go further and offload it to some remote storage and whenever we hit that prefix again we can pull it from that storage back into HBM and keep processing. So for that specific purpose we developed a component called KV block manager that is responsible exactly for that managing KV caches of loading them to different uh tiers of storage and pulling them back when when needed. when we are talking about serving production grade uh inference typically you want to to use Kubernetes for that that's where most of the production grade workloads are being served in Kubernetes you have a number of challenges so first of all you want to make sure that whenever you are decide whenever you're deciding to assign hardware to your uh workers you are topology ology where in this specific case you want to make sure that prefill and decode workers they have high throughput communication interconnect between each other so it can send KV cache relatively quickly because otherwise it will become bottleneck and you might see certain hidden performance on the performance you want to make sure whenever you're scheduling you schedule things so that at least so in this case we have front end prefill and decode uh components. We need to make sure that at least one component, one replica for each component is getting started. It has the cluster has enough GPUs and CPU resources and nodes to to schedule it. Otherwise, it will just be in a hung state where you have one front end, one per field and there is no GPU capacity for the code and it you you can get stuck forever. In Kubernetes is called gang scheduling problem. So you want to schedule everything or nothing. So in order to solve that problem in Kubernetes we developed a bunch of operators topology where scattering and gang scheduling is provided by grove and we have Dynamo operator that essentially allows you to easily transition from local testing to data center scale. There is so many degrees of freedom when you try to decide what is the optimal configuration. So imagine you have a like prefill decode disagregation and you can choose many different tensor parallelism for your prefill worker for your decode worker. You can choose between different ratio between number of GPUs you dedicate to prefill or decode. So this is like a problem that we we we want to guide somehow customers uh we don't want to we want to remove the guesswork from from disagregated service setup and EI configurator is the tool which is part of Dynamo that solves exactly that problem. Essentially what it what it does it takes as an input your target model architecture in this case it's uh quen 32B your GPU budget how many GPUs you can afford to serve it your projected um input sequence length and output sequence length and then it can generate uh this para curve. So we see here internal kind of mode two curves one for aggregated mode one for disagregated mode and essentially the axises here are x is uh tokens per second per user the interactivity the same plot that we saw before and y is throughput per second per GPU and then you can choose a point on this plot and AI configurial will give you the exact configuration you need to render and apply to your kubernetes cluster and in this specific case you can see that by just applying disagregation let's say we take a point at 90 tokens per second per user you might get up to 3x uh improvement in throughput so this is this tool is available in both CLI and the web- based user interface um and it solves very important problem given my requirements what is the configuration and it tells you exactly that so in Dynamo we are uh exposing uh different metrics. So monitoring stack is pretty mature. It's uh it's compatible with Prometheus and Grafana. It works with them out of the box. Uh you can easily extend those metrics and add u some custom stuff that you need if you want to. And here's the recap of the Dynamo architecture in the ecosystem in general. So essentially the top going from top to bottom. It has an API server exposes compatible with open AI API. Router is KV aware. So it makes routing decisions based on what uh worker contains which KV cache. There is disagregating serving mode enabled if you want to. So when the the inference is split into prefill and decode data transfer is happening with the help of the nixel library which is like low latency communication library which is part of Dynamo suite. uh SLO based planner is constantly observing metrics in making autoscaling decisions and rebalancing or reshuffling GPU resources between prefill and the code for Kubernetes related stuff like topology where scheduling scheduling risk growth and Dynamo operator and KV cache manager responsible for KV cache floating. So, Dynamo is fully open source. You can get it from GitHub, build from source or use pre-built containers and like other artifacts. There's available road map u tutorials, technical documentation, getting starting guides and everything else. There is a bunch of pre-optimized recipes uh for the most popular models like llama, GPDSS, deepcarban and others. uh targeting specific uh GPUs. Um so feel free to try it out, play with this and if you are interested in data scale infra optimizations, this is the right thing to look into. Thank you.