Unlocking Developer Productivity across CPU and GPU with MAX

Chris Lattner, Modular18:33 · Jul 2024 · 6,838 views
Thumbnail for Unlocking Developer Productivity across CPU and GPU with MAX Watch on YouTube
TL;DR
  1. 1

    MAX is designed to reduce the fragmentation involved in deploying AI models across different runtimes, hardware targets, and programming languages.

  2. 2

    MAX lets Python and PyTorch developers use a unified stack for CPU and GPU inference, while advanced developers can write custom GPU kernels.

  3. 3

    Mojo extends Python with systems-level performance so teams can write performance-sensitive code without switching to C++, Rust, or CUDA C++.

Summary

Chris Lattner argues that AI engineers are slowed by a fragmented production stack. Models, inference runtimes, preprocessing, data loading, orchestration, and hardware support often require different frameworks and languages. MAX is Modular's attempt to unify this work around Python and PyTorch, with support for CPU and GPU inference in one stack. Lattner describes MAX's existing PyTorch support, native APIs for advanced LLM work, quantization on CPUs, GPU kernel programming, and automatic fusion. He presents Mojo as the language underneath MAX, combining Python's ecosystem with compiler and systems features. Mojo is intended for code that needs predictable performance, such as tokenization, while remaining familiar to Python developers. The talk is also an argument about control. Running models locally can help teams control data, integrate security requirements, customize models, reduce inference costs, and use different hardware. Lattner is direct that GPU support is difficult and says its official launch was planned for September 2024.

Key ideas
01:18

AI production work is fragmented across too many frameworks and languages

Lattner says teams increasingly need more than an inference endpoint. They may need control of their data, internal security integration, model customization, lower inference costs, or support for less common hardware. The stack has also expanded from PyTorch and TensorFlow to ONNX, TensorRT, model-specific runtimes, and many other production technologies. These systems often do not fit together cleanly. Lattner's objection is that this fragmentation slows the move from research and demos into products. AI engineers must absorb new models and optimizations every week while also dealing with hardware and the cost of scaling inference.

05:24

MAX is built around existing Python and PyTorch workflows

Modular's stated approach is to bring several technologies into one stack while allowing developers to keep using familiar tools. MAX is presented as both an AI framework and a separate managed-services business, with the talk focused on the framework. Lattner describes it as a way to deploy PyTorch and work with generative AI, especially inference. Developers can bring existing PyTorch models and choose among paths involving ONNX, TorchScript, and torch.compile. Developers who need more control can use native APIs for work such as KV caches and paged attention. The goal is to give Python developers more capability without asking them to abandon their current ecosystem.

08:48

MAX replaces separate CPU and GPU implementation layers with one stack

Lattner says CPU and GPU work often feels different because a simple device switch can hide two separate implementations. One path may rely on Intel MKL while another relies on CUDA and NVIDIA libraries. MAX attempts to replace that lower layer, including matrix multiplication, fused attention layers, and graph components, so the implementations can work together more predictably. Developers who do not want to write kernels can use an auto-fusing compiler. Advanced developers can write custom kernels directly against MAX and still access CUDA capabilities. Lattner presents this as a way to combine ease of use with lower-level control.

12:14

MAX measures GPU performance against vendor implementations

Lattner says performance comparisons should use the vendor's best implementations rather than an easy target. In the example he gives, MAX is compared with cuBLAS and cuBLASLt on matrix multiplication, a core neural-network operation. He says MAX meets and beats both in the data shown, with a claimed advantage of around 30 percent in that comparison. He also gives a CPU example involving INT4 and INT6 quantization, which he says is five times faster than llama.cpp on cloud CPUs. His broader point is that a unified programming model has to produce useful performance, since latency and inference cost are part of the reason teams choose specialized hardware.

13:29

Mojo was created to build the lower layers of MAX

Modular built the Mojo programming language because Lattner says the company needed to rebuild its AI stack from the bottom up. Mojo is intended to extend Python into systems and accelerator programming. It keeps Python's ecosystem, libraries, packaging, and familiar style, then adds compiler features needed for high performance. Lattner says developers can download the toolchain, use it with Visual Studio Code, and run it on Linux, Mac, and Windows. Developers who only want MAX can ignore Mojo. Those who need more control can use it instead of moving into C, C++, or Rust.

14:29

Mojo makes ordinary performance-sensitive code possible in a Python-like language

Lattner distinguishes Mojo from a slightly faster version of Python. He describes it as a system that targets hardware-level performance while retaining a Python-like feel. His example comes from Llama 3 tokenization and includes linked-list traversal, conditional statements, and loops. These are ordinary programming patterns that Python developers generally avoid when they need speed. Lattner says Mojo can reduce the cost of Python performance work and give researchers more ability to modify the full stack. For managers, he frames the benefit as a more coherent team structure because developers do not all need separate C++, Rust, and Python expertise.

17:05

MAX targets teams that need control beyond an API endpoint

Lattner closes by describing the situations where MAX may fit. He names control over data, integration with company security, customization, lower costs, and portability across hardware. MAX is available as a free download, while Modular also offers help with production deployment on Kubernetes and SageMaker. He says GPU support is difficult and was still being developed, with an official launch planned for September 2024 and early access available through the company's Discord. The ending is practical rather than sweeping: developers can try the framework, inspect open-source code on GitHub, and join the community.

"If you care about control over your data, you want to integrate into your security, you want customization, save money, portability across hardware, then you need to get onto something else."17:05
Who should watch
  • You are deploying models beyond hosted endpoints and need control over data, security, model customization, or hardware placement.
  • Your team uses Python and PyTorch but spends time moving performance-sensitive work into C++, Rust, or CUDA.
  • You are evaluating CPU inference, GPU programming, or a single programming model for an end-to-end AI pipeline.