A vLLM mostantól Google Cloud TPU-kon is futtatható

A vLLM modelljei változtatás nélkül futhatnak TPU-kon

A vLLM október 20-án mutatta be a vllm-torchtpu csomagot, amely a vLLM futtatását teszi lehetővé Google Cloud TPU-kon. A háttér a Google TorchTPU technológiájára épül, és a vLLM-ben már meglévő modellmeghatározásokat használja. Ennek köszönhetően a vLLM-ben működő modellek többsége a modellkód módosítása nélkül szolgálható ki TPU-n.

A projekt az Inferact és a Google mérnöki együttműködéséből született. A készítők célja, hogy az új modellek TPU-ra portolása a meglévő vLLM-kód újrahasznosításával történjen, miközben nyílt forráskódú környezetben javítják a TPU-s következtetés teljesítményét.

A felhasználói élmény is a vLLM megszokott működését követi. A parancssori felület, az OpenAI-kompatibilis szerver, az ütemező és a KV-gyorsítótár kezelője változatlanul az upstream vLLM része marad.

A TPU-specifikus működés a vLLM alatt kap helyet

A vllm-torchtpu külső vLLM-platformbővítményként működik. Telepítés után a vLLM automatikusan felismeri a vllm.platform_plugins belépési ponton keresztül. A modelleket a vLLM szabványos modellfuttatója kezeli, majd a rendszer a torch.compile(backend="tpu") használatával XLA-ra fordítja őket.

Minden TPU-eszköz saját munkafolyamatot kap, ugyanúgy, ahogy a vLLM GPU-k esetén működik. Az eszközök közötti kommunikációt a torch.distributed biztosítja. A TPU-ra szabott kód főként JAX és Pallas kernelekből áll, amelyeket egyedi műveletként regisztrálnak, majd a fordított gráfból hívnak meg.

A csomag telepítése a pip install vllm-torchtpu paranccsal történik. Ehhez Python 3.12 és Cloud TPU VM szükséges. A telepítő a bővítmény mellett megfelelő vLLM-, torch-, torch_tpu- és libtpu-verziókat is telepít. Az első indításkor a rendszer lefordítja a beállított alakzatokhoz tartozó TPU-gráfokat, ami néhány percig tarthat. A lefordított elemek újraindítások közötti megőrzésére a VLLM_XLA_CACHE_PATH környezeti változó használható.

Kiterjedt párhuzamosítás és gyorsítótár-kezelés

A vllm-torchtpu mind a vLLM hat párhuzamosítási dimenzióját támogatja egyetlen gépen és többgépes TPU-szeleteken is. Ezek a tenzorpárhuzamosítás, az adatok párhuzamosítása, a csővezeték-párhuzamosítás, a szakértői párhuzamosítás, valamint a prefillszakasz és a dekódolás kontextusának párhuzamosítása. A beállítások kombinálhatók is.

A fejlesztők különösen a többfordulós ügynöki és multimodális forgalomra összpontosítottak, ahol a kontextus folyamatosan nőhet, a meghívott eszközök kimenete pedig hosszú lehet. Hosszú kontextusoknál a PCP és a DCP a prefillszakaszt, illetve a KV-gyorsítótárat is több eszköz között oszthatja fel.

A rendszer támogatja a KV-gyorsítótár gazdagépi DRAM-ba történő, másolás nélküli kihelyezését. Így a hosszú, többfordulós munkamenetek prefix-gyorsítótára nincs kizárólag a HBM kapacitására korlátozva. A prefillszakasz és a dekódolás külön TPU-gazdagépekre is szétválasztható, a KV-gyorsítótár pedig az átvitel során újraosztható.

A TPU v7x az ajánlott hardver

A vllm-torchtpu a vLLM által támogatott modellek többségét képes kiszolgálni. TPU-specifikus kernel-, kvantálási és párhuzamosítási munkát végeztek a Qwen, DeepSeek, Kimi, GLM, Llama, GPT-OSS és Gemma modellcsaládokhoz. A rendszerben elérhető a spekulatív dekódolás, a prefix-gyorsítótárazás, a több LoRA, a strukturált kimenet és a multimodális bemenet is.

Az ajánlott hardver a TPU v7x, más néven Ironwood, amelyet minden pull request esetén CI-tesztekkel ellenőriznek. A v6e, v5e és v5p TPU-kísérleti támogatásúnak számítanak. A projekt fejlesztői azt szeretnék elérni, hogy az új nyílt modellek a Qwen, DeepSeek, GLM, Kimi, Gemma és GPT-OSS családokból már megjelenésük napján fussanak TPU-n.

A nyilvános ütemtervben szerepel egy olyan automatizált CI-folyamat is, amely az új nyílt súlyú modelleket 24 órán belül felveszi, teljesítményét megméri, majd közzéteszi a támogatási mátrixokat. A fejlesztők közösségi TPU-kapacitási programot is indítanak, amelyben a vllm-torchtpu építői ingyenes TPU-időre jelentkezhetnek, és az Inferact vLLM-karbantartóitól tervezési, illetve pull request-áttekintési segítséget kaphatnak.