Presentamos vLLM TorchTPU

Hoy publicamos vLLM TorchTPU (vllm-torchtpu), un nuevo backend para ejecutar vLLM en las Cloud TPUs de Google. Está construido sobre la pila TorchTPU de Google y ejecuta las definiciones de modelos que ya existen en vLLM, por lo que la mayoría de los modelos que funcionan hoy en vLLM pueden servirse en TPU sin cambios en el código del modelo.

Este trabajo surge de la colaboración de ingeniería entre Inferact y Google que anunciamos el mes pasado. Nuestros objetivos para el proyecto son:

1. Facilitar la incorporación de nuevos modelos a TPU reutilizando el código de modelos existente de vLLM.
2. Impulsar el rendimiento de inferencia en TPU en código abierto, con kernels de Pallas para attention, Mixture-of-Experts y matmul cuantizado.
3. Mantener la experiencia de vLLM intacta: la misma CLI, el servidor compatible con OpenAI, el scheduler y el gestor de KV cache que ya usas en otro hardware.

Cómo funciona

vllm-torchtpu es un plugin de plataforma de vLLM fuera del árbol (out-of-tree). Una vez instalado, vLLM lo detecta automáticamente a través del vllm.platform_plugins entry point.

Los modelos se ejecutan mediante el model runner estándar de vLLM y se compilan con torch.compile(backend="tpu"), que hace el lowering a XLA. Cada dispositivo TPU obtiene su propio proceso de trabajo, de la misma manera que vLLM se ejecuta en GPUs, y la comunicación entre dispositivos pasa por torch.distributed. Esto significa que el scheduler, los workers, el gestor de KV cache y el servidor de API son todos vLLM upstream sin modificar. El trabajo específico de TPU se encuentra por debajo de ellos, principalmente en un conjunto de kernels JAX/Pallas registrados como custom ops e invocados desde dentro del grafo compilado.

Primeros pasos

1. Instalar

   pip install vllm-torchtpu

   Esto requiere Python 3.12 en una Cloud TPU VM, e instala un vLLM coincidente, torch, torch_tpuy libtpu junto con el plugin. La guía de instalación también cubre uv, compilaciones desde el código fuente y Docker.

2. Servir un modelo

   vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \
     --tensor-parallel-size 8 \
     --enable-expert-parallel \
     --kv-cache-dtype fp8 \
     --max-model-len 16384 \
     --async-scheduling

   El primer arranque compila los grafos de TPU para las formas configuradas, lo que puede tardar unos minutos. Configura VLLM_XLA_CACHE_PATH para reutilizar los artefactos compilados entre reinicios (vea Caché de compilación y calentamiento (warmup)).

La decodificación especulativa, el almacenamiento en caché de prefijos, multi-LoRA, las salidas estructuradas y las entradas multimodales están todas habilitadas con las opciones habituales de vLLM. Consulte la guía de inicio rápido para más ejemplos, incluida la inferencia por lotes sin conexión, los modelos multimodales y multi-LoRA.

Kernels

vllm-torchtpu utiliza kernels Pallas personalizados para operaciones complejas que se benefician de un ajuste manual en TPU. Algunos ejemplos:

• Ragged Paged Attention v3 para modelos MHA/GQA
• Multi-head Latent Attention (MLA) densa y dispersa para DeepSeek, Kimi y GLM
• Gated DeltaNet y Kimi Delta Attention para modelos híbridos de atención lineal
• Mixture-of-Experts fusionado con paralelismo de expertos
• Matmuls cuantizados FP8, INT8, MXFP4 y NVFP4

Escalado horizontal

Nos hemos centrado especialmente en las cargas de trabajo que hoy ejercen más presión sobre los sistemas de servicio, en particular el tráfico agéntico de múltiples turnos y el multimodal, donde los contextos siguen creciendo y las salidas de las herramientas pueden ser largas.

Paralelismo

Las seis dimensiones de paralelismo de vLLM están soportadas, tanto en slices de TPU de un solo host como de múltiples hosts, y pueden combinarse:

Para modelos MoE grandes, normalmente obtenemos el mejor rendimiento con atención en paralelismo de datos y paralelismo de expertos (por ejemplo, DP=8, TP=1 con EP en un solo host v7x). Esto mantiene la atención local en cada dispositivo y distribuye los expertos por todo el slice. Para contextos largos, PCP y DCP dividen el prefill y la caché KV entre los dispositivos. La guía de paralelismo tiene más detalles y configuraciones de ejemplo.

Descarga (offloading) de la caché KV y servicio disagregado

vllm-torchtpu usa TPU Sync, una biblioteca de código abierto en C++ que puede leer directamente los búferes de los dispositivos TPU fuera del runtime de Python. Mediante la interfaz de conector KV de vLLM, esto nos proporciona:

• Offloading de la caché KV sin copias a la DRAM del host, de modo que la caché de prefijos para sesiones largas de múltiples turnos no está limitada por la HBM.
• Servicio disagregado de prefill/decode entre hosts TPU con una baja sobrecarga de transferencia KV. El prefill y el decode pueden usar diseños distintos, por ejemplo prefill en PCP=4, TP=2 y decode en DP=8, TP=1, y la caché KV se refragmenta durante la transferencia.

Consulta la guía de servicio disagregado para configurarlo.

Decodificación especulativa y caché de prefijos

La decodificación especulativa admite MTP, EAGLE-3, DFlash, DSpark, n-gram y decodificación de sufijos (guía). La caché automática de prefijos también funciona con modelos híbridos como Qwen3.5 y Kimi K3, donde el estado de atención lineal se almacena en caché junto con la caché KV.

Modelos, características y próximos pasos

vllm-torchtpu puede servir la mayoría de los modelos compatibles con vLLM, con trabajo específico para TPU en kernels, cuantización y paralelismo para las familias Qwen, DeepSeek, Kimi, GLM, Llama, GPT-OSS y Gemma. TPU v7x (Ironwood) es el hardware recomendado y se prueba en CI en cada pull request; v6e, v5e y v5p son experimentales.

La página Supported Models & Features tiene la lista completa de checkpoints y características validados, con instrucciones para habilitar cada uno. Las matrices de soporte nocturnas registran los resultados de corrección y benchmarks de TPU v7x procedentes de CI.

¿Qué sigue?

Queremos que los nuevos modelos abiertos de las grandes familias (Qwen, DeepSeek, GLM, Kimi, Gemma, GPT-OSS) funcionen en TPU el mismo día de su lanzamiento. Algunas de las novedades, del hoja de ruta pública:

• Un pipeline de CI automatizado que incorpore, evalúe y publique matrices de soporte para nuevos modelos de pesos abiertos dentro de las 24 horas posteriores a su lanzamiento.
• Enrutamiento consciente de prefijos, caché y sesiones, y enrutamiento por niveles de memoria TPU de llm-d.
• Interfaces de PyTorch para kernels de Tokamax , compartidos con otros proyectos del ecosistema TorchTPU.
• Servicio de modelos de difusión mediante vllm-omni.

¡Las contribuciones son bienvenidas!

Capacidad TPU para contribuyentes

También estamos abriendo un fondo de capacidad TPU para contribuciones de la comunidad. Si estás desarrollando sobre vllm-torchtpu, puedes solicitar tiempo de TPU gratuito, y los mantenedores principales de vLLM en Inferact te ayudarán con el diseño y la revisión de PR.

Solicita aquí acceso a TPU como contribuyente.

Pruébalo

• Guía de inicio rápido
• Configuración de Cloud TPU
• Recetas de serving de TPU v7x (Ironwood)
• Recetas de vLLM TorchTPU en Ironwood
• Despliegue en Kubernetes
• Guía de contribución

Si tienes preguntas o quieres participar, búscanos en #sig-tpu en el vLLM Developer Slack.