Gros plan sur une puce de calcul spécialisée avec des motifs de circuits complexes sur un établi sombre.

Tokamax : bibliothèque de noyaux optimisés pour les accélérateurs IA

Par Rédaction IP73 dans MLOps et infrastructure IA le

Tokamax est une bibliothèque de noyaux personnalisés qui apporte des implémentations hautement optimisées d’opérateurs essentiels pour les modèles d’intelligence artificielle. Elle cible spécifiquement les accélérateurs matériels récents, incluant les GPU NVIDIA et les TPU de Google. Cette collection de noyaux couvre des blocs fondamentaux comme la multiplication matricielle groupée, utilisée dans les architectures mixture-of-experts, ainsi que diverses variantes d’attention sophistiquées, telles que l’attention paginée (paged attention) ou l’attention DeltaNet.

La bibliothèque est conçue pour évoluer, afin de supporter les derniers modèles open-source sur le matériel disponible. Pour aider les développeurs à écrire, maintenir et déployer ces noyaux, Tokamax fournit des infrastructures comme l’auto-optimisation (autotuning) et des outils de benchmarking à faible surcharge. Les noyaux sont habituellement construits sur JAX et Pallas, mais l’équipe propose aussi des API PyTorch natives pour les TPU, ce qui simplifie l’adoption. Les utilisateurs peuvent installer la dernière version via pip install -U tokamax ou récupérer la version en développement directement depuis le dépôt GitHub.

L’optimisation matérielle est clairement devenue un point critique. Cette bibliothèque attaque ce problème en fournissant des briques de calcul très fines pour les opérations les plus gourmandes.

Illustration de flux de données segmentées traversant un chemin en hexagones, évoquant le parallélisme contextuel.

L’extension des fenêtres de contexte sur les TPU

Cette quête d’efficacité prend tout son sens face aux limites des contextes. L’allongement des séquences est devenu une nécessité pour les modèles de pointe, leur permettant d’ingérer des dépôts de code entiers ou de raisonner sur des tâches complexes. Google, via son blog de développeurs, annonce d’ailleurs le support de l’entraînement en contexte long dans MaxText, son cadre d’entraînement open-source JAX sur Google Cloud TPU. MaxText permet cela en soutenant quatre stratégies de parallélisme contextuel sur TPU : All-gather, A2A (DeepSpeed Ulysses), P2P (Ring Attention), et leur combinaison hybride, USP.

Ces stratégies permettent de répartir une séquence longue sur plusieurs puces, évitant ainsi qu’une seule puce ne doive stocker toutes les activations. L’utilisation du noyau Splash Attention de Tokamax est ici déterminante. Ce noyau permet de sauter les blocs masqués et de gérer les séquences empaquetées sans avoir besoin de remplissage (padding). Les résultats sont significatifs : un modèle Llama 3.1 8B affiche une utilisation des FLOPs de 70 % sur un TPU 5p pour des séquences allant jusqu’à 4 millions de tokens. Sur TPU 7, le contexte peut être étendu jusqu’à 14 millions de tokens, maintenant 50 % d’utilisation des FLOPs sous 6 millions de tokens.

Architectures avancées et convergence des systèmes

Si Tokamax optimise le niveau du noyau, d’autres annonces montrent l’application de cette optimisation dans des systèmes plus larges. Google étend d’ailleurs son portefeuille de TPU de huitième génération avec le TPU 8tx et le TPU 8i Max, visant le calcul confidentiel et l’inférence gourmande en mémoire. Le TPU 8i Max, par exemple, intègre deux puces TPU 8i pour doubler la capacité de calcul et la bande passante mémoire. Parallèlement, les notes de version de MaxText détaillent l’intégration d’architectures très spécifiques, comme GLM-5.3-Flash ou Qwen3.8-Flash-Next. Ces modèles incorporent des concepts complexes tels que les Hyper-Connections ou les Gated DeltaNet.

Il est évident que le travail ne se limite plus à faire tourner des modèles, mais à optimiser chaque étape du chemin de calcul. Les méthodes de parallélisme contextuel, qu’elles soient implémentées via les noyaux de Tokamax ou par des schémas comme le USP, montrent que la compétition se joue sur la réduction de la communication inter-puces tout en maximisant le calcul local. Il faudra suivre comment ces bibliothèques de noyaux interagiront avec les systèmes de parallélisme plus vastes pour que les modèles atteignent effectivement des fenêtres contextuelles de l’ordre du million de tokens sans nécessiter des ressources hors de portée.

Sources

Partager :FacebookXLinkedInBlueskyE-mail

Articles récents

Prochains événements

Tous les événements →