Running Gemma 4 E2B on a TPU with TorchTPU

Gemma 4 E2B runs on a single v5e-1 TPU chip in a Colab notebook and initially produces only five tokens per second when the Hugging Face generate API is used. Adding a compiled decode step raises the throughput to fifty tokens per second, and a ring-cache implementation makes the full 131 072-token context feasible on the same hardware. The notebooks also compare these results with a JAX engine running on the identical chip, showing where TorchTPU narrows the performance gap.

The article walks through five Colab notebooks, each building on the previous one. All notebooks execute on a v5e-1 runtime (one TPU v5e chip with 16 GB of high-bandwidth memory). The model checkpoint is the QAT-trained Gemma 4 E2B repack stored at xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct-text.

The JAX reference engine lives at xbill9/tpu-jax (commit 4b9f8e9). Versions on the Colab image are torch 2.14.1+cpu, torch_tpu 0.2.0, libtpu 0.0.50, jax 0.11.2, compressed-tensors 0.19.0, and Transformers 5.20.1/5.20.2. Every notebook runs in a fresh Python process, so compiled programs do not carry over between runs.

Readers will come away with a clear picture of three things:

1. How TorchTPU’s execution modes differ in raw speed and compile cost.
2. Why a single compiled decode step is the most effective way to boost token-per-second throughput.
3. How a ring-cache lets the model use its maximum context length without losing the speed gains.

TorchTPU modes and raw performance

TorchTPU offers three ways to run PyTorch code on a TPU: strict eager (each operation executes as soon as it is called), defer-and-fuse (operations are collected and compiled in groups with TORCH_TPU_DEFER_AND_FUSE=1), and full compile (the entire function is compiled via torch.compile(backend="tpu")).

Notebook 01 measures a 275 GFLOP stack of matrix multiplies in each mode. The step times are 0.15 ms for eager, 2.69 ms for fuse (1.72 × faster than eager), and 3.00 ms for compile (1.52 × faster than eager).

All three produce identical checksums, confirming numerical parity. On the same tiny operation stack, compile is 98.69 × faster than strict eager and defer-and-fuse is 1.62 × faster. The first call to a compiled function costs between 115 ms and 341 ms; subsequent calls drop to about 1 ms, so every new input shape that appears in a decode loop forces a recompilation, which can dominate latency if shapes vary frequently.

Compiled decode step boosts token throughput

Notebook 02 loads the repack-formatted checkpoint and runs generation with the standard Hugging Face pipeline. The baseline speed is five tokens per second for a 19-token prompt followed by 64 greedy tokens. Notebook 03 adds a single compiled decode step by wrapping the generation loop with torch.compile(backend="tpu"). The same prompt now yields fifty tokens per second, a ten-fold increase.

Notebooks 03 through 05 all produce the same 56-token answer, confirming that the speed gain does not alter the model’s output. The performance jump comes from moving the decode loop out of strict eager mode; compiling the loop fuses many small TPU kernels into a single program that runs entirely on the device, eliminating Python overhead and reducing host-to-device traffic.

Ring cache enables the full 131 072-token context

Even with the compiled decode step, the model remains limited to the default context window. Notebook 05 introduces a ring-cache that streams token embeddings across the TPU’s 16 GB HBM. The cache works by keeping a circular buffer of past key/value pairs; when the buffer fills, the oldest entries are overwritten, allowing the model to attend to up to 131 072 tokens, the maximum supported by the checkpoint.

Running the same 19-token prompt and 64 greedy tokens with the ring-cache reproduces the full answer while preserving the fifty-tokens-per-second speed observed after compilation. The author emphasizes that the ring-cache does not change the numerical result; it merely extends the amount of context the model can consider.

How the PyTorch results compare with JAX

The JAX engine (xbill9/tpu-jax) serves as a performance reference. On the same chip and checkpoint, JAX is 14.8 × faster than the best pure PyTorch approach in notebook 03 and 2.6 × faster than the compiled decode step when the 4-bit weight table is used in notebook 04.

Despite the gap, the combination of compilation and ring-cache brings the PyTorch implementation close to JAX in functional capability: both can now handle the full 131 072-token context and produce identical answers for the tested prompts.

Practical steps to reproduce the results

1. Clone the repository and open the notebooks from the colab/README.md badge.
2. Change the runtime type to v5e-1 TPU (Runtime → Change runtime type).
3. Run all cells in each notebook sequentially; no Hugging Face token is required.
4. For the compiled decode step, wrap the generation loop with torch.compile(backend="tpu").
5. To enable the ring-cache, import the helper from notebook 05 and activate the circular buffer before the decode loop.

Running the notebooks from the command line is also possible via the Colab CLI, which verifies that each notebook executes cleanly on a real v5e-1 chip.

The takeaway

A single compiled decode step is the most effective lever for raising raw PyTorch generation from five to fifty tokens per second on a v5e-1 TPU. Adding a ring-cache then makes the model’s full 131 072-token context available without sacrificing that throughput. Use torch.compile for the decode loop and enable the ring-cache to extract the maximum performance from Gemma 4 E2B on TorchTPU.