tpu-inference is a hardware plugin for vLLM, designed to enable efficient inference of large language models (LLMs) on Google Cloud TPUs. It unifies JAX and PyTorch under a single lowering path, allowing PyTorch model definitions to run performantly on TPUs without additional code changes, while also extending native support to JAX. The library aims to push TPU hardware performance limits and retain vLLM's standardized user experience. It is actively maintained by the vLLM project and Google Cloud, with releases tied to vLLM development.
pip install vllm-tpuVerified import paths — ran on the pinned version, not inferred.
This quickstart verifies the successful installation and configuration of `vllm-tpu` in a Python environment. It checks the versions of `vLLM` and `tpu-inference`, confirms the detected vLLM platform, and lists available JAX devices, indicating the presence and readiness of TPU hardware. Actual model serving typically involves running the `vLLM` API server with `enable_tpu_vllm` set to true, requiring a Google Cloud TPU VM.
Provision a Google Cloud TPU VM (e.g., v5e, v6e, v7x) as per the official documentation for optimal performance and stability. For Ironwood (TPU7x), Google Kubernetes Engine (GKE) is required.
For PyTorch, consider refactoring generation to use a manual decode loop with `torch_xla.core.xla_model.mark_step()` for explicit execution. Aim for static shapes (fixed batch/sequence length) and avoid dynamic branching. Converting models to JAX-compatible checkpoints might offer better performance.
Understand your typical request patterns and model sequence lengths. While `vLLM` handles this automatically, being aware of the bucketization can help in performance analysis and potentially in optimizing input lengths or batching strategies.
Plan inference strategies without relying on speculative decoding for TPU deployments with vLLM. Monitor vLLM and tpu-inference release notes for future support.
Implement a manual decode loop and explicitly force execution with `torch_xla.core.xla_model.mark_step()` after each token generation. Optimize for static input shapes and minimize dynamic control flow within the loop. Consider using a JAX-based model if possible.
Ensure `vllm-tpu` (which includes `tpu-inference` as a dependency) is installed in the active Python environment. If using a Google Cloud TPU VM, ensure the installation steps from the official documentation are followed, preferably within a virtual environment. Use `pip install vllm-tpu` or `uv pip install vllm-tpu`.
Ensure your code is executed on a properly configured Google Cloud TPU VM. Verify the TPU runtime environment setup (e.g., `PJRT_DEVICE` environment variable, `libtpu` availability) as per Google Cloud and vLLM-TPU documentation. Access to a TPU VM and sufficient quota are prerequisites.
Use a smaller model or a TPU generation with more memory if available. Employ model quantization or pruning techniques to reduce memory footprint. Consider model partitioning or pipeline parallelism if supported for the specific model and TPU setup, though vLLM handles much of this automatically.