vLLM TorchTPU es un plugin de plataforma para vLLM que permite ejecutar modelos definidos en PyTorch sobre Google Cloud TPUs mediante TorchTPU. Su objetivo principal es ofrecer la misma experiencia de usuario y funcionalidad que en otras plataformas, sin necesidad de modificar el código de los modelos. A diferencia de los backends anteriores basados en JAX, esta herramienta utiliza kernels personalizados de Pallas para operaciones críticas como la atención ragged paged, la mezcla de expertos (MoE) y la multiplicación matricial cuantizada en formatos FP8, INT8 y MXFP4.
El sistema está diseñado para escalar horizontalmente, soportando paralelismo de tensor, de datos, de pipeline y de contexto, tanto en entornos de un solo nodo como en configuraciones multi-nodo. Esto permite una inferencia eficiente mediante un servidor compatible con OpenAI, que soporta funciones avanzadas como la decodificación especulativa, el caché de prefijos, la inferencia multi-LoRA y entradas multimodales. La herramienta está recomendada para el uso de TPUs de la generación v7x (Ironwood), aunque también soporta experimentalmente las generaciones v6e, v5e y v5p.
Para su instalación, se requiere Python 3.12 y se puede obtener mediante pip desde PyPI o a través de contenedores Docker precompilados. El despliegue es sencillo: el usuario solo necesita ejecutar el comando estándar de vLLM para iniciar el servidor. La primera ejecución compila los gráficos de TPU, un proceso que puede tardar varios minutos, por lo que se recomienda configurar un caché de compilación para persistir los artefactos entre reinicios. Esta solución es la ruta recomendada para servir definiciones de modelos en PyTorch en hardware TPU, ofreciendo un rendimiento optimizado y una integración fluida con la infraestructura existente de vLLM.