Today we are excited to announce the launch of XLA Serving Engine (XSE), a general-purpose model export and inference runtime for the OpenXLA ecosystem, open-sourced under the Apache 2.0 license.
We created XLA Serving Engine to enable a “Train in Python, Serve in C++” workflow. Most ML engineers develop models using Python (with frameworks like JAX and PyTorch); however, some production serving environments do not allow Python due to operational and deployment complexity. XLA Serving Engine offers a solution by providing an API for engineers to export their ML models to an “XLA SavedModel” containing the StableHLO representation of their model’s function, which can then be loaded and executed by the XLA Serving Runtime.
This approach allows users to serve models (typically non-LLMs) directly in a C++ runtime, which enables high throughput at any scale and optimizes for low latency by removing the overhead of the Python interpreter. In our studies, we found that for short requests, XSE delivers 2–5x the throughput of the best Python server we could build, and 8–36x that of a straightforward one, while using 30–70% less CPU time per prediction and a quarter of the memory.
Last year, we began deploying core components of XLA Serving Engine as part of Google’s internal load-bearing, production serving infrastructure. Today, XSE serves over 6 million queries per second and powers some of Google's highest-volume online inference pipelines, including YouTube Shorts, Google Ads, Search, and parts of Gemini Omni, with plans to onboard more customers internally.
Having successfully scaled XSE to support the high-throughput workloads of our internal teams, we are now ready to bring this core piece of infrastructure to the wider ecosystem of hardware partners and Cloud TPU users. We designed XLA Serving Engine to be self-contained and modular, so users can deploy it as a standalone inference engine or plug it into their more complex serving solutions.
Core production-ready features
- Python-Free Serving: XSE enables a “Train in Python, Serve in C++” workflow via the intermediary compiled XLA SavedModel serialization format. In our studies, we found that XSE delivers 2–5x the throughput of a pure Python server.
- Low-latency execution: The XLA Serving Runtime wraps PJRT directly and optimizes the model using the XLA compiler. This eliminates Python runtime overhead and the global interpreter lock. XSE is C++ native, which means one process is capable of using all eight cores.
- Orbax Checkpointing Integration: XLA SavedModel stores JAX model computations as StableHLO and uses Orbax checkpoint to store the weights.
- Flexible Scale: Supports both single-host and multi-host serving by leveraging PJRT.
- Performance Parity: XSE preserves the performance of the model’s forward function (including kernel optimizations). Hand off your model to XLA Serving Engine with confidence that performance will be identical to its training environment.
- Docker Packages: We provide a suite of Docker images that are tested on CPU, GPU, and TPU and expose simple gRPC APIs for calling your model.
Customer testimonials and partner use cases
Our Private Preview partners have been actively leveraging XLA Serving Engine to bring their deep learning stacks to production. Here is a selection of private preview customers (alphabetized) to showcase the diversity of uses:
Arm: Performance Parity on Arm Neoverse platforms
Our work with Arm has proven the production capabilities of JAX and XLA for training and serving on Arm-based CPUs:
“As AI inference scales, software optimization is critical to delivering performance and efficiency across diverse workloads. Through our collaboration with Google and contributions to OpenXLA, we’ve optimized the XLA CPU backend and serving runtime for vision, embedding and recommendation workloads on Arm Neoverse-based platforms, including Google Cloud Axion.”
— Ian Bratt, Vice President of ML Technology and Fellow, Arm
Intel: Native JAX on Xeon server CPUs
We have also been partnering with Intel to showcase the power of Intel’s optimized compute:
“Our partnership with Google showcases the XLA Serving Engine (XSE) on Intel® Xeon® processors, delivering native, high-performance JAX inference accelerated by Intel® AMX. With XSE, enterprises get out-of-the-box performance across Vision, Recommendation, and LLM workloads — and the flexibility to serve them efficiently wherever their business runs, from cloud to edge to on-prem.”
— Bill Pearson, Vice President of Data Center and AI software, Intel
Moloco: Scaling advertising and recommenders
“To support our demand-side platform (DSP) processing over 650 billion daily inferences and peaking at 10.5 million QPS, we needed an inference runtime that could match our massive scale. XLA Serving Engine is the high-performance, C++-based JAX serving solution we’ve been waiting for that keeps us framework agnostic and complements our Cloud TPU Trillium investments.”
— Kunal Kukreja, Staff Machine Learning Engineer, Moloco
What’s next
- Torch Support: We plan to integrate with TorchTPU, allowing PyTorch users to export to an XLA SavedModel to serve their models on TPUs in C++.
- Ahead-of-Time (AOT) Compilation: We are working on rolling out AOT compilation to our internal users before bringing this feature to external users. AOT compilation will reduce cold-start time by offering users the ability to compile their models offline (on CPU) for their target serving hardware.
- C++ Data Preprocessing: We are offering a C++ based preprocessor for users who want to combine their data processing with their model’s graph in C++.
Limitations
XSE is a lightweight, one-shot inference engine. For large autoregressive language models with long context lengths, we recommend popular projects such as vLLM or SGLang, which benefit from LLM-specific memory management techniques.
Get started today
Ready to streamline your JAX production pipeline?
- Check out the XLA Serving Engine quick-start instructions at github.com/openxla/xla-serving. You can find the list of available Docker images here.
- File an issue or contribute pull requests at github.com/openxla/xla-serving.
Acknowledgements
Thank you to everyone who contributed to XLA Serving Engine, including but not limited to the following individuals:
Ashish Agarwal, MD Faijul Amin, Carlos Araya, Isha Arkatkar, Sagun B, Ashok Bhat, Xevi Miró Bruix, Chris Chan, Ray Chen, Chen Cheng, Hyun Min Choi, Gauri Deshpande, Tiffany Dong, Sayce Falk, Happy Gao, Vaibhav Gurunathan, Michael Hudgins, He Jiang, Leon Kennedy, Young Seok (Tony) Kim, Penporn Koanantakool, Kunal Kukreja, Younghee Kwon, Shawn Lu, Xevi Miró Bruix, Sachin Muradi, César Naranjo, Pat Notz, Adam O’Brien, Olivia Qin, AG Ramesh, Shubha Rao, Mayank Raunak, Crefeda Rodrigues, Gassan Salama, Oleg Shaldybin, Brian Seung, Juanli Shen, Prairie Sims, Monica Song, Anne Stern, Quoc Truong, Armando Ugalde Velasco, Jess Wiser, Siqiao (James) Wu, Qinglan Xiang, Yunjie Xu, Jeffrey Yoo, Hyungsuk Yoon, and Wangyuan Zhang.
