Mojo + PyTorch: A Simpler, Faster Path To Custom Kernels
summary
A new workflow uses Mojo to write PyTorch custom kernels without C++, CUDA, or custom build systems and exposes them via a per-op custom op loader or a Max graph API that enables operator fusion via fused tensors. Examples include a grayscale kernel using a parallelizing for_each and a MatMul hand-tuned to match cuBLAS within a few percent, with benchmarks showing 1–15% faster than an equivalent Triton kernel on H100/B200/MI300. A team of ~10 maintains kernels across seven GPU architectures and three vendors and supports CPU, NVIDIA, AMD (Apple silicon in nightlies).
AI-written summary · about
chapters (9)
- 0:00Introduction
- 0:40The case for Mojo
- 2:50Performance vs. ergonomics
- 3:49The problem with MAX
- 7:06Mojo and MAX benefits
- 9:23Writing a Mojo kernel
- 12:24Performance benchmarks
- 14:49Interfacing with PyTorch
- 18:23Conclusion and Q&A
related talks
10mLifting PT to New Heights with MAX and Mojo
24mInside MoE Optimization: A Profiler-Guided Tour of torch.compile and vLLM
19mAI Kernel Generation: What's working, what's not, what's next
06m⚡ Speedy but Sturdy PyTorch Optimizers
10mCombo Kernels: Horizontal Fusion Optimizat...
28mMojo & MAX - A modern AI-first SW Development ecosystem
25mLeveraging AI Ecosystem for Portable and Sustainable GPU Kernels in HPC
08mGPU Programming with Triton Kernels