Sorry, hi. Uh, my name is Mega, and I am a software engineer at Databricks. Um, today I'll be talking about some of the

work that we have been doing at Databricks related to improving LLM inference performance. So, yeah, let's get started. So this is the

outline of the talk. Um, I'll start with just a brief introduction of our LLM serving products, why we continue

using LLMs, some of the key improvements that came with the latest release, um, also focusing on some of the

optimizations that we have upstreamed to VLLM and also some other projects related to FP8 and speculative decoding.

Hopefully, these 30 minutes will be useful before you guys can head to happy hour. Um, so yeah, um, starting with what Databricks

has to offer and why we care about LLM inference performance. So, Foundational Models API is one of our products, um,

and it has two modes for different use cases. So the first is per-token LLM serving. Um, this is where

any user can access and query any open-source model or any proprietary models, including DBX, which is, which was one of

our own mixture-of-experts model trained in-house. Um, and it has very simple pricing for input and output

tokens. This is one of the products that are very widely used by a lot of our users, especially developers. And the

second one that we have is provisioned throughput. Um, this is again where you can use any popular open-source models

or even bring your own. You can also fine-tune your models on our platform and then serve them with this provisioned throughput as an end-to-end experience. And what we will provide is the entire infrastructure where you can pick

minimum and maximum throughput under some latency constraints, and we will provide the hardware that is running on and also autoscale it for you up or

down based on your unique traffic patterns with guaranteed throughput. Yeah, so um, our inference engine is a mixture of backends and we

have uh, integrated VMs um to support many of our serving use cases. So why do we continue uh using VMs? Some of the

key features are highlighted in this slide. Um, first, it's flexible; a lot of uh, broad range of model supports exist on

VMs. As soon as you have like a new open-source model, you will instantly find it launched on the same day with VMs. Uh, we

ourselves upstream the support uh for DBRx um when it at the time of the launch, and it was very easy to add a

custom model. Um, it also has very efficient memory management; continuous batching for improved throughput. Page attention, as you all might be aware, was one of

the key optimizations that was introduced by VMs uh that enables continuous batching. And lastly, it has a

very active open-source community engagement. Some of the people who are also here in the audience um that

continue pushing the boundaries for LLM inference performance. So uh, taking a bird's-eye view of how VM's architecture looks like,

um, on a very high level, uh, the core of VM has an LLM engine which runs an infinite loop in the background, and all

the infinite loop—uh, happening inside that loop is that there is one step that is iteratively generating a batch,

executing it, and then processing it. So inside every step in VM, you have a scheduler which prepares the input

tensors to run this on the GPU. So it allocates KV cache for any new requests, um, it updates the status of the

sequences as pre-fill or decode, and it mixes all the batches together and then finally ships this off and broadcasts it

to the workers, which then runs its own model shard of weights. Um, so this—all of this is happening on the CPU, and now

These input tensors are then run on the GPU, involving many kernel launches.  Basically, a CUDA graph replays all these forward passes on the GPU, and then the sampler generates

the next token ID. This is then shipped back to the CPU where output processing happens. Inside that,

you have detokenization (for stopping sequences), removal of finished sequences from the batch, and finally, streaming the output tokens back to the user. So, a lot is

happening in one decoding step, but the same thing keeps happening in an infinite loop. We ran profiling to see how much time the GPU is idle in each decoding

step. As you can see, only the forward pass uses the GPU. We ran some

Torch traces for Llama 3.1 70B on A100s with FP8 per-tensor

quantization. At batch size 1, 2.6% of the time the GPU is idle. As

you increase the batch size, it becomes 14%, and further increases result in 25%. Clearly, this is very

inefficient, with much room for improvement.  This was the state before our latest release (0.6.0+), which included many performance upgrades—a 2.7x

throughput improvement and 5x latency reduction. This resulted from three key changes driven by the open-source community—

a cross-collaboration across different teams.

Unlike the snail developers in this meme, the community was pretty fast at shipping all of these features.

In less than a month.  But the next few slides, I'm going to focus more on the last

optimization, which was a joint collaboration between Neural Magic and Databricks. Asynchronous output processing is the key idea behind

asynchronous output processing. You delay the output processing of step I to step I+1. So if you see

this diagram, the top shows the state before asynchronous output processing was enabled. You have scheduled forward

and output. But with asynchronous output processing, while the model executor is running

for step I, we're simultaneously processing the outputs from the previous [Music] step, step I-1. So

what happens is this entire large white bubble gets reduced, so your per-token output latency is reduced

because you're reducing GPU idle time by concurrently running your output processor with the model executor. However, there are some

tradeoffs because this delays the output processing. The time to first token will be slightly increased by a few

milliseconds. Also, this approach assumes that all sequences from the previous step will be scheduled for the

next step because the output processing has been delayed. So you might end up decoding one extra token per request,

but overall, the benefits of reduced time per output token and improved throughput

make this a worthwhile tradeoff for most of our use cases. So, while this is a very

simple idea in theory, implementing it had some challenges, especially because of the complexity of VMs.

Architecture and, um, so some of the details I mentioned here are... so this particular asynchronous output processing is compatible only when you

enable CUDA graphs for the VM. Now, um, the VM has integrated CUDA graphs very nicely, in the sense that, uh, you could... CUDA graphs

essentially, the power of CUDA graphs is that you can reduce all the kernel launch overheads, and you can also fuse

operations together and get very high, uh, get very low latency on just the model executor part. So while that is

happening, um, the CPU was always running ahead of the GPU. So inside a model executor, the CPU has already lined up a bunch of

operations. So the way we implemented asynchronous output processing was in a callback, as a callback function, um, so

we basically line up right after the CUDA graph is replaying on the model executor, and we, uh, call the output processing as

a callback function. So, and this way we achieve concurrency. Um, also we tried implementing a very seamless

switching of scheduler states. So the main idea is like if, uh, the scheduler can easily switch from an asynchronous to

synchronous state. So if you have—especially this is, uh, important for when you have requests with beam search

parameters or sampling parameters with best-of-N greater than one. So in such cases, asynchronous output processing is not compatible. So, um, right now what

would happen is that the scheduler, if it sees any request with beam search, it will automatically switch to a

synchronous state. So you can—you will default back to the original behavior of the VM. Um, yeah, so some results of

asynchronous output processing is that we ran, uh, again, 317B on A100s with, um, FP8 per-tensor

quantization. And the speedup that we see on prefill-heavy workloads is... so in this graph, what you

See, the blue line is the baseline, and the red line is with the output, asynchronous output processing.  The x-axis

has throughput; the y-axis has latency. So you see lower latency for the same throughput.

Essentially, there's a 1.01 to 1.09% improvement—a 9x speedup—for prefill-heavy

workloads.  This is basically for 3500 input and 300 output

tokens. So you see, at most, for... basically, the advantage here lies at high batch sizes because the CPU overhead

grows linearly with the batch size. So you see a much higher speedup (the green line) at higher batch sizes, higher requests

per second. Similarly, for decode-heavy workloads—this is where I send requests of, like, 100 input and 300 output

tokens—the speedup is larger. So you get almost a 14% speedup with the

asynchronous output processing at high requests per second.  Another optimization I'd like to talk about—this was led by

AnyScale and Neural Magic—is multi-step scheduling. This brought about a huge improvement in latency

and throughput. The core idea behind multi-step scheduling is that you schedule multiple decode passes in one large step, and what this essentially

does is it minimizes the cost of scheduling and post-processing over *n* steps. Remember the graph

that we saw before, which had scheduled forward and output processing? So, with multi-step, what you do

is you schedule *n* steps, then you do a forward pass for *n* steps, and then you do output processing of *n* steps.

Right.  But the disadvantage is that, if you just have multi-step scheduling, the output tokens are streamed back to the user

Only one out of eight steps. So let's assume you run your multi-step schedule for eight steps; you're only streaming

back without asynchronous output processing. You will be streaming back the output tokens one out of every eight steps, but now, uh, the asynchronous output processing

works quite nicely with multi-step scheduling, and you end up even seeing further boosts to the performance, so

essentially, as soon as the output tokens are ready, the asynchronous output processor will stream it back to the user. Yeah, so here is some—here are some

results when you have all uh these two features together. So in this graph, I again look at—we look at the 3.1, um, 70B

model, FP8 per-tensor quantization on A100s, and the blue line is the baseline

where you have neither multi-step nor asynchronous output processing. Then the red line is when you just enable the asynchronous output

processor, and you get a 9% speedup over the baseline, especially as you can see the advantage of asynchronous output

processing is only at high RPS, right? Because the CPU overhead in that case only grows linearly with the batch size.

For uh, the yellow line is when you enable multi-steps, and this one is for eight steps. So for multi-step, you get an

overall improvement even at low RPS because it is actually even helping at um low batch sizes and trying to amortize

the cost of the scheduler. So that's why you see this um lower latency even with just multi-step, but if you again add on to that

asynchronous output processor, you will get again an 8% speedup over multi-step. So overall, uh, going from the blue

line to the green line, you get a 1.23x speedup over the baseline, which

Databricks' vLLM Optimization for Cost-Effective LLM Inference | Ray Summit 2024

Welcome to the Anyscale YouTube channel! Join us as we explore Ray, an open-source framework designed to accelerate the development and deployment of machine learning applications. From beginner tutorials to advanced insights, our videos cater to all skill levels. Stay up-to-date with the latest advancements and hear from experts through interviews, panel discussions, and keynotes.

Databricks' vLLM Optimization for Cost-Effective LLM Inference | Ray Summit 2024

Databricks' vLLM Optimization for Cost-Effective LLM Inference | Ray Summit 2024

Scripts