TorchTPU：在 Google TPU 上运行 PyTorch 模型与算子

google-pytorch/torch_tpu

GitHub原文发布未知首次收录 2026-10-24上架 2026-10-24 07:01

ELI5：它是什么，解决什么问题

TorchTPU 为 Google TPU 提供 PyTorch 后端，包含面向 TPU 的自定义 ATen 算子实现、用于优化 XLA 计算执行的编译缓存，以及分布式训练（数据并行、张量并行）和模型推理示例。项目还支持与 torch.compile() 集成进行图模式执行，并提供调试与性能基准工具。

# 推理与部署 # 分布式训练 # TorchTPU

核心能力与快速开始

★ 701 · C++

能做什么: 提供 PyTorch 的 TPU 后端；实现多种 PyTorch 操作的 TPU ATen 内核；通过编译缓存优化 XLA 计算执行；支持数据并行与张量并行分布式训练；集成 torch.compile() 图模式执行；附带 Llama、Qwen、ResNet、DLRM、minGPT、UNet 等模型示例；提供调试和基准测试工具。
适用场景: 在 TPU 设备上运行 PyTorch 模型与算子；开展 TPU 上的分布式训练；进行模型推理；借助示例模型和 Marimo 交互式笔记本学习 TPU 上的 PyTorch 使用。
项目形态: 开发框架
部署方式: 需在 TPU 环境（README 提到 v6e GCE 实例最便捷）运行；可通过 pip 安装 torch_tpu==0.2.0 或使用 Bazel 从源码构建 wheel。
主要依赖: Python 3.11+；构建依赖 Bazel/Bazelisk；可从源码构建 wheel；README 中关键 Python 依赖标注为 TBD。
使用限制: 主要面向 TPU 环境，未说明在非 TPU 硬件上的可用性；依赖与部署成本未详细说明；未说明性能基准数据。
许可证: NOASSERTION
最近推送: 2026-10-24 06:58

官方入门项目文档核对 README 依据 ↗

资料检查：2026-10-24 07:00 · 卡片整理：2026-10-24 07:00
卡片由 AI 根据仓库资料整理，未进行安装或运行验证

相似资源

torch-neuronx：AWS Trainium 的原生 PyTorch 后端，支持 eager 执行、torch.compile、分布式训练与 NKI 自定义内核

torch-neuronx 是 AWS Trainium 的原生 PyTorch 后端，通过 PrivateUse1 机制在 PyTorch 中新增 neuron 设备，让用户可以用熟悉的 PyTorch API 把模型和输入迁移到 neuron 设备上训练或推理。项目默认以 eager 模式逐算子执行，也可通过 torch.compile 的 neuron 后端编译模型或模型区域以提升性能；支持 torch.distributed、DTensor 与函数式集合通信，覆盖数据并行、张量并行和完全分片并行；并支持在 eager 模式和 torch.compile 编译图内调用 NKI 自定义内核。项目当前为 Beta，需要 Trn2 或 Trn3 实例及对应 Python 版本，并列出静态形状、连续张量、float64 降级等已知限制。

DeepInfer：面向 DeepSeek 模型的高性能大规模推理框架

DeepInfer 是 DeepSeek 用于支撑其大规模生产服务的高性能推理框架，本次开源版本提供离线推理能力。它支持 DeepSeek-V4.1-Flash 模型，包含优化内核、Attention DP / 均衡 CP、MoE EP + EPLB 等并行策略，以及 DSpark 与置信度调度验证、SWA 有界重放、异步调度、批不变推理、确定性采样和结构化输出等特性，可运行在 NVIDIA SM100 GPU 与华为昇腾 950 NPU 上。完整推理服务所需的 KV 缓存系统、VE/Prefill/Decode 分离以及数据并行副本间的负载均衡未包含在本版本中。

Yunshu：面向 Apple Silicon 的本地 LLM/VLM 推理引擎

Yunshu 是一个基于 MLX、在 Apple Silicon 上本地运行的 LLM 与 VLM 推理引擎，单进程提供 OpenAI、Anthropic 和 Ollama 兼容 API。它支持无损的推测解码（DFlash2、MTP、prompt-copy drafting）、面向混合模型的前缀缓存（RAM、SSD 及可选存储层级）、结构化输出、多模态输入、Web 搜索与 MCP、编码代理（Claude Code、Codex、opencode）接入，以及包含密钥配额、主机遥测和评测的本地 Web 控制台。首个完整调优模型为 Qwen3.8-27B。

Genesis vLLM Patches：面向 Ampere GPU 的 int8 端到端推理补丁集

该项目是 vLLM 0.29.0 的个人运行时补丁分支，核心思路是在没有 fp8 张量核心的 Ampere（sm_86）硬件上把 int8 作为贯穿各阶段的「计算」格式。补丁覆盖 W4A8 Marlin GEMM、int8 per-token-head KV cache、手写 PTX 整数注意力解码内核、int8 TP all-reduce 以及 int4 lm_head 等，目标是让张量从 RMSNorm 到注意力全程不离开整数域。仓库基于特定机器（2× RTX 3090、TP=2、一个特定量化检查点）调优，附带大量实测对比数据与未解决问题记录，处于活跃开发状态，并非成品。

这是基于来源内容整理的摘要。原文链接失效后，本页摘要和讨论仍会保留。