vLLM 与 SGLang 现在能通过 PyTorch 原生跑在 TPU 上，要记住的数字是 2.01 美分

2026 年 10 月 20 日，Google 发布了 vLLM（与 Inferact 合作）与 SGLang（与 RadixArk 合作）的原生 PyTorch 后端，两者都构建在 TorchTPU 之上。真正的新闻不是「这两个引擎能跑在 TPU 上」——它们通过基于 JAX 的后端早就能跑。真正的新闻是：开发者反馈希望在 PyTorch 里定制和调试自己的服务栈，现在每个框架的 worker 模型、调度器、KV 缓存管理器和 API 服务器都原样保留，TPU 实现放在下面。官方公布的基准是 Qwen 3.5 397B-A17B 在 Ironwood TPU 上每百万 token 成本 2.01 美分。这个数字适不适用于你的负载是另一回事，而这件事在迁移之前就能测出来。

vLLM 与 SGLang 现在能通过 PyTorch 原生

一、TorchTPU 到底是什么，去掉宣传话术

TorchTPU 把 TPU 变成原生 PyTorch 设备。张量就是 tpu 设备上的标准 PyTorch 张量，没有包装类型；代码可以用 eager 模式跑，也可以用 torch.compile 通过 XLA 编译；集合通信走 torch.distributed。因为每个设备跑自己的 PyTorch 进程，两个引擎都保留各自框架的机制，只是把 TPU 实现放在下面。多数模型沿用上游标准 PyTorch 定义，TPU 相关工作集中在 kernel 和层实现上，用 Pallas 写的自定义 kernel 直接作为自定义算子被 PyTorch 在编译图里调用。最后这个设计决策才是该关心的：它意味着 TPU 优化不会分叉你的模型代码，上游模型更新也不会变成合并冲突。

# Get a server up from the official repos, then read their README for
# the current flags -- the engines are moving weekly.
set -euo pipefail
git clone https://github.com/vllm-project/vllm-torchtpu.git
cd vllm-torchtpu
python -m venv .venv && . .venv/bin/activate
pip install -U pip
pip install -e .

# Serve an OpenAI-compatible endpoint. Model choice must be one of the
# supported families (Qwen, DeepSeek, GLM, Kimi, GPT-OSS, Gemma).
export VLLM_TARGET_DEVICE=tpu
vllm serve Qwen/Qwen3.5-8B-Instruct --port 8000 --max-model-len 32768

# Enable host-side KV offload if you expect long prefixes that outlive HBM.
# Flag names live in the repo's docs -- check them, do not copy from blogs.

二、对 Agent 负载真正重要的功能

Google 自己给出的「为什么这次发布重要」的解释是对的：Agent 会反复重发不断膨胀的多轮对话，所以要高效服务它们，关键是找到已缓存的前缀而不是重新计算。两个引擎都把可复用前缀放在 HBM 里，开启主机下沉后，把更老的前缀挪到主机内存作为第二层。对首个版本来说，支持列表完整得少见：包含混合模型的自动前缀缓存；经 TPU Sync 把 KV 缓存下沉到主机内存和 Google Cloud Managed Lustre；工具调用与结构化输出；MTP、EAGLE-3、DFlash、DSpark 和 n-gram 草稿的投机解码；带专家并行的数据并行注意力；prefill 与 decode 的上下文并行；FP8 KV 缓存；跨主机分离式 prefill/decode；视觉语言模型与嵌入模型；以及 FP8、FP4、INT8、INT4 权重。

# Does prefix caching actually pay on YOUR traffic? Measure cold vs warm TTFT.
# OpenAI-compatible endpoint, streaming on, so the first chunk is the signal.
import json, time, urllib.request

ENDPOINT = "http://127.0.0.1:8000/v1/chat/completions"

def ttft(messages, model="Qwen/Qwen3.5-8B-Instruct"):
    body = json.dumps({"model": model, "messages": messages,
                       "stream": True, "max_tokens": 64}).encode()
    req = urllib.request.Request(ENDPOINT, data=body,
                                 headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    with urllib.request.urlopen(req, timeout=180) as r:
        for line in r:
            if line.startswith(b"data: ") and b"[DONE]" not in line:
                return time.perf_counter() - t0
    return float("nan")

prefix = [{"role": "system", "content": "You are a build log analyst. " * 300}]
cold = [prefix[0], {"role": "user", "content": "Summarize run %d in 3 lines." % i}
        for i in range(1)]
warm = [prefix[0], {"role": "user", "content": "Summarize run %d in 3 lines." % i}

c = ttft(cold)
w = ttft(warm)                     # identical prefix, second call
print(f"cold_ttft={c*1000:.0f}ms warm_ttft={w*1000:.0f}ms speedup={c/w:.2f}x")

# If warm is not meaningfully faster, prefix caching is off, or your traffic
# never repeats a prefix. Both are fixable; the number you cannot fix is none.

三、TPU Sync 是那件没人会写文章讲的事

上面这些功能里有两个会把 KV 缓存搬离计算它的设备。分离式服务把 prefill 放一台主机、decode 放另一台，于是每个请求的 KV 缓存都要在两者之间传输。分层下沉把放不进 HBM 的那部分推到主机 DRAM 和二级存储，之后有请求复用前缀时再拉回来。两条路径都会直接加到延迟上——除非传输与计算重叠。两个引擎都用 TPU Sync 做这件事：一个开源 C++ 库，直接读 PyTorch 张量背后的设备缓冲区，因此主机侧不需要额外的暂存拷贝，传输在自己的线程上、在 Python 之外运行，prefill 和 decode 继续。在 vLLM 里它通过 KV connector 接口接入；在 SGLang 里它既是分离式传输后端，也是 HiCache 的 I/O 后端。

# Turn your own traffic into dollars per million tokens, not vibes.
# Inputs: measured prompt/completion tokens per request and your hardware cost.
python3 - <<'PY'
import json

HW_COST_PER_HOUR = 4.20          # your actual rate, from the invoice
REQUESTS = [
    # prompt_tokens, completion_tokens, seconds
    (12500, 480, 2.9),
    (12510, 460, 1.1),           # warm prefix, same conversation expanded
    (800,   300, 0.7),
]

total_s = sum(r[2] for r in REQUESTS)
p_tok = sum(r[0] for r in REQUESTS)
c_tok = sum(r[1] for r in REQUESTS)

cost = total_s / 3600 * HW_COST_PER_HOUR
per_m = cost / ((p_tok + c_tok) / 1_000_000)

print(f"requests={len(REQUESTS)} prompt={p_tok} completion={c_tok}")
print(f"wall_seconds={total_s:.1f} cost_usd={cost:.4f} per_million_tokens=${per_m:.4f}")

# Compare like with like: same model, same sequence shape, same batch policy.
# A cents-per-token figure you cannot reproduce is a marketing number.
PY

四、那些数字，以及它们没告诉你的

基准是公开的，通过 SemiAnalysis 的 InferenceX 项目发布——这一点让数字可被核查。对 Qwen 3.5 397B-A17B、8k 输入 1k 输出、FP8 精度，Google 报告的总体最低每百万 token 成本是 2.01 美分，Ironwood TPU 在帕累托前沿的大部分区段领先，并在 AgentX（该项目专门评估 Agent 长上下文与多轮负载的赛道）上刷新了 SOTA。另外，针对 2.8 万亿参数的 Kimi K3，一个 Pallas megakernel 把整个 decode 步骤作为单个 kernel 执行，在 Ironwood 上对单个用户跑到每秒 700+ token。需要保留的注意事项是所有推理基准共有的：那是特定模型在特定序列形状下的数字。你实际的每 token 成本取决于提示长度分布、缓存命中率，以及你是否真能把前缀保持住。

# Disaggregated serving only pays if the KV transfer overlaps compute.
# Sanity check that prefill and decode hosts are actually talking to each
# other, and that transfer time is not serialized into your TTFT.
ssh prefill-host 'ss -tnp | grep -E ":7100|tpu-sync" | head'
ssh decode-host  'ss -tnp | grep -E ":7100|tpu-sync" | head'

# Then measure the boundary directly: short prompt (transfer-light) versus
# long prompt sharing a cached prefix (transfer-heavy). If the long-prompt
# case degrades superlinearly, the transfer is on the critical path.
python3 traffic_shape.py --prompt-mix short,long-cached,long-cold --n 60

# The thing to watch is not average latency. It is p99 for long-cached
# prompts, because that is the shape agents actually generate.

五、把它当成本问题来测，而不是框架选型问题

要避免的错误是：拿一个全新的单轮提示测一下每秒 token 数，然后据此做决定。Agent 会重发上下文，那才是你的负载。所以要测两件事：冷前缀下的首 token 时间，以及同样长度热前缀下的首 token 时间，然后算出你真实流量带来的有效缓存命中率。如果命中率很低，前缀缓存就变成「用内存换不到任何东西」的成本项，2.01 美分也不是你的数字。下面的脚本刻意与框架无关：一个从官方仓库起服务，一个在 OpenAI 兼容端点上对比冷/热首 token 时间，还有一个把你自己的 token 计数换算成每百万成本，好拿来和你现在跑的东西比。

# Before moving a workload, prove the hit rate is real over a day of traffic.
# Parse your gateway logs; every field here should already exist.
python3 - <<'PY'
import json, collections

hits = collections.Counter()
prompt_tokens = 0
records = 0

with open("/var/log/llm-gateway.jsonl") as fh:
    for line in fh:
        try:
            r = json.loads(line)
        except json.JSONDecodeError:
            continue
        records += 1
        prompt_tokens += r.get("usage", {}).get("prompt_tokens", 0)
        # Most gateways expose cached tokens; name varies by vendor.
        cached = (r.get("usage", {}).get("cached_tokens")
                  or r.get("usage", {}).get("prompt_tokens_details", {}).get("cached_tokens")
                  or 0)
        hits["cached" if cached else "cold"] += 1

total = sum(hits.values()) or 1
print(f"requests={records} prompt_tokens={prompt_tokens}")
print(f"cached_requests={hits['cached']} share={hits['cached']/total:.1%}")

# Rule of thumb: below roughly a third, prefix caching is a memory cost
# with no throughput payoff, and you should fix the traffic shape first.
PY

六、这替换不了什么

基于 JAX 的两个运行时 vLLM-Jax 与 SGLang-JAX 仍会继续维护和支持，所以现有部署不会被人从脚下抽走。当前模型覆盖 Qwen、DeepSeek、GLM、Kimi、GPT-OSS 与 Gemma 4，官方把「新开源模型的 Day-0 支持」表述为目标而非保证。更硬的约束不在软件：如果你的负载长在别的加速器厂商上，那么在你并不拥有的硬件上更便宜的每 token 数字只是谈资。诚实的读法是：对那批「因为工具链原因才拒绝 TPU」的团队（也就是绝大多数），TPU 推理刚刚变得可考虑。在再读一份基准之前，先拿自己的流量跑一遍成本模型。