TorchTPU, XLA, JAX: how Google attacks NVIDIA's software lock-in

Why the real 2026 battle is not a FLOPS war

AI compute gets endless commentary, but its software layers rarely get explained. People talk about the chips (Ironwood, Blackwell, Trainium) as if the silicon were self-sufficient. The playing field has shifted: what decides an accelerator’s adoption today is no longer its peak FLOPS FLOPS Floating-Point Operations Per Second. A raw throughput metric for floating-point compute, in tera or peta. For LLM inference it is rarely the limiting factor: memory bandwidth almost always comes first. , it is the friction a developer hits getting their existing PyTorch code to run on it.

This friction has a name and an address. The default path of AI development runs through CUDA CUDA Compute Unified Device Architecture. NVIDIA's GPU computing platform: language, compiler and libraries (cuBLAS, cuDNN). Its software ecosystem is the main lock-in against alternatives like ROCm; at runtime, its context also reserves an incompressible slice of VRAM. : the code stack, the kernels, the tooling, the muscle memory of debugging at 2 a.m., all of it native to NVIDIA. It is a software moat, not a hardware one, and it is exactly what Google set out to dismantle. The 2026 shift is that Google stopped answering only with silicon and started answering with a stack: a framework developers love (PyTorch), a compiler that acts as the pivot ( XLA XLA Accelerated Linear Algebra. An ML compiler that takes the whole computation graph instead of running each operation one after another, and rewrites it for the target (TPU, GPU, CPU). Its signature transform is fusion: merging several operations into a single kernel so intermediate values are never written back to memory. ), a serving engine (vLLM) and, at the top, a project named TorchTPU whose ambition fits in one line from the team: “it should feel like PyTorch”.

This piece takes the problem by that friction. It speaks to two readers at once: the ML engineer who wants the detail (execution model, compilation path, install commands, benchmark numbers) and the reader who wants the whole picture: who depends on whom, and why it reshuffles the deck against NVIDIA. Each brick is explained first for what it does, then for the place it holds in the edifice. We start with the keystone, because nothing stands without it: the compiler.

XLA: the compiler that pivots the whole edifice

XLA (Accelerated Linear Algebra) is a compiler for machine learning. To see what it changes, you have to see how a framework runs without it. PyTorch, by default, handles each operation one after another: a kernel for the add, a kernel for the multiply, a kernel for the reduction. Each kernel reads its inputs from memory, computes, and writes its result back to memory, where the next kernel will read it again. On an accelerator, these memory round-trips cost far more than the computation itself.

XLA does not receive an operation: it receives the whole compute graph, and rewrites it for the target machine. Its signature transformation is called fusion. On the canonical example from the OpenXLA documentation (a computation of the form reduce_sum(x + y * z)), three separate kernels launch without compilation: one for the multiply, one for the add, one for the reduction. XLA melts them into a single kernel, and that is the detail that decides everything: it never writes the intermediate values (y*z, then x + y*z) to memory. It makes them “stream” straight from one step to the next, keeping them in the GPU registers. The documentation is explicit and worth quoting word for word:

XLA’s operation comes down to a chain. A framework hands it a graph expressed in StableHLO StableHLO High Level Operations, in its stable, versioned form. An operation set that acts as the portability layer between frameworks (PyTorch, JAX, TensorFlow) and the XLA compiler. One StableHLO graph can target several architectures. , the versioned operation set that serves as the portability layer between frameworks and the compiler. XLA first applies target-independent optimization passes (common-subexpression elimination, fusion, partitioning for parallelism), then hardware-specific passes, and finally emits a binary optimized for TPU, GPU or CPU. It is this independence from the target that makes XLA portable: the same StableHLO graph can target several architectures.

Hold onto this point, it is the key to the whole piece. XLA was originally developed inside TensorFlow; it is now the heart of the open-source OpenXLA project, and it supports three framework frontends: TensorFlow, JAX and PyTorch. The TPU TPU Tensor Processing Unit. Google's ML acceleration ASIC. Not a lone chip but a network: chips linked directly by the ICI interconnect in a torus topology, each combining TensorCores (dense matmul) and SparseCores (embeddings, collectives). It only knows how to talk to the XLA compiler. (Tensor Processing Unit), for its part, only knows how to speak to XLA. So anything that wants to run on a TPU, JAX code as much as PyTorch code, ends up, one way or another, going through XLA. The question “how do you run PyTorch on a TPU” is therefore, at bottom, the question “how do you cleanly bring a PyTorch graph to XLA”. The rest of the piece describes the competing paths to get there.

JAX: the framework the TPU was designed for

JAX JAX Google's functional ML framework. It composes four transforms (grad, jit, vmap, pmap) over NumPy-like code, using XLA for compilation and SPMD parallelism. Co-designed with the TPU. is Google’s ML framework. For the non-specialist reader, it is “NumPy on steroids”: the same array API, augmented with automatic differentiation and compilation for GPU/TPU. For the engineer, it is a functional-programming library whose strength is the composability of four transformations: grad computes a function’s gradient (the building block of all learning); jit compiles the function through XLA; vmap automatically vectorizes a function written for a single vector so it operates on a batch, with no Python loop; pmap parallelizes it across several accelerators under an SPMD SPMD Single Program, Multiple Data. A parallelism model where every rank runs exactly the same program over different shards of data. It is the model the TPU is built for; its divergent counterpart, MPMD (one rank doing extra logging, say), used to break the optimization. model (Single Program, Multiple Data).

The internal mechanism is worth unfolding, because it explains why JAX and the TPU get along so well. JAX does not “see” your values: it traces them. When you apply jit, JAX runs your function with symbolic values (tracers) to reconstruct the sequence of operations as a jaxpr, its functional intermediate representation, which it then lowers to StableHLO for XLA. The Python code describes a computation rather than executing it step by step. It is a deliberately lazy model, and it is exactly the kind of whole graph XLA knows how to optimize.

Why have TPUs historically been optimized for JAX? Because the two were tuned together. Reuters’ December 2025 reporting says it plainly: for years, TPUs were tuned first for JAX, Google’s favored internal framework, which created a “bottleneck” (the term is in quotes in the article) for customers already built on PyTorch. The production JAX stack rests on Pathways, the orchestration layer that lets a single JAX client drive computations across thousands of chips, beyond the boundaries of one TPU pod. It is the system that trains Gemini internally, and labs like Anthropic, xAI and Apple use JAX as one of their foundation-model building tools.

The problem, from Google’s vantage point, is the other side of the scale. PyTorch dominates research and, increasingly, production. According to the Linux Foundation’s Shaping the Future of Generative AI report, relayed by the PyTorch Foundation, PyTorch leads model training with 63% adoption; the same foundation puts the share of AI research implementations that rest on it at over 70%. It is this imbalance that Google is trying to neutralize: silicon cut for JAX, facing a code world that writes PyTorch.

PyTorch/XLA: the bridge that already exists today

Before TorchTPU, running PyTorch on a TPU went (and still goes) through PyTorch/XLA, the torch_xla package. It is a public product, installable right now (pip install torch==2.8.0 'torch_xla[tpu]==2.8.0'), unlike TorchTPU. Understanding its execution model is understanding the step TorchTPU wants to remove.

PyTorch/XLA rests on lazy tensors. When you write z = x + y on an XLA tensor, nothing runs: the operation is recorded into a symbolic graph. The actual computation happens only at a synchronization point, typically the call to xm.mark_step(), which “cuts” the accumulated graph, compiles it and runs it on the TPU. HuggingFace’s documentation captures the contrast well:

import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()              # an XLA device, not CUDA
x = torch.randn(2048, 2048, device=device)

for step in range(100):
    y = x @ x                         # nothing runs: the op is recorded
    x = torch.tanh(y)                 # the graph grows, still nothing on the TPU
    xm.mark_step()                    # here: cut the graph, compile, execute

# A print(x) BEFORE mark_step would force a device -> host materialization
# and break tracing: the model's ergonomic Achilles' heel.

It is powerful (letting XLA see a whole graph lets it fuse and optimize), but it is also the Achilles’ heel. Any operation that demands a concrete value (a print, a log, a checkpoint) forces a materialization and a device→host transfer that breaks tracing. And if the shapes change at every iteration (dynamic shapes), XLA recompiles a new graph each time: a performance sinkhole, avoided by padding to fixed shapes. The historical frustration shows through in a researcher’s comment on Hacker News about training research models on the old bridge: “a mess of undocumented behavior and bugs (silently hanging after 8 hours of training!)”.

Two version details, because they matter in production. Since the 2.7 release (official post dated May 2025, the repository’s C++11 notice carrying the date of March 18, 2025), C++11 ABI builds became the default; Google no longer ships pre-C++11 ABI wheels, and these new wheels improve tracing performance: on a Mixtral 8x7B on v5p-256, the repository documents 39% MFU versus 33% before. Support spans Python 3.9 to 3.13, with 2.9 nightlies available.

The point that matters most for what follows is written in black and white in the official pytorch/xla repository: “Once TorchTPU is public it will replace PyTorch/XLA.” PyTorch/XLA is not a competitor to TorchTPU; it is its predecessor, destined to be replaced.

TorchTPU: making the TPU “behave like PyTorch”

Here is the heart of the piece. TorchTPU is Google’s new engineering stack for running PyTorch natively on a TPU. Its history is short: an internal project first revealed by Reuters on December 17, 2025 (the codename had leaked), then detailed officially on April 7, 2026 on the Google Developers Blog. The guiding principle is that a developer should be able to take an existing PyTorch script, switch its device to tpu, and launch the training loop without touching a line of business logic.

The TPU should be an obvious choice for any PyTorch user to target. […] Getting access through PyTorch has always been difficult. We are changing that this year.

Lee Howes, Engineering Lead, TorchTPU, Google (remarks reported by The Stack)

Eager First: the philosophical break

Where PyTorch/XLA imposed the lazy model, TorchTPU establishes an “Eager First” philosophy. Technically, it builds on PyTorch’s PrivateUse1 interface, the extension point meant for wiring in a new device type. The integration goes deep enough that you handle real torch.Tensor objects that happen to live on a TPU. The team’s phrasing: “No subclasses, no wrappers; just ordinary, familiar PyTorch Tensors on a TPU.” No subclass, no wrapper: an ordinary PyTorch tensor that responds like one on CUDA.

import torch

# TorchTPU: PrivateUse1 interface, real torch.Tensor living on the TPU
x = torch.randn(2048, 2048, device="tpu")
y = x @ x                 # eager execution, like on CUDA
print(y.sum())            # the value is available, no mark_step

# Custom kernel: decorate a JAX function, TorchTPU wires it via Pallas
@torch_tpu.pallas.custom_jax_kernel
def fused_softmax(x):
    ...

Three eager modes cover the development cycle. Debug Eager runs one operation at a time, with CPU synchronization after each step. Slow, but precious for hunting shape mismatches, NaNs, out-of-memory crashes. Strict Eager dispatches op by op but asynchronously, to reproduce the default PyTorch experience. Fused Eager is the breakthrough: by automatic reflection over the operation stream, TorchTPU fuses steps on the fly into denser blocks before handing them to the TPU. The claimed gain is “a 50% to 100+% performance increase over Strict Eager, with no setup required”, no user configuration. The three modes share a Compilation Cache, single-host or persistent multi-host.

Static compilation: Dynamo → StableHLO → XLA

For peak performance, TorchTPU integrates natively with torch.compile, and the path it takes is a deliberate architectural choice. FX graph capture through Torch Dynamo, then, instead of routing to Torch Inductor (PyTorch’s default backend), use of XLA as the backend compiler, through a translation layer that maps PyTorch operators directly to StableHLO. Google’s justification: XLA is “rigorously battle-tested for TPU topologies” and knows how to natively optimize the overlap between dense compute and collective communications across the ICI ICI Inter-Chip Interconnect. The proprietary link that connects TPU chips directly to one another, in a 2D/3D torus topology, without going through the host network. It is what makes a pod of thousands of TPUs behave like a single machine, Google's equivalent of NVLink. (Inter-Chip Interconnect). In other words, Google is not reinventing a TPU backend for PyTorch: it is rewiring PyTorch onto the TPU backend matured over years for JAX.

The MPMD problem PyTorch/XLA could not solve

The real architectural progress hides here, and it is less visible than the eager numbers. PyTorch/XLA supported only pure SPMD: all ranks run exactly the same code. But real PyTorch code diverges slightly from one rank to another: rank 0 typically does a bit of extra logging. This divergence, MPMD (Multiple Program, Multiple Data), broke TPU optimization, which is cut for SPMD. TorchTPU is architected to support these divergent executions by isolating the communication primitives where necessary, “at minimal cost”, while preserving the global view XLA needs. The concrete consequence: PyTorch’s DDP, FSDPv2 and DTensor distributed APIs work “out of the box”.

Portability does not erase the hardware, and Google gives the classic trap as an example. Many models hard-code an attention head dimension at 64, whereas current TPUs hit their peak matrix-multiply ( matmul matmul Matrix multiplication, the dominant operation in a neural network's layers. When we say a GPU is computing, it is doing matmul about 90% of the time. ) efficiency at 128 or 256. TorchTPU recommends a staged flow: first get a correct execution, then refactor suboptimal architectures or inject custom kernels. The fact that your code runs does not mean it runs well: hardware awareness stays your job.

vLLM: the serving layer

Everything above concerns training and general execution. vLLM occupies a distinct layer: production LLM inference serving. Born in 2023 at UC Berkeley’s Sky Computing Lab, it is the reference open-source inference engine, and its coupling with the TPU stack became, in late 2025, a signal about the real direction Google is taking.

Its central innovation, PagedAttention PagedAttention An algorithm introduced by vLLM that manages the KV cache like an operating system's virtual memory: in pages, without requiring a contiguous block per request. It eliminates internal and external fragmentation, letting a server handle 2 to 4× more concurrent requests on the same VRAM. , attacks the memory waste of serving. To generate text, a transformer re-reads at every token the KV cache KV cache The stored key and value vectors an LLM has already computed for every token it has processed. It avoids recomputing attention over the whole history, at the cost of a memory footprint that grows with the context length. (the stored key and value tensors for the whole sequence); naive serving pre-reserves a contiguous block sized for the maximum length, and wastes most of the reserved VRAM. PagedAttention applies the idea of an operating system’s virtual memory: the KV cache is split into blocks (16 tokens per block by default, configurable), addressed through a table and stored non-contiguously. Fragmentation collapses. The broader story, offloading and routing included, is in our piece on the KV cache as the central object of serving.

Second pillar: continuous batching continuous batching Iteration-level scheduling: adding and removing requests from the batch at every generation step, instead of waiting for a whole batch to finish. Formalized by Orca (OSDI 2022), popularized by vLLM. It multiplies an inference server's throughput by 2 to 4× under heavy concurrency. . Instead of waiting for the whole batch to finish, vLLM works at the granularity of the iteration: as soon as a request completes, its KV blocks are freed and the next request in the queue is inserted at the following decode step. The accelerator stays saturated. The concept comes from Orca (OSDI 2022); the most-cited figure (up to 23× throughput) is Anyscale’s measurement against HuggingFace’s naive static batching, not a general gain, and it has to be presented as such.

This is where it connects to the thesis of this piece. vLLM is multi-backend: NVIDIA, AMD, Intel (CPU, GPU, Gaudi), IBM Power, AWS Trainium/Inferentia via plugin, and TPU. In October 2025, vLLM TPU was rebuilt around tpu-inference, a hardware plugin that unifies JAX and PyTorch under a single JAX→XLA lowering path. A model defined in PyTorch runs on a TPU with no code change, lowered through Torchax (a PyTorch→JAX bridge, now in its own repository google/torchax), while JAX is natively supported. The vLLM blog states it unambiguously:

The documented gains are real: on the previous path (PyTorch/XLA), the team had obtained +3.6× throughput on Llama 3.1-8B on v6e-1 and +2.1× on Llama 3.1-70B on v6e-8; the tpu-inference rebuild reaches “nearly 5x” the February 2025 prototype. A limit documented by Google Cloud frames the usage: on GKE, vLLM on TPU is single-host; multi-host is not supported, which for 400B+ models points to JetStream, Google’s multi-host TPU inference engine.

Note the detail that resonates with the rest of the piece. Even to serve PyTorch, the real path on a TPU runs through JAX→XLA via Torchax. On the serving side, the destination is not “PyTorch on TPU”: it is always XLA, reached through JAX.

How it all fits together: the mental model

Here is the edifice, top to bottom. This is the diagram to keep in mind.

Layer 1: the frameworks, what the developer writes. Two worlds: PyTorch (imperative, eager, dominant) and JAX (functional, lazy, co-designed with the TPU). Competitors at the API level, they converge on the same compiler.

Layer 2: the compiler, the pivot. XLA, fed with StableHLO. JAX gets there through jit/jaxpr. PyTorch gets there through three historically distinct paths: the old PyTorch/XLA bridge (lazy tensors, mark_step); the new TorchTPU (eager-first, torch.compile → Dynamo → StableHLO → XLA, MPMD); for serving, Torchax (PyTorch → JAX → XLA) inside vLLM.

Layer 3: the silicon, where it runs. The TPU, which is not an isolated chip but a network: chips linked by the ICI in a 2D/3D torus topology, each combining TensorCores Tensor Cores Hardware units specialized in low-precision matrix multiplication, introduced by NVIDIA with Volta (2017). Each generation adds supported formats: FP16 → FP8 (Hopper) → FP6/FP4 (Blackwell). They run the bulk of the inference compute. (dense matmul) and SparseCores (embeddings, gather/scatter, collectives). XLA can also target GPU and CPU: that is where portability comes from.

Around them, the satellite tools, by tier of need. TorchTitan: the native PyTorch platform for large-scale training (composable parallelism through DTensor and DeviceMesh), a TorchTPU integration target for 2026. Pathways: the JAX-side orchestrator that extends an XLA computation beyond a pod. Helion: a “PyTorch with tiles” Python DSL that automates kernel autotuning. Pallas / Mosaic: the JAX custom-kernel language, which lowers through Mosaic onto the TPU, the way to extract the last drop of performance where XLA cannot yet optimize. MaxText / JetStream: the reference JAX implementation for LLM training, and the multi-host TPU inference engine.

The line to remember: PyTorch and JAX fight over the developer; XLA reconciles them on the silicon; vLLM serves the result. TorchTPU is the project that turns “PyTorch can technically run on a TPU” into “PyTorch runs on a TPU like it does at home”.

NVIDIA’s real moat is not the silicon

NVIDIA does not dominate because its chips work some inaccessible magic. They are excellent, full stop. It dominates because the default path runs through CUDA, and that path holds a decade’s head start: tooling, optimized kernels, network effect. Over 6.5 million developers use NVIDIA software (figure released in August 2025), more than 6 million of them on CUDA. It is a software moat, and it is what TorchTPU targets, not the peak FLOPS.

The numbers behind the moat are eloquent. In fiscal 2026 (closed at the end of January 2026), NVIDIA posted roughly $194 billion in datacenter revenue, up 68% year over year, about 90% of its total revenue (~$216 billion), with a non-GAAP gross margin on the order of 71% for the year (75% in Q4 alone). Networking revenue alone (NVLink, AI Ethernet) is near $30 billion. On market share, IDC puts NVIDIA around 81% of datacenter AI chips, its closest rival, AMD, holding only about 10%: an analyst figure, to be treated as such. As for revenue visibility, it was Jensen Huang, at GTC in October 2025, who quantified it (“the first technology company in history to have visibility into half a trillion dollars”) at roughly $500 billion of Blackwell + Rubin from early 2025 to the end of 2026, a figure since raised to some $1 trillion through the end of 2027.

It is this wall that the software friction sustains, and that is why the fight plays out at the compiler level, not the die. For the exact nature of this lock-in and how it compares to ROCm, see our piece on CUDA vs ROCm in 2026.

Google’s answer: portability as a weapon

Google’s idea is to turn the AI chip into a commodity by removing the switching cost. Three levers, all already engaged in 2026.

The first: make the TPU “PyTorch-friendly” with TorchTPU, so PyTorch shops no longer have to rewrite their stack in JAX. The second: ally with Meta, PyTorch’s creator and keeper. Reuters reports that Google works closely with Meta and is considering open-sourcing parts of TorchTPU to speed adoption. Meta, one of NVIDIA’s biggest customers (~$72 billion in 2025 capex), has a direct interest in lowering its costs and strengthening its negotiating leverage; the reporting indicates it explored renting TPU capacity as early as 2026 and deploying them in its own datacenters from 2027. The third: multi-stack portability (JAX, PyTorch/XLA, ONNX) that lets the same model train on NVIDIA, Google or AMD with minimal changes.

The business context is explosive. Sundar Pichai stated, on Alphabet’s Q3 2025 earnings, that Google Cloud had signed more deals worth over $1 billion in the first nine months of 2025 than in the previous two years combined. And above all Anthropic: the October 23, 2025 agreement grants access to “up to one million TPU chips… well over a gigawatt of capacity coming online in 2026”, valued at “tens of billions of dollars”. Thomas Kurian, CEO of Google Cloud, justified it this way:

Anthropic’s choice to significantly expand its usage of TPUs reflects the strong price-performance and efficiency its teams have seen with TPUs for several years.

Thomas Kurian, CEO, Google Cloud

The agreement was extended in April 2026: a Google investment in Anthropic of up to $40 billion (announced April 24, “up to”), and a next-generation capacity agreement through Broadcom: 3.5 GW per Broadcom’s regulatory filing, “up to 5 GW” in Google’s phrasing. A new development, finally: Google started selling TPUs for delivery to a restricted circle of customers in their own datacenters, where historically you could only rent them on Google Cloud.

Two chips for two jobs: the 8th generation

At Google Cloud Next, on April 22, 2026, Google unveiled an 8th-generation TPU split for the first time into two distinct chips. (Not to be confused: Ironwood, the 7th generation, had reached general availability back in late 2025; it is the 8t/8i split that is the April 2026 announcement.) The rationale: training and inference have diverged into two workload profiles, and a single silicon serves both poorly at once, the same phase-specialization logic seen on the NVIDIA side with Vera Rubin’s disaggregation.

The TPU 8t, for training, bets on scale. Native FP4 FP4 A 4-bit floating-point format, the 2026 frontier of high-throughput inference. Four times less memory than FP16, but a very narrow dynamic range: it only holds up with fine-grained scaling through block formats (MXFP4, NVFP4). that doubles MXU throughput, an Axion ARM host, and above all a fabric named Virgo that links over 134,000 chips with up to 47 Pb/s of non-blocking bisection bandwidth, for over 1,700 ExaFlops at near-linear scaling, enough to exceed one million chips in a training cluster, driven by JAX + Pathways. The claimed gain is up to 2.7× performance per dollar vs Ironwood (Google’s figure; several outlets cite 2.8× or even 2.85×, which are not the primary number). The superpod counts 9,600 chips, and Google’s consumer-facing post advances 121 ExaFlops and 2 PB of HBM shared per superpod.

The TPU 8i, for inference and reasoning, bets on latency. 288 GB of HBM and above all 384 MB of on-chip SRAM (×3 vs the previous generation) to house a larger KV cache as close as possible to the compute. The radical novelty is topological: dropping the 3D torus for Boardfly, which brings a 1,024-chip configuration’s network diameter down from 16 hops to 7 (−56%), for latency improved “up to 50%”. And a new fixed block, the CAE (Collectives Acceleration Engine): two TensorCores and one CAE per chip replace Ironwood’s four SparseCores, and cut on-chip collectives latency “by 5x”, decisive for autoregressive decoding and chain-of-thought. Claimed gain: up to 80% performance per dollar vs Ironwood in inference. The two chips share bare-metal access and up to 2× perf/watt, and both “run native PyTorch via TorchTPU in preview”. The 8th generation therefore arrives with TorchTPU as a front-line software argument.

A structural limit to keep in mind: the TPU was long rentable only on Google Cloud, never buyable standalone. NVIDIA GPUs run everywhere: multiple clouds, on-prem, workstations, edge. Google has started to correct this by selling TPUs to a few customers, but for a multi-cloud or sovereign strategy, the deployment-flexibility gap stays a decisive factor. A sign that Google is not trying to purely replace NVIDIA: it also offers NVIDIA Vera Rubin NVL72 instances on its own Virgo fabric.

How to choose for your situation

You are a PyTorch-first shop on GPU evaluating the TPU to cut the bill. Migrate nothing in preview. The public, stable path stays PyTorch/XLA (torch_xla[tpu]), not TorchTPU. Run an honest A/B test (the same PyTorch model on a GPU cluster and on a TPU slice) and compare sustained throughput, energy and bill. Watch two signals that change the game: the release of TorchTPU’s public GitHub repository with the flip to GA, and the validation of linear scaling at Pod size. Until these two milestones are cleared, TorchTPU is a credible promise, not a production foundation.

You do LLM serving and want the best cost per token. vLLM is the default option, multi-backend. On TPU, start from tpu-inference (single-host on GKE); for 400B+ multi-host models, switch to JetStream. Keep your model definitions in PyTorch: Torchax lowers them with no code change. The threshold that justifies seriously evaluating the TPU for inference: an MoE or reasoning workload at high concurrency with a low latency target, exactly the profile the 8i (CAE, big SRAM) was designed for.

You are starting a very large-scale research project on TPU. JAX + MaxText + Pathways stays the path of least resistance and the most mature: it is the stack that trains Gemini. Adopt PyTorch-on-TPU only if your code and talent capital is already massively PyTorch.

You have a multi-cloud, on-prem or sovereign constraint. The NVIDIA GPU stays the safe choice as long as TPUs are not widely buyable outside Google Cloud. Keep your pipelines portable (Keras 3, ONNX, framework-agnostic models) to switch between GPU and TPU as prices and availability evolve.

Conclusion

TorchTPU will not bring down the CUDA wall on its own. What it changes is the height of that wall: it turns “PyTorch can technically run on a TPU” into “PyTorch runs on a TPU like it does at home”. The day an infrastructure lead can say “we move significant workloads without heroic engineering”, the negotiating leverage flips. That is precisely what Google is trying to make true.

What remains is the question that decides everything, and that no spec will settle before the first GA deployments. Will TorchTPU’s maturity be enough to migrate the PyTorch-first shops? Or will JAX/Pathways stay the path of least resistance for whoever wants the maximum out of the TPU? One clue leans for TorchTPU: teams like Liquid AI moved from JAX to PyTorch in 2023 for ecosystem reasons, not code, proof that framework lock-in is more surmountable than people think. But one clue leans against: even vLLM TPU, to serve PyTorch, lowers it to JAX→XLA via Torchax. At least on the serving side, the real destination stays XLA reached through JAX. The real question for 2027 is therefore no longer “can the TPU run PyTorch”: it is how many developers will cross the bridge once it is stable, and which of them will find they were, without knowing it, already on the JAX side.

Sources and method

This piece draws on verification work closed on June 9, 2026. Every figure in it is, unless noted otherwise, a verified fact from a primary source.

Software primary sources:

• OpenXLA: architecture and tf2xla. Definition of XLA, fusion (“Fusion is XLA’s single most important optimization”, verbatim), the reduce_sum(x + y*z) example, StableHLO as the portability layer, TF/JAX/PyTorch frontends.
• JAX: key concepts. grad/jit/vmap/pmap, tracing → jaxpr → StableHLO.
• PyTorch/XLA: pytorch/xla repository (install command, C++11 ABI, “Once TorchTPU is public it will replace PyTorch/XLA”, verbatim) and 2.7 release post. The lazy/mark_step model and the “CPU and CUDA tensors… are lazy” quote come from the HuggingFace PyTorch/XLA blog. The “mess of undocumented behavior and bugs” comment is from user in-silico on Hacker News.
• TorchTPU: Google Developers Blog, April 7, 2026. “it should feel like PyTorch”, PrivateUse1, three eager modes, “50% to 100+%”, Dynamo→StableHLO→XLA, MPMD, roadmap, preview. The long Lee Howes quote (“The TPU should be an obvious choice…”) is reported by The Stack, not by the official post, which contains no direct quotes.
• vLLM: tpu-inference blog, October 16, 2025 (verbatim fallback quote, +3.6×/+2.1×, “nearly 5x”); PagedAttention (design doc) and SOSP 2023 paper; Torchax on google/torchax; multi-host/JetStream on Google Cloud GKE.

Hardware and business primary sources:

• TPU 8t/8i: Google Cloud technical deep-dive (April 22, 2026) and the associated blog.google post: all the Table 1 specs, Virgo, Boardfly, CAE, and the “up to” figures (2.7× training, 80% inference, 5× collectives, 50% latency, 56% diameter). Ironwood: Google blog (April 2025), GA reached in late 2025.
• NVIDIA: FY2026 results (8-K / NVIDIA Newsroom) for datacenter revenue of ~$194 billion ($193.7 billion at the filing), +68%, ~90% of total revenue; the “half a trillion” visibility is Jensen Huang’s at GTC (October 2025). The ~81% market share (IDC) is an analyst estimate.
• Deals: Anthropic/Google announcement (October 23, 2025, PR Newswire) (verbatim Kurian quote, “up to one million TPU chips”); the April 2026 extension (Google investment “up to $40B”, 3.5 GW capacity per the Broadcom filing / “up to 5 GW” per Google); Reuters (December 2025) for the Meta collaboration and the contemplated open-sourcing of TorchTPU. Liquid AI (JAX→PyTorch 2023): the founder’s public testimony.

Method notes and points to handle with care:

• Vendor “up to” figures. All 8t/8i gains (perf/$, latency, perf/watt) are Google claims phrased as “up to”, not verified by an independent benchmark. Same for Fused Eager’s +50 to 100%.
• Consistency notes. NVIDIA’s 75% gross margin is the Q4 figure, not the full year (~71% non-GAAP over the year). vLLM’s “23×” is a maximum against naive static batching (Anyscale measurement), not a general gain; the concept itself comes from Orca.
• Secondary figures, outside Google’s primary source: the “Sunfish/Zebrafish” codenames, the Broadcom (training)/MediaTek (inference) split, the TSMC 2 nm node and a “2027” GA timeline do not appear in the official deep-dive (Google announces a GA “later this year”, i.e. 2026). The official post itself carries an internal inconsistency on the Boardfly chip count (“up to 1,152” vs “up to 1,024 active chips”).
• Currencies. Amounts in this article are in US dollars, as reported by the cited sources (SEC filings, analyst notes, press releases). The French edition converts to euros at an indicative 1 USD = 0.92 EUR. Direct quotes (“half a trillion dollars”, “tens of billions of dollars”) are reproduced in their original currency.
• Preview API. The TorchTPU code snippets (device="tpu", @torch_tpu.pallas.custom_jax_kernel) are illustrative: the API is in preview and may change before GA.
• Header image. Google TPU v3 board photographed by Zinskauf, Wikimedia Commons, under CC BY-SA 4.0 license, cropped for the header (illustrative image of an earlier TPU generation).

A French version of this article, published on June 9, 2026, is the original.