The Core Update
Google just rolled out direct PyTorch support for vLLM and SGLang on their Tensor Processing Units (TPUs). This means you can now run large language model inference with these popular engines using PyTorch natively on Google's specialized hardware. Previously, TPU support for vLLM and SGLang relied on a JAX-based backend. This new update, developed with Inferact for vLLM and RadixArk for SGLang, delivers a more native PyTorch experience directly on TPU infrastructure.Official Source: Google Announcement
Technical Impact & Mechanism
The primary friction point was PyTorch developers' need to customize and debug their serving stacks on TPUs. The prior JAX-based path added complexity. Google addressed this with a new initiative calledTorchTPU.
TorchTPU transforms TPUs into native PyTorch devices. This is a significant shift. Tensors now reside directly on a tpu device, requiring no wrapper types. Your PyTorch code can execute eagerly. Alternatively, compile it using torch.compile via XLA for performance. Distributed operations flow through torch.distributed as expected.
Under the hood, vLLM and SGLang retain their existing worker models, schedulers, KV cache managers, and API servers. Google integrated TPU-specific implementations beneath these components. This platform-plugin architecture means much of the original framework code is reused. TPU optimizations, such as custom Pallas kernels, can evolve independently. These Pallas kernels are invoked as custom operators directly from PyTorch within the compiled graph, enabling fine-tuned performance.
Example:
import torch
# Before: Indirect or JAX path for vLLM/SGLang on TPU
# Now: Direct PyTorch tensor on TPU, no special wrappers
device = torch.device("tpu") # TPU is now a first-class PyTorch device
# Example tensor creation on TPU
tpu_tensor = torch.randn(1024, 1024, device=device)
# Standard PyTorch operations run natively on TPU
result = tpu_tensor @ tpu_tensor.T
# Integrate a model for compiled execution
# from vllm import LLM
# my_llm_model = LLM("meta-llama/Llama-2-7b-hf", device=device)
# compiled_model = torch.compile(my_llm_model)
# output = compiled_model.generate(...)
print(f"Tensor device: {tpu_tensor.device}")
Action Plan for Developers & Businesses
- Evaluate Existing Workloads: If you currently run vLLM or SGLang on a JAX-TPU setup, assess the benefits of migrating to the native PyTorch path.
- Update PyTorch Environment: Ensure your PyTorch installation is compatible with
TorchTPUfor direct TPU hardware access. - Optimize with Pallas Kernels: Investigate and utilize custom Pallas kernels for specific model layers to maximize inference performance on TPUs.
- Plan for Easier Debugging: Capitalize on the improved debugging and customization capabilities for your LLM serving stack on TPUs.
Need help architecting your next-gen AI inference stack? Exploring cloud cost optimization for LLMs? I build scalable systems.
Explore Case Studies & Work | Contact Waleed