Présentation de vLLM TorchTPU

Aujourd'hui, nous publions vLLM TorchTPU (vllm-torchtpu), un nouveau backend pour exécuter vLLM sur les Google Cloud TPUs. Il repose sur la pile TorchTPU de Google et exécute les définitions de modèles déjà présentes dans vLLM, de sorte que la plupart des modèles qui fonctionnent aujourd'hui dans vLLM peuvent être servis sur TPU sans modification du code des modèles.

Ce travail est issu du partenariat d'ingénierie entre Inferact et Google annoncé le mois dernier. Nos objectifs pour ce projet sont de :

1. Faciliter l'apport de nouveaux modèles sur TPU en réutilisant le code de modèles existant de vLLM.
2. Pousser les performances d'inférence TPU en open source, avec des kernels Pallas pour l'attention, le Mixture-of-Experts et le matmul quantisé.
3. Conserver l'expérience vLLM identique : le même CLI, le même serveur compatible OpenAI, le même scheduler et le même gestionnaire de KV cache que vous utilisez sur d'autres matériels.

Fonctionnement

vllm-torchtpu est un plugin de plateforme vLLM hors arbre. Une fois installé, vLLM le détecte automatiquement via le vllm.platform_plugins entry point.

Les modèles s'exécutent via le model runner standard de vLLM et sont compilés avec torch.compile(backend="tpu"), qui génère du code pour XLA. Chaque périphérique TPU dispose de son propre processus worker, de la même manière que vLLM s'exécute sur GPU, et la communication entre les périphériques passe par torch.distributed. Cela signifie que le scheduler, les workers, le gestionnaire de KV cache et le serveur API sont tous du vLLM amont non modifié. Le travail spécifique au TPU se situe en dessous, principalement dans un ensemble de kernels JAX/Pallas enregistrés comme custom ops et appelés depuis le graphe compilé.

Premiers pas

1. Installation

   pip install vllm-torchtpu

   Cela nécessite Python 3.12 sur une VM Cloud TPU, et installe un vLLM correspondant, torch, torch_tpu, et libtpu aux côtés du plugin. Le guide d'installation couvre également uv, les builds depuis les sources et Docker.

2. Servir un modèle

   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

   Le premier démarrage compile les graphes TPU pour les formes configurées, ce qui peut prendre quelques minutes. Définissez VLLM_XLA_CACHE_PATH pour réutiliser les artefacts compilés entre les redémarrages (voir Compilation & cache de préchauffage).

Le décodage spéculatif, le cache de préfixe, le multi-LoRA, les sorties structurées et les entrées multimodales sont tous activés avec les habituels flags de vLLM. Voir le quickstart pour plus d'exemples, notamment l'inférence par lots hors ligne, les modèles multimodaux et le multi-LoRA.

Kernels

vllm-torchtpu utilise des kernels Pallas personnalisés pour les opérations complexes qui bénéficient d'un réglage manuel sur TPU. Quelques exemples :

• Ragged Paged Attention v3 pour les modèles MHA/GQA
• Multi-head Latent Attention (MLA) dense et creuse pour DeepSeek, Kimi et GLM
• Gated DeltaNet et Kimi Delta Attention pour les modèles hybrides à attention linéaire
• Mixture-of-Experts fusionné avec parallélisme d'experts
• matmuls quantisés FP8, INT8, MXFP4 et NVFP4

Passage à l'échelle

Nous nous sommes fortement concentrés sur les charges de travail qui exercent aujourd'hui le plus de pression sur les systèmes de service, en particulier le trafic agentique multi-tours et multimodal, où les contextes continuent de croître et où les sorties des outils peuvent être longues.

Parallélisme

Les six dimensions de parallélisme de vLLM sont prises en charge, aussi bien sur des tranches TPU à hôte unique qu'à plusieurs hôtes, et elles peuvent être combinées :

Pour les grands modèles MoE, on obtient généralement le meilleur débit avec l'attention en parallélisme de données et le parallélisme d'experts (par exemple, DP=8, TP=1 avec EP sur un seul hôte v7x). Cela maintient l'attention locale à chaque dispositif et répartit les experts sur l'ensemble de la tranche. Pour les longs contextes, PCP et DCP divisent le préremplissage et le cache KV entre les dispositifs. Le guide de parallélisme contient plus de détails et des exemples de configurations.

Déchargement du cache KV et service désagrégé

vllm-torchtpu utilise TPU Sync, une bibliothèque C++ open source capable de lire directement les tampons des dispositifs TPU, en dehors du runtime Python. Via l'interface de connecteur KV de vLLM, cela nous donne :

• Déchargement du cache KV sans copie vers la DRAM de l'hôte, afin que le cache de préfixe pour les longues sessions multi-tours ne soit pas limité par la HBM.
• Service désagrégé préremplissage/décodage entre hôtes TPU avec un faible surcoût de transfert KV. Le préremplissage et le décodage peuvent utiliser des dispositions différentes, par exemple le préremplissage sur PCP=4, TP=2 et le décodage sur DP=8, TP=1, et le cache KV est redivisé pendant le transfert.

Consultez le guide du service désagrégé pour le mettre en place.

Décodage spéculatif et mise en cache de préfixes

Le décodage spéculatif prend en charge MTP, EAGLE-3, DFlash, DSpark, n-gram et le décodage par suffixe (guide). La mise en cache automatique des préfixes fonctionne également pour les modèles hybrides comme Qwen3.5 et Kimi K3, où l'état d'attention linéaire est mis en cache avec le cache KV.

Modèles, fonctionnalités et perspectives

vllm-torchtpu peut servir la plupart des modèles pris en charge par vLLM, avec un travail TPU spécifique sur les noyaux, la quantification et le parallélisme pour les familles Qwen, DeepSeek, Kimi, GLM, Llama, GPT-OSS et Gemma. Le TPU v7x (Ironwood) est le matériel recommandé et est testé en CI à chaque pull request ; le v6e, le v5e et le v5p sont expérimentaux.

La page Supported Models & Features contient la liste complète des points de contrôle et fonctionnalités validés, avec la manière d'activer chacun d'eux. Les matrices de support nocturnes suivent la correction et les résultats de benchmarks du TPU v7x issus de la CI.

Et maintenant ?

Nous voulons que les nouveaux modèles ouverts des grandes familles (Qwen, DeepSeek, GLM, Kimi, Gemma, GPT-OSS) fonctionnent sur TPU le jour même de leur sortie. Une partie de ce qui arrive, issue de la feuille de route publique:

• Un pipeline CI automatisé qui intègre, évalue (benchmark) et publie des matrices de compatibilité pour les nouveaux modèles à poids ouverts dans les 24 heures suivant leur sortie.
• Routage conscient des préfixes, du cache et des sessions, et routage par niveaux de mémoire TPU de llm-d.
• Interfaces PyTorch pour noyaux Tokamax , partagés avec d'autres projets de l'écosystème TorchTPU.
• Le service de modèles de diffusion via vllm-omni.

Les contributions sont les bienvenues !

Capacité TPU pour les contributeurs

Nous ouvrons également un pool de capacité TPU pour contributions communautaires. Si vous développez sur vllm-torchtpu, vous pouvez demander du temps TPU gratuit, et les mainteneurs principaux de vLLM chez Inferact vous aideront pour la conception et la revue des PR.

Demandez ici l'accès TPU contributeur.

Essayez-le

• Démarrage rapide
• Configuration de Cloud TPU
• Recettes de service TPU v7x (Ironwood)
• Recettes vLLM TorchTPU sur Ironwood
• Déploiement Kubernetes
• Guide de contribution

Si vous avez des questions ou souhaitez participer, retrouvez-nous sur #sig-tpu sur le Slack des développeurs vLLM.