Un kernel CUDA Top-K accélère le RAG agentic sur le GPU
Un développeur a conçu un noyau CUDA personnalisé optimisant la récupération de similarité vectorielle pour les systèmes d'IA agents. Cette innovation vise à éliminer un goulot d'étranglement majeur : les transferts inutiles de données entre le processeur et la carte graphique. Dans les architectures RAG standard, chaque requête d'un agent nécessite une recherche de similarité. Le processus habituel consiste à extraire l'embedding de la requête vers la mémoire vive, à effectuer le calcul de similarité sur le CPU, puis à renvoyer les résultats sur le GPU. Cette navette constante sur le bus PCIe ralentit considérablement le pipeline, annulant les gains de puissance de calcul du processeur graphique. La solution proposée consiste à héberger intégralement le corpus de référence dans la mémoire VRAM du GPU. Un noyau personnalisé de 343 lignes calcule les produits scalaires, sélectionne les k résultats les plus pertinents et fusionne les partiels entièrement sur place. Seuls l'embedding de la requête et les k indices finaux traversent le bus PCIe. Cette approche réduit drastiquement la latence et libère les ressources CPU. Les tests réalisés sur une carte NVIDIA GTX 1080 démontrent une accélération allant jusqu'à 8,6 fois par rapport aux méthodes basées sur le CPU, et ce sur des corpus allant de 10 000 à 1 million d'entrées. Pour des valeurs de k comprises entre 8 et 32, le noyau GPU surpasse systématiquement les solutions classiques. Au-delà de 100 résultats, la version actuelle, qui privilégie la lisibilité du code et la vérification exacte des résultats à l'optimisation extrême, voit ses performances dépassées par le CPU. L'auteur note explicitement cette limite et annonce qu'une version utilisant une sélection par tournoi optimisée pour les warps GPU sera développée. Au-delà des performances, cette implémentation soulève un parallèle intéressant avec les télécommunications. Le mécanisme de sélection des meilleurs vecteurs rappelle de près le choix des faisceaux radio dans les stations de base 5G, où le terminal doit identifier rapidement les signaux les plus forts sans surcharger le processeur central. Cette architecture open source met en lumière une réalité souvent négligée dans le développement d'agents IA : la latence réseau interne peut devenir plus critique que la puissance de calcul elle-même. Ce travail s'inscrit dans une série dédiée à l'optimisation des infrastructures d'inférence IA. Après avoir éliminé les calculs de préremplissage redondants et permis le partage temporel de la carte graphique entre multiples micro-agents, cette étape consolide le traitement en mémoire GPU. La phase suivante portera sur la persistance de l'état des agents entre les requêtes, afin d'éviter les réinitialisations à froid qui pénalisent les raisonnements multi-sauts. Le code complet et la procédure de benchmark sont disponibles publiquement, permettant aux chercheurs et ingénieurs de reproduire les résultats et d'adapter l'approche à leurs propres déploiements. Cette démonstration confirme que la maîtrise des déplacements de données est aussi cruciale que l'optimisation algorithmique pour faire passer les architectures IA de la démonstration à la production.
