Lightning Talk: Pluggable PyTorch LLM Inference Architecture With VLL... Yahav Biran & Maen Suleiman

In this talk, Yahav from the Anapuna ML team discusses the challenges of running machine learning models on hardware that isn't GPU-based. He introduces the Torch Neuron, a tool that connects with PyTorch to facilitate the integration of diverse accelerators without the need for separate codebases. The speaker highlights the functionality of Torch Neuron, which allows for seamless switching between different devices while maintaining hardware-agnostic code. Additionally, he explains the two execution modes it supports: eager mode for training scenarios and torch.compile for optimized serving. The importance of enabling performance portability across various hardware vendors and the efficient management of model operations through scheduling and execution paths are emphasized. This presentation showcases a significant advancement in machine learning model deployment and efficiency.

Hello everyone. Sorry for the delay. My name is Yahav and I am from the Anapuna ML team. Uh so before we start, I want to have a a quick show of hands. Uh who uh who is ever tried running on hardware that was not GPU? Okay. So K keep your heads up hands up if it was not if it was not the hardware itself and if it

was like the failure was uh during the the software. Okay. I have a friend and Okay. And how many people just move back to GPUs? Okay. Okay. So so in Amazon what we did we were trying to solve that problem. we heard customers uh having those types of issues and so uh and we did one of the areas that we're going to focus today is the VLM

part that includes text as well as uh VLM as a a content understanding. We are doing uh more work on covering other other modalities. As as PyTorchbased LLM serving matures, we're seeing a more fundamental h shift. Uh the challenge isn't just getting the model to run. It's integrate it's integrating diverse sets of accelerators efficiently without uh separate incompatible uh code bases. Um for those who uh use

the our accelerators and other accelerators then GPU known uh knows the the the the mo the the methods of forking VLM do everything from the beginning and that of friction. So so what we've built and this is available in beta today and we built the torch neuron. It hooks into the pietorch through the private use one uh device backend mechanism. This is the same extension other nonvidia

uh accelerators use like XPU and MPS. And this is the official uh PyTorch way to register new devices. What it means h practically is that trinium just show up as a device neuron sitting right alongside to uh your uh other favorite hardware options. Once you register the device, we register the device in the torch neuron. uh the device uh uh the device the PyTorch standard dispatcher just

and uh use that uh back end. And the cool thing here is that if you write a code that is agnostic to hardware and there is no uh available uh h hardware as the back end as you specify the back end it will fail over to uh to uh supposed to this is how the the how it will look like basically like I said from tuga to

uh to cuda to neuron and If you use the torch. Accelerators, you don't even need to do that because the device uh will uh discovered automatically. Okay. So, torch neuron supports uh two execution mode. One is eager and one is torch.compile. The eager mode is the is the options that researchers are using and this is more valid for training h scenarios post- training scenarios when you want

to uh to run specific sections and torch.compile Compile is for the uh for serving where you want the compiler and uh Nikki what I'm going to talk about is to basically to see to h to group more operation fuse them so they will run more optimally. Thank you. Great. So, so Nikki is the is the the kernel have I'm going to start with the left side. This

is the most trivial one. You just specify the you just specify the call the code that that's it. The compiler will handle everything. Compiler might use tensor tensor and other operations, right? But you don't have a lot of control into of the operation. Moving to the right side, you can see um you can see that we can explicitly set when to load the data and do the

add and store it back to the HBM and the the HBM is the the device uh memory and the compiler will look will like tensor operation but but the idea is that the compiler decides what will be uh running and how and then on the very uh right side uh We offer the ISA option where you can actually um control the you can explicitly say which the

calling which DMA to call and the tensor instruction. Here I'm defining how it runs on the hardware because uh because I know the the dimensions of the tensors right so this is the the level of uh of uh flexibility that these are the tradeoff that you have trading vers trading versus simplicity and control. I want to mention that this is different than the GPU kernels like Triton.

This is a single thread approach compared to a multi-thread approach with CUDA. We can talk more have more uh more h questions. Okay. So I'm going to skip out. This is uh how you scope how you scope the operation. H basically you can have a method and you can scope the entire model. Right? Okay. think part here of of what we what we've done. I'm talking about

the the the inference aspect but this is also including the the the trinium. We just wanted to reduce a little bit uh the complexity of this uh of this uh h graph. But essentially what we did we implemented right so we intercepted the torch dynamo the torch neuron x will uh implement is implementing the torch dynamo and basically torch dynamo will create the the fx graph and

from that point it will go through it will go through this uh this process out fx graph if it's in training it's going to be also backward and here we are going to have the mlir which will create an ML dialect which will create our HLO the HLO and eventually reduce to neuron IR and then uh created the the neuron exe basically runs the operations that I

showed before. Okay. So now we have that after we have that that that background how how PyTorch native will will will use the VLM plug-in basically the VLM uses the entry point and it it detect the the the platform automatically using the mechanism that I mentioned before and um and that's how it will look like right so on the on the very top you have your uh

VLM M that's the um the upstream VLM right you have the API server the model register etc. And we implemented the VLM uh neuron plugin that basically implement each part of that thing. Our goal is to make that plug-in as slim but there are cases where we need to uh we need to to implement that one minute. Thank you. So the scheduler for example we use uh

we use buckets right so let's say that you run a a model with a th you compile it with,024 as a sequence length and there is a request that is that we don't want to use the entire section and pad it because it's going to be wasteful and for that we're going to intercept that to extend that scheduleuler and we're going to slice it for so this

is just an example and then from there uh so once that the mode another thing is that we are imple implementing the uh the runner interface the ray h executor. So basically all this all this uh h um the path that I showed the execution path happen natively and these are the other uh uh the other uh components. So let's just summarize pluggable inference isn't just about

portability it's about performance portability. Second is the architecture provides clean abstraction boundaries right uh that keeps the code maintainable and then the third this approach enables ecosystem growth not just uh h for AWS but specifically to AWS h for more hardware vendor to join I was on time okay I'm going to stay for questions after that if you uh if you would like we have uh another

session where I'm going to talk to you about uh distributed AI. Just thank you.