A Google natív PyTorch-motort indít a vLLM és az SGLang mellé TPU-khoz

PyTorch-alapú futtatás TPU-kon

A Google Developers Blog október 20-i bejelentése szerint a vLLM-hez az Inferacttal, az SGLanghoz pedig a RadixArkkal együttműködésben készítették el az új natív PyTorch-backendeket. A két motor a TorchTPU-ra épül, amely a TPU-t natív PyTorch-eszközként teszi elérhetővé.

A Google szerint a fejlesztők standard PyTorch-tenzorokkal dolgozhatnak, a kód pedig közvetlenül vagy a torch.compile segítségével, XLA-fordítással futtatható. A kommunikációs műveletek a torch.distributed rendszeren keresztül zajlanak. A vLLM és az SGLang saját munkafolyamat-modelljét, ütemezőjét, KV-gyorsítótár-kezelőjét és API-szerverét megtartja, miközben a TPU-specifikus megvalósítások ezek alatt kapnak helyet.

A megközelítés a meglévő keretrendszerek minél nagyobb részét újrahasznosítja. A modellek többsége a vLLM vagy az SGLang szabványos, upstream PyTorch-definícióit használja, a TPU-optimalizálásokat pedig a kernelek és a rétegimplementációk tartalmazzák.

Ügynöki, multimodális és nagy méretű modellek

A támogatott funkciók között szerepel az automatikus prefixgyorsítótár, a KV-gyorsítótár gazdagépi memóriába és a Google Cloud Managed Lustre tárhelyére történő áthelyezése, valamint az eszközhívás és a strukturált kimenet. Az interaktív alkalmazásokhoz spekulatív dekódolást is kínálnak többek között MTP, EAGLE-3, DFlash, DSpark és n-gram vázlatkészítőkkel.

A motorok támogatják a nagy áteresztőképességű kötegelt feldolgozást, a hosszú kontextust, a külön gazdagépeken végzett előfeldolgozást és dekódolást, továbbá a kép- és videóbemenetet kezelő látás-nyelvi modelleket és a beágyazási modelleket. A nagy modellekhez tensor-, adat-, expert-, pipeline- és context-parallelizmust, valamint FP8, FP4, INT8 és INT4 ellenőrzőpontokat sorol a Google.

A TPU Sync nyílt forráskódú C++ könyvtárként kezeli a KV-gyorsítótár mozgatását és kezelését. A Google szerint a könyvtár közvetlenül olvassa a PyTorch-tenzorok mögötti eszközpuffereket, így nincs szükség további gazdagépi köztes másolatra. Az átvitel külön szálakon futhat, miközben az előfeldolgozás és a dekódolás folytatódik.

Éles használat és nyilvános mérések

A Google közlése szerint mindkét motor nagy léptékben, előzetes ügyfeleknél már éles környezetben fut. A teljesítményt a SemiAnalysis InferenceX közösségi projektben teszik közzé. A Qwen 3.5 Ironwood TPU-n, vagyis TPU 7x-en végzett kezdeti eredményeire építve a Google mérnökei 2,01 centes költséget értek el egymillió tokenenként magas konkurenciájú konfigurációkban.

A vállalat az AgentX ügynöki, hosszú kontextusú és többfordulós feladatokra létrehozott pályáján is új eredményeket publikált. Állítása szerint az Ironwood TPU a Pareto-front nagy részén vezet. A 2,8 billió paraméteres Kimi K3 modellhez készített, Pallas-alapú megakernel egyetlen kernelként futtatja a teljes dekódolási lépést, és a Google közlése szerint egy felhasználónak másodpercenként több mint 700 tokent szolgál ki Ironwood TPU-n.

Az SGLang-TorchTPU a Google szerint már a Metánál is élesben fut. A vLLM közben közösségi TPU-kapacitásalapot indít, amely ingyenes TPU-számítási kapacitást biztosít azoknak a fejlesztőknek, akik aktívan hozzájárulnak a két motor fejlesztéséhez. A Google a jövőben további teljesítmény-optimalizálásokat, fejlettebb funkciókat és új nyílt modellek natív TPU-támogatását tervezi.