Lightning Talk: Building a PyTorch‑native VLLM Plugin for IBM Spyre - Thomas Parnell & Thomas Ortner

This talk delves into the integration of the Spire device with the VLLM plugin, which enhances performance and efficiency on the IBM C platform, a crucial technology that handles a substantial volume of global transactions. The speaker discusses the evolution of the VLLM Spire next plugin, highlighting its current capabilities, which are maintained by a vibrant community of contributors. They explain the unique design choices made to incorporate Spire, including the use of a third-party model repository and specific adaptations to leverage the latest advancements in PyTorch, such as Torch Compile. A demonstration showcases how the plugin utilizes Spire to efficiently execute model layers while addressing the complexities that arise from handling different device decompositions, thus paving the way for more effective model training and deployment. Throughout the presentation, the speaker emphasizes the importance of collaboration within the community and invites contributions to the ongoing development of this open-source project.

Welcome everybody. I have the great pleasure of taking the talk before me a step further and actually dive into how one can use the Spire device in a VLM plugin. And as we heard before, so Spire in principle is a high-performance and energy-efficient add-on to IBM C platform. And actually for those who may not platform is is actually it's running 70% of the world's transactions by volume.

So it's a quite powerful device and quite well utilized. So although this this talk is not about Spire itself, it's taking Spire and using it in VLM as you will see. So this is actually the second generation if you wish plugin. We have been developing this plugin already. We have open-sourced it actually beginning of last year. It's a community project, a community maintains project and it has

actually about 550 uh commits and and about 40 contributors to this project. And actually it's project and it supports a lot of features that you would expect from VLM. However, there's a little bit of a catch here. So this plugin is somewhat weird as you might see in a sense that we don't use much of VLM VLM upstream as we would like to use. In particular, we

have our own Spire and we have the model code that doesn't really come from VLM upstream, but there is a third-party repository called the foundation model stack where we pull the model code from. And there's also a lot of other details that the plugin overwrites as opposed to the upstream This was a design choice that we did back then when we started to develop this model because

first of all, VLM back then didn't have really good support for torch and in particular not torch compile. And also Spire was only loosely integrated into the PyTorch ecosystem. As you might have seen in the last talk, the second change. So, now Spire is very well using this private use one device, as well as the first point is has also changed quite dramatically quite dramatically because VLLM

is now actually a power user of VLLM. And actually at the position to request features from PyTorch. good position. And this is also what you can see in the snippet here, which was a slide from Richard from last year's PyTorch conference, which just reiterates this point. As we also saw in the talk before, so I'm not going to repeat everything here, but TorchSpire is essentially a device

listed by this private use one device, and you can use it like so. Basically, you just say that your tensor should be on this device Spire, and the tensor will be placed on Spire. So, now how will we use this, and what can can we do better in order to have this VLLM Spire next developing. First one is that we can now reuse much more from upstream

VLLM that we could use before. So, for example, the modeling codes, the runner, and the worker, all of those parts can largely be overtaken from upstream VLLM. Furthermore, we can utilize quite a lot of the power from Torch Compile itself. Spire is natively in the PyTorch ecosystem, we can use all the features that VLLM uses from Torch Compile, basically. This includes also testing. So, we the our

aim here with this next plugin is to reuse as much as possible from Spire upstream from VLLM upstream, excuse me. We'll show you what our current development stage is, and what we followed. In a sense, we followed an iterative approach, which means that we started from the CPU backends in VLLM, and then iteratively wrapped or integrated every single layer to be supported on Spire. So, what you

see here on the left is a typical throw model. So, the user starts its input processing on the CPU device. So, all the blue boxes here are CPU. All the red boxes are on Spire. So, it processing the input in CPU, then it hits the model with the CPU device, and the individual layers that you see here in red are actually executed on the Spire device. So,

what happens? The tensor comes to the embedding, then gets converted to the Spire device, the operation happens, and then the tensor goes back to CPU. However, this has the This current approach has the downsides that there is the need for having an entry device, in this case the CPU, as well as an out-of-tree device, in this case the on the same stack, {quotes} and {quotes}, which might

come with its own challenges. As you heard before, it might happen that the out-of-tree device has a specific set of decompositions and specific set of lowerings, and a specific set of fusion passes, and there may be many more specific details of the out-of-tree device. our PyTorch is a bit lacking behind in those regards, and this goes back actually to the question that we had One issue that

we found is that the entry device, for example CPU, has a different set and a conflicting set of decompositions than the Spire device would support. And therefore, one solution that we is to have a thread-safe way of switching between the different decomposition tables, so switching between a decomposition table that's relevant for CPU and one that's relevant for Spire. And actually, the nice thing because of this tight

integration of Spire to PyTorch is that we can directly port this to upstream PyTorch, and this is what we did, and it's being there in 2.12. Just didn't make it for 2.11. Now, I just briefly wanted to dive a bit into the layer wrappings of how we do this. So, the VLM model is typically composed of two types of layers, but the majority of it. One is

what's called a custom op what's called a pluggable layer. And both of those mechanisms allow quite an easy way of registering out of three implementations. And what this means is, as you can see here at the bottom, you take an operation like the RMS norm, and all you do is you say register OOT, and then you specific implementation for this RMS norm operation, and it just works

out of the box. In addition to this, there's also this new VLM IR on the horizon, which even further simplifies this registration. This IR separates the semantics from the actual implementation, and then also the dis happens much more naturally in the sense that if if you use this operation, and the inputs are such that are not by Spire, you will fall back to the CPU variant. If

they are supported by Spire, you will be using this new out of three Spire implementation. Just to reiterate our goal here for this VLM Spire next is to stick as closely as possible to VLM upstream in order to consume all the features that the community adds to it. On the right-hand side, a very important piece is also the attention backends. So, you might know that VLM has

a vast variety of attention backends that's currently supported, and we are the Spire, we also have our own attention backends implemented. Again, it's very easy to do in VLM. You just create a new class, inherit from the right base classes, and then you're ready to go. What I'll show you in a second is a demonstration of using this VLM Spire next plugin in a demo. And what

you'll see is we query one of our IBM's model Granite 3.3 8 billion with a question what PyTorch is. And if this plays Oops. One second. Now it plays. So what you will see is that now VLM Spire start that we designed and then on the right you will see the activity trace of Spire. So you see now that some some the device is now executing something.

It's busy. And eventually it will come back with the answer to our prompt which is PyTorch is a open-source machine learning library which is in this case. This just goes to demonstrate that we have a first working prototype of this Spire next plugin. To summarize, we wanted to share our experiences and some community messages with all of you because we feel that's important to also give back

to the community the feedback that what we've learned. The first and good message VLM has a lot of entry points for customizations and they just work pretty well. So this goes for exchanging custom operations, exchanging pluggable layers, as well as attention back-end. So this this just works in a sense. Where some work is needed is on the testing infrastructure actually. So what we upstream tests are actually

hardcoded and bound to either CPU devices or CUDA devices. So some work might be needed here in order to relax this constraint a bit and support better out of three Furthermore, as I mentioned before there's this VLM IR on the horizon and we believe it's it's a great thing to do and we should push it forward because it simplifies how OOT devices can have their own custom

implementations. And finally, recently there was also an announcement about the model runner V2. Although this is yet in flux, so there is nothing cast in stone yet, it would still be good to have an eye on out of that whatever happens to the model runner or any part of the upstream VLLM would also be would also allow support for out of tree devices because with all the

new devices coming coming up as you heard before from AWS or from several other companies, it would be good to have VLLM support this. I'm at the end. Like I mentioned before, this is an all open source project and never so we have both both repositories published so the VLLM Spire next as well as the torch Spire and we welcome contributions and feel free to check out

our GitHub page. Thanks.