AI 推理术语表
软件栈

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 仓库。