PyTorch on TPU: Hands On With Colab

A developer benchmarked running Gemma 4 E2B, Google's quantization-aware-trained 4-bit model, in PyTorch on a single Google Cloud TPU v5e chip via the TorchTPU backend, publishing reproducible Colab notebooks. A compiled decode step lifted throughput from about 5 tokens per second under Hugging Face generate to roughly 31-33 tokens per second, and a 4-bit per-layer table pushed it to about 50 tokens per second, while a ring cache enabled a 131,328-token context at about 38.7 tokens per second. A JAX engine serving the same checkpoint on the same chip remained faster, at roughly 131 tokens per second on short prompts.

This article provides a step by step guide to running Gemma 4 E2B, Google's quantization-aware trained QAT model, in PyTorch on a Google Cloud TPU with TorchTPU. Every step is a Colab notebook that runs on a Colab v5e-1 runtime, one TPU v5e chip with 16 GB of high-bandwidth memory HBM . The measures are what TorchTPU does with PyTorch code, how far Hugging Face Transformers on TorchTPU trails a JAX engine on the same chip, what closes the gap, and how much context fits. https://github.com/xbill9/torchtpu-dev/tree/main/colab https://github.com/xbill9/torchtpu-dev/tree/main/colab One compiled decode step takes PyTorch from 5 to 50 tokens per second, and a ring cache takes it to the full context. Same chip, same checkpoint, same 19-token prompt, 64 greedy tokens 32 in notebook 02 : | One Colab v5e-1 chip | Decode tok/s | Context | Notebook | |---|---|---|---| | Hugging Face generate | 5.0 to 5.3 | short prompt | 04 | | Best way before the fix defer-and-fuse | 8.8 to 9.1 | short prompt | 02, 03 | | Compiled decode step | 31.3 to 33.4 | 256 tokens | 04 | | Compiled step, 4-bit per-layer table | 50.2 to 50.4 | 256 tokens | 04 | | Ring cache | 38.6 to 38.7 | 131,328 tokens | 05 | | JAX engine, bf16 weights | 130.8 to 132.9 | short prompt | 03, 04 | JAX was 14.8x the best PyTorch way in notebook 03 and is 2.6x the compiled step with the 4-bit table in notebook 04. On that prompt every way in notebooks 03 to 05 writes the same 56-token answer. TorchTPU torch tpu is a PyTorch backend for TPUs: tensors live on a tpu device, and the same Python code runs eagerly, fused, or under torch.compile . The question for a PyTorch user is how much of the TPU a plain Transformers model gets, measured against a JAX engine that serves the same checkpoint on the same chip. The model is Gemma 4 E2B in the W4A16 repack xbill9/gemma-4-E2B-it-qat-q4 0-w4a16-ct-text : Google's QAT 4-bit weights in the compressed-tensors format. Its twin -text-emb4 stores the vocabulary tables 4-bit as well. The JAX engine is xbill9/tpu-jax at commit 4b9f8e9 . Every number below comes from an executed notebook in colab/runs/ . Versions on the Colab image: torch tpu 0.2.0, torch 2.14.1+cpu, libtpu 0.0.50, jax 0.11.2, compressed-tensors 0.19.0, Transformers 5.20.1 notebooks 02 and 03 and 5.20.2 04 and 05 . Each way runs in a fresh Python process, so none inherits another's compiled programs. Each notebook opens in Colab from its badge in the repository's colab/README.md . Pick Runtime → Change runtime type → v5e-1 TPU , then Runtime → Run all . No Hugging Face token is needed. From a terminal, the Colab CLI runs the same notebooks 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 05 e2b ring cache php OK: 05 e2b ring cache ran clean on v5e1 - runs/2026-10-23/05 e2b ring cache output.ipynb 1/1 notebooks verified Notebook 01 runs one stack of matrix multiplies, 275 GFLOP per step, in each mode: mode first call s step ms TFLOP/s checksum eager 0.15 2.91 94.6 359.0 fuse 2.69 1.68 163.2 359.0 compile 3.00 1.91 143.6 359.0 fuse vs eager: step 1.72x faster, checksum differs by 0.0e+00 compile vs eager: step 1.52x faster, checksum differs by 0.0e+00 Strict eager runs each operation as it is called. Defer-and-fuse TORCH TPU DEFER AND FUSE=1 collects operations and compiles them in groups. torch.compile backend="tpu" compiles a whole function. All three give the same checksum. On a stack of tiny operations, compile is 98.69x strict eager and defer-and-fuse 1.62x. A compiled function compiles once per input shape. Notebook 01 feeds four lengths: raw: 4 compiled graph s first pass ms: 341.2, 115.7, 116.7, 143.8 second pass ms: 1.1, 1.0, 0.8, 0.8 bucketed: 2 compiled graph s first pass ms: 120.5, 1.0, 1.0, 143.7 second pass ms: 0.9, 0.9, 0.8, 0.9 Padding to 128-token buckets gives 2 programs instead of 4, with a largest difference of 0 against the unpadded results. A first call costs 115 to 341 ms and a warm one about 1 ms, so every new shape in a decode loop is a new compile. Reading one value back to the host with .item each step costs 1.62x under defer-and-fuse. Notebook 02 loads the repack. Two settings decide whether the model reaches the TPU at all: claim expected today tensor names matched as stored / mapped 0 / all but tied 0 / 540 tensors on TPU after model.to 'tpu' all 0 tensors on TPU with device map='tpu' all 1098 The repack's tensors are named for the multimodal model model.language model. and the text model looks for model. , so none of the 540 match until key mapping={r"^model\.language model\.": "model."} maps them. model.to 'tpu' leaves every tensor on the host; device map="tpu" in from pretrained puts all 1,098 on the chip. Notebook 02 then runs generate four ways, 32 greedy tokens: way first call s warm call s warm tok/s TPU memory GB eager 16.21 6.99 4.6 9.85 fuse 359.49 3.50 9.1 10.83 compressed 27.80 6.85 4.7 9.85 buffer 20.66 4.83 6.6 9.83 All four write identical tokens. Defer-and-fuse is the fastest warm way at 9.1 tokens per second, and its first call takes 359 s, because generate grows the sequence by one token per step and every length is a new shape. A compiled step over a fixed buffer runs at 6.6. Notebook 03 runs the tpu-jax engine and PyTorch side by side, on the same checkpoint, prompt and libtpu : engine load s first 64 s warm 1 s warm 64 s decode tok/s memory GB jax w4a16 7.3 39.50 0.014 0.625 103.1 6.62 jax bf16 16.1 23.39 0.010 0.492 130.8 9.30 torch eager 21.2 20.98 0.203 13.531 4.7 9.87 torch fuse 42.8 658.66 0.127 7.268 8.8 11.91 torch buffer 20.0 11.57 0.167 10.113 6.3 9.83 Decode speed is 63 tokens over the warm 64-token time minus the warm 1-token time, so the prompt drops out. JAX with 4-bit weights is 11.7x the fastest PyTorch way and JAX with bf16 weights 14.8x. The two stacks compute the same model. Over 83 positions the mean KL divergence between JAX bf16 and PyTorch eager is 0.001073, with the same top token at every position. All five engines write the same 56-token answer. Notebook 04 starts with memory. Loading the repack with its weights left 4-bit, and watching the chip: after load: 6.56 GB packed weights present: True after forward pass 1: 9.83 GB after forward pass 2: 9.83 GB packed weights present: False weight dtype afterwards: torch.bfloat16 Transformers' compressed-tensors integration unpacks every layer to bf16 at the first forward pass and keeps the result. After that the model is the bf16 model, so every way in this article computes in bf16, JAX bf16 included. Keeping a table 4-bit on the chip takes code that reads the packed values itself Step 9 . A compiled step with fullgraph=True raises at its first graph break and names the line: modules with a wrapped forward: 542 break: Data pointer comparison File "/usr/local/lib/python3.13/dist-packages/compressed tensors/offload/utils.py", line 45, in send tensors - if tensor.data ptr == value.data ptr : compressed-tensors wraps the forward of all 542 modules so each module's inputs move to its device first, and torch.compile cannot trace a comparison of memory addresses, so it cuts the program at every one. On one chip every tensor is already on the TPU, and the package's own remove dispatch takes the wrappers off: python from compressed tensors.offload.dispatch import remove dispatch remove dispatch model 542 wrapped modules before, 0 after The step is one token in, the next token out, over Hugging Face's StaticCache , whose tensors never change shape. Transformers' sliding-window cache layer counts its length in a Python integer, which a compiled step treats as a constant and recompiles on. While the cache is no longer than the 512-token window, the plain static layer's methods give the same result with a counter on the chip: for method in "update", "get seq length", "get mask sizes" : setattr cu.StaticSlidingWindowLayer, method, getattr cu.StaticLayer, method cache = StaticCache config=model.config, max cache len=256 def step : return model input ids=tok, past key values=cache, use cache=True, logits to keep=1 .logits :, -1 .argmax -1 step = torch.compile step, backend="tpu", dynamic=False The seven ways of notebook 04, two runs, 64 tokens: | Way | tok/s | x generate | Graphs / breaks | |---|---|---|---| | generate | 5.0 to 5.3 | 1.0 | | | Static cache, strict eager | 4.2 to 4.5 | 0.8 to 0.9 | 0 / 0 | | Static cache, defer-and-fuse | 10.6 to 11.0 | 2.0 to 2.2 | 0 / 0 | | Compiled step, wrappers on | 6.4 to 6.6 | 1.2 to 1.3 | 24 / 60 | | Compiled step 🥉 | 31.3 to 33.4 | 5.9 to 6.7 | 1 / 0 | | Compiled step, 4-bit table 🥈 | 50.2 to 50.4 | 9.4 to 10.1 | 1 / 0 | | JAX bf16 🥇 | 131.2 to 132.9 | 25.0 to 26.2 | | With the wrappers removed the step compiles once and every later token reuses it. Its first 64-token call, compiling included, takes 36.5 to 38.3 s. Gemma 4's per-layer embedding table is the largest tensor in E2B. -text-emb4 stores it 4-bit, and a small module looks up packed rows and unpacks only those: python def forward self, ids : p, s = self.packed ids , self.scale ids q = p.unsqueeze -1 self.shifts & 15 .to torch.float32 - 8 w = q.view ids.shape, s.shape -1 , -1 s.unsqueeze -1 .to torch.float32 .to torch.bfloat16 return w.view ids.shape, -1 torch.tensor self.embed scale, dtype=torch.bfloat16 It frees 3.4 GB 9.82 GB to 6.44 GB after load and makes the compiled step 1.51x to 1.60x faster arithmetic over the two runs, evidence/arithmetic.txt , writing the same tokens: Google's QAT checkpoint already stores that table on a 4-bit grid. A server that streams its reply reads each token back to the host as it is made. Read back that way, the compiled step runs at 31.8 to 36.5 tokens per second and the 4-bit-table step at 32.4 to 37.2, so each token's round trip to the host sets the pace and the 4-bit table's lead disappears. Collect tokens on the chip and read them at the end where the client allows it. Notebook 05 replaces Transformers' cache with the layout a 4 GB laptop GPU engine uses to hold E2B's whole context gemma4-dev https://github.com/xbill9/gemma4-dev , local-pytorch-1650ti-2b-q4w4a16 . Only 15 of the 35 layers compute keys and values: full layers: 3 33024 slots head size 512 sliding layers: 12 768 slots head size 256 shared layers: 20 reuses an earlier layer's head size sliding window: 512 positions cache for a 33,024-token context: 0.212 GB cache for a 131,328-token context: 0.816 GB Full-attention layers keep one slot per position. Sliding layers look back 512 positions, so each keeps a ring of 512 + 256 slots that records which position each slot holds: room for the window plus one 256-token piece of the prompt, so writing a piece never overwrites a key the same piece still needs. The other 20 layers reuse the last full and last sliding layer's keys and values. Prompts are read in fixed 256-token pieces, so the model compiles to two programs, one for a piece and one for a decode step, whatever the prompt's length. The ring replaces Gemma4TextAttention.forward ; its attention is plain PyTorch with a mask: php def attend self, q, pos : scores = torch.einsum "nhd,kd- hnk", q, self.k .float kpos, qpos = self.kpos None, : , pos :, None keep = kpos = 0 & kpos <= qpos if self.window: keep = keep & kpos qpos - self.window weights = torch.softmax scores.masked fill ~keep None , float "-inf" , dim=-1 .to q.dtype return torch.einsum "hnk,kd- nhd", weights, self.v On the short prompt the ring writes the same 64 tokens as notebook 04's step: way first call s tok/s streamed graphs breaks notebook 04 step, 256 slots 47.8 1 0 ring, 33,024 92.1 45.2 33.3 2 0 ring, 131,328 90.7 38.7 30.7 2 0 ring, 33,024: first 64 of 64 tokens identical to notebook 04's step the answer ends at token 56 ring, 131,328: first 64 of 64 tokens identical to notebook 04's step the answer ends at token 56 Notebook 05 times decode from the first token to the 64th, collected at the end, which puts notebook 04's step at 47.8 to 47.9. The long prompts fill the context with numbered records, each with a random six-digit code, and ask for the code of the record a quarter of the way in. They come from the same generator, seed and order as the laptop engine's long-context run, so its answers apply unchanged: context prompt read s read tok/s decode answer found record laptop answer same 33024 61 0.1 1218 45.2 915965 yes 1 915965 yes 33024 4081 0.7 5978 45.2 610395 yes 33 610395 yes 33024 32761 5.0 6601 45.2 319550 no 723 319550 yes 131328 32761 9.9 3305 38.5 390185 yes 272 390185 yes 131328 130981 39.8 3288 38.1 567239 no 91 567239 yes graphs after every prompt: 2, 0 and 2, 0 programs, breaks Across the runs of notebook 05 the 130,981-token prompt reads in 35.5 to 39.8 s. The laptop engine, on a GTX 1650 Ti, read it in 2,024 s and decoded at 21.8 tokens per second there. All five answers match the laptop engine's, the two wrong codes included, and llama.cpp gave the same three answers on the three prompts it was given: finding one record among thousands of near-identical ones is at the edge of what E2B does at 32k tokens and beyond. Every step attends over every slot and masks the empty ones, so a step is the same program at every position. Decode runs at 45.1 to 45.3 tokens per second at every depth with a 33,024-token context and 38.1 to 38.7 with 131,328. Reading follows the same rule: the same 32,761-token prompt reads at 50% to 56% of the speed with the larger context. Size the context to the work. Asked for 4,352 slots, Transformers' StaticCache gives the sliding layers 512 whatever the length: 3 layers of StaticLayer 4352 slots 12 layers of StaticSlidingWindowLayer 512 slots The sliding layer allocates min sliding window, max cache len slots, and Step 8's fix writes each token at its position, so a 4,081-token prompt writes past the end of those 512 slots. Run on the chip, that recipe halted the TPU with a fatal error RuntimeUnexpectedCoreHalt , XLA error E0200 the first time a result was read back. Beyond the window the cache needs its own layout, which is the ring. way after load GB with cache GB peak GB notebook 04 step, 256 slots 6.44 16.67 ring, 33,024 6.44 6.65 16.91 ring, 131,328 6.44 7.26 16.91 The full 131,328-token context costs 0.82 GB beside 6.44 GB of model. The peak, 16.6 to 16.9 GB in every compiled way, comes from compiling: read right after the first call, it was 16.67 GB with a 33,024-token ring and 16.64 GB with a 131,328-token one. That is close to the chip's 16 GiB 17.2 GB , so a larger model meets the limit at compile time first. | One v5e chip, Gemma 4 E2B | Decode tok/s | First call | Longest context run | HBM after load | |---|---|---|---|---| | generate , defer-and-fuse | 8.8 to 9.1 | 359 s | 19-token prompt | 9.82 GB | | Compiled step + 4-bit table 🥈 | 50.2 to 50.4 | 36.5 to 37.3 s | 256-token cache | 6.44 GB | | Ring cache, 131,328 🥉 | 38.6 to 38.7 | 91 to 95 s | 130,981-token prompt | 6.44 + 0.82 GB | | JAX bf16 🥇 | 130.8 to 132.9 | 23.4 to 23.5 s | 19-token prompt | 9.30 GB in use | For short chats, notebook 04's compiled step with the 4-bit per-layer table: a static cache, the sliding-window fix, the compressed-tensors wrappers removed and one torch.compile call. It pays under 40 s once and then decodes at 50 tokens per second, about 10x generate , writing the same tokens. For anything past 512 tokens, notebook 05's ring cache: two programs, compiled once in 91 to 98 s, then the full 131,072-token context at 38 to 39 tokens per second, or 45 with a 33,024-token context. Both are plain PyTorch over Transformers' own Gemma 4 code, with no custom kernels. For the most tokens per second on this chip, the JAX engine, 2.6x the fastest PyTorch way. The goal of this article was to find how much of a TPU v5e chip PyTorch gets on Gemma 4 E2B through TorchTPU and Hugging Face Transformers, measured against a JAX engine on the same chip. The key to the solution was one compiled decode step: a static cache, compressed-tensors' module wrappers removed, and a ring cache for long context. The results were: generate , and at 50.2 to 50.4 with the per-layer table kept 4-bit Scope: Gemma 4 E2B-it QAT only, as the W4A16 repack with its 4-bit-tables twin, text only, one Colab v5e-1 runtime per run on 2026-10-21 notebooks 01 to 03 and 2026-10-23 04 and 05 , greedy decoding, one sequence at a time, one short prompt and five long ones; notebook 04 ran twice and notebook 05's figures cover three runs and one earlier run of the same code. JAX speeds are the tpu-jax engine at 4b9f8e9 . The laptop figures are from gemma4-dev's GTX 1650 Ti run on the same prompts, a different engine on different hardware. Batching, vLLM or SGLang on TorchTPU, and a serving endpoint were not run. The strategy for running PyTorch on a TPU was validated with an incremental step by step approach.