vLLM e SGLang agora rodam nativo em TPU o que muda pra quem usa PyTorch

Bora entender rapidinho

Olha só, se você já tentou rodar inferência em TPU com PyTorch, sabe a dor. Ou você ia pro JAX, ou ficava na vontade. Pois bem, o Google acabou de anunciar backends PyTorch nativos pro vLLM (feito com a Inferact) e pro SGLang (com a RadixArk), rodando em cima de uma camada nova chamada TorchTPU.

A ideia é simples e poderosa: você traz suas definições de modelo PyTorch, seu fluxo de torch.compile, seus coletivos torch.distributed — e ganha aceleração em TPU sem reescrever sua stack de serving em JAX.

Sabe aquele movimento que a gente viu com a React Foundation indo pra Linux Foundation? É o mesmo padrão: governança aberta + alinhamento upstream é o que faz projeto de infraestrutura sobreviver de verdade.

Mas o que tem de novo aqui, e onde estão as pegadinhas? Vamos lá! 🚀

Como o TorchTPU funciona por baixo

O TorchTPU transforma a TPU em um dispositivo PyTorch nativo. Sem tipos wrapper de tensor. Os tensores ficam num device tpu, o código roda eager ou compila via torch.compile com XLA, e os coletivos passam pelo torch.distributed.

import torch
import torch_tpu  # registra o backend do device 'tpu'

# PyTorch padrão — sem wrapper, sem classe de tensor especial
device = torch.device("tpu")
x = torch.randn(4096, 4096, device=device)
y = torch.randn(4096, 4096, device=device)

# Roda eager ou compila via XLA
z = torch.compile(lambda a, b: a @ b)(x, y)
print(z.device)  # tpu

A escolha arquitetural mais importante: cada device roda seu próprio processo PyTorch. Os dois engines mantêm o modelo de worker, scheduler, gerenciador de KV cache e API server nativos do framework — com implementações TPU por baixo. Kernels customizados em Pallas são chamados direto do PyTorch como operadores customizados dentro do grafo compilado. Esse design de plugin de plataforma deixa as otimizações de TPU andarem no próprio ritmo, sem fork do framework upstream.

O rolê do KV Cache

Cargas agênticas quebram o serving ingênuo. Agentes reenviam conversas multi-turno que crescem, então recalcular o prefixo é puro desperdício. Os dois engines mantêm prefixos reutilizáveis na HBM e, com host offload ligado, jogam prefixos mais antigos pra memória do host como segundo tier.

Duas features tiram o KV cache do device que o computou:

• Serving desagregado — prefill e decode rodam em hosts TPU separados, então o KV cache viaja do host de prefill pro host de decode.
• Offload de KV em tiers — o overflow vai pra DRAM do host e armazenamento secundário via Google Cloud Managed Lustre, e volta quando um request posterior reusa o prefixo.

De qualquer forma, a transferência precisa sobrepor a computação ou vai direto pra latência. Os dois engines resolvem isso com o TPU Sync, uma lib C++ open-source que lê os buffers do device por trás dos tensores PyTorch direto — sem cópia de staging no host — e roda transferências em threads próprias fora do Python enquanto prefill e decode continuam. O vLLM-TorchTPU expõe isso via interface de KV connector do vLLM; o SGLang-TorchTPU usa como backend de transferência de desagregação e como backend de I/O pro HiCache.

Benchmarks e as pegadinhas honestas

O Google benchmarkou publicamente pelo projeto SemiAnalysis InferenceX, e o número de destaque é real: 2,01 centavos por milhão de tokens no Qwen 3.5 rodando em Ironwood TPU (TPU 7x) — o menor custo por milhão de tokens geral na submissão deles. Também publicaram resultados SOTA no AgentX, a trilha de long-context e multi-turno, onde o Ironwood lidera na maior parte da fronteira de Pareto.

Um megakernel Pallas pro Kimi K3 de 2,8T de parâmetros roda o passo de decode inteiro como um kernel só, servindo 700+ tokens/seg pra um único usuário no Ironwood.

Comparação rapidinha

Onde isso complica

• TorchTPU é novo. A história dos tensores sem wrapper é limpa, mas o ecossistema de kernels Pallas é jovem. Se seu modelo tem uma op esquisita, você pode acabar escrevendo o kernel na mão.
• Host offload não é de graça. Offload de KV em tiers pra DRAM/Lustre só ajuda se a transferência sobrepõe a computação. Se sua carga é sensível a latência e de contexto curto, o tier extra pode atrapalhar mais do que ajudar.
• JAX não vai a lugar nenhum. vLLM-JAX e SGLang-JAX continuam mantidos e suportados. Se seu time já roda JAX, migrar pro PyTorch nativo é movimento lateral, não upgrade.
• Contexto do benchmark importa. O número de 2,01¢ é num modelo específico, hardware específico e formato de carga específico. Seu resultado vai variar — especialmente entre distribuições prefill-heavy e decode-heavy.

Se você também tá de olho no lado de feature flags da plataforma Vercel, o deep dive do Vercel Flags GA é uma leitura paralela útil sobre como abstrações de plataforma amadurecem.

Vale a pena pra você?

Se você serve modelos abertos em escala e vem pagando o imposto do JAX — seja em fricção de tooling ou em gap de skill do time — isso é uma mudança genuinamente significativa. O caminho PyTorch-nativo significa que suas definições de modelo, seu fluxo de debug e sua intuição de torch.compile tudo carrega junto.

Se você roda cargas menores ou já tá fundo no JAX, a conta é menos óbvia. O número de 2,01¢/M-token é impressionante, mas é benchmark de fronteira, não garantia.

Próximos passos:

1. Checa os projetos upstream no GitHub do vLLM-TorchTPU e SGLang-TorchTPU — lê o issue tracker antes de commitar.
2. Se você contribui pro vLLM, vale aplicar pro pool de TPU comunitário — compute grátis pra contribuidores ativos é raridade.
3. Benchmarks o formato da sua carga antes de migrar. Distribuições prefill-heavy e decode-heavy se comportam bem diferente sob KV offload.

Fonte: anúncio do Google Developers Blog