Optimizing CPU LLM Inference in PyTorch: Lessons From VLLM - Crefeda Rodrigues & Fadi Arafeh

This talk focuses on optimizing large model inference on CPU within the PyTorch ecosystem, specifically targeting VLM on Arm architecture. The speakers, Krupafila and Fadi from Arm, discuss their journey from facing various challenges such as build issues and performance bottlenecks to achieving significant throughput improvements. They describe the software stack for VLM on Arm, highlighting key libraries like oneDNN, Arm Compute Library, and OpenBLAS. They detail several optimization strategies including the replacement of the default allocator with MiMalloc, enhancements in memory management, and custom kernel optimizations for attention and linear layers. Through collaboration and targeted improvements, they managed to achieve a sixfold increase in performance, showcasing the importance of deep integration and system-level optimizations.

Hello everyone. Welcome for to welcome for our talk. My name is Krupafila. And hello, I'm Fadi. We are from Arm in a group based in Manchester. We primarily look at optimizing ML inference on Arm server scale CPUs. So today's talk is going to be about optimizing LM inference on CPU in the PyTorch ecosystem, mainly looking at VLM. We will look at VLM CPU path and mainly talk

about our journey from making where we started working with VLM where things just didn't work to making it perform well on Arm. We focus on key optimizations at the system and kernel level and we will talk about the lessons we learned along the way. So before we get to the optimization story, let's have a look at what the VLM software stack looks like on Arm. At the

high level, you have your models, so your Llamas, your Gemma, that you want to run through a serving engine like VLM. And VLM is built on top of PyTorch with which has its different acceleration backends. Here are a couple of backends that we work on. So oneDNN is an open source library that is part of the UXL Foundation. It contains performant routines for ML workloads. We optimize

our Arm assembly routines from compute library and integrated them into oneDNN. Another library that we work with is Arm Compute Library which we have used primarily to optimize our quantized path, looking into int4 quantization. And finally, OpenBLAS which is the BLAS backend of PyTorch. So our team works across the stack right from the serving engine and down to hand optimized kernels. And you will see this in

the entire talk, the journey that we took. And at the low level, you have your Arm silicon where your routines leverage Arm ISA features like Neon and SVE to accelerate LLM inference on CPU. So let's just look at some of the practical deployments of CPU in the inference space. So the CPU can act as a AI head node, mainly as an orchestration layer, coordinating AI agents as

well as coordinating work for the accelerators. The CPU is also a general purpose computer, so it can be used for general compute as well as AI specific inferences. Existing server deployments already have a lot of CPU infrastructure, so why not leverage them for low cost CPU inference? And finally, during off-peak batching when your accelerator is busy running big LLMs, we can use the CPU to run low

cost low batch inferencing. So as I mentioned, the state of VLM when we started working with it on Arm CPU was in a pretty bad shape. First of all, we wanted to just run the pip package and that was not available for non-CUDA builds. So then we decided, okay, let's just build it from source. But there we ran into a bunch of build issues as well as

mysterious crashes. Then once we got the build sorted, we look at, oh, can we actually run with VLM features? Well, they didn't work out of the box for us. Once we sorted that and we said, okay, let's run an LLM with the key runtime flags, we found out that performance was not good on CPU and that was due to low CPU utilization. Once we sorted all of

these problems, we then ran into accuracy issues. So yeah, a lot of issues from when we started working and we addressed all of them, so you can see it's all of these issues are closed. And we had a long journey just to make VLM work out of the box for CPU and just run. Thank you. And okay, now that the issues are fixed and the stack runs,

let's look at performance. And for this, we introduce a throughput benchmark with Llama 3.1 8B running on a single socket and 96 core Arm Neoverse V2. And when we first ran this benchmark right after we got started with this, we got a throughput that was way too low for what we expected and that was a bit surprising to us because 80% of the model runtime is basically

spent in layers that are dispatched to very highly optimized GEMM implementations. And in this case, the GEMM implementations that were being utilized, if you benchmark them stand alone, they basically reach 80 to 90% of the system maximum flops. So they're good. The GEMMs were not necessarily the problem here. And that basically hinted to us that there's a lot of low-hanging fruit to be gained at the system

level. the baseline, this becomes the baseline, the point in time when we started. This plot is going to stay alive with us throughout this presentation and we're going to show you how each optimization moved the needle and by how much relative to the baseline throughput wise. Now given that LLMs create intense multi-threaded memory pressure on CPUs, it made sense for us to start by looking at the

memory allocator. And in this case, we looked at the PyTorch allocator and we noticed that it's using GLIBC malloc, right? Which was sub-optimal because malloc does not provide efficient reuse of large tensor allocations. So these This basically shows us loads of page faults and generally slow and noisy benchmarks, What we need here for LLM sort of workloads is a caching allocator which provides efficient reuse of large

tensor buffers and one that scales well under multi-threaded workloads. And luckily, MiMalloc addresses both of those needs for us, right? And basically, once we enabled MiMalloc as the default PyTorch allocator on Arm, we see a 133% throughput uplift. And what this teaches us is that no matter how performant your GEMM is, if your allocator is sub-optimal, you're leaving a lot of performance on the table. Now the

next bottleneck showed up when we tried to scale out high number of threads, right? The benchmarks that we've been showing are with 96 cores, which is a lot of cores, right? And what we What we were seeing here is in this case, more cores was basically hitting or regressing performance. And to understand what was going on, we had to profile. And basically, once you profile a layer

like paged attention, you see that 74% of your time is spent on this compute dynamic next. It's basically a routine in the libgomp runtime we ship with PyTorch, right? And the problem here is that this is not specific to paged attention. It applies to other ops across the stack as well. And given that 74% is not an acceptable time to spend at an op called gomp iter

dynamic next, we had to understand what this op does. So we look at the source code, right? And the op is not really doing much. It's just basically distributing loop iterations across threads at runtime. Each thread basically comes, it takes a it its next chunk of work by incrementing a shared pointer atomically, right? And this is key. So the only important thing that's happening here is this

atomic fetch add. And if you look down to the assembly, you see something that looks similar to this. And this is basically a read-modify-write loop with retries under contention. not something you'd want to have on your hot path, right? This is a nasty thing to have on the hot path because basically what happens with a high number of threads is that you're going to have a lot

of threads colliding on the atomic update. And once this happens, the loop will retry and basically you're losing performance. And what's even more interesting is that we benchmarked this on an Arm Neoverse Arm core. And Arm Neoverse cores come with the Arm large system extension which basically has an atomic instruction in hardware that is 100 times more efficient than this loop and can basically replace it. And

the problem here was that the runtime that we were shipping with PyTorch was just not built to utilize instruction or extension on capable CPUs. And once you fix that, you get 9% higher throughput, right? bottlenecks can hide in unexpected places like a very simple atomic add op in the and that sometimes you can only really see the missed hardware opportunity once you dive all the way down

to the assembly as we just did. The next thing is integration overhead. And to help you understand what this means, I'm going to show you the BFloat16 linear layer code path on arm. So, you start with the LLM. You go down through PyTorch. PyTorch has all sorts of lap wrappers. You go through ATen, is one of them. Then, you go to oneDNN, all the way down through

arm compute library. This way the micro gem kernels lives. And all the way down to the silicon. And that's a very deep stack. Right? And with stacks as deep as this, we end up paying two sorts of what we call taxes, right? The first tax is that things take a very long time to propagate across the stack. Something happening in compute library takes time to to propagate

all the way to the LLM. But, the next thing here is that libraries are opinionated, right? And with a deep stack, a lot of opinionated libraries, um you're basically going to have a lot of API friction, which hurts An example of that that's relevant here is PyTorch. For example, PyTorch likes to store weights in plain format. And once you have when you have a linear layer, every

time before the linear layer, it would basically repack the weights into a gem friendly format right before doing the matmul. This is problematic because packing is not free, Now, this is solved a little bit with Torch Compile and graph freezing, but at this time we still needed a fast eager path, which allows us to run models that at the time did not compile out of the box.

So, what do we do about that? Okay. So, the pragmatic the pragmatic thing here for us is to basically leverage an existing path that the community has added between the LLM and oneDNN. What we did here is basically extend the existing interface connecting the LLM and oneDNN to to support the arm BFloat16 matmuls. And once we did that, now we can basically pre-pack our weights once at

model load time, and later on when we want to run the actual linear layer, we just run the matmul going directly through oneDNN. This eliminated the redundant repacks, and it also eliminated some of the overhead that we were paying due to the deep stack. And once you do that, you get 16% higher throughput. And the lesson from this is that sometimes the bottleneck is neither the runtime

or the kernel. It basically lies the abstraction boundaries between the layers. Overall, up to this point with these system level improvements, we almost three xed our throughput. And what's important to note here is that this three x um throughput gain was not gated by a single heroic kernel, right? We got here by optimizing how the whole system interacted in general. And with that, I'm going to leave

you with Grafida, who will tell you about one heroic kernel built around this idea. So, yeah. Let's now look at the model itself, focusing on some of the kernel optimizations. So, the attention is a key backbone in all LLM models. And the page attention uh on CPU was contributed by the ecosystem. This, when we started working with it, was primarily optimized for x86. It supported VLM features

like chunk prefill mode, as well as other attentions like sliding window attention that's used in GPT-OSS. However, this was not optimized on arm. So, let's take a look at some of the hotspots in the attention. So, the attention is primarily two gems and a softmax. So, we focused on the gem, which is the QKV gem. We accelerated them using custom micro kernels that implemented arm instructions optimized

instructions for floating point and BF16. So, these are the fused multiply accumulate instructions, FM LA and BFM MLA. Then, we pivoted to the softmax, uh which occupied 30% of the runtime. The key issue in softmax is the exponential function. We replaced the slow standard exponential with a vectorized implementation, uh vectorized third-degree polynomial approximation. And with these two key optimizations, we were four x faster on the page

attention Now, that translated to 12% improvement in the Llama benchmark that we've been seeing. So, once we sorted the attention, the next key bottleneck is the linear layers. And a key optimization strategy, especially for CPU, is to quantize them. We want to do quantization to reduce memory bandwidth, but also make the linear layers go fast. To start off with, we looked at int8 quantization, mainly to help

reduce memory bandwidth, as well as leverage the arm i8mm instructions, which is two x faster than running with BF16. We added those kernels into oneDNN. And here we really benefited from the ecosystem. Uh we had Fujitsu that contributed to the 256-bit kernels in oneDNN. We had Intel that did the plumbing for int8 kernels in general in VLM. And finally, we contributed to the 128-bit kernel inside oneDNN.

And as you can see, we also were able to reach within 99% of BF16 accuracy on MLPerf Llama benchmark, as well as a really good bump in throughput performance on our Llama benchmark with around 78% improvement. To take the quantization story even further, we look at info, which halves the memory use of int8, but also helps us run bigger models on CPU. Uh to leverage that, in

fact, uh on the info path, we do not have any native uh info instructions. So, in order to uh address that, we looked at unpacking info weights into int8 and leveraging the i8mm instructions that I mentioned before. All our info kernels are in the arm Clady path, which we integrated uh and made use of. And also with the info quantized model, we were able to reach within

99% BF16 accuracy on the same MLPerf benchmark. That also translated to another 6% performance boost on our throughput So, we've come a long way from where things didn't work out of box. We addressed those. We then looked at system level overheads, fixed them, and then we looked at kernel level optimizations. And with all of these key strategies, we were able to get six x performance boost from

when we started. Thank you. And now, what happens if you want to run a model that's much larger, and you just can't get it to serve within your latency budget, right? And usually the answer is simple, right? You just run on a bigger machine with a much higher number of cores. But, it's not as simple as that, right? Because usually, with high number of cores, you start

seeing multiple sockets and multiple NUMA nodes. And basically, cross-socket or remote NUMA access is very expensive. So, in this case, this is the system that will be working with. Um for this part of the presentation, we have two sockets, um 96 cores per socket. And this shows a heat map for core-to-core latency. And as you can see from the heat map, the latency shoots up as soon

as you start going outside your socket. practically means when you serve, we're going to take a new latency-oriented where we have a much bigger model, right? Queen 2.5 M32B parameter models. we're going to report the decrease in time per output token, i.e. latency, compared to a baseline. And here, our baseline is basically the single socket um time per output token. So, now that you start spanning requests

across two sockets, you basically get worse latency, and that's coming from the cross-socket remote NUMA um access Sweet. So, we doubled the compute, we made it twice more expensive, and we ended up with a worse latency. What do we do about that? Well, the solution is to basically keep execution um local whenever possible. If you're interested in throughput, you can basically have model replicas and run one

model replica per socket or NUMA node. But, in this case, we're interested in latency. And the best strategy here is to basically partition the model across sockets. And one way to do that is through tensor parallelism. And in this case, each um NUMA node, socket, or rank basically operates on its own um weight shard, which we know how to do because we already accelerated the single NUMA

case um throughout this presentation. But, at some point, these ranks would need to communicate to share their results. And one example of that is the all reduce um up here that we're showing in the tensor parallel um in Llama style MLP. So, the interesting bit that we haven't covered yet is how the communication happens between the sockets. So, let's see how it used to happen on the

arm paths. So, on the arm path we used to go through torch distributed. The communication used to happen through glue, all reduce um was um a ring um basically all reduce implementation through generic message passing um that happens over TCP sockets. Right? And this is problematic. Because here we have two ranks on the same machine, yet we're paying the overhead as if the ranks were on different

machines by communicating through TCP. Basically, that shows up as 50% of the inference time is basically just spent in weights, and this is not good. And overall, now once we start sharding our weights properly with tensor parallel, if you run across two sockets and use this um communication path, your latency is still worse off than the single socket baseline. The main problem here is, as we said

before, we don't want to communicate through TCP because that does not match the system topology. The system topology is two ranks on the same machine. Now, luckily, VLLM comes with its own custom shared memory communicator where um x86 oriented implementation was contributed by the community. We enabled that for arm, and the communication here happens through mapped shared memory. And with this, we see minimal cross numa cross

numa traffic, very lightweight synchronization, and most importantly, the communication path matches the system topology because we're now communicating through shared memory. And overall, once you do this and you're on with tensor parallel equal to, you see latency going down by 43%, which is basically 88% efficiency compared to a system where communication is free, i.e. the perfect um case. And this is a long way from where we

started, right? We started with, you know, correctness issues, crashes, poor And throughout this um presentation, we shared our journey on how we moved on from that to being able to support out of the box low precision inference across um sockets. Right? And this has been achieved through a lot of ecosystem collaboration in open source, um a lot of kernel level, system level, and integration level happening all

over the place in arm compute library, oneDNN, arm compute library AI, openBLAS, PyTorch, VLLM. And it will only get better. Thank you. >> [applause] >> I'm not sure if we have time for questions. Can I have one? Please go. This problem was that the instruction for atomic adds. How how did you find it? How did you find the problem with atomic instruction? Yeah, so basically it's exactly

the trace that I showed on the slides. We basically profiled the op. We found that op as the bottleneck. And then we went down to the assembly of that op. We saw that loop, and we could tell that what this loop is doing, and we know that we have an arm instruction that, you know, like should be set here in instead of the loop, which is a

lot more efficient. Thank you. Thanks.