Google open-sourced XLA Serving Engine (XSE) 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 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_modelPrebuilt 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 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 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.
