Abstract:

This presentation details the strategic integration of a native JAX backend into vLLM, a high-performance LLM inference engine traditionally based on PyTorch. The primary goal is to provide optimized performance on Google's TPUs, where JAX is a predominant framework, while maintaining seamless support for the vast PyTorch ecosystem. The core of this initiative is a new interoperability framework, torch.ax, which enables graph-level integration between PyTorch and JAX. This eliminates performance-degrading graph breaks and allows for a unified codebase. The speakers explain how torch.ax works by re-implementing PyTorch operations in JAX, facilitating a gradual optimization path where models can run immediately ("Day 0" support) and then be progressively enhanced with native JAX components for maximum performance. The presentation also introduces LMD, a Kubernetes-native inference stack that uses vLLM as its engine to orchestrate complex serving patterns like disaggregated serving and smart routing for modern LLM applications.

Unifying PyTorch and JAX in vLLM for High-Performance TPU Inference

• 00:00:05 - Announcing JAX Support in vLLM: The talk introduces a major new feature for the vLLM inference engine: a native JAX backend. The primary motivation is to unlock the best possible performance on TPUs, where JAX is a common framework. This integration is made possible by a new tool called torch.ax.
• 00:02:23 - The Goal: A Unified Backend for All Models: The team aims to support four distinct user paths with a single, unified code path: (1) existing PyTorch models, (2) custom PyTorch models, (3) new, highly-optimized JAX models, and (4) PyTorch models running on the JAX backend for immediate "Day 0" support.
• 00:08:22 - Integrating JAX into vLLM's Architecture: The proposed design reuses vLLM's core scheduler but introduces a new JAX TPU Worker. This worker will leverage JAX's SPMD (Single Program, Multiple Data) capabilities and include a collection of highly optimized model primitives, such as paged attention written in JAX.
• 00:16:45 - The Magic of torch.ax: Graph-Level Interoperability: The key technology, torch.ax, provides "graph-level" interoperability between PyTorch and JAX. Unlike simpler methods that cause performance-killing "graph breaks," torch.ax allows JAX and PyTorch operations to exist within the same computational graph, ensuring no performance is lost.
• 00:21:39 - How torch.ax Works: torch.ax works by creating a PyTorch tensor subclass that internally holds a JAX array. It uses a torch.dispatch hook to intercept PyTorch operations as they are called and replaces them with their equivalent JAX implementations. This process allows a PyTorch model to be traced into a clean, efficient JAX graph.
• 00:30:07 - A Path for Gradual Optimization: torch.ax enables a flexible workflow. A new PyTorch model can be run immediately on the JAX backend. Later, developers can replace performance-critical parts of the model (e.g., the attention mechanism) with native JAX code while leaving the rest in PyTorch, gradually improving performance without a complete rewrite.
• 00:35:11 - Introducing LMD for Advanced Orchestration: The talk introduces LMD (LLM Native Development), a Kubernetes-native inference stack. It is designed to handle modern, complex LLM serving patterns that stress systems in different ways (e.g., compute-bound prefill vs. memory-bound decode).
• 00:37:04 - LMD's Key Features: LMD acts as an orchestrator on top of vLLM (its engine) to enable advanced features like disaggregated serving (using separate nodes for prefill and decode), smart routing, and hierarchical KV caching. It is designed to be an open-source, hardware-agnostic solution for production-grade LLM serving.
• 00:43:22 - The Key Advantage Over PyTorch XLA: Unlike a traditional PyTorch XLA approach which would require maintaining two separate, parallel code paths (one for PyTorch, one for JAX), torch.ax allows for a single, unified backend. It translates PyTorch code into JAX primitives, effectively consolidating the development and maintenance effort.