vLLM TorchTPU: většina modelů má běžet na Cloud TPU bez změny kódu

Týmy Inferact a Google představily vLLM TorchTPU, samostatně instalovaný backend pro provoz vLLM na Google Cloud TPU. Jejich hlavní tvrzení je užší než „všechny modely fungují bez úprav“: většina modelových definic, které už podporuje vLLM, má běžet na TPU bez změny samotného modelového kódu. Konkrétní checkpoint a kombinaci funkcí je přesto nutné ověřit v matici podpory.

Co plugin zachovává a co nahrazuje

Balíček vllm-torchtpu se registruje jako platformní plugin a vLLM jej po instalaci automaticky načte. Model dál prochází standardním runnerem vLLM, ale kompiluje se pomocí torch.compile(backend="tpu") do XLA. Každé TPU zařízení dostane vlastní worker proces a komunikace mezi zařízeními používá torch.distributed.

Plánovač, API server, správa workerů a KV cache tak zůstávají součástí upstream vLLM. TPU specifická vrstva leží níž: projekt registruje vlastní JAX/Pallas operace pro attention, Mixture-of-Experts a kvantizované maticové násobení. Oznámení popisuje podporu FP8, INT8, MXFP4 i NVFP4, ale nepublikuje benchmark, ze kterého by šlo odvodit obecné zrychlení proti GPU nebo jinému TPU backendu.

Instalace a první spuštění

Na Cloud TPU VM s Pythonem 3.12 se plugin instaluje příkazem:

pip install vllm-torchtpu

Balíček přidá odpovídající verze vLLM, torch, torch_tpu a libtpu. Oznámení ukazuje spuštění modelu Qwen3-Coder na osmi TPU zařízeních:

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

První start kompiluje grafy pro zvolené tvary a může trvat několik minut. Proměnná VLLM_XLA_CACHE_PATH umožňuje uložit kompilační artefakty a použít je při dalším spuštění.

Podpora má několik úrovní

Projekt uvádí optimalizace pro rodiny Qwen, DeepSeek, Kimi, GLM, Llama, GPT-OSS a Gemma. To však neznamená stejnou úroveň ověření. Tabulka podporovaných modelů a funkcí rozlišuje test při každém pull requestu, noční end-to-end běh, komponentový test a pouhý příklad či recept.

Doporučenou platformou je TPU v7x (Ironwood), na níž běží průběžné CI. Starší v6e, v5e a v5p označuje oznámení jako experimentální. Stejně je potřeba ověřovat jednotlivé kombinace checkpointu, kvantizace, paralelismu, prefix cache, spekulativního dekódování nebo multimodálních vstupů; chybějící kombinace v tabulce nemusí být nefunkční, ale nebyla daným způsobem validována.

TorchTPU tedy snižuje množství kódu, které je nutné přepsat při přechodu z GPU backendu vLLM. Neodstraňuje ale potřebu zkontrolovat podporu konkrétního modelu, zahřívací kompilaci ani výkon na cílové TPU konfiguraci.