Le code qui colle à la carte graphique passe en MIT

La bibliothèque FlashKDA de Moonshot AI double la vitesse de prefill sur GPU Hopper. Une intégration transparente désormais sous licence open source MIT.

Un modèle de langage passe l'essentiel de son temps de calcul dans quelques opérations très répétées, dont le mécanisme d'attention. Les exécuter vite suppose du code écrit au plus près de la carte graphique, ce qu'on appelle des kernels. Moonshot AI diffuse les siens sous le nom de FlashKDA, une implémentation bâtie sur CUTLASS des kernels de Kimi Delta Attention, le mécanisme d'attention maison qui équipe ses modèles.

Le gain annoncé porte sur une phase précise. Sur une carte H20, l'entreprise revendique un prefill, c'est-à-dire l'ingestion initiale du contexte avant que le modèle ne commence à répondre, accéléré de 1,72 à 2,22 fois. La comparaison se fait contre l'implémentation de référence de flash-linear-attention, donc contre un autre code exécutant le même mécanisme, et non contre l'attention classique. Des relevés pour la génération GB200 accompagnent le dépôt.

L'intégration se veut indolore. Une fois installée, la bibliothèque se substitue automatiquement au chemin habituel de flash-linear-attention sans modification du code appelant, une variable d'environnement permettant de revenir à l'implémentation en Triton. Un mode de journalisation indique si la substitution a bien eu lieu, ou pour quelle raison elle a été écartée. Les tests fournis vérifient l'égalité exacte des sorties avec l'implémentation de référence.

Les contraintes sont posées sans détour. L'ensemble réclame des cartes de génération Hopper ou plus récente, CUDA 12.9 et PyTorch 2.4 au minimum, et n'accepte pour l'instant que des dimensions de clé et de valeur fixées à 128. Le code est diffusé sous licence MIT.