Scaling Smarter: Technical Dive Into How Databricks Optimizes Model Serving

Key Takeaways

All right. Thank you everyone for coming. I hope your Data and AI Summit is going well. Sorry for the delay. And some uh we were waiting for the moderator. Um so what what I'll be talking to you today about is foundation model serving and specifically how we optimize performance for it. Um just to give you like an initial introduction. Most Databricks has something called Mosaic AI uh model serving. It's this unified deployment and governance platform for deploying model endpoints, agents, all these these fun things you can play with. Um it's it's secure. It's highly available. It's natively integrated into all of the Databricks um the uh pipelines and so on. And I would I would encourage you to go check it out. But I'm not really going to talk a lot about this. I'm going to focus on the foundation models aspect of it. Which is the large language models and our other genera- generative AI models that we we host ourselves and optimize ourselves. So why should you care about these? Um we've seen a lot of rapid growth recently in usage for these. So personally, I'm not a big believer in AGI ASI coming anytime soon, but there is a lot of emerging patterns we've seen for enterprise and data analysis use cases where these models actually do really really well. And so this is not the forum to share our usage numbers or something, but suffice to say like we've seen exponential growth in the usage of these models. And if you've been at the other keynotes and other talks here, you'll see a lot of higher level products that are being built. And if you drill down deep enough, there's an LLM sitting under them. So things like the bricks, the agent bricks, AI functions, there's a whole host of other products that we're building. At the same time, what we're trying to do is we're trying to stay at the forefront of performance for these open weights models. Um and if you look back where if you look at the some of the workloads that we care about, that we see our customers care about. There is a big gap between us and sort of out of the box DIY deployments of ELM or something or other competitors that we're not going to name on the slides, but are one some of the leading competitors that we have in the field. There's also a lot of uh velocity. We've been improving performance rapidly over the last 6 months or something like 3x in some of our benchmarks and some of our models. I I think there's another 2 to 3x to go in the next couple of quarters, so there's a lot of improvement happening. Um So, we're improving. Growth use um usage is also growing rapidly, so we're trying to keep up with that usage. Um The rough agenda of the talk today is I'm going to talk a little bit about the challenges, what makes these gen AI models special as opposed to, you know, just serving other classical AI models and so on. I'm going to give you a quick uh primer on LLM inference just so everybody's on the same page. Uh then we're going to go a little deep, talk about uh optimizations that we worked on, and then come back a little bit and sort of talk about how we think about balancing performance with quality. Like you can go very, very fast if you don't actually care about the quality of the text you generate. So, we spend a lot of time thinking about that balance. Um Okay, so let's start with challenges. Um Anything that is sort of challenging in a classical AI model kind of gets amplified with these large language models that we're dealing with. And the the first thing is they're large, it's right in the name, and they're getting a lot larger. Like the the size of these models the open weights models has been growing rapidly over the last 6 months or so. And when you look at the closed proprietary models, they're much bigger than what's out there um in many way like the the larger open AI models and Anthropic models and so on. Just the physics of moving these bytes between GPUs, CPUs, disks, downloading them as you auto scale, like it just makes things hard. Um and so one of the things just to get the economics to be competitive, to get the performance to be competitive, you have to compress these models. Uh there's lots of ways to compress them, and a lot of the basic ways actually lead to significant degradations in model quality. So, you have to do this very carefully. Um and the other thing is like you have to serve these off hardware accelerators. Like there's a lot of ways you can serve them off CPUs, you tend to make compromises in either speed or quality. You can't really build a serving system that way. It in in situations where like if you think of a classical AI model like a classifier where you can scale up relatively easy easily to hundreds of thousands of QPS on CPUs, we're trying to scale one of these LLMs to tens of QPS, you you absolutely need very expensive, very high-end hardware accelerators. And finally, the bot there's multiple bottlenecks in the system. Uh and you really have to do these cross-layer optimizations. You have to up co-optimize the CPU code with the GPU code. You have to co-optimize the design of your distributed system with how the servers work and where the bottlenecks in the servers are. Um And so, all of these challenges make setting up one of these LLM uh inference systems fairly hard. There's also another aspect of it, which is there are very strong economies of scale in serving these models. If you just take something like a llama 4 model, serve it on your own on your own private H100, it can be 10 times as expensive depending on your workload than something like Cloud or something. So, the models themselves have a very high fixed cost. So, you sort of need to amortize it across users. You need to like have um multi-tenant system that can really exploit the fact that um the hardware is pretty powerful if you can batch users together. Um and and once you start scaling up, it unlocks a whole host of new optimization opportunities. Like you can have differentiated servers, which is a prefill decode disaggregation. You can uh invest in techniques like low-rank adaptation or like parameter efficient fine-tuning and stuff to sort concentrate different fine-tuned models on on the same server. Um And then but there's the other the flip side of it. When when you're building these these serving platforms, simplicity is very challenging, right? If you just look at the inference engines out there, the models out there, there's tens of models, hundreds of flags, very nuanced interactions between different features on these. And if if you're thinking, I'm going to launch I I want my own LLM, it's very hard to account for all that, get it right. There's a lot of complexity you have to pay for up front. So, our our perspective has been we just want to build something that just works from the perspective of the user. You do an SQL query, it just runs, it auto scales, performance is great. You you talk to an endpoint, it works. You don't have to think about what the quantization scheme is, what the hardware is. Like all these things are sort of tuned to be the best given the sort of workloads that we see on Databricks. Um So, yeah. So, but underneath it all, I think we have to start with a very strong investment in our core platform. Um so, performance is super important for LLMs, but at scale, you cannot be performant unless you're reliable. You cannot be performant unless you're highly available. And so, all of that underpinning all of model model serving is a very strong infrastructure foundation. And then what we do is we integrate the best-in-class open source innovations where you know, like anything that comes out, we're like, okay, if this is faster than what we have, let's just pull it in, let's use it. But then we find a lot of gaps in in open source and other things that we have access to, and we we do a lot of deep research, a lot of deep engineering to overcome and go beyond what what is achievable out there. And so, that that is kind of what I'm going to spend the bulk of this talk talking about. Um But before you know, I can go deep, I need to just give a quick overview of LLM inference. So, at the heart of it, you probably heard about this, LLMs are autoregressive. What that basically means is you just do one token at a time. So, if you're trying to generate a string of text, as far as the LLM is concerned, it's a sequence of numbers, tokens. And you start with your prompt, and it generates one token, feeds it back into the model, generates the next, feeds it back, generates the next, and just goes on until it generates an end token, or it runs out of whatever token budget it has. That That is the basic process. Now, there's There's some important implications there. Because it's variable length, you can't predict how long this is going to be. You can't do smart load balancing and stuff up front. It has to be reactive because you don't know how much work a prompt is going to take. Um it also has implications around how you do um um well, yeah, I'll talk a little bit more about that. Um Another thing that comes out of this autoregressive nature is you're running the same model multiple times, right? Most LLMs are basically stacks of transformer blocks, and when people talk about it, they'll give you like a scary diagram like this sometimes. There's lots of little components. Um but base The The things to take away from this are the model itself is basically a stack of the same module stacked up like 20 layers, 30 layers, 80 layers, depending on the model, right? And you're doing it in a loop. And there are things you can cache within those blocks. So, this is when you hear about the KB cache. There are There There are vectors that you compute on the first pass that you can just reuse in the next pass instead of recomputing them. And what that gives rise to is like you have to maintain uh sometimes very large KB cache to achieve good performance. Um what that also means is because when you start a prompt, you have to compute everything, and then the next time, you only have to compute one token's uh characteristics. And the next time, just one more. So, it gives this split of prefill phases and decode phases, and each of them has very different bottlenecks in terms of the hardware. Like prefill is bottlenecked on the compute capacity of your hardware accelerators. Decode is bottlenecked on basically memory bandwidth. And that gives rise to like different optimizations you can do. Uh gives rise to people splitting up these two things into different servers, so that one is optimized for one case, the other one is optimized for the other case. Um What What most servers This is sort of table stakes. Like if you're running any server right now, LLM server in production, you kind of have to batch requests. And the reason for that is um on modern hardware, the hardware capacity of um is is way beyond what you need for one one prompt, unless it's a very long context prompt. And so you want to be able to batch multiple requests, keep the hardware full, fully utilized. Uh and then you've also got these random arrival processes, especially in a multi-tenant system, right? Your requests are coming in, you don't know when they're coming in. So as soon as a request comes in, you want to be able to integrate it into the batch. So that's dynamic batching, right? Like you want to be able to just add requests into the batch, grow the batch dynamically as requests arrive. Then there is something which most systems do now, which is continuous batching, which is you you want to do Okay, so you're running decode, a new request arrives, you immediately want to put its new tokens into that batch. Another one arrives, you want to put it So there's no clear separation between prefill and decode steps anymore. You're running these mixed batches as you need to, and it creates um additional complications. And features you may have heard of, like chunk prefill and stuff, are sort of like managing how complicated this process gets. Um So yeah, so that is roughly all you need to know about LLM inference to make like reasoned uh judgments about like what different optimizations and stuff you should pursue. Um And then you need to know a little bit about GPUs, right? So uh The main thing you need to know about GPUs is they're a hardware accelerator that sits next to your CPU. Uh the CPU submits work to it asynchronously. So, it's just submitting like do this, do this, do this, do this, like run this matrix multiplication, run this uh non-linear function transformation on it. And as long as the GP as long as the CPU is able to feed those tasks fast enough to the GPU, uh it's fine. The GPU just keeps 100% utilized. It just keeps doing that work. It's pipelined. Um If if the CPU is too slow, you get bubbles. And that's kind of what I spend most of my time trying to figure out how to get rid of very subtle bubbles within the GPU. Um There's also things like there's bottlenecks because the GPU and the CPU usually don't share memory. And so, you have to be very careful about how you're moving data and weights around. And if you do it wrong, it can have a pretty big impact on your performance. Um Okay. So, that's sort of introduction 101 to you know, what's what's important in the space. Let's get get into more of like deeper techniques of of what what we do. Um but first, you know, like why why do we care about performance? You know, with Databricks, we could just launch a container or you could just launch a container with VL 11 and call it a day. It's kind of good enough performance. Uh um We we kind of at Databricks, you know, we pride ourselves as user obsession. So, we really want these responses to be fast. Um but then there's also this other side. I said demand is growing exponentially. Hardware and cost cannot grow exponentially. So, we have to make the system much much more efficient to absorb all that growing demand. And that has to be beyond what we have seen coming out of sort of open source and so on. Um then there's also product co-optimization. Like we do a we spend a lot of time on the performance team co-optimizing the performance in conjunction with the with our product and research teams so that you know, from your perspective as a user, you are like, "Oh, this is an agent brick. I'm going to hit an optimized function and you don't have to worry about like all the months of optimization that went under the hood to like make that 1/10 as expensive as using GPT and just as fast or faster. Um So, okay. So, when when I think about performance, right? There's nothing Sometimes there's moonshots, but most of the time it's this very um it's this very predictable loop where you profile something, you find the bottlenecks, you optimize them, you go back and you fix it and you go again and you find new bottlenecks and that is kind of how we approach it. So, we are we are working on some sort of more speculative things where we're trying bigger ideas, different model architectures, like a lot of stuff, you know, because we have a research team. But for the most part it's like just stack wins on top of each other. So, you take techniques and you like if you improve scheduling, you get a 2x improvement. Then on top of that you go in, okay, I've got better scheduler that is better at sending work to the GPU. I'm going to make the GPU more efficient and that's another 1 and 1/2 x for instance. Then maybe okay, let's do less work, right? And that's quantization, so you work with smaller numbers, the GPU has can compute them faster. And then the third one is, okay, let's share um hardware between users and that's another 3x or so improvement in uh performance per dollar. And so all of these if you stack them all together, you're sort of getting in the 10x range, right? But if you look at any of these individual things, it's like, uh okay, this is a few PRs or somebody somebody spent a couple of works on uh weeks on this. Um So, yeah, so one of the things I will say is if you look at the last generation of um LLM inference engines like VLM V0 and so on and the newer ones which are much faster, sometimes 2x faster, right? The main difference has been this. They have optimized CPU overhead relative um to uh where they were. In the past, what used to happen a lot is the CPU was doing complicating scheduling decisions, copying tensors, like all these things, and you pay 100 microseconds here, 100 microseconds here, a millisecond there, and you've eaten up a lot of time that your GPU could be doing work. Um, and now that we've started optimizing this, especially because these these engines tend to be written Python because they're based on torch, we we are seeing a lot of wins. And so, this is an example, I'm sorry for all of you in the back, but this is an example of comparing, you know, a naive way of doing this where you just have sort of a single-threaded way you're you compute something, you're like, "Hey, what is my batch that I should run on the GPU?" Submit it to the GPU. GPU comes back with, "Here's some logits." You pick a token, you feed it back into the system, and you go again, right? The GPU is is okay. It's it's kind of utilized, but there's a lot of dead space if you look at the bottom a stream, right? Um, but what you can do is you can build your system to be pipelined. Use multiple threads, use multiple processes, do work out a CPU-heavy work in in non-blocking ways, and then submit it to the GPU as efficiently as possible. And when you do that, you can sort of bring these uh, the those GPU execution blocks closer together and get rid of bubbles, and that has a big difference on how much throughput you can get out of the system. And so, this this kind of looks like if if you all have used torch traces and stuff, this kind of looks like what you would see if you're optimizing a system. You're basically uh, bringing the GPU blocks closer together. Um, that is only part of the picture. Um, every time you ask the GPU to do work, a lot of the time there's a dish scheduling overhead, uh, just the way the work is submitted, especially on NVIDIA. Um, and so, there's ways to deal with that. There's CUDA graphs where you're basically trying to cache the operation graphs on the GPU side. They're pretty finicky, but when they work, they're really powerful. There's newer stuff like BDL, which is sort of getting beyond that. It lets kernels within themselves stagger and and execute earlier, start execution earlier. Um There's also something that we have found particularly effective. If you look within these models, you often find parallel paths. Like these are sub graphs that don't depend on each other. They have So the model splits off into two paths and then it converges back together. And if if you do some careful engineering, you can basically make those two paths execute concurrent in parallel on the GPU. Uh this isn't as simple as just, you know, writing some torch code that splits them into streams most of the time. You have to do like kernel optimizations and stuff. But when when this happens, the best case is you can completely hide the execution time of one path. So that that's like a free free lunch in some Well, you have to do a lot of work to get it working. But yes, it gives you a big speed up when it works. And so we've done a lot of this to optimize say Lama 4 or for BaFFT. Um The thing that you will hear a lot about is when people talk about performance optimization is this idea of fusing and writing custom kernels and stuff. And basically a kernel is is just, you know, a a task that executes on a GPU. And the fewer of them you have, generally it's better because starting one of these tasks every time you you pay some cost in starting it up, bringing it down. If you start it up a task to say add one number and then you start up another task to apply a non-linearity to the result of that, right? That can be done more efficiently if you just write some a new task that adds one, applies the non-linearity, then writes it out. You skip intermediate writes, you save scheduling overheads. So that's basically the idea behind fusion. And it's it's pretty powerful. It works fairly well. Uh it requires a lot of GPU expertise to get this right. I'll talk a little bit about that in a second. Um you can go very far in this direction. And um when you when these slides are posted, you can look at that link or Google it. It's some research out of Stanford. Uh and we know other companies done this. This idea of mega kernels, you take the whole model and you just write raw GPU code to execute it. There's no abstraction like clean abstractions and stuff in head and torch abstractions. You just write the C++ code. It just does the whole model. And you know, this one in theory is uh in their blog post is 1,000 tokens per second plus four llama 1B, which is beyond significantly beyond where we are. Um so, it's it's very powerful. It's just its impact is limited to lower batch sizes and so on. So, we don't really personally invest in it. Uh but it is interesting. Um beyond fusion, fusion is like the basic the most basic kernel optimizations you can do. Beyond that, there's a lot of opportunities just write better algorithms. Um and you can if you spend any time online, you'll find a lot of versions. So, the flash attention is probably, you know, one of the most um uh well-known ones where they optimize the attention operation and then they optimize it further in in version three for the newer capabilities of Nvidia's uh the new hardware. Uh there's fast sampling kernels that we use. Um you can also just take the existing kernels and you can tune them in smart ways for the models and the hardware that you care about. Which is where we kind of have my advantage because for us, we're not trying to solve the general problem that Yell Lamb or Nvidia or so on are trying to solve. We're we have some core models that we care about. We have some core workloads that we care about. And so, we can really make those go fast even if we have to make trade-offs with with other workloads that don't care about. Um But, like I said, custom kernels aren't easy. And I have a lot of war stories I can talk to you about about very popular kernels and stuff having very subtle bugs. Uh this is one of them. Um so, we spent weeks trying to figure out why one of our servers was sporadically crashing. Uh and it turns out that it was a bug in a very well-known sampling kernel. Um And it was only triggered when the final layer of the model produced uh one-hot vector, which is pretty hard to reproduce just based on the inputs. Uh and when it did that, there's a loop within the kernel that turned into an infinite loop because they were doing an exact inequality with two floating-point numbers, which had just happened to be lower precision than they needed for that loop to uh terminate. And so, this was something that, you know, like these are subtle, very subtle bugs that crop up. And this is a very complicated kernel. They got the really the the complicated part of the algorithm really they nailed that. Uh but this, you know, like little bugs like this do crop up. We've seen lots of small issues like this. Um but switching gears a little bit, um I talked a little bit about model compression being necessary, right? These models are very big. Um Why is it necessary? It's necessary because, partly because of the competitive landscape, right? Like if you're going to host Llama 4, for instance, right? It is up against Anthropic's Claude, it's up against other model providers. And within that landscape, you absolutely have to compress this model to get the within the economic um like the dollars per token that you need to be in and the right performance. Um And when you compare you can compress the weights, that gets rid of a lot of the memory bottlenecks. Um you can improve do the compute in eight bits as well, which gives you additional flops, basically, because you got twice as many um matrix multiplications, but roughly you can do. Um Um We have not You can go further. There's lots of people that go down to 4-bits. Theoretically, I think the latest research is you know, about 3-bits is where you start losing in an ideal world model quality. Uh but we've we found enough challenges even going down to 4-bits and the speed-ups aren't really worth it. So, we we we tend to stay you know, balance it out at 8-bits. Um if you do this wrong, it's very easy to slow down the process. Like you're going for better performance. Uh you compress weights. If you don't compress use compressed compute, you would basically are upscaling downscaling so much and it it just wastes enough time that your system becomes slower. Uh what's much more troubling is it's very easy to degrade model quality. There's like at least three or four common um very common quantization schemes that folks use. Uh and you just flip those flags in something like vLM, it's very easy to end up with something like this. So, this is a a real example from Llama 8B running in I believe TensorRT FP8. Uh and you just go to it, you say, "Hey, tell me a very long story." And it's fine for like a thousand tokens and then it just devolves into gibberish. Um and if you're running a basic vibe check on this system or something, you're not going to see this. Because what's happening is errors accumulating, accumulating, accumulating and you get to a long enough context and it just throws like you you hit a edge condition the model basically just falls over. Um this is not You can go into playground, you can put this prompt into our 8B. This is not what it does. It tells you the actual story because we did the work to make the FP8 quantization work. Um So, yes. So, this is the the the lesson I would like Well, I'd like to leave you with is you need to quantize with care. We think about looking at quantization schemes that quote unquote, are quality neutral. Um we carefully select them. We look at layer activations, a large spread of uh benchmarks, and so on. Um And we've developed new kernels to actually maintain quality, which do hybrid computation, like some of it is in higher precision, some of it is in lower precision. So, you kind of thread the needle between getting better performance, but not degrading quality. Um So, So, that's quantization. Um Remember calling back to something I said earlier, which is there are very strong economies of scale here. Like you want to concentrate a lot of servers, a lot of users on heavyweight hardware. Like it's better for all of us to share H100 than each of us to get like the our own A10, roughly, which is a much uh all of us would get better performance. Now, that starts to break down when folks have fine-tuned models, right? You you can't it's different weights. And so, we spent a lot of time figuring out how to uh share uh share servers with fine-tuned models. And this is sort of this this technology underpins AI the agent bricks. Um Now, I'm not going to spend a lot of time talking about it, but uh just want to give a high-level We use this idea of training smaller adapters that live with the models. It's a pretty common technique. Uh but when we started to implement it, we found that the the the more com- the popular ways of doing it actually caused quality issues. And so, we spent a lot of time with our research team finding a version of it that maintained the quality that we needed, uh which was very slow out of the box. So, we spent a lot of time with the performance team trying to optimize it with things like those parallel bots, and so on, and custom kernels. Um There are There have been a couple of talks uh at the conference that go much deeper into this. So, I would point to them, and you know, when the slides and stuff are released uh videos are released, I would encourage you to look at them if you're interested in this. Um There's What I will say is You know, I don't have enough time to talk about all the optimization work that we do. I will say I have a very firm belief there's a lot more headroom. Um and some of it is data-driven, some of it is intuitive. Um I think we have another 2x, 3x, so on to go. Um And beyond that, Genie is a very fast-evolving space, right? Like you're getting new models every month, uh new techniques, new hardware coming online. So, what you as users you should expect is the cost and performance of these models to just improve. Quarter over quarter, year over year, uh you should expect significant improvements. Um But as performance is increasing, it's very important to pay attention to quality. Uh quality is a lot more important than performance. Uh Some inference runtimes that are out there are very, very good at producing gibberish very, very fast. Uh And we have used them in production. We've had to deal with some of these issues. Um It's hard to detect. It's hard to Like it's hard to detect during releases if you don't have a good release process. It's hard to debug. Um Sometimes it's obvious, you know, like the gibberish example I shared with you. Bad FP8 quantization, you'll you'll start by checks start failing, you start seeing coding mishaps and so on. Um Other times it's it's much less obvious. So, this is This is a bug. It's a real bug from a Llama 4 implementation. Um where it's applying the non-linearity on the wrong side of a matrix multiplication. So, that's a pretty big difference in the math. And this is in a core component. Now, the way this manifests is the human eval score goes down by five points. Which is Okay. I I don't really know how to debug that. You dig into what's happening, there's one line difference in this whole generation that it has. Uh and nine most of the generations are not affected. Um and so when especially when you start getting to these larger models, they're surprisingly resilient. And so small bugs are not going to show up in live checks and end evals. You have to do much more thorough testing and evaluation. And so we've built this system for very strict quality control over the last few quarters and the idea that every release we have, whether it you know, we're releasing something that's open source or inbuilt or something, has to go through a gauntlet of tests. It has to go through a range of academic evals where we then have to analyze the differences between the release and the prior release based on also correct the statistical differences in these evals, then look at the layer act like very deep statistics on layer activations and stuff. And then just you know, the usual stress tests, evals, unit tests, like all that other stuff. Um I raise this because like for for folks that are running at scale, the there's a a very great thriving LLM inference ecosystem. There's amazing work coming out of that. But if you want to operate at scale, you have to you know, have a foundation of basically very good CICD behind it because you can't just take the releases and put them into production. There's a lot that can go wrong and even if you have standard CICD, you have to in extend it with this more gen AI specific version of it. Um Okay, so that kind of brings us to the end of the talk 2 minutes before uh it's supposed to end. Um what I want you to take away from this is that on on the Databricks model serving team, we're focused on delivering this high-quality, high performance experience and we're also work doing a lot of deep work to continuously improve that. Um We're very excited about you know what our users are building with these things. Foundation model endpoints they underpin a lot of our features. Um as a company we're sort of more on the agnostic end of like hey if you want to use Anthropic well yeah we'll help you use that. It's more about data intelligence right? It's less about we're not a GenAI company we're we're not we're not really pushing our own models. Um but in a lot of cases these open weights models are the right tool especially when they're fine tuned and so we have to make them sort of fly. Um If yeah like if you want to try out how well all the stuff I talked to you today about is working um you know go into Agent Bricks click optimize on an endpoint see see how well it does. Um run some AI functions in in SQL. Um play around with the foundation models in in the AI playground and excited to see uh what you all do with this. And yeah that's the end of my talk.