← All AI Engineer talks

AI Engineer World's Fair 2024

Unlocking Developer Productivity across CPU and GPU with MAX

About this talk

Modular co-founder Chris Lattner argues that fragmented inference frameworks impede secure, customizable production AI. He introduces MAX and Mojo as a unified, Python-oriented stack intended to simplify model deployment and CPU/GPU programming, drawing on established compiler and machine-learning infrastructure while targeting higher developer productivity and hardware performance.

Chapters

  1. 0:00Why production AI needs deployment control
  2. 1:27Inference-framework fragmentation and engineering challenges
  3. 3:21Modular's unified AI stack and Python-first approach
  4. 9:17Making GPU programming as accessible as CPU programming
  5. 13:48How Mojo powers MAX and targets hardware-level performance

Talk transcript

  1. 0:00

    [upbeat music] All right. Good morning, everyone. I'm here to talk to you about Modular and accelerating the pace of AI.

  2. 0:19

    Uh, you know what gen AI is. I'm not gonna tell you all about this. Let me tell you one of the things I think is really cool about it, and very different than certain other technologies, is that it's super easy to deploy.

  3. 0:30

    There's lots of great endpoints out there. There's a lot of good implementations, a lot of ways to make it super easy to build a prototype and get, get going very quickly.

  4. 0:39

    But despite all the availability of all these different endpoints, sometimes you do have other needs. Sometimes you might wanna go, you know, and control your data instead of sending your data to somebody else.

  5. 0:51

    Sometimes you might wanna integrate it into your own security 'cause you're-- got your critical company data in your model and you don't wanna fine-tune it somewhere else. Sometimes you wanna customize the model.

  6. 1:00

    Like, there's research happening all the time, right? A lot of things in building proprietary m-models that work best for your use cases, uh, can make your applications even better.

  7. 1:09

    And of course, the inference endpoints are expensive, and so sometimes you wanna save money. Sometimes there's hardware out there that's really interesting, and you wanna explore out from the, the mainstream, and you wanna go do this.

  8. 1:21

    And if you care about any of these things, what you need to do is you need to go beyond the endpoint.

  9. 1:27

    And so how do you do that? Well, if you have-- many of you have explored this, I'm sure. The answer has shifted. It used to be that we had things like PyTorch and TensorFlow and Caffe and things like this, but as inference became more important, the world shifted.

  10. 1:41

    First we got ONNX, TensorRT, things like this. And today we have an explosion of these different frameworks, some of which are specific to one model. And that's cool if you care about that one model, but if you have many different things you wanna deploy and you wanna work with, it's very frustrating to have to switch between all,

  11. 1:57

    all these different technologies. And of course, it's not just the model. You all know there's this gigantic array of different technologies that get used to build real-world things in production.

  12. 2:07

    And of course, no-none of these are really actually designed for gen AI.

  13. 2:11

    So my, my, my concern about this, my, my objection to the status quo is that this fragmentation slows down getting the, the research and the in-in-innovations coming into gen AI into your products.

  14. 2:22

    And I think we've seen so many of these demos. Last year was really the year of the gen AI demo. But still, we're struggling to get gen AI into products in an economical and, and, uh, good way.

  15. 2:33

    And so whose fault is it? Well, m- is it our fault? Like, many of you are AI engineers. If you don't, let's sympathize with the plight of the AI engineer because y'all, these folks that are building this have new models and optimizations coming out every week, right?

  16. 2:52

    Every product needs to be enhanced with gen AI. This is not, like, one thing. We're getting dumped on, and there's so much to do, and we can't even keep up.

  17. 2:58

    There's no time to deal with new hardware and all the other exciting new features. And of course, once you get something that actually works, the costs end up making it very difficult to scale these things because getting, getting things into production means suddenly you're paying on a per unit basis.

  18. 3:13

    So it's not the AI engineer's fault. We should look at the concerns and look at the challenges faced here. And so I think that we need a new, new approach, right?

  19. 3:21

    We've learned so much. Let's look at what we need to do. How do we solve and improve the world here? This is what Modular is about. And so I'll give you a quick intro of what we're doing and kind of our approach on this.

  20. 3:33

    First of all, who are we? Modular is a fairly young company. We've been around for a couple of years. Um, we, we have brought together some of the world's experts that built all of these things.

  21. 3:43

    And so we've built TensorFlow and PyTorch. We built compilers like LLVM and MLIR and XLA and all, all of these different things. And so what, what I can say about that is that we learned a lot, and I apologize [laughs] because we know why it is so frustrating to use all these things.

  22. 4:01

    But, but really it's-- it was a, you know, the world looked very different five years ago. Gen AI didn't exist. It's, it's understandable. We, we tried really hard, but, but we have learned.

  23. 4:11

    And so what our goal is, is to make it so you can own your AI, you can own your data, you can control your product, you can deploy where you want to.

  24. 4:20

    You can do this and it-- make it much easier than the current systems work today. And so how? Well, what we're doing is really going back to the basics.

  25. 4:31

    We're gonna-- we're bringing together the best in class technologies into one stack, not one solution per model.

  26. 4:38

    Our goal is to lift Python developers, PyTorch users. This is where the entire industry is, and so we wanna work with existing people. We're not trying to like say, "Hey, ditch everything you know and try something new."

  27. 4:50

    We wanna gradually teach and give folks new tools so they can be superpowers, so they can have superpowers. And finally, uh, so I spent a lot of time at Apple.

  28. 4:59

    Like, I want things to just work. Like, you wanna build on top of infrastructure. You do not wanna have to be experts in the infrastructure. And this is the way all of this stuff should work, and unfortunately, it's just not the case today in AI.

  29. 5:10

    And so at Modular, we're building this technology called MAX. I'll explain super fast what this is. Um, MAX is two things. One is an AI framework, which I'll spend a bunch of time about.

  30. 5:22

    Uh, the AI framework is free, widely available. We'll talk about it today. The other is our managed services. This is how Modular makes money, very traditional. We're not gonna spend a lot of time talking about that today.

  31. 5:33

    And so if you dive into this AI framework, well, it's... We see it as two things. It's the best way to deploy PyTorch. It's also the best way to do gen AI.

  32. 5:43

    And both halves of this are really important. And MAX is currently very focused on inference. And so these are areas where, uh, PyTorch is challenging at times. This is where gen AI is driving us crazy with cost and complexity.

  33. 5:57

    And so really focusing on this problem is something that we are, uh, we're all about.

  34. 6:02

    The other thing, as I said before, is Python. So we natively speak Python. That is where the entire world is. We also have other options, including C++, which we'll talk about later.

  35. 6:12

    So how do we approach this? Well, as I said, we work with PyTorch out of the box. You can bring your models, your model works. We can talk to the v- wide array of PyTorchy things like ONNX and, uh, TorchScript and TorchCompile and, like, all this stuff, and so you can pick your path and, and that's all

  36. 6:27

    good. Uh, if you wanna go deeper, you can use native APIs. Native APIs are great if you w- if you speak the language of KV caches and page attention and things like this, and you care about pushing the state-of-the-art of LLM and other gen AI techniques.

  37. 6:40

    That's very cool. And also, um, MAX is very different in that it r- really rebuilds a ton of the stack, which I don't have time to talk about. But, um, we do not build on top of cuDNN and the, the NVIDIA libraries and on top of the Intel libraries.

  38. 6:56

    We replace all that with a single consistent stack, which is really different approach, and I'll talk about what that means later. And so what you get is you get a whole bunch of technology that you don't have to worry about.

  39. 7:07

    And so again, as a next-generation technology, you get a lot of fancy compiler technologies, runtimes, high-performance kernels, like, all this stuff, uh, in the box, and you don't have to worry about it, which is really the point.

  40. 7:20

    Now why would you use MAX? So it's, it's a AI framework. You have one, right? W- And so there are lots of different reasons why people might want to use an alternative thing.

  41. 7:30

    Um, for example, developer velocity, your team being more productive. That's actually incredibly important, particularly if you're pushing state-of-the-art. But it's also very hard to quantify. And so I'll do the same thing that, you know, kind of people generally do, is go and talk about the quantifiable thing, which is performance.

  42. 7:46

    And so I'll give you one example of this. Um, we just shipped a release that has our int four, int six K fancy quantization approach. Um, this is actually five X faster than Llama.cpp.

  43. 8:00

    And so if you're using Llama.cpp today on, in cloud CPUs, this is actually a pretty big deal, and five X can have a pretty big impact on, um, you know, the actual perceived latency of your product and performance and cost characteristics.

  44. 8:15

    And the way this is possible is, again, this combination of really crazy compiler and technology and other stuff underneath the covers. But the fact that you don't have to care about that is actually pretty nice.

  45. 8:27

    It's also pretty nice that this isn't just one model. This is... [chuckles] You know, we have this make it easy to do int four technology, and then we demonstrate it with a model that people are very familiar with.

  46. 8:38

    And so if you care about this kind of stuff, this is, this is actually pretty interesting, and it's a next-generation approach to a lot of the things that are very familiar, but it's also done in a generalizable way.

  47. 8:48

    Now, CPUs are cool, and so, I mean, so far we've been talking about CPUs, but GPUs are also cool. And [chuckles] what I would say and what I've seen is that the, uh, uh, CPUs and AI are, are kind of well understood, but GPUs are where most of the pain is.

  48. 9:03

    And so I'll talk just a little bit about our approach on this.

  49. 9:07

    And so first, before I tell you what we're doing, let me tell you our dream. And this is, this is not a small ambition. This is kind of a crazy dream.

  50. 9:17

    Imagine a world where you can program a GPU as easily as you can program a CPU in Python.

  51. 9:25

    Okay. Not C++, in Python. That, that, that, that is, that, that is a very different thing than the world is today.

  52. 9:34

    Imagine a world in which you can actually get better utilization from the GPUs you're already paying for. I don't know your workload, but you're probably somewhere between thirty percent, maybe fifty percent utilization, which means you're paying for, like, two to three times the amount of GPU that you should be, right?

  53. 9:49

    And that, that is understandable given the technology today, but that's not great for lots of obvious reasons. Imagine a world where you have the full power of CUDA. So you don't have to say there's a powerful thing and there's an easy-to-use thing.

  54. 10:05

    You can have one technology stack that scales.

  55. 10:08

    Well, this is something that is really hard. [chuckles] This is something where, you know, NVIDIA has a lot of very good software people, and they've been working on this for fifteen years.

  56. 10:16

    Um, but I don't know about you, I don't run fifteen-year software on my cell phone. Like, it doesn't [chuckles] run BlackBerry software either. And I think that it's time to really rethink this technology stack and push the world forward, and that's what we're trying to do.

  57. 10:29

    And so how does it work? Well, you know, it's just like PyTorch. You use one line of code and switch out CPU to GPU.

  58. 10:36

    Ha ha. We've all seen this, right? This, this doesn't say anything. I actually hate this kind of a, a demo, um, because the way this is usually implemented is by having a big fork at the top of two completely different technology stacks.

  59. 10:49

    One built on top of Intel MKL, one built on top of CUDA. And so as a consequence, nothing actually works the same except for the, the thing on the slide.

  60. 10:58

    And so what Modular's done here is we've gone down and said, "Let's replace that entire layer of technology. Let's replace the matrix multiplications. Let's replace the fused attention layers.

  61. 11:08

    Let's replace the graph thingies. Let's replace all this kind of stuff and make it work super easily, super predictably, and let's make it all stitch together." And yeah, it looks fine on a slide, but the slide is missing the point.

  62. 11:21

    So if you are an advanced developer, and so many of you don't wanna know about this, and that's cool. If you are an advanced developer, like I said, you get the full power of CUDA.

  63. 11:30

    And so if you want, you can go write custom kernels directly against MAX, and that's, that's great. And for advanced developers, which I'm not gonna dive too deeply into, it's way easier to use than things like the, uh, Triton language and things like this, and it has good developer tools and it has all the things that you'd

  64. 11:48

    expect from a world-class implementation of GPU programming technology. Um, for people who don't wanna write kernels, you also get a very fancy auto-fusing compiler and things like this. And so you get good performance for the normal cases without having to write the hand fuse kernels, which is, again, a major usability improvement Now, you know, it's cool, like,

  65. 12:08

    there's a lot of, lot of things out there that promise to be easy, but what about performance, right? A lot of the reason to use a GPU in the first place is about performance.

  66. 12:17

    And so one of the things I think is pretty cool, and one of the things that's very important to Modular, is that we're not comparing against those standards. We're comparing against the vendor's best.

  67. 12:26

    In this case, NVIDIA, they're experts in their architecture. And so if you go look at, again, there's a million ways to measure things, a micro benchmark. Go look at the core operation within a neural network, matrix multiplication.

  68. 12:41

    This is the most important thing for a wide variety of workloads, and again, one set of data. But we compare against CuBLAS, the hard-coded thing, and then also against Cutlass, the more programmable C++y thing.

  69. 12:55

    And so, MAX is, uh, meeting and beating both of these, you know, by just a little bit. I mean, it's... You know, it depends on your bar, and data's complicated.

  70. 13:04

    But you know, if you're winning by 30%, 30% is actually a pretty big deal given the amount of cost, the amount of complexity, the amount of effort that goes into these kinds of things.

  71. 13:14

    And so I've talked a lot about the what, but I haven't talked about the how. And so the how is actually a very important part of this, and I'll just give you a sample on this.

  72. 13:23

    So we are crazy enough that we [laughs] decided to go rebuild the world's first AI stack from the bottom up for gen AI. And as part of doing that, what we realized is we had to go even deeper, and so we built a programming language.

  73. 13:38

    We have a new programming language that's called Mojo. And so the thing about Mojo is, if you don't wanna know about Mojo, you don't have to use Mojo. [laughs] You can just use MAX, it's fine.

  74. 13:48

    But we had to build Mojo in order to build MAX. And I'll tell you just a couple of things about this. Our goal is that Mojo's the best way to extend Python,

  75. 13:58

    and that means that you can get out of C, C++ and Rust. And so what is it as a programming language? It's a full-- It's Pythonic, so it looks like Python.

  76. 14:06

    It feels like Python. Everything you know about Python comes over, and you can not have to retrain everything, which is a really big deal. You get a full tool chain.

  77. 14:14

    You can download it on your computer. You can use Visual Studio Code. It's open source, available on Linux, Mac, Windows. Two hundred thousand people, twenty thousand people in Discord.

  78. 14:23

    It's, it's really cool. Would love for you to go check it out if you're interested in this. But what is Mojo? [laughs]

  79. 14:31

    Like, what, what, what actually is it? Fine, there, there's a programming language thing going on. Well, what we decided is we decided that AI needs two things. It needs everything that's amazing about Python.

  80. 14:42

    This is, in my opinion, the developers. This is the ecosystem. This is the libraries. This is the community. This is even, sorry, the package managing and, like, all the things that people are used to using already.

  81. 14:56

    Those are the things that are great about Python. But what is not great about Python, unfortunately, is its implementation. And so what we've done is we've combined the things that are great about Python with some very fancy, highfalutin compilery stuff, MLIR, all, all this good stuff, that then allows us to build something really special.

  82. 15:15

    And so while it looks like Python, please do forget everything you know about Python, because this is, this is a different beast.

  83. 15:23

    And I'm not gonna give you a full hour-long presentation on Mojo, but I'll give you one example of why it's a different beast, and I'll pull it back to something many of you care about, which is performance.

  84. 15:32

    And all what I'll say is that Mojo's fast.

  85. 15:35

    How fast? Well, it depends. Like, this isn't a s- slightly faster Python. This is a working back from the speed of light of hardware kind of system. And so many people out there have found that it's a thou- a hundred times to a thousand times faster.

  86. 15:47

    In crazy cases, it can be even, even better than that. But the speed is not the point. The point is what it means. And so in Python, for example, you should never write a for loop. [laughs]

  87. 16:00

    Python is not designed for writing for loops if you care about performance, at least. In Mojo, you can go write, uh, code that does arbitrary things. This is an example pulled from our Llama Three, written in Mojo, that does tokenization using a standard algorithm.

  88. 16:15

    It's chasing linked lists, has if statements, for loops. Like, it's just normal code, and it's Python. I mean, it feels like Python, and that, that is really the point.

  89. 16:23

    And so for you, the benefit of Mojo is, first of all, you can ignore it if you don't wanna care about it. But if you do, you don't have to learn C, C++.

  90. 16:32

    You have lower cost by default versus Python because performance is cost. It means that as a researcher, if you use this, you can actually have full stack hackability. And if you're a manager, it means that you don't have to have people that know Rust on your team, and C++, and things like this.

  91. 16:48

    You can have a much more coherent engineering structure where you're able to scale into the problem no matter where it is. And so if you wanna see something super polarizing, go check the Modular blog, and we'll explain how it's actually faster than Rust, which many people consider to be the gold standard, even though it's, again, a fifteen-year-old

  92. 17:03

    language. So I have to wrap things up. They'll get mad at me if I go over. Um, the, the, the thing that I, I'm here to say is that many of you may want to go beyond the AP- the API, and they're fantastic.

  93. 17:15

    There's amazing technology out there. I'm very excited about them too. But if you care about control over your data, you want to integrate into your, your security, you want customization, you wanna save money, you want portability across hardware, then you need to get onto something else.

  94. 17:28

    And so if you're interested in these things, then MAX can be very interesting to you. MAX is free. You can download it today. It's, it's totally available. Go nuts.

  95. 17:37

    Uh, we didn't talk about production or deployment or things like this, but if you wanna do that, we can also help. We support production deployment on Kubernetes, SageMaker, and we can make it super easy for you.

  96. 17:47

    Our GPU support, like I said, is actually really hard. We're working really hard on this. We wanna do this right, and so it'll launch officially in September. If you join our Discord, you can get early access, and we'd be very happy to work with you ahead of that too.

  97. 18:01

    Uh, we're cranking out new stuff all the time. And so if you are interested in learning more, you can check out modular.com, find us on GitHub, a lot of this is open source, and join our Discord.

  98. 18:12

    Thank you, everyone. [upbeat music]