Le moteur de service SGLang vient d’intégrer un plugin pour les Google TPUs, annonçant son support natif de l’écosystème PyTorch sur ces accélérateurs. Ce développement, selon github.com, permet d’exécuter des modèles de langage et de vision-langage en utilisant l’API compatible OpenAI, tout en conservant les flux de travail PyTorch habituels.

Ce plugin, SGLang-TorchTPU, exploite torch_tpu et les noyaux Pallas pour adapter SGLang au TPU7x (Ironwood). Les développeurs peuvent ainsi configurer diverses topologies de service. On passe du déploiement sur une seule machine hôte avec quatre puces à des configurations multi-hôtes, incluant le service désagrégé de pré-remplissage et de décodage.

L’efficacité à l’échelle est au cœur de cette intégration. Elle repose sur la combinaison de techniques comme le prefix caching, le déchargement HiCache KV vers la mémoire hôte, ainsi que le parallélisme de tenseur, d’expert, de données et de pipeline. Les performances sont encore affinées par l’intégration du décodage spéculatif, explorant des approches telles qu’EAGLE3 ou DFlash.

L’optimisation au niveau du matériel

Cette mise à jour ne se contente pas d’être une simple portabilité. Elle est intrinsèquement liée aux capacités matérielles des TPUs. Les TPUs, comparés aux GPU, sont des machines intrinsèquement séquentielles dotées d’un registre vectoriel très large, ce qui impose des contraintes spécifiques sur la manière de structurer les calculs. Les noyaux Pallas gèrent cette complexité.

Les chercheurs ont montré comment cette architecture permet d’optimiser des modèles comme Muse Spark Open, un modèle mixture-of-experts (MoE) de 816 milliards de paramètres. En fusionnant les 62 couches de décodage en un unique programme Pallas, chaque étape de décodage ou de vérification devient un flux continu d’accès aux poids depuis la mémoire HBM, minimisant les arrêts qui freinent la vitesse.

Une analyse de PyTorch indique que cette approche a permis de diviser par deux le temps par jeton de sortie pour une requête concurrente, bien que ce gain se réduise avec l’augmentation de la concurrence. De plus, la gestion de la mémoire, notamment grâce au L2 HiCache sur l’hôte, permet de réduire l’utilisation de la HBM par puce, améliorant l’efficacité globale.

Modularité pour modèles complexes

La capacité à séparer les composants est un atout majeur pour les modèles de vision-langage. Les sources indiquent que l’encodeur visuel peut être exécuté sur un serveur dédié, envoyant ensuite les plongements d’images (embeddings) à un serveur de langage seul. Cette désagrégation permet à chaque partie du système de choisir sa propre taille de TPU et son hôte, offrant une souplesse inédite dans le déploiement.

Les résultats montrent que cette séparation améliore la latence. Pour un modèle de 235B-A22B, le temps de premier jeton (TTFT) reste bas même à cinq requêtes par seconde, une performance qui se dégrade significativement si l’encodeur est hébergé sur le même hôte.

Pour mettre en place cette architecture, on peut lancer séparément un serveur d’encodeur en mode encoder-only et un serveur de langage en mode language-only, qui se connecte à l’encodeur via des URL. Cela assure une indépendance totale entre les deux fonctions. Les recettes de modèles sont disponibles pour plusieurs architectures, incluant des versions de Gemma et DeepSeek.

Sources

• GitHub – sgl-project/sglang-torchtpu: SGLang’s TorchTPU Plugin, github.com
• Writing TPU kernels with Pallas, docs.jax.dev
• Optimizing Muse Spark Open on TPU v7 with SGLang-TorchTPU, PyTorch, 10 novembre 2026
• Meta’s Muse Ember, Muse Spark 1.3, and Muse Wonder land on Google Cloud, Crypto Briefing, 9 novembre 2026