Rust 跑 Gemma 4：速度减半，内存却省下一半

核心洞察

这篇文章最有意思的地方在于，它证明了一件事：你可以在完全不碰 JAX 的情况下，用 Rust 写的模型跑在 TPU 上。速度只有 Python 编译版的一半，但内存占用不到一半。如果你正在纠结要不要为了 TPU 支持把整个推理栈迁到 JAX，这篇实测数据值得看一眼。

---

本文是一份手把手教程：在 Colab 的 v5e-1 运行时（单颗 TPU v5e 芯片）上，用 Rust 写的模型配合 TorchTPUPyTorch的TPU后端（torch_tpu），允许PyTorch模型在TPU上运行。 跑 Gemma 4 E2BGoogle 开源的小型大语言模型，本教程中被部署为推理服务的模型本体，权重挂载路径为 /mnt/models/gemma-4-E2B-it。，每一步运算都落在 TPU 上。一个 Rust crate 负责读取 checkpoint、解包 4-bit 权重、跑 35 层网络和环形缓存，并驱动解码循环。Python 负责托管后端并启动引擎。全程不涉及 JAX。

代码在这里：https://github.com/xbill9/torchtpu-dev/tree/main/colab/rust

Rust 写出的 token 和编译版 Python 完全一致，流式速度大约是后者的一半，峰值内存不到后者的一半。 同一颗芯片、同一个 checkpoint、同一个 19-token 提示词、64 个贪心 token、33,024 token 上下文：

两种 Rust 方式在全部 64 个 token 和三个长提示词答案上都和编译版 Python 模型一致。

---

核心结论

1. 在单颗 Colab v5e-1 TPU 芯片上，用 Rust 编写的 Gemma 4 E2B 模型（经 TorchTPU 后端）生成的 64 个 token 与编译版 Python 模型完全一致，长提示词答案也 3/3 一致，全程不涉及 JAX。

2. 性能对比：Rust 延迟融合模式流式速度约 14.7–15.7 tok/s，是 Python torch.compile 版（33.4–33.5 tok/s）的约 0.45 倍；但 Rust 峰值内存仅 7.75 GB，不到编译版 16.91 GB 的一半。

3. 设置环境变量 TORCH_TPU_DEFER_AND_FUSE=1 可将 Rust 模式从严格 eager 的约 4.8 tok/s 提升至约 15 tok/s 流式速度（约 3.2 倍），且 token 输出完全不变，crate 代码无需修改。

4. tch 0.26 只能针对 PyTorch 2.13.0 构建（2.14.1 上编译报错），而编译版 Python 环形缓存需要 PyTorch 2.14.1；在 2.13.0 上运行 torch.compile 会把模型切成 65 个程序并导致从第 21 个 token 起回复偏离。

5. Rust 引擎在 CPU 上与 Hugging Face Transformers 对拍 64/64 token 一致，KL 散度仅 1.9e-5，验证了权重解包与模型实现的正确性。

到底在测什么

TorchTPU（torch_tpu）是 PyTorch 的 TPU 后端。这里要回答的问题是：一个模型能从 Python 里搬出去多远，还能继续在它上面跑。Rust 方式跑在 PyTorch 2.13.0 上，Python 基线跑在 2.14.1 上。notebook 会在两者之间切换版本，原因在第 3 步会讲。

模型是 Gemma 4 E2B 的 W4A16 重打包版本 xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text，逐层表来自它的 -text-emb4 孪生版本。Python 基线来自 notebook 05，环形缓存在 PyTorch on TPU: Hands On With Colab 这篇文章里有介绍。

Colab 镜像上的版本：Python 3.13.18、torch_tpu 0.2.0、libtpu 0.0.50、Transformers 5.20.2、compressed-tensors 0.19.0、rustc 1.99.1。每种方式都在全新的 Python 进程里跑，不会继承别的进程编译好的程序。

---

第 1 步：跑 notebook

在 Colab 里打开 notebook 06，选 运行时 → 更改运行时类型 → v5e-1 TPU，然后 运行时 → 全部运行。不需要 Hugging Face token。notebook 会在虚拟机上装 Rust 并编译 crate：

Finished `release` profile [optimized] target(s) in 1m 44s
rustc 1.99.1 (58cf8f9c3 2026-10-13)
Rust installed and crate built in 122 s

从终端跑的话，Colab CLI 会在真实的 v5e-1 上执行并保留执行后的副本：

git clone https://github.com/xbill9/torchtpu-dev
cd torchtpu-dev/colab
make setup
python3 tools/verify_on_tpu.py 06_e2b_rust_on_torchtpu

OK: 06_e2b_rust_on_torchtpu ran clean on v5e1 -> runs/2026-11-01/06_e2b_rust_on_torchtpu_output.ipynb

1/1 notebooks verified

---

第 2 步：Rust 怎么碰到芯片

TorchTPU 0.2.0 没有提供受支持的 C++ API，而且它链接的是 libtorch_python，所以必须由 Python 托管进程。导入 torch_tpu 会把 TorchTPU 的 kernel 注册到 libtorch 的 dispatcher 上。tch crate 调用的就是同一个进程里加载的那份 libtorch，所以 Rust 对 tpu tensor 的调用会落到 TorchTPU 的 kernel 上。

桥接层是一个 PyO3Rust与Python互操作的绑定库，用于把Rust crate编译成原生Python模块torchtpu_rs。 扩展。PyO3 0.30 把 crate 构建成原生 Python 模块 torchtpu_rs，把 Rust 函数和结构体变成 Python 函数和类。tch 的 python-extension feature 负责在 Python torch.Tensor 和 tch::Tensor 之间转换，靠的是两个小助手函数 take 和 give。一个函数解包 Python tensor、跑 ATen 调用、再把结果交回去：

#[pyfunction]
fn matmul_sum(py: Python<'_>, x: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
    let x = take(x)?;
    give(py, x.matmul(&x).sum(Kind::Float))
}

在芯片上跑：

Rust matmul then sum: 64.0 on tpu:0
built in Rust, moved: [0.0, 1.0, 2.0, 3.0, 4.0] on tpu:0
bf16 kept: torch.bfloat16

结果留在 tpu:0 上，bf16 精度也没丢。

---

提示：tch 没有 TPU 设备

tch::Device 里没有 tpu 这个变体。type_as 会跟随已经在芯片上的 tensor，所以 crate 在主机上构建，再靠"照葫芦画瓢"的方式搬到设备上：

/// 不点名设备地从主机搬到设备：tch 没有 `tpu`，但 `type_as` 会跟随 tensor。
fn on(t: &Tensor, like: &Tensor) -> Tensor {
    t.type_as(&like.to_kind(t.kind()))
}

Python 传进来一个 torch.zeros(1, device="tpu")，之后所有 tensor 都跟着它走。同理，别对这些 tensor 调 Tensor::device()。

---

第 3 步：把 PyTorch 钉在 tch 生成的版本上

tch 0.26 是针对 PyTorch 2.13.0 生成的，它的绑定在 2.14.1 上编译不过：

'qr' is not a member of 'torch'

所以 crate 要对 2.13.0 构建（TorchTPU 0.2.0 支持这个版本），用 LIBTORCH_USE_PYTORCH=1 让构建指向已安装的 PyTorch：

[dependencies]
tch = { version = "0.26.0", features = ["python-extension"] }
torch-sys = { version = "0.26.0", features = ["python-extension"] }
pyo3 = { version = "0.30", features = ["extension-module"] }
safetensors = "0.8.0"
memmap2 = "0.9.12"

---

第 4 步：Rust 里的模型

crate 大约 580 行。weights.rs 用内存映射读 safetensors 文件，把 compressed-tensors用于压缩和量化模型权重的库，本教程用其4-bit格式并在Rust中解包为bf16。 的 4-bit 格式解包成 bf16。gemma.rs 装模型：35 层，其中 15 层有自己的 key 和 value，20 层复用它们，和 Python 环形缓存的布局一样。

环形缓存的 attention 会把空槽位和窗口外的槽位都 mask 掉，所以每一步在每个位置都是同一个程序：

fn attend(&self, q: &Tensor, pos: &Tensor) -> Tensor {
    let scores = q.matmul(&self.k.tr()).to_kind(Kind::Float);
    let (kpos, qpos) = (self.kpos.unsqueeze(0), pos.unsqueeze(1));
    let mut keep = kpos.ge(0).logical_and(&kpos.less_equal_tensor(&qpos));
    if let Some(w) = self.window {
        keep = keep.logical_and(&kpos.greater_tensor(&(&qpos - w)));
    }
    let weights = scores.masked_fill(&keep.logical_not().unsqueeze(1), f64::NEG_INFINITY);
    // softmax，然后乘 value

Python 用一个目录和一个芯片上的 tensor 启动它：

import torch_tpu
import torchtpu_rs

engine = torchtpu_rs.Engine(MODEL_DIR, EMB4_DIR, torch.zeros(1, device="tpu"), 33024)
tokens, read_s, decode_s = engine.generate(ids, 64, 1)

最后一个参数控制 token 多久回一次主机：1 是流式，8 是每八个一批读回。

PyO3 让 Engine 变成一个 Python 类。Rust 模型藏在 #[pyclass] 后面，方法收发的是普通 Python 列表和数字，所以 token id 过边界时根本不涉及 tensor：

#[pyclass(unsendable)]
struct Engine(gemma::Gemma);

#[pymethods]
impl Engine {
    // new、layout 和 logits 省略
    #[pyo3(signature = (ids, max_new, read_every=1))]
    fn generate(&mut self, ids: Vec<i64>, max_new: i64, read_every: i64) -> (Vec<i64>, f64, f64) {
        tch::no_grad(|| self.0.generate(&ids, max_new, read_every))
    }
}

#[pymodule]
fn torchtpu_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // 两个测试函数省略
    m.add_class::<Engine>()?;
    Ok(())
}

unsendable 让 PyO3 拒绝从创建引擎之外的线程调用它。

---

第 5 步：和 Transformers 对答案

在花任何 TPU 时间之前，Rust 引擎先在 CPU 上和 Hugging Face Transformers 并排跑，PyTorch 2.13.0，用真实 checkpoint：

load 6.8 (3, 12, 20, 1024)
top5 rust [818, 102500, 4420, 16520, 18318] [25.75, 18.875, 18.75, 18.5, 18.125]
top5 hf   [818, 102500, 4420, 16520, 18318] [25.625, 18.75, 18.625, 18.25, 18.0]
max abs diff 0.375 KL(hf||rs) 1.8714923498919234e-05
same 64 of 64
first diff None

前五名候选完全一样，KL 散度 1.9e-5，64 个 token 全对。布局那行读出来是 3 个全注意力环、12 个滑动环、20 个共享层。

---

第 6 步：芯片上的三种方式

notebook 06 的最终运行：

way                       load s  first s  graphs  breaks  streamed  in batches  peak GB
Rust fused, torch 2.13      23.1     83.5       -       -      15.1        15.2     7.75
Rust eager, torch 2.13      23.8     22.8       -       -       4.8         4.8     8.62
Python compile, 2.14.1      25.9    101.1       2       0      33.5        45.1    16.91
Rust fused, torch 2.13: first 64 of 64 tokens identical to Python compile on 2.14.1
Rust eager, torch 2.13: first 64 of 64 tokens identical to Python compile on 2.14.1

延迟融合是 Rust 的快模式。 torch.compile 追踪的是 Python，看不见 Rust。TORCH_TPU_DEFER_AND_FUSE=1 在 ATen 之下工作，crate 一行都不用改。它把 Rust 从 4.8 tok/s 拉到 15.1 tok/s 流式，大约 3.2 倍，token 完全一样。这次运行里 Rust 融合版的流式速度是编译版 Python 单步的 0.45 倍，批量 Python 达到 Rust 的 2.97 倍（算术在 evidence/arithmetic.txt 里）。

编译让 Python 付出更多内存代价：16.91 GB 峰值对 Rust 融合版的 7.75 GB，是 2.18 倍，已经接近芯片的 16 GiB。

---

提示：这个模型别在 2.13 上跑 torch.compile

编译版 Python 环形缓存需要 2.14.1。在 2.13.0 上同样的源码行为不一样：

Python compile, 2.13        25.6     34.6      65      64      15.9        16.2    11.26
Python compile, 2.13: first 21 of 64 tokens identical to Python compile on 2.14.1

它被切成 65 个程序、64 次图断裂，从第 21 个 token 开始回复就不一样了。换 Transformers 和 compressed-tensors 版本没用，换 PyTorch 版本有用。原因没测。2.13.0 上的 Rust 方式不受影响，和 2.14.1 一致。

---

第 7 步：长提示词

notebook 05 的提示词用编号记录填满上下文，每条记录带一个随机码，然后问其中一条：

way                       prompt  read s read tok/s  decode   answer  found
Rust fused, torch 2.13        61     0.1        776    14.6   915965    yes
Rust fused, torch 2.13      4081     1.0       4172    14.8   610395    yes
Rust fused, torch 2.13     32761     8.0       4116    15.6   319550     no
Rust eager, torch 2.13     32761    28.5       1148     4.8   319550     no
Python compile, 2.14.1     32761     5.0       6555    45.1   319550     no

Rust 融合版读 32,761 token 的提示词用了 7.6 到 8.0 秒，编译版环形缓存 5.0 秒，Rust eager 大约 29 秒。三种方式给出同样的三个答案。第三个在所有方式里都是错的，和 Python 环形缓存给出的码一样，所以这个限制属于模型在这个深度上的能力。解码速度在每个深度都保持稳定，因为每一步都要对同样大小的缓存的所有槽位做 attention。

---

对比一览

---

那到底选哪个

要速度，留在这颗芯片上的编译版 Python 环形缓存，PyTorch 2.14.1：批量 45 tok/s，流式 33 tok/s。

要 Rust 写模型，设 TORCH_TPU_DEFER_AND_FUSE=1：流式大约 15 tok/s，token 和答案都一样，峰值内存不到一半。TorchTPU 把这个变量标为实验性，可能会变。

严格 eager 留着调试用，大约 5 tok/s。

---

小结

这篇文章的目标是在 TPU 上跑 Gemma 4 E2B，模型用 Rust 写，走 TorchTPU，不碰 JAX。方案的关键是 tch 和 TorchTPU 注册 kernel 的那份 libtorch 共享同一个进程，所以 Rust 对 tpu tensor 的 ATen 调用能落到 TorchTPU 的 kernel 上，PyO3 再把 Rust 引擎作为模块暴露给 Python。Colab v5e-1 上的结果：

• Rust 在 64/64 个 token 和 3/3 个长提示词答案上和编译版 Python 一致
• Rust 在 CPU 上和 Transformers 64/64 个 token 一致，KL 散度 1.9e-5
• 一个环境变量把 Rust 从大约 5 tok/s 拉到大约 15 tok/s 流式
• Rust 峰值 7.75 GB，编译版 16.91 GB
• Rust 融合版流式速度是编译版 Python 的 0.45 倍，因为 torch.compile 追踪不了 Rust
• tch 0.26 只能对 PyTorch 2.13.0 构建
• Python 仍然托管进程、构建提示词、做分词
• PyTorch 2.13.0 上的 torch.compile 把这个模型切成 65 个程序并改变了回复

范围说明：只测了 Gemma 4 E2B-it QAT，纯文本，每次运行一颗 Colab v5e-1，日期 2026-11-01，贪心解码，一次一条序列，四个提示词，33,024 token 上下文。Rust 的区间覆盖三次运行，编译版 Python 覆盖两次。最早的 Rust 运行用 Transformers 5.21.0 和 compressed-tensors 0.20.1，其余用 5.20.2 和 0.19.0。Rust 方式跑在 PyTorch 2.13.0 上，Python 基线跑在 2.14.1 上，差一个版本。131,328 token 上下文、服务端点、Rust Colab kernel、以及 Rust 融合版和编译版之间时间花在哪，都没跑。最后的 crate 格式化提交在 CPU 上构建并检查过，TPU 运行用的是它之前的源码。

在 TorchTPU 上跑 Rust 模型的策略，是通过一步一步增量验证出来的。

---

参考

• Rust crate：https://github.com/xbill9/torchtpu-dev/tree/main/colab/rust
• Notebook 06：https://github.com/xbill9/torchtpu-dev/blob/main/colab/notebooks/06_e2b_rust_on_torchtpu.ipynb
• 执行记录：https://github.com/xbill9/torchtpu-dev/tree/main/colab/runs/2026-11-01
• PyTorch on TPU: Hands On With Colab：https://dev.to/gde/pytorch-on-tpu-hands-on-with-colab-5b5g
• tch：https://github.com/LaurentMazare/tch-rs
• PyO3：https://github.com/PyO3/pyo3
• W4A16 重打包：https://huggingface.co/xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text
• 4-bit 表重打包：https://huggingface.co/xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text-emb4
• Colab CLI：https://github.com/googlecolab/google-colab-cli

常见问题（FAQ）

不用JAX也能在TPU上跑模型吗？

可以。本文实测用Rust写模型配合TorchTPU，在单颗TPU v5e上跑Gemma 4 E2B，全程不涉及JAX，Python只负责托管后端和启动引擎，所有运算都落在TPU上。

Rust版和Python编译版速度内存差多少？

同一颗芯片、同一checkpoint下，Rust流式速度约14.7到15.7 tok/s，约为Python编译版33.4 tok/s的一半；但峰值内存仅7.75 GB，不到Python版16.91 GB的一半。

Rust版生成的token和Python一致吗？

一致。两种Rust方式在全部64个贪心token和三个长提示词答案上都与编译版Python模型完全相同，说明把模型搬出Python后输出质量没有损失。