{"slug": "google-open-sources-xla-serving-engine-serve-jax-models-at-36x-throughput", "title": "Google Open-Sources XLA Serving Engine: Serve JAX Models at 36x Throughput", "summary": "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.", "body_md": "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.\n\n## Why C++ Serving Wins on Cost\n\nPython-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.\n\nGoogle’s internal production numbers — not synthetic benchmarks — show what this unlocks:\n\n- 2–5x throughput over their best Python server\n- 8–36x over a straightforward Python setup\n- 30–70% less CPU per prediction\n- 4x less memory footprint\n\nThese 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.\n\n## The Workflow: Export Once, Serve Anywhere\n\nThe [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.\n\n```\n# Install the export package\npip install xla-saved-model\n\n# Start the gRPC server (CUDA 12 GPU)\ndocker run -p 9000:9000 \\\n  us-docker.pkg.dev/xla-serving-engine/xla-serving/server:cuda12-latest \\\n  --model_path=/path/to/xla_saved_model\n```\n\nPrebuilt 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.\n\n## Not a JAX-Only Story\n\nThe 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.\n\nThis 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.\n\n## Know When to Stop: XSE vs vLLM vs SGLang\n\nGoogle 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.\n\n| Workload | Recommended Engine | \n|---|---|\n| Recommender / ranker | XLA Serving Engine | \n| Vision / multimodal encoder | XLA Serving Engine | \n| Encoder-style LM (BERT) | XLA Serving Engine | \n| General LLM serving | vLLM | \n| RAG / multi-turn chat | SGLang | \n\nGetting 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.\n\n## The OpenXLA Play\n\nXSE 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.\n\nXSE 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).", "url": "https://wpnews.pro/news/google-open-sources-xla-serving-engine-serve-jax-models-at-36x-throughput", "canonical_source": "https://byteiota.com/xla-serving-engine-google-open-source/", "published_at": "2026-11-06 18:06:48+00:00", "updated_at": "2026-11-06 18:22:54.623056+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "ai-infrastructure", "ai-tools", "developer-tools"], "entities": ["Google", "XLA Serving Engine", "OpenXLA", "JAX", "PyTorch", "torchax", "vLLM", "SGLang"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/google-open-sources-xla-serving-engine-serve-jax-models-at-36x-throughput", "markdown": "https://wpnews.pro/news/google-open-sources-xla-serving-engine-serve-jax-models-at-36x-throughput.md", "text": "https://wpnews.pro/news/google-open-sources-xla-serving-engine-serve-jax-models-at-36x-throughput.txt", "jsonld": "https://wpnews.pro/news/google-open-sources-xla-serving-engine-serve-jax-models-at-36x-throughput.jsonld"}}