谷歌与 Meta 联手打造 PyTorch 原生 TPU 推理引擎

SGLang-TorchTPU 的核心目标是让 TPU 成为 PyTorch 生态中的一等公民。当前绝大多数模型以 PyTorch 编写并在其生态中部署，推理服务栈依赖 torch.compile、torch.distributed 等工具链。TorchTPU 由谷歌构建，使 TPU 能够作为原生 PyTorch 设备运行，支持 eager 模式和通过 XLA 编译两种路径。SGLang-TorchTPU 在此基础上，将 SGLang 的完整服务特性带到 TPU，并叠加了针对 TPU 优化的 Pallas 内核。

在架构设计上，每个 TPU 设备运行独立的 PyTorch 进程，SGLang 的调度器、RadixAttention 和 API 服务器保持不变，TPU 相关工作在其下层完成。KV 缓存存储于为 Pallas ragged paged attention 设计的分页内存池中，默认采用 FP8 格式，注意力计算和 MoE 分组矩阵乘法均以 PyTorch 自定义算子的形式嵌入 torch.compile 计算图。整套实现以 sglang-tpu 插件形式发布，用户只需在启动时传入 --device tpu 参数即可加载。

KV 缓存的跨设备传输依托谷歌开源的 TPU Sync 库实现。在 KV 缓存卸载场景中，KV 页面被移至 DMA 固定的主机内存池；在预填充与解码分离部署场景中，解码端在预填充每步提交后逐块拉取 KV 数据，从而将最终需要传输的数据量降至最低。当两侧使用不同的张量并行规模时，KV 数据会在传输过程中完成重新分片，例如从 16 个预填充分片重组为 8 个解码分片，有效降低了异构部署的通信开销。

功能覆盖方面，SGLang-TorchTPU 在发布时已支持 OpenAI 兼容 API、连续批处理、分块预填充、结构化输出、RadixAttention 前缀缓存、EAGLE-3 等多种推测解码方案、数据并行注意力与专家并行、FP8 KV 缓存、上下文并行以及预填充解码分离部署等核心特性。多模态方面支持 Qwen3-VL、Kimi MoonViT 和 Gemma 4 视觉模型。超大规模模型方面，支持跨主机张量并行与专家并行，以及 FP8、MXFP4 和 INT4 权重格式。

在可靠性保障机制上，项目采用测试驱动的支持矩阵：所有标注为已支持的模型和功能均与真实 TPU 硬件上的 CI 测试绑定，支持矩阵由测试注册记录自动生成，确保文档声明与实际运行状态一致。目前 CI 覆盖单主机和多主机 TPU v7x 上的 15 个模型，共注册 193 项测试，涵盖内核正确性、端到端服务精度（GSM8K、GPQA、AIME、MMLU-Pro 等标准评测）、推测解码、前缀缓存、多主机执行及 Agent 工具调用工作负载等场景，每晚自动回归运行。

在强化学习训练方向，SGLang-TorchTPU 可作为 RL 训练的 rollout 引擎，与 TorchTitan 训练框架通过 Miles 框架协同工作，实现端到端的 TPU 上 RL 训练闭环。谷歌 AI 推理技术负责人 Qi Zhou 表示，随着 Agent 工作负载的增长，推理效率愈发关键，此次合作为开发者提供了在 TPU 上运行主流开源模型的生产级引擎，同时兼顾 Cloud TPU 的成本优势。Meta PyTorch-TPU 负责人 Zongwei Zhou 则确认该引擎已在 Meta 生产环境中运行，并高度评价其与现有 AI 基础设施的无缝集成能力。

要点

• SGLang-TorchTPU 由 RadixArk、谷歌与 Meta 三方联合开发，已正式开源，支持同一套 SGLang 引擎在 GPU 和 TPU 上无缝切换部署。
• 首批支持 Muse Spark Open、DeepSeek V4、Kimi K3、Gemma 4 等九款主流开源大模型，并已在 Meta 生产环境中实际运行。
• KV 缓存传输依托谷歌开源的 TPU Sync 库，支持预填充解码分离部署及跨不同张量并行规模的 KV 重新分片。
• 项目采用 CI 驱动的支持矩阵机制，193 项 TPU 测试每晚自动回归，确保文档声明与实际运行状态严格对应。
• 该引擎可与 TorchTitan 通过 Miles 框架配合，作为强化学习训练的 rollout 引擎，覆盖训练与推理的完整闭环。