vLLM y SGLang ya corren nativo en TPU lo que todo dev de PyTorch debe saber

¡Hola Devs! Vamos a darle

Si alguna vez intentaste correr inferencia en TPU con PyTorch, sabes el dolor. O te ibas a JAX, o te quedabas con las ganas. Pues bueno, Google acaba de anunciar backends PyTorch nativos para vLLM (hecho con Inferact) y SGLang (con RadixArk), corriendo sobre una capa nueva llamada TorchTPU.

La idea es simple y potente: traes tus definiciones de modelo PyTorch, tu flujo de torch.compile, tus colectivos torch.distributed — y ganas aceleración en TPU sin reescribir tu stack de serving en JAX.

¿Te acuerdas del movimiento que vimos con la React Foundation pasando a Linux Foundation? Es el mismo patrón: gobernanza abierta + alineación upstream es lo que hace que un proyecto de infraestructura sobreviva de verdad.

Pero, ¿qué hay de nuevo aquí y dónde están las trampas? ¡Vamos a ver! 🚀

Cómo funciona TorchTPU por dentro

TorchTPU convierte la TPU en un dispositivo PyTorch nativo. Sin tipos wrapper de tensor. Los tensores viven en un device tpu, el código corre eager o compila vía torch.compile con XLA, y los colectivos pasan por torch.distributed.

import torch
import torch_tpu  # registra el backend del device 'tpu'

# PyTorch estándar — sin wrapper, sin clase de tensor especial
device = torch.device("tpu")
x = torch.randn(4096, 4096, device=device)
y = torch.randn(4096, 4096, device=device)

# Corre eager o compila vía XLA
z = torch.compile(lambda a, b: a @ b)(x, y)
print(z.device)  # tpu

La decisión arquitectónica que más importa: cada device corre su propio proceso PyTorch. Los dos engines mantienen el modelo de worker, scheduler, gestor de KV cache y API server nativos del framework — con implementaciones TPU por debajo. Los kernels customizados en Pallas se llaman directo desde PyTorch como operadores customizados dentro del grafo compilado. Este diseño de plugin de plataforma deja que las optimizaciones de TPU avancen a su propio ritmo, sin forkear el framework upstream.

El tema del KV Cache

Las cargas agénticas rompen el serving ingenuo. Los agentes reenvían conversaciones multi-turno que crecen, así que recalcular el prefijo es puro desperdicio. Los dos engines mantienen prefijos reutilizables en HBM y, con host offload activado, mandan prefijos viejos a memoria del host como segundo tier.

Dos features sacan el KV cache del device que lo computó:

• Serving desagregado — prefill y decode corren en hosts TPU separados, así que el KV cache viaja del host de prefill al host de decode.
• Offload de KV en tiers — el overflow va a DRAM del host y almacenamiento secundario vía Google Cloud Managed Lustre, y regresa cuando un request posterior reusa el prefijo.

De cualquier forma, la transferencia tiene que solaparse con la computación o se va directo a la latencia. Los dos engines resuelven esto con TPU Sync, una lib C++ open-source que lee los buffers del device detrás de los tensores PyTorch directo — sin copia de staging en el host — y corre transferencias en sus propios threads fuera de Python mientras prefill y decode siguen. vLLM-TorchTPU lo expone vía la interfaz de KV connector de vLLM; SGLang-TorchTPU lo usa como backend de transferencia de desagregación y como backend de I/O para HiCache.

Benchmarks y las advertencias honestas

Google hizo benchmark público a través del proyecto SemiAnalysis InferenceX, y el número estrella es real: 2.01 centavos por millón de tokens en Qwen 3.5 corriendo en Ironwood TPU (TPU 7x) — el menor costo por millón de tokens general en su submission. También publicaron resultados SOTA en AgentX, la pista de long-context y multi-turno, donde Ironwood lidera en la mayor parte de la frontera de Pareto.

Un megakernel Pallas para Kimi K3 de 2.8T de parámetros corre el paso de decode entero como un solo kernel, sirviendo 700+ tokens/seg a un solo usuario en Ironwood.

Comparación rápida

Dónde se complica

• TorchTPU es nuevo. La historia de tensores sin wrapper es limpia, pero el ecosistema de kernels Pallas es joven. Si tu modelo tiene una op rara, puede que termines escribiendo el kernel tú.
• Host offload no es gratis. Offload de KV en tiers a DRAM/Lustre solo ayuda si la transferencia se solapa con la computación. Si tu carga es sensible a latencia y de contexto corto, el tier extra puede estorbar más que ayudar.
• JAX no se va a ningún lado. vLLM-JAX y SGLang-JAX siguen mantenidos y soportados. Si tu equipo ya corre JAX, migrar a PyTorch nativo es movimiento lateral, no upgrade.
• El contexto del benchmark importa. El número de 2.01¢ es en un modelo específico, hardware específico y forma de carga específica. Tu resultado va a variar — sobre todo entre distribuciones prefill-heavy y decode-heavy.

Si también andas checando el lado de feature flags de la plataforma Vercel, el deep dive de Vercel Flags GA es una lectura paralela útil sobre cómo maduran las abstracciones de plataforma.

¿Te conviene?

Si sirves modelos abiertos a escala y has estado pagando el impuesto de JAX — ya sea en fricción de tooling o en gap de skills del equipo — esto es un cambio genuinamente significativo. El camino PyTorch-nativo significa que tus definiciones de modelo, tu flujo de debug y tu intuición de torch.compile todo se traslada.

Si corres cargas más pequeñas o ya estás metido en JAX, la cuenta es menos obvia. El número de 2.01¢/M-token es impresionante, pero es benchmark de frontera, no garantía.

Próximos pasos:

1. Checa los proyectos upstream en GitHub de vLLM-TorchTPU y SGLang-TorchTPU — lee el issue tracker antes de commitear.
2. Si contribuyes a vLLM, vale la pena aplicar al pool de TPU comunitario — compute gratis para contribuidores activos es rareza.
3. Benchmarks la forma de tu carga antes de migrar. Distribuciones prefill-heavy y decode-heavy se comportan muy distinto bajo KV offload.

Fuente: anuncio del Google Developers Blog