Google Open-Sources XLA Serving Engine: Serve JAX Models at 36x Throughput

Google open-sourced the XLA Serving Engine (XSE) on November 5, a C++ inference runtime that already serves YouTube Shorts, Google Ads, Search and parts of Gemini Omni at over 6 million queries per second. Google reports internal production gains of 2–5x throughput over its best Python server and 8–36x over a straightforward Python setup, with 30–70% less CPU per prediction and 4x less memory footprint, for recommenders, rankers, vision and multimodal encoders, encoder-style language models and short fixed-horizon decoding. Models export once from JAX, or from PyTorch via torchax, to the XLA SavedModel format and run in prebuilt CPU, CUDA 12, CUDA 13 and Cloud TPU server images; Google directs autoregressive LLM serving with paged KV caches to vLLM and SGLang instead.

Google open-sourced XLA Serving Engine XSE https://opensource.googleblog.com/2026/11/supercharging-xla-inference-xla-serving-engine.html on November 5 — the C++ inference runtime that already drives YouTube Shorts, Google Ads, Search, and parts of Gemini Omni at over 6 million queries per second. For ML teams caught between sluggish Python serving stacks and the complexity of ONNX conversions, XSE offers a cleaner path: export once from JAX or PyTorch via torchax , and serve in C++ with no Python in the request path. Why C++ Serving Wins on Cost Python-based inference servers have a ceiling. The GIL serializes concurrent inference calls; a multi-gigabyte Python environment loads on every container restart; and even optimized Flask or FastAPI servers burn CPU on framework overhead. XSE eliminates all three. Google’s internal production numbers — not synthetic benchmarks — show what this unlocks: - 2–5x throughput over their best Python server - 8–36x over a straightforward Python setup - 30–70% less CPU per prediction - 4x less memory footprint These gains apply to the workloads XSE is designed for: recommenders, rankers, vision and multimodal encoders, encoder-style language models think BERT-scale , and short fixed-horizon decoding tasks. If your production stack looks anything like Google Ads or YouTube Shorts, those numbers are directly relevant. The Workflow: Export Once, Serve Anywhere The GitHub quick start https://github.com/openxla/xla-serving at openxla/xla-serving is straightforward. Install the export package, export your JAX model to an XLA SavedModel, and pull a prebuilt server image. Install the export package pip install xla-saved-model Start the gRPC server CUDA 12 GPU docker run -p 9000:9000 \ us-docker.pkg.dev/xla-serving-engine/xla-serving/server:cuda12-latest \ --model path=/path/to/xla saved model Prebuilt images cover CPU, CUDA 12, CUDA 13, and Cloud TPU. The XLA SavedModel format packages the model’s StableHLO computation alongside an Orbax weight checkpoint — no ONNX conversion, no TensorFlow dependency, and full support for custom XLA calls and Pallas kernels. The official docs https://xla-serving.readthedocs.io cover the full export API. Not a JAX-Only Story The common reaction to any JAX announcement is: “cool, but I’m on PyTorch.” XSE anticipated this. The torchax library bridges PyTorch operations to their JAX equivalents transparently — take any HuggingFace PyTorch model, run it through torchax, and export to XLA SavedModel without touching model code. The result lands in the same C++ serving runtime as a native JAX model. This matters because the recommenders and vision encoders most likely to benefit from XSE’s throughput gains are often PyTorch models. The JAX-only limitation no longer holds. Know When to Stop: XSE vs vLLM vs SGLang Google is direct about where XSE does not belong: autoregressive LLM serving with paged KV caches and continuous batching. For those workloads, the recommendation is vLLM or vllm-jax / vllm-torchtpu on TPUs and SGLang. XSE is a one-shot inference engine — no continuous batching scheduler, no speculative decoding, no KV cache management. | Workload | Recommended Engine | |---|---| | Recommender / ranker | XLA Serving Engine | | Vision / multimodal encoder | XLA Serving Engine | | Encoder-style LM BERT | XLA Serving Engine | | General LLM serving | vLLM | | RAG / multi-turn chat | SGLang | Getting this boundary wrong is expensive. Deploy XSE on a Llama-style workload expecting chat throughput and you will be disappointed. That is not a failure of the tool — it is a workload mismatch. The OpenXLA Play XSE is the production serving layer of Google’s OpenXLA https://openxla.org stack: XLA compiler → StableHLO IR → XSE runtime . By open-sourcing the entire stack, Google is positioning StableHLO as the universal ML intermediate representation — competing with ONNX on hardware breadth and winning on native TPU, Pallas kernel, and multi-device sharding support. XSE shipped at v0.0.13 on release day. This is not a research dump — it has been running inside Google’s production infrastructure for years before the open-source release, and the version history reflects active, ongoing development. If your stack serves encoders or recommenders at any meaningful scale, it is worth a look. Start at github.com/openxla/xla-serving https://github.com/openxla/xla-serving .