SGLang把服务栈搬上TPU：同一引擎兼顾GPU与TPU，Meta已用于生产
核心概要
SGLang团队与Google、Meta合作发布开源后端SGLang-TorchTPU，让同一套SGLang服务引擎在GPU和TPU上运行。此前TPU并非原生PyTorch设备，服务栈难以直接复用PyTorch生态；新后端把调度器、RadixAttention和API服务器原样保留，TPU相关工作放在下层。该后端已在Meta生产环境运行，并作为Miles强化学习流程的rollout引擎与TorchTitan配合。
深度剖析
SGLang的完整服务栈现在可以跑在Google Cloud TPU上，且与GPU共用同一个引擎，部署时可按工作负载和成本选择硬件。 此前TPU不是原生PyTorch设备，以PyTorch为中心的服务栈无法直接沿用torch.compile和torch.distributed这套生态。 该后端已在Meta生产环境运行，Meta工程师还贡献了编译与分离式服务工作，并用生产负载从一开始塑造引擎。
每个TPU设备运行自己的PyTorch进程，因此SGLang的调度器、RadixAttention和API服务器无需改动，TPU相关工作全部位于其下层。 TorchTPU把TPU变成可即时执行或经XLA编译的原生PyTorch设备，SGLang-TorchTPU在其上叠加TPU调优的Pallas内核。 KV缓存放在为Pallas不规则分页注意力布局的分页池中，默认FP8 KV，注意力与MoE分组矩阵乘的Pallas内核以PyTorch自定义算子形式在torch.compile图中调用，整体作为sglang-tpu插件在传入--device tpu时加载。
同一套PyTorch栈也覆盖训练侧：SGLang-TorchTPU生成rollout，TorchTitan训练，Miles在两者之间同步权重。 此前强化学习训练同样面临GPU或TPU的选择，训练器一侧未必能沿用同一栈。 在单台4芯片TPU v7x主机上以GRPO训练Qwen3-4B、数据为DAPO-Math-17k，与6块H200 GPU跑同一配方和超参、各两个随机种子对照，奖励与回复长度曲线在种子间波动范围内一致，AIME-24在评测噪声内。
启示与展望
已展示的范围是TPU v7x上的服务与单主机RL训练：CI覆盖单主机与多主机TPU v7x上的15个模型、193项TPU测试，Miles TPU训练处于预览阶段。面向需要在PyTorch流程中部署或训练模型的工程师与研究者。发布方计划继续增加模型并优化已支持模型，先从DeepSeek V4.1 Flash和Reflection Beam开始。
Miles TPU训练目前是小规模实验，仅演示单主机上的端到端RL流程，早期学习行为与GPU基线相近；后续计划包括更大稠密模型与MoE模型、多主机训练、经ICI的更快权重传输、缩小与GPU的步时差距、更长运行的对照以及训练与推理之间逐位一致的log-prob。部分最新模型与功能仍标为实验性，需通过测试后才进入支持层级。