Gemma 4 in Rust on TorchTPU: A Native TPU Rig With PyO3 and MCP

8.5 relevance

Running Gemma 4 on TPU with Rust and PyO3 is cutting-edge AI infrastructure work directly relevant to the reader.

AI/ML dev.to

Answer the board on AI readiness

Arrive with evidence, not opinion. See how your foundation benchmarks against your peers, get a clear read, and leave with a ranked place to start.

Why It Matters

For a platform engineer focused on AI/ML agent orchestration and cloud infrastructure, this demonstrates a practical pattern for running custom inference engines on TPUs with Rust's performance benefits, using MCP servers for provisioning and PyO3 for Python-Rust interop, which could inform your own agent-based deployment pipelines.

Key Takeaways

Evaluate using Rust with a minimal ATen C++ shim (instead of tch) to build TPU-native inference engines that boot fast and integrate with Python orchestration via PyO3.

Gemma 4 in Rust on TorchTPU: A Native TPU Rig With PyO3 and MCP

Summary

A developer built a native TPU inference rig for Gemma 4 E2B using Rust on a Cloud TPU v5e-1 VM, with all 35 layers, ring cache, and decode loop written in Rust and reaching TorchTPU through a 154-line C++ shim over ATen instead of the tch crate. The setup boots in 4 minutes 29 seconds from Python 3.14 to serving an OpenAI-compatible endpoint, decodes at a median 12.20 tokens per second, and matches Hugging Face Transformers on 64 of 64 greedy tokens on CPU. A Claude Code MCP server provisions the TPU, and the VM builds the Rust crate as a PyO3 extension on boot, with Python hosting the HTTP server and process.

Author

xbill