软件栈
SGLang-JAX
也称为 SGLang-JAX、sgl-jax、SGL-torchtpu
先用大白话
SGLang-JAX 是 SGLang 服务引擎面向 TPU 的 JAX 原生版本,独立于基于 PyTorch 的 SGLang 运行时。
技术定义
SGLang-JAX 是一个 JAX 原生的服务引擎,把 SGLang 风格的调度和 prefix caching 与 JAX 模型实现及 TPU 专用 kernel 结合,而不是把 PyTorch 版 SGLang 运行时翻译到 TPU 上。
工程细节
vLLM 通过 TorchAX 翻译登陆 TPU,SGLang 则选择了另一条路,用 JAX 重新实现了引擎。这样可以直接使用成熟的 TPU 原语,代价是维护一套与上游 SGLang 分叉的第二代码库。Google 和 RadixArk 已宣布基于 TorchTPU 的 PyTorch 原生替代方案 SGL-torchtpu,它能让 TPU 上的 SGLang 重新共享上游代码路径。
为什么重要
两个引擎采用两种不同的 TPU 策略,说明生态尚未定型。TorchTPU 的目标是让两者收敛:vLLM 和 SGLang 都获得原生 PyTorch 设备,Pallas kernel 工作可以共享,模型的 day-0 支持也能与 NVIDIA、AMD 同步登陆 TPU。
如何在 InferenceX 中解读
TPU InferenceX 预览的结果使用 TorchTPU vLLM。TorchTPU SGLang 开源后会补上 TPU 上的 SGLang 结果,届时 SemiAnalysis 会把 TPU 基准测试迁入公开的 InferenceX 仓库。