Les modèles de décoder autoregressifs, au cœur de nombreuses applications d’intelligence artificielle générative, produisent des tokens un par un, une approche qui semble inefficace au premier abord. Face à la demande croissante en IA générative, des efforts ingénieux sont déployés pour optimiser ces modèles, que ce soit via des noyaux CUDA personnalisés, des CUDA Graphs, des accélérateurs dédiés ou des techniques comme l’échantillonnage spéculatif. Même une amélioration marginale de la latence ou du coût est considérée comme un gain significatif. Cet article présente une technique d’optimisation de la génération de tokens dans PyTorch en utilisant l’interleaving de flux CUDA. Bien que simple à implémenter, cette méthode cible un goulot d’étranglement souvent négligé et peut entraîner des gains de performance notables. L’expérience repose sur un modèle GPT-2 simple de Hugging Face, exécuté sur une GPU NVIDIA L40S avec PyTorch 2.10.0. Le code présenté est à but pédagogique et ne doit pas être utilisé en production sans benchmarking préalable. L’optimisation cible les charges d’inférence natives PyTorch, courantes en développement, bien que pour les environnements de production, des bibliothèques dédiées comme vLLM ou NVIDIA TensorRT-LLM offrent généralement de meilleures performances. Le modèle de base affiche une complexité temporelle en O(N²) : le temps de génération quadruple lorsque la longueur de séquence double, tandis que la fragmentation mémoire CUDA reste élevée, résultant d’allocations de tenseurs croissantes. Cette fragmentation nuit à l’efficacité mémoire et peut provoquer des erreurs. La première optimisation consiste en un cache KV (Key-Value), qui stocke les activations intermédiaires pour éviter de recalculer toute la séquence à chaque étape, réduisant la complexité à O(N). Ce gain est spectaculaire, surtout pour des séquences longues. Cependant, la fragmentation mémoire persiste. Deux solutions sont explorées : les allocations mémoire expandables (via PYTORCH_ALLOC_CONF="expandable_segments:True") et le StaticCache de Hugging Face, qui pré-alloue un espace fixe pour le cache. Le StaticCache diminue fortement la fragmentation, mais introduit un coût computationnel supplémentaire, car l’attention s’applique à toute la longueur maximale, même pour les tokens non pertinents. Pour des séquences courtes, cela peut ralentir le modèle ; pour des séquences longues (200 et 400 tokens), il améliore les performances de 9 à 10 %. La compilation du modèle avec torch.compile amplifie ces gains, en particulier avec un cache statique, car elle permet des optimisations JIT sur des graphes de calcul à taille fixe. Toutefois, le véritable goulot d’étranglement réside dans la vérification de fin de séquence (EOS) à chaque étape. L’appel stop_gpu.item() provoque une synchronisation bloquante entre CPU et GPU, entraînant des périodes d’idling du GPU — un problème critique en haute performance. Pour résoudre cela, une technique d’interleaving de flux CUDA est proposée. Deux flux sont utilisés en mode « ping-pong » : pendant que le GPU calcule le token i, la CPU prépare le token i+1 dans un autre flux, tout en vérifiant la condition d’arrêt du token i-1. Grâce à des synchronisations non bloquantes et à une gestion explicite des flux, le GPU reste actif tout au long du processus. Le profilage avec Nsight™ Systems montre une utilisation constante du GPU, sans idling, et une amélioration significative du débit. Les résultats montrent que cette méthode, combinée au StaticCache et à torch.compile, peut améliorer la performance jusqu’à 5 fois par rapport à la version de base. L’efficacité dépend fortement du rapport entre le temps de chargement des noyaux et le temps de calcul, avec un gain maximal pour de petites tailles de batch. Toutefois, l’utilisation de flux CUDA comporte des risques : erreurs CUDA, corruption de données ou augmentation de la fragmentation mémoire si les synchronisations ne sont pas correctement gérées. Cette technique doit donc être appliquée avec prudence, après un benchmark adapté à l’environnement. En résumé, l’interleaving de flux CUDA est une optimisation puissante pour masquer les latences liées aux synchronisations hôte-périphérique dans les générateurs autoregressifs PyTorch. En combinant cette méthode avec le StaticCache, la compilation et une gestion rigoureuse des flux, il est possible d’atteindre des performances proches du maximum théorique sur les modèles de décoder.