pyo3

Gemma 4 in Rust on TorchTPU: A Native TPU Rig With PyO3 and MCP --- title: "Gemma 4 in Rust on TorchTPU: A Native TPU Rig With PyO3 and MCP" published: true description: "Gemma 4 E2B with the model in Rust, served from a TPU v5e-1 VM through PyO3 and a C++ shim over ATen, with no tch and no pinned PyTorch, an OpenAI endpoint, an MCP server and a local chat UI." tags: rust, pytorch, tpu, mcp cover_image: raw.githubusercontent.com --- This article provides a step by step guide to serving Gemma 4 E2B from a Cloud TPU v5e-1 VM with the model written in Rust and every operation run by TorchTPU. A Claude Code MCP server provisions the TPU, the VM builds the Rust crate on boot and serves an OpenAI-compatible endpoint, and a local shim gives it a chat page and a second MCP server for demos. Python hosts the process and the HTTP server; the 35 layers, the ring cache and the decode loop are Rust. github.com/xbill9/torchtpu-dev/tree/main/tpu-pyto… The Rust model reaches TorchTPU through a 154-line C++ shim over ATen, so it builds against whichever PyTorch is installed: here Python 3.14, PyTorch 2.14.1 and a nightly Rust. It matches Hugging Face Transformers on 64 of 64 greedy tokens on CPU, boots on a fresh TPU VM in 4 minutes 29 seconds from installing Python to starting the service, and decodes at a median 12.20 tokens per second on the TPU. !The local chat page streaming from the Rust model on the TPU, with the build and versions in the side panel --- What Is Being Built? Gemma 4 in Rust on TorchTPU: Hands On With Colab ran the same Rust engine in a Colab notebook. This rig takes it to a TPU VM of its own and keeps it there: | Layer | Where | What | | :--- | :--- | :--- | | Capacity | Claude Code | An MCP server that creates a v5e-1 queued resource and reads its boot stages | | Model | TPU VM | torchtpu_rs, a Rust crate built on the VM as a PyO3 extension | | Kernels | TPU VM | ATen calls through PyTorch's dispatcher to TorchTPU | | API | TPU VM, port 8000 | rust_openai_server.py, OpenAI-compatible, streaming | | Demo | This laptop | A shim on 127.0.0.1:8015 with a chat page, and a demo MCP server | The checkpoint is the W4A16 repack xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text with the int4 vocabulary tables from its -text-emb4 twin. The directory name follows the rig naming scheme: tpu (Cloud TPU API), pytorchrust (the forward pass is Rust over PyTorch's ATen), v5e1, 2b, and q4w4a16emb4 for the encoding. Versions on the VM, read from its /v1/engine: Python 3.14.10, PyTorch 2.14.1, torch_tpu 0.2.0, libtpu 0.0.51, Transformers 5.20.2, rustc 1.101.0-nightly. --- Step 1 — Install the Rig The rig is a Claude Code plugin with its own MCP server and skill: git clone https://github.com/xbill9/torchtpu-dev cd torchtpu-dev/tpu-pytorchrust-v5e1-2b-q4w4a16emb4 ./project-setup.sh . --project <gcp-project-id> Every setting lives in tpu.env: the zone, the two checkpoints, the Python and PyTorch versions the VM installs, the cache length and the firewall tags. PyTorch and TorchTPU come from public PyPI: PYTHON_BIN=python3.14 TORCH_PIP_SPEC=torch==2.14.1 torch_tpu==0.2.0 transformers==5.20.2 compressed-tensors==0.19.0 'jinja2>=3.1' huggingface_hub fastapi uvicorn SERVE_CONTEXT=8192 SERVE_DEFER_AND_FUSE=1 Python 3.14 is the newest version torch_tpu 0.2.0 ships wheels for. The Rust toolchain is named in the crate's rust-toolchain.toml, channel = "nightly", and the crate is on edition 2024. --- Step 2 — Reach TorchTPU Without tch TorchTPU registers its kernels on libtorch's dispatcher and ships no C++ API of its own, so Rust reaches it through ATen. The tch crate wraps every ATen operation in generated C++ that compiles against PyTorch 2.13.0 only. The model uses a few dozen operations, so the crate wraps those itself, in rust/src/aten.cpp: #define RET(expr) TRY(return box(expr);) return nullptr; extern "C" { Tensor* aten_matmul(Tensor* a, Tensor* b) { RET(at::matmul(*a, *b)) } Tensor* aten_linear(Tensor* x, Tensor* w) { RET(at::linear(*x, *w)) } Tensor* aten_softmax(Tensor* a, int64_t d, int k) { RET(at::softmax(*a, d, kind(k))) } Tensor* aten_index_copy_(Tensor* t, int64_t d, Tensor* idx, Tensor* src) { RET(t->index_copy_(d, *idx, *src)) } // torch.Tensor <-> at::Tensor, through libtorch_python. Both directions share storage. Tensor* aten_from_python(PyObject* o) { RET(THPVariable_Unpack(o)) } PyObject* aten_to_python(Tensor* t) { TRY(return THPVariable_Wrap(*t);) return nullptr; } } rust/src/aten.rs gives them tch's method names, so the model code reads as before. build.rs asks the installed torch for its headers, libraries, C++ ABI and version, and compiles the shim against them with C++20, which PyTorch 2.14's headers require. The module records what built it: >>> torchtpu_rs.BUILD {'torch': '2.14.1+cpu', 'rustc': 'rustc 1.101.0-nightly (3e92e1a2d 2026-11-01)', 'python': '3.14.7', 'crate': '0.1.0'} A module built for one PyTorch refuses to load under another, because the C++ ABI differs: import torchtpu_rs compares torch.__version__ with the version it was built against and raises ImportError on a mismatch. --- Step 3 — The PyO3 Bindings Every PyO3 item is in rust/src/bindings.rs; the model and weight code never see a Python object. rust/PYO3.md lists each crossing, and rust/python/torchtpu_rs.pyi is the type stub. Four rules shape the code. Token ids are lists, tensors share storage. decode(ids: list[int], ...) returns list[int]; generated tokens never cross as tensors. The one tensor that crosses on every load is like, a torch.zeros(1, device="tpu") that names the device: Rust builds each weight on the host and moves it with type_as(like). Long calls release the GIL; the callback takes it back. The server streams through a callback that runs once per token batch: let done = py.detach(|| { aten::no_grad(|| { gemma.decode(&ids, max_new, &stop, read_every, |batch| { let Some(cb) = &on_tokens else { return true }; Python::attach(|py| { let call = PyList::new(py, batch).and_then(|l| cb.call1(py, (l,))); match call.and_then(|r| r.is_truthy(py)) { Ok(go_on) => go_on, Err(e) => { failed = Some(e); false } } }) }) }) }); A false return ends the decode with finish="cancelled", and an exception in the callback is raised from decode. Generation stops at the checkpoint's eos_token_id list, [1, 106, 50]. The engine is #[pyclass(unsendable)]. It owns TPU tensors and one ring cache, so PyO3 checks the calling thread on every access. A call from another thread raises pyo3_runtime.PanicException, which derives from BaseException. The server builds the engine on one dedicated thread and queues every request to it: class EngineWorker(threading.Thread): """Owns the one `torchtpu_rs.Engine`. Nothing outside this thread may touch it.""" Errors a caller can cause are ValueError. An empty prompt, max_new below 1, or a prompt longer than the cache are checked in bindings.rs before the model's own asserts. tests/test_rust_bindings.py runs all of these on CPU against the real checkpoint: test_bad_input_is_value_error ... ok test_callback_exception_propagates ... ok test_cancel_and_length ... ok test_decode_stops_and_streams ... ok test_gil_released_during_decode ... ok test_module_surface ... ok test_tensor_functions ... ok test_unsendable ... ok Ran 8 tests in 247.568s OK --- Step 4 — Check It Against Transformers Before any TPU time, the Rust module and Hugging Face Transformers 5.20.2 run the same prompt on CPU, both on PyTorch 2.14.1, one process per side: top5 rust [818, 8409, 102500, 4420, 16520] hf [818, 102500, 8409, 4420, 16520] KL(hf||rs) 1.078e-05; same 64 of 64; first diff None The five top candidates are the same set, the KL divergence of the first token's distribution is 1.1e-5, and all 64 greedy tokens match. --- Step 5 — Provision the TPU From Claude Code, the rig's MCP server creates a flex-start v5e-1 queued resource in us-west4-a. The startup script travels in the instance metadata with the crate and the server packed inside it, 47 KiB rendered, so the VM needs no SSH to get its code: > create_tpu_queued_resource 🚀 Queued Resource gemma4-rust-qr creation initiated in us-west4-a, serving `xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text-emb4`. A TPU VM has no serial-console command, so the startup script writes each stage to a guest attribute and get_rig_progress reads it back without SSH. The boot on 2026-11-02: STAGE: python g++ (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0 STAGE: torch torch 2.14.1+cpu STAGE: rust STAGE: crate rustc 1.101.0-nightly (3e92e1a2d 2026-11-01) Finished `release` profile [optimized] target(s) in 21.03s STAGE: weights STAGE: smoke torchtpu_rs 0.1.0 matmul_sum 512.0 on tpu:0 TPU smoke OK STAGE: unit STAGE: serving From installing Python to starting the service took 4 minutes 29 seconds. The smoke test is a Rust call on a tensor on the TPU, through the shim, checked for the right value and a result on tpu:0. The service then loads the weights in 6.8 seconds and runs a warmup decode in 18.8 seconds before /health answers. --- 🔎 Tip: Pin jinja2>=3.1 on a TPU VM The TPU image's own jinja2 3.0.3 sits in /usr/lib/python3/dist-packages, which is on the path of the deadsnakes Python. Transformers' chat templates need 3.1 or later and refuse to render with the older one: ImportError: apply_chat_template requires jinja2>=3.1.0 to be installed. Your version is 3.0.3. 'jinja2>=3.1' in the pip spec puts a newer copy in /usr/local/lib/python3.14/dist-packages, which comes first on the path. --- Step 6 — The OpenAI Endpoint The VM serves /v1/chat/completions and /v1/completions, streamed or not, plus /v1/models, /health and /metrics. /v1/engine reports what is loaded and how it is bound: "torch_tpu_defer_and_fuse": "1", "load_seconds": 6.83, "warmup_seconds": 18.8, "layout": {"full_rings": 3, "sliding_rings": 12, "shared_layers": 20, "context": 8192}, "engine_class": "Engine — #[pyclass(unsendable)], owned by thread torchtpu-rs-engine" Decoding is greedy. Each non-streamed response carries the engine's own seconds for reading the prompt and for decoding, and system_fingerprint reads torchtpu_rs-greedy. --- Step 7 — Measure It The same prompt and length as a vLLM run on the same chip type, Write a short paragraph about the ocean. at 100 tokens, five requests one at a time, rate from the engine's own timer: request 1: 83 tokens, finish stop, read 0.088 s, decode 6.651 s, 12.33 tok/s, client 6.83 s request 2: 83 tokens, finish stop, read 0.086 s, decode 6.720 s, 12.20 tok/s, client 6.89 s request 3: 83 tokens, finish stop, read 0.086 s, decode 6.734 s, 12.18 tok/s, client 6.91 s request 4: 83 tokens, finish stop, read 0.087 s, decode 6.762 s, 12.13 tok/s, client 6.94 s request 5: 83 tokens, finish stop, read 0.083 s, decode 6.561 s, 12.50 tok/s, client 6.74 s median decode tok/s 12.20, min 12.13, max 12.50 Streaming eight prompts through the shim from a laptop, run_rust_benchmark measured a first token in 0.175 to 0.256 seconds and 11.5 to 12.6 tokens per second as the client saw them. --- Step 8 — Demo It From a Laptop demo/install.sh installs two things on the laptop: a shim on 127.0.0.1:8015 as a systemd user service, and a demo MCP server registered with Claude Code: demo/install.sh --project <gcp-project-id> > endpoint_status 🟢 local http://127.0.0.1:8015/v1 → http://136.85.219.123:8000 - TPU /health: 200 - queued resource `gemma4-rust-qr` in us-west4-a, project aisprint-491218 The shim finds the VM's address with gcloud and looks it up again when the VM stops answering, so a new flex-start VM needs no reconfiguration. It passes every request through, streams included, so any OpenAI client works against http://127.0.0.1:8015/v1. Its root is a chat page: replies stream, the conversation carries over, and Stop cancels the decode on the TPU. A client that disconnected after 10 streamed chunks moved the server's cancel counter from 1 to 2, and the engine went idle. The demo MCP server, torchtpu-rs-demo, has five tools, all through the shim: > ask_gemma_rust "What is a TPU? Two sentences." A TPU stands for **Tensor Processing Unit**, which is a specialized type of integrated circuit designed to accelerate machine learning workloads. ... --- xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text-emb4 · finish `stop` · 17 prompt + 55 completion tokens · prompt read 0.088 s · decode 12.2 tok/s on the TPU · 4.59 s round trip from here stream_gemma_rust, engine_info, endpoint_status and serving_metrics cover streaming, the build, the shim's view of the VM and the server's counters. --- Compare and Contrast One TPU v5e chip, Gemma 4 E2B, decode tokens per second for one request at a time. Cost per million output tokens is computed from each rate and the v5e-1 flex-start list price, USD 0.60 an hour (arithmetic, evidence/arithmetic.txt): | One v5e chip | Decode tok/s | Model code | Stack | USD per M tokens | | :--- | ---: | :--- | :--- | ---: | | vLLM 0.30.0 🥇 | 144.1 | Python | vLLM TPU, bf16 checkpoint | 1.16 | | Python torch.compile, Colab 🥈 | 33.5 streamed, 45.1 batched | Python | TorchTPU, PyTorch 2.14.1 | | | Rust, tch, Colab 🥉 | 15.1 streamed | Rust | TorchTPU defer-and-fuse, PyTorch 2.13.0 | | | Rust, ATen shim, this rig | 12.20 | Rust | TorchTPU defer-and-fuse, PyTorch 2.14.1 | 13.66 | The vLLM figure is from the gcp-inference demo on 2026-10-12, read from vLLM's own metrics. vLLM decodes 11.8x faster than this rig on the same chip type. The compiled Python step decodes 2.75x faster when streaming. torch.compile traces Python and cannot see a Rust model, so the Rust ways run TorchTPU's defer-and-fuse mode, which fuses operations below ATen. This rig streams at 0.81x the Colab tch run. The two differ in PyTorch version, ATen layer, cache length (8,192 tokens against 33,024), a Python callback per streamed token, and where the timer sits; which of those accounts for the gap is not measured. --- So, Which One? For throughput on a v5e, serve with vLLM: 144.1 tokens per second on one request and USD 1.16 per million tokens. For a model written in Rust on TorchTPU, this rig is the working shape: a crate that builds against the installed PyTorch, PyO3 bindings that keep the event loop free while the TPU works, and an OpenAI endpoint any client can use. It decodes at about 12 tokens per second. For a demo, the shim gives one local address for a browser, an OpenAI client and Claude Code, whichever VM is up. --- Summary The goal of this article was to serve Gemma 4 E2B from a TPU VM with the model in Rust, through TorchTPU and without JAX, and to operate it from Claude Code. The key to the solution was a C++ shim over the few dozen ATen operations the model uses, compiled against the installed PyTorch, with PyO3 bindings that release the GIL during decoding and confine the engine to one thread. The results on one TPU v5e chip were: - 🟢 Rust matches Transformers on CPU on 64 of 64 tokens, with a KL divergence of 1.1e-5, on PyTorch 2.14.1 - 🟢 No pinned PyTorch: Python 3.14, PyTorch 2.14.1 and nightly Rust on the VM - 🟢 4 minutes 29 seconds from installing Python to starting the service, with the crate built on the VM in 21 seconds - 🟢 OpenAI endpoint, streaming, cancel on disconnect, an MCP server for the TPU and one for demos, and a local chat page - ⚠️ 12.20 tokens per second median decode, 11.8x slower than vLLM on the same chip type - ⚠️ One request at a time and greedy decoding only; no tool-call parsing - ❌ No compile path from Rust: torch.compile traces Python, and TorchTPU 0.2.0 has no C++ API Scope: Gemma 4 E2B-it QAT W4A16 with int4 vocabulary tables, text only, one flex-start v5e-1 queued resource in us-west4-a on 2026-11-02, an 8,192-token cache, greedy decoding, one request at a time. The vLLM and torch.compile figures come from other runs on the same chip type (gcp-inference on 2026-10-12, Colab notebook 06 on 2026-11-01) with other checkpoints and cache lengths. Batched host reads, a longer cache and concurrent requests were not run on this rig. The parity check and the binding tests ran on a laptop CPU. The strategy for serving a Rust model on TorchTPU was validated with an incremental step by step approach. --- References - The rig: github.com/xbill9/torchtpu-dev/tree/main/tpu-pyto… - PyO3 bindings: github.com/xbill9/torchtpu-dev/blob/main/tpu-pyto… - Demo shim, chat page and MCP server: github.com/xbill9/torchtpu-dev/tree/main/tpu-pyto… - Gemma 4 in Rust on TorchTPU, Hands On With Colab: dev.to - gcp-inference: xbill9/gcp-inference - PyO3: PyO3/pyo3 - W4A16 repack: huggingface.co - 4-bit-tables repack: huggingface.co

Gemma 4 in Rust on TorchTPU: Hands On With Colab --- title: "Gemma 4 in Rust on TorchTPU: Hands On With Colab" published: true description: "A Rust crate reads the checkpoint, runs all 35 layers of Gemma 4 E2B and drives the decode loop on one TPU v5e chip in Colab, with every operation a TorchTPU kernel and the same tokens as compiled Python." tags: rust, pytorch, tpu, gemma cover_image: raw.githubusercontent.com --- This article provides a step by step guide to running Gemma 4 E2B with the model written in Rust and every operation executed by TorchTPU, on a Colab v5e-1 runtime: one TPU v5e chip. A Rust crate reads the checkpoint, unpacks the 4-bit weights, runs the 35 layers and the ring cache, and drives the decode loop. Python hosts the backend and starts the engine. No JAX is involved. github.com/xbill9/torchtpu-dev/tree/main/colab/ru… Rust writes the same tokens as compiled Python at about half the streaming speed and less than half the peak memory. Same chip, same checkpoint, same 19-token prompt, 64 greedy tokens, 33,024-token context: | One Colab v5e-1 chip | Streamed tok/s | Batched tok/s | First call | Peak memory | | :--- | ---: | ---: | ---: | ---: | | Rust, defer-and-fuse | 14.7 to 15.7 | 15.2 to 16.0 | 82 to 84 s | 7.75 GB | | Rust, strict eager | 4.8 to 5.0 | 4.8 | 23 to 25 s | 8.62 GB | | Python torch.compile | 33.4 to 33.5 | 45.1 to 45.2 | 98 to 101 s | 16.91 GB | Both Rust ways match the compiled Python model on all 64 tokens and on all three long-prompt answers. --- What Is Being Measured? TorchTPU (torch_tpu) is a PyTorch backend for TPUs. The question here is how far a model can move out of Python and still run on it. The Rust ways run on PyTorch 2.13.0 and the Python baseline on 2.14.1. The notebook swaps versions between them, for a reason Step 3 covers. The model is Gemma 4 E2B in the W4A16 repack xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text, with the per-layer table from its -text-emb4 twin. The Python baseline comes from notebook 05, the ring cache in PyTorch on TPU: Hands On With Colab. Versions on the Colab image: 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. Each way runs in a fresh Python process, so none inherits another's compiled programs. --- Step 1 — Run the Notebook Open notebook 06 in Colab, pick Runtime → Change runtime type → v5e-1 TPU, then Runtime → Run all. No Hugging Face token is needed. The notebook installs Rust and builds the crate on the VM: 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 From a terminal, the Colab CLI runs it on a real v5e-1 and keeps the executed copy: 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 --- Step 2 — How Rust Reaches the Chip TorchTPU 0.2.0 ships no supported C++ API, and it links libtorch_python, so Python has to host the process. Importing torch_tpu registers TorchTPU's kernels on libtorch's dispatcher. The tch crate calls that same libtorch, loaded once in the same process, so a Rust call on a tpu tensor lands on the TorchTPU kernel. The bridge is a PyO3 extension. PyO3 0.30 builds the crate as a native Python module, torchtpu_rs, and turns Rust functions and structs into Python functions and classes. tch's python-extension feature converts a Python torch.Tensor to a tch::Tensor and back, in two small helpers, take and give. A function unpacks a Python tensor, runs ATen calls and hands the result back: #[pyfunction] fn matmul_sum(py: Python<'_>, x: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> { let x = take(x)?; give(py, x.matmul(&x).sum(Kind::Float)) } Run on the chip: 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 The result stays on tpu:0, and bf16 survives the trip. --- 🔎 Tip: tch Has No TPU Device tch::Device has no tpu variant. type_as follows any tensor already on the chip, so the crate builds on the host and moves by example: /// Host to device without naming the device: tch has no `tpu`, but `type_as` follows the tensor. fn on(t: &Tensor, like: &Tensor) -> Tensor { t.type_as(&like.to_kind(t.kind())) } Python passes in torch.zeros(1, device="tpu") and every tensor follows it. Avoid Tensor::device() on these tensors for the same reason. --- Step 3 — Pin PyTorch to What tch Generates tch 0.26 is generated for PyTorch 2.13.0, and its bindings fail to compile on 2.14.1: 'qr' is not a member of 'torch' So the crate builds against 2.13.0, which TorchTPU 0.2.0 supports, with LIBTORCH_USE_PYTORCH=1 pointing the build at the installed 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" --- Step 4 — The Model in Rust The crate is about 580 lines. weights.rs memory-maps the safetensors files and unpacks the compressed-tensors 4-bit format to bf16. gemma.rs holds the model: 35 layers, of which 15 own keys and values and 20 reuse them, the same layout as the Python ring. The ring's attention masks every slot that is empty or out of the window, so a step is the same program at every position: 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, then the values Python starts it with a directory and a tensor on the chip: 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) The last argument sets how often tokens come back to the host: 1 streams, 8 reads in batches of eight. PyO3 is what makes Engine a Python class. The Rust model sits behind a #[pyclass], and its methods take and return plain Python lists and numbers, so token ids cross the boundary with no tensor at all: #[pyclass(unsendable)] struct Engine(gemma::Gemma); #[pymethods] impl Engine { // new, layout and logits omitted #[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<()> { // the two test functions omitted m.add_class::<Engine>()?; Ok(()) } unsendable makes PyO3 refuse calls on the engine from any thread but the one that created it. --- Step 5 — Check It Against Transformers Before any TPU time, the Rust engine runs on CPU beside Hugging Face Transformers, on PyTorch 2.13.0 with the real 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 The top five candidates are the same, the KL divergence is 1.9e-5, and all 64 tokens match. The layout line reads 3 full-attention rings, 12 sliding rings and 20 shared layers. --- Step 6 — The Three Ways on the Chip Final run of 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 Defer-and-fuse is the fast mode for Rust. torch.compile traces Python, so it cannot see Rust. TORCH_TPU_DEFER_AND_FUSE=1 works below ATen and needs no change to the crate. It takes Rust from 4.8 to 15.1 tokens per second streamed, about 3.2x, at the same tokens. In this run Rust fused streams at 0.45x the compiled Python step and batched Python reaches 2.97x Rust (arithmetic, evidence/arithmetic.txt). Compiling costs Python more memory: its 16.91 GB peak against 7.75 GB for Rust fused is 2.18x, close to the chip's 16 GiB. --- 🔎 Tip: Don't Run torch.compile on 2.13 For This Model The compiled Python ring needs 2.14.1. On 2.13.0 the same source behaves differently: 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 It splits into 65 programs with 64 graph breaks and writes a different reply from token 21. Changing Transformers and compressed-tensors versions made no difference, and changing the PyTorch version did. Its cause is untested. The Rust ways on 2.13.0 are unaffected and match 2.14.1. --- Step 7 — Long Prompts Notebook 05's prompts fill the context with numbered records, each with a random code, and ask about one record: 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 fused reads the 32,761-token prompt in 7.6 to 8.0 s, the compiled ring in 5.0 s and Rust eager in about 29 s. All three ways give the same three answers. The third is wrong in every way, the same code the Python ring gives, so the limit belongs to the model at this depth. Decode speed holds steady at every depth, because each step attends over every slot of a cache of the same size. --- Compare and Contrast | One v5e chip, Gemma 4 E2B | Streamed tok/s | Model code | PyTorch | Peak memory | | :--- | ---: | :--- | :---: | ---: | | Python torch.compile 🥇 | 33.4 to 33.5 | Python | 2.14.1 | 16.91 GB | | Rust, defer-and-fuse 🥈 | 14.7 to 15.7 | Rust | 2.13.0 | 7.75 GB | | Rust, strict eager 🥉 | 4.8 to 5.0 | Rust | 2.13.0 | 8.62 GB | --- So, Which One? For speed on this chip, stay with the compiled Python ring on PyTorch 2.14.1: 45 tokens per second in batches and 33 streamed. For a model in Rust, set TORCH_TPU_DEFER_AND_FUSE=1: about 15 tokens per second streamed, the same tokens and answers, and less than half the peak memory. TorchTPU marks that variable experimental, so it may change. Keep strict eager for debugging, at about 5 tokens per second. --- Summary The goal of this article was to run Gemma 4 E2B on a TPU with the model written in Rust, through TorchTPU and without JAX. The key to the solution was tch sharing the one libtorch that TorchTPU registers its kernels with, so a Rust ATen call on a tpu tensor reaches a TorchTPU kernel, with PyO3 exposing the Rust engine to Python as a module. The Colab v5e-1 results were: - 🟢 Rust matches compiled Python on 64 of 64 tokens and 3 of 3 long-prompt answers - 🟢 Rust matches Transformers on CPU on 64 of 64 tokens, with a KL divergence of 1.9e-5 - 🟢 One environment variable takes Rust from about 5 to about 15 tokens per second streamed - 🟢 Rust peaks at 7.75 GB against 16.91 GB for compile - ⚠️ Rust fused streams at 0.45x compiled Python, because torch.compile cannot trace Rust - ⚠️ tch 0.26 builds against PyTorch 2.13.0 only - ⚠️ Python still hosts the process, builds the prompts and tokenizes - ❌ torch.compile on PyTorch 2.13.0 splits this model into 65 programs and changes the reply Scope: Gemma 4 E2B-it QAT only, text only, one Colab v5e-1 runtime per run on 2026-11-01, greedy decoding, one sequence at a time, four prompts, a 33,024-token context. Rust ranges cover three runs and compiled Python ranges two. The earliest Rust run used Transformers 5.21.0 and compressed-tensors 0.20.1, and the rest 5.20.2 and 0.19.0. The Rust ways ran on PyTorch 2.13.0 and the Python baseline on 2.14.1, one version apart. The 131,328-token context, a serving endpoint, a Rust Colab kernel, and where the time goes between Rust fused and compile were not run. The final crate formatting commit was built and checked on CPU, and the TPU runs used the source before it. The strategy for running a Rust model on TorchTPU was validated with an incremental step by step approach. --- References - Rust crate: github.com/xbill9/torchtpu-dev/tree/main/colab/ru… - Notebook 06: github.com/xbill9/torchtpu-dev/blob/main/colab/no… - Executed run: github.com/xbill9/torchtpu-dev/tree/main/colab/ru… - PyTorch on TPU: Hands On With Colab: dev.to - tch: LaurentMazare/tch-rs - PyO3: PyO3/pyo3 - W4A16 repack: huggingface.co - 4-bit-tables repack: huggingface.co - Colab CLI: googlecolab/google-colab-cli

Why AI Infrastructure Is Moving to Rust ? GitHub just rewrote its Copilot agent runtime in Rust - and you are probably running Rust inside your Python AI code already. So is Rust replacing Python in AI? No: Python stays the language you drive AI with, while Rust replaces the C and C++ engine underneath and the servers around it. In this video: why Python only gives the orders, the memory bug C++ lets through and Rust refuses to compile, four reasons Rust fits AI systems (no garbage collector, real parallel threads, memory layout control, PyO3), real Rust projects in the AI stack, the honest downsides, and who should actually learn Rust. The C++ and Rust snippets are our own examples; the PyO3 snippet is the official example from the PyO3 README (Apache-2.0 / MIT). Copilot performance figures are GitHub's own measurements, excluding model and network time. Chapters: 0:00 Intro 0:13 GitHub's Rust rewrite 0:48 Rust you already run 1:20 Who really does the math 1:42 Two different jobs 2:00 The problem with C++ 2:42 What Rust does instead 3:07 Why Rust fits AI systems 4:13 Rust inside Python 4:35 Rust around AI 4:59 The catch 5:23 Should you learn Rust? Stock footage: Video by Atakan Ozkan on Pexels; Video by Nothing Ahead on Pexels. Sources (checked 2026-10-17): github.blog huggingface/tokenizers safetensors/safetensors openai/tiktoken github.com/pydantic/pydantic/tree/main/pydantic-c… pydantic.dev pola-rs/polars astral-sh/uv astral-sh/ruff qdrant/qdrant lancedb/lancedb huggingface/candle tracel-ai/burn PyO3/pyo3 PyO3/maturin chromium.org microsoft.com docs.python.org doc.rust-lang.org doc.rust-lang.org docs.pytorch.org huggingface.co docs.lancedb.com docs.python.org raw.githubusercontent.com Independent explainers of AI creative tools. Not affiliated with, endorsed by, or sponsored by any company or tool shown. Prices and features change fast - always check the official sites. Narration uses an AI-generated voice. Script and facts are written and checked by the creator. #rust #python #ai #AItools #AIart #AIvideo