Představení vLLM TorchTPU: Nový backend pro Google Cloud TPUs

Tým vLLM oznámil uvedení vLLM TorchTPU, nového backendu určeného pro běh modelů vLLM na Google Cloud TPUs. Tento backend je postaven na Google TorchTPU stacku a umožňuje nasazení modelů bez nutnosti měnit jejich kód, což usnadňuje práci vývojářů.

Jak to funguje

vLLM TorchTPU je plugin, který se automaticky načítá a používá standardní model runner vLLM. Všechny modely, které již fungují v rámci vLLM, mohou být nasazeny na TPU bez jakýchkoli úprav. To znamená, že vývojáři mohou snadno přenést své existující modely na TPU, aniž by museli investovat čas do přepisování kódu.

Backend využívá Pallas jádra pro zlepšení výkonu inference, zejména při operacích jako attention a Mixture-of-Experts. Tato optimalizace může přispět k rychlejšímu zpracování a efektivnějšímu využití TPU.

Instalace a spuštění

Pro instalaci vLLM TorchTPU je třeba použít následující příkaz:

pip install vllm-torchtpu

Tento příkaz vyžaduje Python 3.12 na Cloud TPU VM a nainstaluje potřebné komponenty vLLM, torch, torch_tpu a libtpu. Po úspěšné instalaci lze model spustit pomocí příkazu:

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í spuštění modelu zahrnuje kompilaci TPU grafů, což může trvat několik minut. Pro opětovné použití kompilovaných artefaktů je doporučeno nastavit proměnnou VLLM_XLA_CACHE_PATH.

Praktické omezení

Zatím není jasné, jaké konkrétní modely byly úspěšně testovány na TPU, což je důležité vzít v úvahu při plánování nasazení.