dev.to | 📄 原文链接 | 2026-11-04 收录

在 TorchTPU 上用 Rust 跑 Gemma 4：一块 TPU 上的原生推理

来源：dev.to — 2026-11-02

📋 概述

这篇实操文把 Gemma 4 E2B 搬到一台 Cloud TPU v5e-1 虚拟机上：35 层模型、环形缓存与解码循环全部用 Rust 写，经 154 行 C++ shim 调用 PyTorch 的 ATen，因此不锁死 PyTorch 版本。CPU 上与 Hugging Face 的 64 个贪心 token 全部一致，新机器从装 Python 到服务起来只要 4 分 29 秒，TPU 上解码中位数 12.20 token/秒。

🔑 核心要点

架构分四层：MCP 服务器申请 TPU、Rust crate 编译成 PyO3 扩展、ATen 下发内核、OpenAI 兼容接口对外服务。
「不用 tch、不锁 PyTorch」是关键设计，靠的是那 154 行 C++ shim。
实测版本为 Python 3.14.10、PyTorch 2.14.1、torch_tpu 0.2.0、nightly Rust。
权重是 W4A16 重打包版本，目录名把平台、语言、芯片、规模、量化方式全部编码进去。

💡 金句

“The Rust model reaches TorchTPU through a 154-line C++ shim over ATen, so it builds against whichever PyTorch is installed.”
译：Rust 模型通过一层 154 行的 C++ shim 走 ATen 接到 TorchTPU，所以装哪个版本的 PyTorch 都能编译。

← 返回 dev.to 首页