Running AI Workloads on TPUs: JAX, PyTorch and vLLM Without the Lock-In

On this page

1. 01 · Why Software Decides the TPU Question
2. 02 · The TPU Software Stack in Late 2026
3. 03 · JAX: The Native Path
4. 04 · PyTorch on TPUs Through TorchTPU
5. 05 · Serving on TPUs: vLLM and SGLang
6. 06 · What Ports and What Does Not
7. 07 · An Illustrative ML Workflow on OCI
8. 08 · A Portability Test Plan You Can Start Now
9. 09 · Key Takeaways
10. 10 · Frequently Asked Questions

01 · Why Software Decides the TPU Question

Developer guide · the TPU software stack in late 2026

"Will our code even run on a TPU?"

When Oracle and Google announced on 27 October 2026 that Google TPUs will be available through OCI from calendar year 2027, most coverage focused on hardware: pod sizes, exaflops and the first TPU offering on a hyperscale cloud outside Google. For engineers, the first question is different. A TPU is only useful if your training code, your serving stack and your libraries run on it without a year of rework.

The good news is that the TPU software story in late 2026 is far more open than a few years ago. JAX runs natively, PyTorch has a first-class TPU runtime in TorchTPU, and vLLM and SGLang have TPU backends. The harder news is that anything written directly against CUDA does not run on a TPU at all. This post maps the stack, shows what ports, and gives you a test plan to start before the hardware reaches OCI.

02 · The TPU Software Stack in Late 2026

Answer: The TPU stack has five software layers above the chip: your model code, a framework (JAX or PyTorch via TorchTPU), optionally a serving engine (vLLM or SGLang), the XLA compiler, and the TPU runtime. You write the top layers; Google maintains XLA and the runtime; Oracle will host the hardware on OCI from 2027.

On a GPU, most teams think of CUDA, cuDNN and NCCL sitting under PyTorch. On a TPU, that role is played by XLA, a compiler that turns the framework's computation graph into an optimized TPU program, and by the TPU runtime, which manages devices, memory and communication across a pod. There is no layer where hand-written CUDA can be dropped in.

The figure below shows the stack as a set of layers, with brackets marking who owns each part. The practical point is the boundary between the green group and the blue group. Everything above XLA is code you control, and if that code only calls framework operations it can move between GPU and TPU. Everything below is maintained for you, so you rarely touch it directly.

Read the stack from the top down when you assess a workload. If every dependency in your code stops at the framework layer, porting is mostly configuration. If anything reaches past the framework into CUDA, that piece has no path through XLA and needs a replacement before the workload can run on TPU 8t or TPU 8i.

03 · JAX: The Native Path

Answer: JAX is the framework Google builds its own TPU models with. It compiles functions through XLA, expresses parallelism through sharding annotations rather than device-specific code, and scales the same program from one chip to a full pod. If you are starting a new model and expect to run on TPUs, JAX is the lowest-risk choice.

JAX suits TPUs because it was built around compilation. You write numerical Python, mark functions for just-in-time compilation and describe how arrays are split across a mesh of devices; XLA plans the computation and communication. The same code runs on a laptop CPU, a GPU node or a TPU pod slice.

Google's ecosystem adds higher-level tools such as MaxText, an open-source reference for training large language models in JAX, and Pathways, Google's system for orchestrating very large computations across many chips. You do not need either to start, but both show how Google structures large TPU jobs.

A minimal JAX program is device-agnostic. The snippet below runs unchanged on CPU, GPU or TPU; only the device list printed at the start differs.

import jax
import jax.numpy as jnp

print(jax.devices())          # lists TPU, GPU or CPU devices

@jax.jit                      # compiled by XLA for whatever device is present
def predict(w, x):
    return jnp.tanh(x @ w)

w = jnp.ones((1024, 1024))
x = jnp.ones((8, 1024))
print(predict(w, x).shape)

The real work in JAX on TPUs is not the syntax but the sharding plan: how parameters, activations and optimizer state are split across chips. That plan is where TPU pod topology starts to matter, and it is the part you should prototype early.

04 · PyTorch on TPUs Through TorchTPU

Answer: TorchTPU is Google's PyTorch runtime for TPUs, built together with Meta. It is eager-first, so PyTorch code behaves the way PyTorch developers expect, and it supports torch.compile for compiled performance. For most PyTorch teams it is the path to TPUs without switching frameworks.

For years the biggest barrier to TPU adoption was that most of the industry writes PyTorch. TorchTPU addresses that directly. "Eager-first" means operations run as you call them, which keeps the familiar debugging experience: you can print tensors, step through code and use standard PyTorch idioms. When you want performance, torch.compile captures the model graph and hands it down through the compiler stack to XLA.

That does not mean every PyTorch program runs unchanged. Code that calls .cuda(), reads torch.cuda properties, relies on NCCL-specific behaviour or uses CUDA-only extensions will still fail. Code that selects its device from configuration and uses standard distributed APIs is in a much better position. The cleanup is usually modest, but it must be done and tested.

The choice between JAX and PyTorch usually follows existing code. If your models and pipelines are already in PyTorch, TorchTPU lets you keep them; for new large-scale training aimed at TPUs, JAX is the most mature path. Many organizations will run both.

05 · Serving on TPUs: vLLM and SGLang

Answer: vLLM runs on TPUs through its tpu-inference backend, which serves both JAX and PyTorch models behind the same vLLM interface. Native PyTorch backends for vLLM and SGLang on TPU were announced on 20 October 2026. For inference teams, this means the serving API your applications call can stay the same when the accelerator underneath changes.

Inference is where many enterprises will meet TPUs first: Oracle's announcement names production inference for reasoning, agentic and long-context applications, and Google designed TPU 8i for that work. Google says TPU 8i delivers up to 80% better performance per dollar for inference than Ironwood. Only your own benchmark will show whether that holds for your models.

vLLM's tpu-inference backend unifies JAX and PyTorch model implementations under one serving engine, so applications keep calling the same endpoint while operations teams change the hardware. SGLang is gaining a native PyTorch backend for TPU under the same 20 October 2026 announcement. Both projects move quickly, so check coverage for the exact architectures you serve.

If you already serve through vLLM or SGLang, you are well placed for a TPU inference pilot. If you use a custom or vendor-specific server, evaluate a move to an open engine first; that also pays off on GPUs.

06 · What Ports and What Does Not

Answer: Framework-level code ports with configuration changes. Custom CUDA kernels do not port and must be replaced with framework operations or rewritten as TPU kernels. Triton and CUDA-only libraries need TPU-supported equivalents. Hard-coded CUDA calls need cleaning up. Every path ends with an accuracy and performance benchmark.

The fastest way to estimate porting effort is to walk each component of a workload through a short set of checks. Start with the question that matters most: does this code only use PyTorch or JAX operations? If yes, the port is mostly about device selection, sharding and configuration. If no, find out what it depends on beyond the framework, because that dependency sets your effort level.

The decision flow below shows the checks in order, with a separate serving check on the right and a benchmark at the end of every path. Custom CUDA kernels in the hot path are the most expensive finding, because they must be rewritten and revalidated. Triton kernels and CUDA-only libraries come next. Hard-coded device strings and NCCL-specific code are the cheapest to fix but the most common.

Most workloads land in more than one box. A fine-tuning pipeline that is pure PyTorch, uses one CUDA-only quantization library and serves through vLLM is a good candidate: one library swap plus a backend switch. A codebase full of custom fused kernels is a poor near-term candidate, whatever the hardware economics.

07 · An Illustrative ML Workflow on OCI

Answer: A TPU workflow on OCI will likely look like a GPU workflow with a different accelerator: code and CI build container images, training jobs read datasets and write checkpoints to Object Storage, a model registry holds approved versions, and vLLM serves them to applications. Oracle has not yet published how TPU jobs will be orchestrated on OCI.

It helps to picture where TPUs will sit in an end-to-end pipeline, even before Oracle publishes service details. The figure below is an illustrative workflow built from generic OCI building blocks. Your code repository and CI pipeline build training and serving images and push them to a container registry. Training jobs on TPU pods pull those images, read datasets from Object Storage and write checkpoints back. Approved checkpoints are promoted to a model registry. vLLM, running its tpu-inference backend, loads those weights and serves an HTTPS inference API that applications and agents call. Observability spans both training and serving.

The yellow box is the honest gap. Oracle has not said whether TPU capacity on OCI will be scheduled through Kubernetes, a managed service or something else. Keep the orchestration layer replaceable: containers, Object Storage paths and a registry carry over whatever Oracle announces.

08 · A Portability Test Plan You Can Start Now

Answer: You do not need TPUs on OCI to start. Before 2027, audit dependencies, remove CUDA-specific code, standardise serving on vLLM or SGLang and benchmark on any TPU capacity you can reach, so the OCI decision becomes a pricing and placement question.

The aim is accelerator-ready workloads you can place on NVIDIA, AMD or Google hardware in OCI by cost and availability. Google's Amin Vahdat described the direction as needing "a diverse range of compute", and software portability is how you take advantage of that range.

1. Inventory dependencies. For each workload, list the framework, version, custom kernels, Triton code, CUDA-only libraries and serving engine.
2. Remove hard-coded CUDA. Make device selection and distributed setup configurable. This also helps on AMD Instinct GPUs, which OCI already offers.
3. Standardise serving. Move inference to vLLM or SGLang where you can, so the serving API stays stable across accelerators.
4. Prototype the hard parts. For any custom kernel, test whether a framework operation or compiler fusion gets close enough before you plan a rewrite.
5. Benchmark with your own models. Measure accuracy, throughput, latency and cost per token or per step against your GPU baseline. Vendor figures are a starting point, not a result.
6. Record the result. Put each workload's portability status and benchmark date in your architecture records, ready for the 2027 capacity conversation.

The accelerator is a placement decision. Portable code is what lets you make it on price and performance instead of on what your software can tolerate.

09 · Key Takeaways

1. The TPU stack is open in late 2026. JAX runs natively, PyTorch runs through TorchTPU, and vLLM and SGLang have TPU backends.
2. XLA replaces CUDA. Framework-level code compiles through XLA; anything written directly against CUDA has no path and must be replaced.
3. Serving is the easiest on-ramp. If you already serve through vLLM or SGLang, a TPU inference pilot is mostly a backend switch plus a benchmark.
4. Orchestration on OCI is not announced. Keep containers, Object Storage and a model registry as the stable parts of your pipeline.
5. Start the portability work now. Audit dependencies, remove CUDA-specific code and benchmark, so 2027 becomes a placement decision.

10 · Frequently Asked Questions