Cet article détaille l'optimisation du noyau d'attention 2-simpliciale en utilisant TLX (Triton Low-Level Extensions) pour améliorer l'efficacité matérielle sur les GPU. L'attention 2-simpliciale, qui modélise les interactions entre triples de tokens, est plus complexe que l'attention par produit scalaire standard, mais offre une meilleure efficacité en termes de tokens, particulièrement utile pour les tâches de raisonnement logique et mathématique. Cependant, cette complexité entraîne une croissance cubique de la charge de calcul avec la longueur de la séquence, rendant son utilisation complète peu pratique. Pour pallier ce problème, une fenêtre glissante bidimensionnelle est introduite, limitant le champ d'attention de chaque token de requête à un voisinage local, ce qui réduit considérablement la complexité computationnelle tout en préservant l'expressivité du modèle. L'article présente également TLX, une extension du langage Triton qui permet un contrôle fin sur l'exécution des noyaux GPU, offrant des performances élevées tout en maintenant une productivité de développement. En utilisant TLX, les auteurs ont réécrit le noyau d'attention 2-simpliciale pour atteindre jusqu'à 588 TFLOPs sur les GPU NVIDIA H100, soit une amélioration de 60 % par rapport aux noyaux précédents. Cette optimisation est rendue possible grâce à une conception alignée sur le matériel, incluant une décomposition stratégique des opérations ternaires en opérations binaires pour tirer parti des Tensor Cores. Cette approche permet de transformer l'attention 2-simpliciale en une conception compatible avec les Tensor Cores, tout en conservant l'équivalence mathématique du modèle original.