在机器学习领域,“训练用 Python,服务用 C++”已成为谷歌遵循的技术准则。Python 虽便于数据科学家上手,但其解释器在推理服务管道中会引入额外开销与部署难题。为此,由行业支持的 OpenXLA 项目提供了将机器学习计算编译为不同处理器原生代码的工具,其最新开源组件即为 XSE(XLA Serving Engine,XLA 推理引擎)。

专注高并发短请求

XSE 并非针对大规模大语言模型(LLM)推理——此类场景更适合 vLLM 或 SGLang——而是专为任务导向性强、规模较小的模型设计。这类模型在谷歌内部广泛存在,用于实时结果排序、视频推荐及嵌入处理等决策环节,用户往往无感察觉。

谷歌产品管理总监指出,尽管 LLM 推理占据主流关注,但生产环境中仍有大量涉及排序、推荐和视觉模型的推理工作。这些请求通常极短,却具有极高的每秒查询数(QPS)和严格的延迟预算。若每个请求均由 Python 解释器处理,将成为系统瓶颈甚至导致服务崩溃。

目前,包括 YouTube Shorts、Google Ads 和搜索在内的多项谷歌服务,已通过 XSE 处理超过每秒 600 万次查询。测试显示,针对短请求场景,XSE 的吞吐量是最佳优化 Python 服务器的 2 至 5 倍,内存使用量仅为后者的四分之一。

彻底剥离 Python 运行时

在该架构下,研究人员使用 Python 编写模型代码并借助 Google JAX 辅助计算。谷歌还计划集成 TorchTPU,使 PyTorch 用户能将模型导出为在 TPU 上运行的 C++ 服务格式。导出工具会将计算过程打包为可移植中间表示形式 StableHLO,并将训练权重封装于“XLA SavedModel”中。

部署阶段,C++ 推理运行时利用 XLA 将这些计算编译为目标芯片的原生代码,运行环境中完全不存在 Python 解释器。虽然 XLA SavedModel 支持在 GPU 和 CPU 上运行——谷歌已与 Arm 和英特尔合作开发相关生态系统——但其在自家 TPU 加速器上的表现尤为突出。

Python 的效率瓶颈持续困扰着机器学习社区。除 XSE 外,业界正通过多种方案加速“去 Python 化”进程,如 PyTorch 提供的原生 C++ 运行时 LibTorch,以及微软及其合作伙伴推出的支持 Python 绑定的 ONNX Runtime。此外,Rust 语言也开始在这一领域崭露头角。