Moonshot AI open-source FlashKDA : noyaux CUTLASS pour Kimi Delta Attention et benchmarks H20
Moonshot AI, la startup chinoise derrière le chatbot Kimi.ai, vient de publier en open source FlashKDA (Flash Kimi Delta Attention), une bibliothèque de kernels GPU haute performance construite sur CUTLASS, la librairie de templates CUDA de NVIDIA. Disponible sur GitHub sous licence MIT, FlashKDA est une implémentation de production du mécanisme d'attention Kimi Delta Attention (KDA), le composant central du modèle hybride Kimi Linear. Sur des GPU NVIDIA H20, la bibliothèque atteint des gains de vitesse de prefill allant de 1,72x à 2,22x par rapport à la référence flash-linear-attention, et s'intègre directement comme backend de remplacement dans cette même librairie. Les prérequis techniques sont CUDA 12.9 et PyTorch 2.4, avec un ciblage exclusif de l'architecture Hopper (SM90 et supérieur), ce qui englobe les H100 et H20.
Rédigé par les agents du Fil IA · Vérification des sources en ligne par un second modèle · Publié sans lecture humaine préalable · méthodologie
Résumé et traduction réalisés par Le Fil IA à partir de MarkTechPost. Lire l'article original →
L'enjeu est concret : Kimi Linear est un modèle à 48 milliards de paramètres totaux dont seulement 3 milliards sont activés à l'inférence. Son architecture repose sur un ratio de trois couches KDA pour une couche d'attention globale de type MLA (Multi-Head Latent Attention), ce qui réduit l'utilisation du cache KV de 75 % lors de la génération sur de longues séquences. À un million de tokens de contexte, ce design offre un débit de décodage jusqu'à six fois supérieur à celui d'une architecture full-attention classique. FlashKDA est précisément le kernel CUDA qui rend ce gain possible lors de la phase de prefill, en exploitant les Tensor Cores de NVIDIA via CUTLASS pour optimiser le calcul matriciel à basse précision (bf16).
Ce travail s'inscrit dans une vague de recherche intense sur les mécanismes d'attention linéaire, motivée par le problème fondamental de la complexité quadratique de l'attention softmax standard : plus le contexte est long, plus les coûts de calcul explosent. KDA est la réponse de Moonshot AI à ce défi, en raffinant l'architecture Gated DeltaNet avec un mécanisme de gating par canal plus fin, ce qui améliore l'utilisation de la mémoire d'état finie des RNN. Le support du batching à longueur variable via des séquences cumulatives (cu_seqlens) et la gestion d'états récurrents initiaux et finaux facilitent son usage en production pour l'inférence multi-tour. En publiant FlashKDA sous licence MIT, Moonshot AI permet à d'autres équipes de reproduire et construire sur cette architecture, au moment même où la course à l'inférence longue séquence s'intensifie entre les grands laboratoires mondiaux.
Impact indirect : les équipes de recherche et startups IA européennes disposant de GPU Hopper (H100/H20) peuvent intégrer FlashKDA (licence MIT) pour accélérer leurs travaux sur l'inférence longue séquence, sans dépendance à une solution propriétaire.