Google gooses Python AI performance by... removing Python

"Train in Python, serve in C++" might as well be Google's motto when it comes to much of its ML work.

Python is easy for data scientists to wield, but keeping its interpreter in the serving pipeline can add overhead and deployment headaches. The industry-backed OpenXLA Project provides tools for compiling machine learning computations into native code for different processors.

The latest component of this setup to be open sourced is XSE, short for the XLA Serving Engine.

One-shot inference

XSE is not so much suited for large-scale LLM inference (which would be better served by vLLM or SGLang), but rather for smaller, more sharply task-oriented models, of which there are plenty at Google.

"LLM serving gets most of the attention right now, but a significant amount of production inference is still ranking, recommendation, embedding, and vision models. These tend to be short requests with very high [queries-per-second] and tight latency budgets," Nathan Beach, Google director of product management, wrote in a LinkedIn missive Tuesday.

At Google, XSE already handles more than six million queries per second across services including YouTube Shorts, Google Ads, and Search.

These models aren't there for a leisurely chat. They rank results, recommend videos, and make other decisions on the fly, often without users noticing.

A Python interpreter mulling each request would be a bottleneck, if not the end of Google itself. In Google's tests on short requests, XSE delivered two to five times the throughput of its best optimized Python server, using a quarter of the memory.

No Python here

In this setup, the ML boffins write their model code in Python, using Google's JAX to help with the sums. Google also plans integration with TorchTPU so PyTorch users can export models for C++ serving on TPUs. An export tool packages the model's computations in StableHLO, a portable intermediate representation, alongside its trained weights in an "XLA SavedModel."

Ready to precognate, the C++ serving runtime uses XLA to compile those computations into native code for the target silicon. No Python interpreter to see here.

XLA SavedModels can run on GPUs and CPUs - Google partnered with both Arm and Intel to develop the XLA ecosystem. But Google obviously thinks it'd be a lot cooler if you ran your models on its own line of TPU accelerators.

Python's sluggishness continues to be an issue for the ML community. Beyond XSE, other approaches have been taken, many of which involve jettisoning the Python code as quickly as possible. PyTorch offers a native C++ runtime called LibTorch. Also, Microsoft and its collaborators offer the ONNX Runtime, which supports an open graph format with Python bindings.

Even Python's own reference implementation relies heavily on C, with Rust now getting a foot in the door. (R)