PrivateUse1 后端
也称为 PrivateUse1、PyTorch out-of-tree 后端、自定义设备后端
先用大白话
PrivateUse1 是 PyTorch 允许厂商接入新设备类型的钩子,让张量可以像放在 CUDA 上一样放在 device="tpu" 上。
技术定义
PrivateUse1 是 PyTorch 中预留的 dispatch key 和设备类型,允许 out-of-tree 后端注册自己的张量、内存分配器和算子实现,使框架把新加速器当作一等设备处理。
工程细节
TorchTPU 用 PrivateUse1 暴露一个位于 device="tpu" 上的普通 torch.Tensor,而不是由 JAX 数组支撑的包装对象。PyTorch dispatcher 把 ATen 操作路由到 TPU 后端,后端既可以在 bring-up 和调试时 eager 执行,也可以通过 torch.compile 捕获图。DDP、FSDP2 和 DTensor 等分布式 API 通过同一机制工作。此前的 TorchAX 方案则通过 __torch_dispatch__ 拦截每个操作并翻译成 JAX 调用。
为什么重要
这个后端钩子让"原生"一词有了具体含义。服务引擎可以复用上游的模型代码、调度器、continuous batching 和功能逻辑,不需要 PyTorch 到 JAX 的翻译层,从而降低每个新模型和引擎功能登陆 TPU 的成本。
如何在 InferenceX 中解读
TPU InferenceX 预览的结果来自基于该后端构建的 TorchTPU vLLM 栈。Google 预计在 10 月中旬的 PyTorch Conference 前后结束私测并开源,之后 SemiAnalysis 会把 TPU 基准测试从自己的 fork 迁移到公开的 InferenceX 仓库。