Google open-sources XLA Serving Engine for Python-free ML inference at scale

Google has released the XLA Serving Engine (XSE), an Apache 2.0-licensed inference runtime for the OpenXLA ecosystem that enables a "Train in Python, Serve in C++" workflow for machine learning models. The engine exports models from frameworks like JAX into an "XLA SavedModel" format containing StableHLO representations, which the C++-native XLA Serving Runtime then loads and executes directly via PJRT without Python interpreter overhead. Google reports that XSE achieves 2-5x the throughput of optimized Python servers and 8-36x that of basic implementations, while using 30-70% less CPU time and roughly 25% of the memory.

XSE has been deployed in Google's production infrastructure since 2025 and currently serves over 6 million queries per second across YouTube Shorts, Google Ads, Search, and portions of Gemini Omni. The release includes Docker images tested on CPU, GPU, and TPU with gRPC APIs, supports both single-host and multi-host serving through PJRT, and integrates with Orbax for checkpointing. Google is working to add PyTorch support through TorchTPU integration and ahead-of-time compilation to reduce cold-start latency, though the project explicitly recommends vLLM or SGLang for large autoregressive language models with long contexts.

Architecture and runtime characteristics

• Python-free serving pipeline: XSE exports models to an "XLA SavedModel" containing StableHLO computation graphs and Orbax-stored weights, enabling C++-native serving without Python runtime or GIL constraints.
• PJRT-based execution: The runtime wraps PJRT directly and applies XLA compiler optimizations, allowing a single process to utilize all eight CPU cores and eliminating interpreter-induced latency.
• Multi-backend Docker deployment: Google provides tested Docker images exposing gRPC APIs for CPU, GPU, and TPU inference, with verified partner optimization for Arm Neoverse and Intel Xeon with AMX acceleration.
• Performance parity guarantee: XSE preserves forward-function performance including kernel optimizations, ensuring that model behavior in serving matches the training environment without numerical divergence.
• Scalability model: Supports both single-host and multi-host configurations through PJRT, with internal Google deployments demonstrating sustained operation at 6 million queries per second.

Why it matters

XLA Serving Engine addresses a persistent operational gap in open-source ML infrastructure: the difficulty of deploying Python-trained models in production environments that restrict or penalize interpreter overhead. By open-sourcing a runtime that Google itself relies on for high-volume services, the project gives Cloud TPU users and hardware partners a production-tested alternative to Python-based serving stacks, while the Apache 2.0 license and modular design reduce vendor lock-in concerns. For the broader OpenXLA ecosystem, XSE provides a concrete serving target for StableHLO and strengthens the toolchain around JAX, with planned PyTorch integration potentially expanding its relevance beyond Google's primary framework.