Run inference on Cloud TPU

Serving is the process of deploying a trained machine learning model to a production environment to run inference. TPU v5e and later versions support inference. Meeting latency service level objectives (SLOs) is a priority for serving workloads.

This document discusses serving a model on a single-host TPU. TPU slices with eight or fewer chips have one TPU VM or host and are called single-host TPUs. For information about multi-host inference with PyTorch, see vLLM-TorchTPU and SGLang-TorchTPU. For information about multi-host inference with JAX, see Perform multi-host inference using Pathways.

Set up your TPU

Whichever stack you choose, you need a Google Cloud project with a single-host TPU before you can serve a model. The following steps apply to every combination of library and framework:

  1. Create a project. You need a Google Cloud account and project to use Cloud TPU. For more information, see Set up a Google Cloud project for TPUs.

  2. Check your quota. Ensure that you have sufficient quota for the number of TPU cores that you plan to use for inference. TPU v5e uses separate quotas for training and serving. In the Cloud TPU API, the serving-specific quotas for TPU v5e are as follows:

    • On-demand v5e resources: TPU v5 lite pod cores for serving per project per zone
    • Preemptible v5e resources: Preemptible TPU v5 lite pod cores for serving per project per zone

    For other TPU versions, training and serving workloads use the same quota. For more information, see Cloud TPU quotas.

  3. Create a TPU VM. Choose a TPU version and zone that has serving capacity, and then create a single-host TPU VM. For more information, see TPU regions and zones and Manage TPU resources.

After the TPU VM is running, follow the setup guide for your library and framework in Choose an inference stack.

Choose an inference stack

Cloud TPU supports large language model (LLM) inference through two open-source serving libraries: vLLM and SGLang. Both libraries run on Cloud TPU with either a JAX or a PyTorch backend, so you can choose the library that best fits your workload regardless of the framework that your model is written in.

Choose the library first, and then select the backend that matches your model's framework, as shown in the following table:

Serving library JAX backend PyTorch backend
vLLM vllm-jax vllm-torchtpu
SGLang sglang-jax sglang-torchtpu

Inference with vLLM

vLLM is an open-source library for fast inference and serving of LLMs, built around continuous batching and PagedAttention. Cloud TPU works with vLLM through the vllm-jax plugin for JAX models and the vllm-torchtpu plugin for PyTorch models.

For additional examples of using vLLM to serve a model on TPUs, see the following resources:

vLLM with JAX

The vLLM TPU documentation covers installation, a quickstart, supported models, and tuning options. For more information, see Serve JAX models with vLLM on TPU.

vLLM with PyTorch

The vllm-torchtpu plugin adds TorchTPU kernel and runtime support to vLLM. For more information, see Serve PyTorch models with vLLM on TPU.

Inference with SGLang

SGLang is an open-source serving framework for large language and multimodal models that pairs a high-performance runtime with a structured generation frontend. On Cloud TPU, SGLang runs through the SGLang-JAX backend (sglang-jax) for JAX models and the SGLang-TorchTPU backend (sglang-torchtpu) for PyTorch models.

SGLang with PyTorch

For information about installation, server launch, and supported configurations, see Serve PyTorch models with the SGLang-TorchTPU runtime on TPU.

SGLang with JAX

For information about installation, server launch, and supported configurations, see Serve JAX models with the SGLang-JAX runtime on TPU.

Profile inference performance

After setting up inference, you can use profilers to analyze performance and TPU utilization. For more information about profiling, see the following resources: