CuTeDSL chez Perplexity

Chez Perplexity, nous utilisons le moteur d'exécution maison ROSE (Runtime-Optimized Serving Engine) pour exécuter des modèles allant des plongements aux LLM à des billions de paramètres. Situé en amont de toutes les API, notamment Sonar, Search et Embeddings, notre moteur a été conçu pour permettre aux chercheurs d'explorer les modèles et architectures les plus adaptés, tout en déployant rapidement les meilleurs candidats en production. Pour y parvenir, nous devons être capables d'implémenter, de déployer et d'optimiser l'inférence pour des architectures de pointe à un rythme soutenu. Dans cet article de blogue, nous examinons comment les noyaux CuTeDSL intégrés à notre moteur d'inférence nous permettent de créer les noyaux GPU nécessaires pour amener les modèles à des performances optimales sur les GPU NVIDIA.

Des modèles personnalisés, créés par différentes équipes de Perplexity, sont hébergés en interne sur des GPU NVIDIA Hopper et Blackwell à l'aide du moteur d'inférence ROSE conçu par l'équipe d'inférence IA. ROSE a initialement été développé pour mieux prendre en charge des modèles Llama personnalisés pour le décodage et la classification de texte, avec une interface compatible avec le serveur d'inférence NVIDIA Triton. Au fil du temps, il est devenu un moteur d'inférence à part entière prenant en charge des LLM d'envergure, ainsi que des modèles de classement, de classification, de notation et d'autres modèles basés sur les transformateurs pour un large éventail d'équipes afin de propulser les produits de Perplexity. L'objectif de conception de ROSE est d'être simple : avec le temps, des efforts sont investis pour atteindre des performances optimales, mais la priorité absolue est de s'adapter aux architectures et cas d'utilisation en constante évolution de nos modèles.

L’espace de configuration du noyau, selon les types de données, les dimensions de tête, les constantes à la compilation et les paramètres de lancement.

À un niveau élevé, ROSE se compose d'un moteur qui adapte un modèle à une interface avec laquelle les clients peuvent interagir. Les moteurs sont hautement réutilisables, gérant la planification et le traitement par lots des requêtes, ainsi que la récupération des jetons décodés, des plongements ou des scores des modèles sous-jacents. L'initialisation des appareils, les canaux de communication inter-processus et inter-nœuds, ainsi que le chargement des poids sont tous pris en charge par ces derniers. Pour les LLM, le moteur est responsable du traitement par lots, du découpage (chunking) et de l'échantillonnage. Il est également responsable de l'allocation du stockage KV, pour l'attention complète et linéaire, ainsi que de la correspondance de préfixes. Les moteurs de plongement sont plus simples, car ils effectuent uniquement un traitement par lots en ligne. L'exécution est optimisée de telle sorte que le moteur prépare la série de travaux suivante pendant que le modèle s'exécute en parallèle sur l'accélérateur.

ROSE expose une riche collection de couches personnalisées pour la modélisation, enveloppant les noyaux GPU construits avec CuTeDSL, Triton, CUDA, CUTLASS et cuBLAS. Pour les implémentations de multiplication matricielle et d'attention, nous comptons principalement sur les noyaux NVIDIA de CUTLASS et cuBLAS. Le reste du pipeline, y compris les noyaux de plongement, de norme, de MoE et d'activation, fait partie de la bibliothèque de noyaux de ROSE. Pour prendre au mieux en charge le matériel sur lequel nous exécutons l'inférence, ROSE est doté de ses propres primitives de communication inter-appareils pour le routage, l'acheminement et la combinaison MoE. Bien que les autres noyaux soient relativement simples, ils doivent être hautement spécialisés pour atteindre des performances optimales sur toutes les différentes dimensions cachées et tous les appareils sur lesquels ils s'exécutent. De plus, nous comptons sur la spécialisation au moment de la compilation pour accommoder tous les réglages et ajustements de divers modèles (biais de poids ou absence de biais de poids pour la norme RMS, fonctions d'activation, etc.). Par conséquent, nous avons décidé d'utiliser des DSL pour les implémenter : nous avons initialement construit ROSE avec Triton. Cependant, nous avons depuis adopté CuTeDSL comme environnement principal pour la programmation GPU.

D'un point de vue mathématique, les noyaux d'inférence sont généralement simples. Par conséquent, lors du développement d'un noyau, la complexité ne réside généralement pas dans l'expression du calcul à effectuer, mais dans la manière exacte dont le matériel doit le réaliser. Pour créer les meilleurs noyaux, nous avons besoin d'un langage qui offre les bonnes abstractions sur le matériel sous-jacent. Le bon langage nous permet de spécifier les opérations de haut niveau pour effectuer le calcul, en les spécialisant et en les optimisant agressivement à travers tous les points d'un vaste espace de configuration.

CuTeDSL est un DSL basé sur Python pour les GPU NVIDIA, construit au sommet de l'algèbre de mise en page CuTe et de MLIR. Il compile juste à temps (JIT) en un code PTX hautement optimisé, tout en conservant le contrôle des primitives matérielles de bas niveau. Il vise à offrir les mêmes fonctionnalités et performances que CUTLASS C++, sans les temps de compilation prohibitifs qui nuisent à la vitesse de développement. Parmi les langages que nous avons examinés ou utilisés (notamment CUDA, cuTile, TVM IR et Triton), nous le considérons comme le choix le plus adapté au développement de nos noyaux.

Dans l'ensemble, bien qu'il ne soit pas aussi haut niveau et concis que Triton ou cuTile, nous avons constaté que l'algèbre de mise en page CuTe et ses compositions (exposées via cute.Tensor) constituaient des substituts raisonnables. De plus, CuTeDSL offre un accès sans restriction aux primitives matérielles, nous fournissant le niveau de contrôle requis pour atteindre des performances optimales dans tous les cas limites. Avec CUDA, les opportunités de spécialisation à la compilation étaient limitées, car chaque paramètre supplémentaire ajoutait un coût déraisonnable aux temps de compilation ou de la complexité au système de construction. En revanche, avec un DSL compilé juste à temps, nous pouvons nous spécialiser de manière plus agressive sans sacrifier le rythme d'itération. Contrairement à TVM IR, qui utilisait CUDA comme étape intermédiaire, nous avons trouvé que le pipeline basé sur MLIR qui convertit vers PTX était plus robuste et plus facile à déboguer qu'un pipeline émettant CUDA comme étape intermédiaire. De plus, grâce à l'utilisation d'opérations personnalisées pour envelopper les instructions PTX, nous avons constaté que nous pouvions exploiter pleinement toutes les fonctionnalités du matériel sous-jacent.

L'essentiel du coût de compilation des noyaux CUDA est généralement consacré à la phase d'expansion des modèles (templates) des noyaux, l'optimisation et la compilation à partir de l'étape PTX pouvant être réalisées en quelques secondes. Avec de nombreux paramètres à la compilation, cette étape peut exploser pour atteindre plusieurs minutes pour un noyau non trivial. Ce goulet d'étranglement est accentué par le fait que les opportunités de compilation parallèle n'apparaissent que plus tard, aux niveaux IR et PTX, la gestion des modèles C++ se faisant de manière séquentielle. Bien qu'il existe des astuces pour résoudre ce problème, en déplaçant les instanciations de modèles vers des unités de compilation séparées, elles nécessitent des ajustements importants des systèmes de construction et l'utilisation de macros lourdes. En revanche, les noyaux CuTeDSL compilés juste à temps (JIT) ne nécessitent pas de règles de construction complexes et n'entravent pas les temps de compilation, car le cycle de développement peut se concentrer sur un seul point de l'espace de configuration, tandis que des tests unitaires minutieux ultérieurs peuvent couvrir l'ensemble de l'espace de conception du noyau.

Même si la compilation JIT déplace les coûts du temps de développement vers le temps d'exécution, la surcharge est plus qu'acceptable. Lorsqu'un modèle est déployé, il n'active qu'une fraction des points de l'ensemble de l'espace de configuration, entraînant des surcharges de l'ordre de la seconde lors du démarrage ou de l'échauffement du moteur d'inférence. En revanche, précompiler toutes les configurations d'un noyau pourrait prendre plusieurs minutes. De plus, les surcharges peuvent être davantage atténuées en mettant en cache les artefacts compilés en JIT.

Le coût de compilation élevé de CUDA est accentué par la difficulté de retracer l'origine des erreurs lorsque la répartition et la spécialisation basées sur des modèles sont utilisées, comme c'est le cas avec CuTe et CUTLASS. Avec un DSL, les erreurs peuvent être pointées vers la ligne de code qui les a générées, ce qui facilite grandement le débogage. De plus, il est beaucoup plus simple de mapper le code généré à l'instruction DSL qui lui a donné naissance via des étapes MLIR intermédiaires ou en s'appuyant sur les informations de débogage attachées au code PTX généré. Grâce à ces aspects du DSL, nous nous sommes trouvés beaucoup plus productifs lors du développement de noyaux.

Bien que les noyaux partagent de nombreuses similitudes entre les modèles, ils doivent tout de même être dotés d'une gamme de paramètres de configuration susceptibles d'avoir un impact sur les performances. Les paramètres pertinents incluent :

Dimension cachée / Dimension de tête / Nombre de têtes d'attention

Types de données d'entrée/sortie

Schéma de quantification

Constantes à la compilation (biais, fonctions d'activation, etc.)

L’architecture ROSE, y compris son moteur, son modèle, sa bibliothèque de noyau et ses services.

La spécialisation à la compilation est utile tant pour les performances que pour le débogage. Particulièrement dans un processeur pipeline, comme c'est le cas avec les GPU, une branche inexistante vaut mieux qu'une branche toujours ou jamais prise. De plus, la compilation d'un noyau pour un ensemble spécifique de paramètres de configuration se traduit par un nombre réduit d'instructions PTX/SASS, qui sont plus faciles à analyser et à déboguer lors de l'utilisation de cuda-gdb au cours d'une investigation.

Pour une taille de lot et une configuration données, le trafic attendu vers la DRAM et le nombre d'opérations arithmétiques peuvent être estimés. Par la suite, des décisions optimales peuvent être prises au cas par cas pour l'allocation optimale de la grille, la synchronisation et l'utilisation de certaines fonctionnalités matérielles. CuTeDSL nous permet de définir et de réutiliser des composants pour une spécialisation poussée sans engendrer de surcoût inacceptable au moment de la compilation.

De manière générale, les performances des noyaux d'inférence doivent être optimisées pour deux cas d'utilisation distincts :

Préremplissage (Prefill) : petite taille de lot (peu de blocs, de séquences), nombreux jetons par séquence. Débit à maximiser pour les plongements et les jetons d'entrée des LLM. Latence de l'ordre de dix à quelques centaines de millisecondes.

Décodage (Decode) : tailles de lots de petite à grande, de un (sans décodage spéculatif) à quelques jetons par séquence (spéculation brouillon-cible, prédiction multi-jeton). Latence proche de la latence de lancement du noyau.

Pour GEMM et l'attention, l'étape de préremplissage s'exécute sur un grand nombre de jetons et les performances sont limitées par le débit de calcul de l'appareil. Le décodage implique la même opération sur de petites tailles de lots, avec des temps d'exécution dominés par le coût du chargement des activations et des poids de la DRAM vers les multiprocesseurs de streaming et les cœurs tensoriels, limités par la bande passante mémoire de l'appareil à la place. Même si la taille du lot ou le nombre total de jetons dans un lot est un indicateur relativement clair, le seuil de basculement entre le préremplissage et le décodage est plus subtil car la bande passante mémoire est également déterminée par la dimension cachée du modèle et le type de données de l'activation. Les noyaux sont par conséquent divisés en implémentations de préremplissage et de décodage si des différences algorithmiques ou des attributions de grille/bloc avantagent l'un ou l'autre.

Le paramètre qui diffère universellement entre les noyaux de préremplissage et de décodage est le nombre de groupes de threads (warps) dans un bloc (le nombre de threads, divisé par la taille du groupe de 32 threads par warp). Il est souhaitable de faire de ce paramètre une constante à la compilation, car cela peut aider à l'allocation des registres. Dans CUDA, cela se fait généralement en spécifiant le nombre de threads via __launch_bounds__ et en fournissant le nombre de groupes de threads comme paramètre de modèle :

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

Le nombre de registres disponibles pour tous les threads d'un bloc est une contrainte matérielle, de 64K sur Hopper et Blackwell. Si le nombre de threads est connu à l'avance, le compilateur sait exactement combien de registres il peut attribuer à chaque thread et peut réussir à compiler des noyaux optimisés de manière plus agressive. Cependant, cela a un coût au moment de la compilation, car le noyau doit maintenant être recompilé pour différentes tailles de groupes de threads. Avec CuTeDSL, ce n'est plus un problème, car nous n'avons pas à payer le coût de compilation pour l'ensemble de l'espace de configuration au préalable. Le DSL s'attend également à ce que la taille du bloc soit une expression constante, effectuant implicitement l'optimisation impliquée par __latch_bounds__ dans CUDA.

Outre le fait de rendre le nombre de groupes de threads constant, il est souhaitable de le faire varier entre le préremplissage et le décodage. Pour exploiter pleinement la bande passante mémoire d'un GPU, il ne suffit pas d'attribuer toujours un nombre élevé de groupes de threads par bloc, car la pleine bande passante ne peut être atteinte qu'en utilisant autant de multiprocesseurs de streaming (SM) disponibles que possible. Par conséquent, pour un nombre de jetons donné et des types de données et dimensions cachées connus, nous choisissons d'abord une dimension de vecteur par thread idéale (8 pour bfloat16, 4 pour float32) pour garantir que le code généré utilise les instructions de chargement/stockage de 128 bits (256 bits sur Blackwell). Connaissant la taille du vecteur, nous pouvons calculer le nombre de groupes de threads requis par le noyau. Nous essayons d'abord de maximiser le nombre de SM : s'il y a moins de groupes de threads que de SM, un seul groupe de threads par bloc est choisi. Il s'agit généralement de la variante de décodage du noyau. S'il y a suffisamment de groupes de threads, nous augmentons progressivement leur nombre à 8, 16 ou 32, en évaluant une limite supérieure après laquelle nous ne constatons généralement plus d'amélioration des performances.

Bien que Triton ait effectué certaines de ces optimisations sous le capot, nous avons constaté que nous pouvions obtenir des résultats similaires en utilisant soigneusement l'algèbre de mise en page CuTe, sans être contraints par les lacunes d'un compilateur de haut niveau. Particulièrement lorsque des réductions horizontales sous le niveau du groupe de threads sont impliquées, nous avons constaté que CuTeDSL nous offrait le contrôle granulaire requis pour générer du code optimal.

Dans certaines situations, pour répartir le travail sur le plus grand nombre de blocs possible, cooperative_groups::this_grid().sync() peut être utilisé pour synchroniser tous les blocs participants. Des restrictions s'appliquent : puisqu'un bloc ne peut pas être désordonné ou désalloué d'un SM une fois qu'il commence son exécution, pour garantir une progression constante, le nombre de blocs ne peut pas dépasser le nombre de SM disponibles sur l'appareil, afin de s'assurer qu'à tout moment tous les blocs peuvent mettre à jour et interroger les emplacements de mémoire globale impliqués dans la synchronisation. Cette restriction est appliquée par cudaLaunchCooperativeKernel. Le coût d'une barrière de grille est d'environ 3 us sur Hopper, plus le coût du vidage de toutes les opérations de mémoire précédentes via une barrière de mémoire implicite dans certaines situations. Dans l'ensemble, c'est comparable à la surcharge liée à la division d'un noyau en deux lancements, ce qui conduit à une décision de conception sensible aux performances. Malgré le fait que le lancement dépendant programmatique (PDL) puisse réduire les surcharges de lancement à partir de Hopper, pour les noyaux de décodage, éviter cette surcharge en vaut la peine.

Si l'une des charges de travail nécessite plus de blocs que le nombre de SM, le noyau doit être divisé, en s'appuyant sur la barrière implicite des lancements de noyaux. C'est principalement le cas pour le préremplissage. Si le nombre de blocs idéal pour toutes les charges de travail s'ajuste dans la grille, une barrière de grille est préférée, comme c'est le cas pour les noyaux de décodage fonctionnant sur un plus petit nombre de jetons.

Dans CuTeDSL, ces décisions peuvent être prises dans le code de lancement côté hôte et les noyaux peuvent être composés et entrelacés avec des barrières de grille selon les besoins, comme l'illustre l'exemple ci-dessous. Avec une telle stratégie hybride, les noyaux de décodage peuvent être optimisés pour atteindre des latences minimales et ceux de préremplissage peuvent l'être pour un débit maximal. Contrairement aux solutions basées sur CUDA, cette logique peut être définie plus succinctement et avec moins de surcoût au moment de la compilation.

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

Dans le cas de nos noyaux MoE all-to-all, comme détaillé ci-dessous, le fait d'avoir la flexibilité de créer des noyaux avec ou sans barrières de grille et de les optimiser à la compilation nous a aidés à obtenir des gains de performance substantiels grâce à CuTeDSL.

Le premier noyau que nous avons construit dans CuTeDSL était un noyau de norme RMS spécialisé, optimisé pour les petites dimensions. La norme QK, utilisée dans la couche d'auto-attention des récents modèles Qwen, est appliquée par tête, sur 64, 128 ou 256 éléments. Contrairement à la norme RMS régulière, qui fonctionne sur des dimensions cachées dépassant les milliers, nécessitant une réduction horizontale à la fois à travers la mémoire partagée et les registres, la norme QK ne nécessite une réduction qu'entre les groupes de threads, avec la possibilité de regrouper plusieurs têtes dans un seul groupe de threads. C'est devenu un cas d'utilisation où nous pouvions optimiser de manière agressive à la compilation pour trouver les meilleures attributions de jetons aux groupes de threads pour toutes les dimensions de tête et tous les types de données.

La norme RMS est une opération assez simple, bien qu'elle nécessite une somme horizontale pour calculer la norme au carré d'un vecteur :

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

Puisque les têtes d'attention sont généralement petites et que les opérations par thread sont fortement vectorisées, un seul groupe de threads peut calculer la norme au carré. Afin de maximiser la bande passante mémoire, chaque thread doit traiter 128 bits d'entrée, ce qui se traduit par 8 valeurs bfloat16. Par conséquent, un demi-groupe de threads (half-warp) suffit pour traiter une tête, et la norme au carré peut être déterminée à l'aide d'une réduction horizontale exécutée par shuffle_sync_bfly. En partant de la dimension de la tête et du type de données, nous pouvons déterminer le nombre de têtes par groupe de threads et spécifier les arguments ainsi que le nombre de tours de réduction au moment de la compilation.

Nous pouvons ensuite construire un noyau spécialisé avec relativement peu de lignes de code, en nous appuyant sur des chargements/stockages vectorisés automatiquement et en utilisant des opérations de bas niveau spécifiques à l'architecture pour la réduction du groupe de threads. Le noyau détermine sur quelle tête (Q ou K) travailler et identifie les poids correspondants à appliquer. Il charge un vecteur entier de la tête et effectue d'abord une somme locale. Ensuite, en fonction du nombre de threads par tête, la norme au carré est calculée. Enfin, la norme RMS, automatiquement vectorisée, est déterminée et réécrite en mémoire. Les paramètres spécifiques au modèle, tels que WEIGHT_BIAS et EPS, sont repliés en constantes.

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

Par rapport à un noyau de norme RMS classique, cette variante spécialisée offre un débit et une latence de 2 à 3 fois supérieurs pour le préremplissage et le décodage.

La répartition et la combinaison MoE réorganisent les jetons d'entrée vers les experts auxquels ils sont acheminés, effectuant une accumulation pondérée après que les noyaux GEMM MoE ont exécuté les multiplications matricielles pertinentes. Les performances revêtent un intérêt particulier, car divers produits reposent sur des modèles MoE de tailles variées. Selon les types de déploiement, nous distinguons différents noyaux à optimiser :

TP=1 : Noyau local pour un seul GPU. La répartition réorganise les jetons vers les experts pertinents et la combinaison les accumule. Utilisé principalement pour les déploiements de Qwen3 30B sur H100.

DP=1, TP ≠ 1 : Noyau mononeud, parallèle par tenseur. L'entrée est répliquée, la répartition sélectionne les jetons pour le rang actuel, et la combinaison les accumule à travers tous les rangs. Accumulation via NVLink ou NVLS. Utilisé pour les grands préremplisseurs.

DP ≠ 1, NVLink : Noyau parallèle par données avec parallélisme par tenseur optionnel. Chaque rang envoie des jetons à chaque autre rang hébergeant des experts spécifiques. Les informations de routage sont agrégées, les jetons sont transférés via NVIDIA NVLink. Utilisé pour les déploiements de données parallèles H200 ou GB200.

DP ≠ 1, InfiniBand : Noyau pair à pair parallèle par données et parallèle par tenseur fonctionnant à la fois sur NVLink et NVIDIA Quantum InfiniBand. Utilisé pour les déploiements Hopper multi-nœuds.

Le noyau local compte le nombre de jetons attribués à chaque expert sur la base de la table de routage et attribue un index unique à chaque jeton dans la liste des jetons appartenant à chaque expert. Après l'étape de routage, les jetons sont copiés dans des tampons facultativement remplis (padded) pour être transmis à un noyau GEMM MoE. En utilisant les emplacements des jetons, le noyau de combinaison effectue ensuite une accumulation dans les registres, écrivant les jetons finaux dans un tampon de sortie.

La phase de répartition varie entre le préremplissage et le décodage : pour un petit nombre de jetons, il y a peu d'informations de routage à lire, nous les traitons donc dans un seul bloc et signalons le reste de la grille, des blocs individuels copiant des morceaux de jetons. L'attribution des jetons est déterminée en incrémentant atomiquement des compteurs dans la mémoire partagée. Étant donné que le préremplissage doit traiter davantage d'informations de routage, le routage est réparti sur plusieurs blocs dans un noyau séparé. Les indices sont d'abord attribués localement à l'aide de compteurs dans la mémoire partagée, puis agrégés sur tous les blocs de la mémoire globale pour déterminer la position de chaque jeton. La grille de ce noyau est limitée au nombre de SM afin de pouvoir utiliser des barrières à l'échelle de la grille. Un noyau subséquent effectue les copies, lancées sur un nombre beaucoup plus grand de blocs et de groupes de threads pour exploiter pleinement la bande passante mémoire disponible.

Le noyau de combinaison est relativement simple, car il effectue une accumulation vectorisée, en s'appuyant sur les informations de routage et sur les décalages (offsets) calculés par la répartition. CuTeDSL s'est avéré précieux ici, car le calcul pouvait être exprimé par de simples tenseurs CuTe qui se compilent en un code efficace pour tous les types de données et dimensions cachées.

Dans un déploiement parallèle par tenseur, les informations de routage et les jetons d'entrée sont répliqués. Le noyau de répartition local peut être entièrement réutilisé avec une légère modification pour router et sélectionner uniquement les jetons attribués aux experts situés sur le rang actuel. Avec quelques paramètres et une spécialisation à la compilation, les noyaux de répartition locaux sont réutilisés.

Étant donné que le nombre de rangs est généralement inférieur au nombre d'experts vers lesquels un jeton est routé, dans le cas moyen, on peut supposer que chaque jeton a au moins une de ses répliques routée vers chaque rang. Dans cette situation, le noyau de combinaison local peut être réutilisé pour agrèger les copies locales d'un jeton, suivi d'une réduction globale (all-reduce) pour sommer à travers les rangs participants. Dans le cas du noyau de préremplissage, nous lançons des noyaux séparés, mais le décodage fusionne les deux opérations en un seul lancement.

Pour la réduction globale, nous utilisons NVLink simple si seulement deux rangs participent, passant à NVLS lorsque 4 rangs ou plus sont impliqués pour une augmentation de 10 à 20 % du débit. Grâce au DSL, nous pouvons partager la plus grande partie de la logique de réduction et des échanges de tampons entre pairs entre le MoE et les noyaux de réduction globale réguliers. C'est une situation où un contrôle étroit sur le matériel est bénéfique, car les opérations DSL sur multimem.ld_reduce et multimem.st sont requises pour NVLS. L'implémentation de la réduction elle-même suit largement l'exemple all_reduce des tests NCCL.

Les noyaux parallèles par données s'appuient sur NVLink pour agréger les jetons pour plusieurs appareils. Dans la phase de répartition, avant d'échanger des jetons, les pairs acheminent d'abord localement pour identifier le décalage cible sur chaque rang pour chaque jeton qu'ils envoient. Par la suite, ils commencent à copier les jetons dans des tampons de réception alloués pour chaque émetteur afin d'éviter les conflits. Des écritures de type libération-stockage (release-store) par jeton sont utilisées pour signaler l'achèvement. Du côté du récepteur, les informations de routage sont agrégées pour pouvoir condenser tous les jetons dans un seul tampon pour le GEMM MoE. Les récepteurs interrogent l'indicateur écrit par les pairs avec un chargement-acquisition (acquire-load) pour attendre que les écritures se terminent.

Du côté de la combinaison, les jetons sont copiés via NVLink vers les rangs source, avec un indicateur par jeton indiquant la réception. Le noyau de combinaison est similaire aux variantes locale et mononeud, avec de légères différences d'adressage et une logique de synchronisation supplémentaire.

Nous avons construit une bibliothèque d'enveloppes PTX pour exposer des charges et des stockages de type libération-acquisition portés au niveau du système et de l'appareil afin d'implémenter ce noyau. En améliorant la granularité des indicateurs de synchronisation et en allouant plus de blocs pour envoyer et recevoir des jetons via NVLink, nous avons obtenu des améliorations multipliées par 2 par rapport à nos noyaux CUDA précédents.

Nos noyaux pair à pair ont atteint par le passé des performances de pointe sur les GPU Hopper couplés à des adaptateurs InfiniBand ConnectX-7. Nous avons depuis porté les noyaux de combinaison vers CuTeDSL pour de nouveaux gains de performance. Étant donné que les versions CUDA s'appuyaient déjà fortement sur la spécialisation à la compilation, elles se sont avérées être les goulets d'étranglement de compilation dans notre référentiel, nous obligeant à omettre certains paramètres. En migrant vers le DSL compilé en JIT, nous nous sommes spécialisés pour un plus grand nombre de paramètres et avons ajusté les allocations de grille, introduisant de nouvelles différences entre le préremplissage et le décodage pour un gain de latence d'environ 3 %.

Du point de vue des noyaux ROSE, la prise en charge de Blackwell se divise en deux catégories : les noyaux agnostiques et les noyaux spécifiques à l'architecture. Les noyaux spécifiques à l'architecture, qui s'appuient sur des cœurs tensoriels, nécessitent un travail explicite d'adaptation. Cependant, grâce à l'utilisation du DSL, les noyaux agnostiques s'adaptent de manière triviale.

Un GPU Blackwell possède jusqu'à 160 SM, contre 132 pour Hopper. Pour la plupart des noyaux, il s'agit d'une constante à la compilation et la logique du code de lancement côté hôte en tient compte pour déterminer le noyau optimal (préremplissage ou décodage) à exécuter. De plus, Blackwell introduit également des instructions de chargement de 256 bits, qui peuvent être facilement émises en augmentant la largeur du vecteur (dernière dimension des tenseurs CuTe) pour améliorer l'utilisation de la bande passante mémoire.

Blackwell introduit une prise en charge matérielle pour un nombre important de schémas de quantification (MXFP8, MXFP4, NVFP4). La prise en charge de ceux-ci a été implémentée principalement en ajustant les noyaux existants pour la mise à l'échelle par bloc FP8. La prise en charge du JIT est cruciale ici, car de nouveaux modes peuvent être pris en charge sans faire exploder l'espace de configuration à la compilation. Par le passé, nous avons trouvé que Triton peinait à exprimer les réductions horizontales sur un petit nombre d'éléments, sur la moitié ou le quart d'un groupe de threads. Un contrôle plus fin des réductions était crucial pour atteindre des performances optimales lors du calcul d'une valeur maximale absolue sur 16 ou 32 éléments pour NVFP4 et MXFP8/MXFP4, respectivement.

En tant que langage relativement nouveau, les LLM ont eu peu d'exposition au code CuTeDSL lors de leur pré-entraînement. De plus, le style de notre code actuel, avec des noyaux optimisés pour les activations et les normes, était légèrement différent des exemples existants CUTLASS ou Quack, qui étaient fortement optimisés pour le TMA et les instructions de multiplication matricielle spécifiques à l'architecture. Atteindre un point où les LLM pourraient aider à l'écriture de noyaux a été un processus graduel.

Nous avons construit manuellement un corpus initial de noyaux et de fonctions utilitaires, suivant des modèles cohérents, avant de pouvoir nous fier à des agents pour la conversion des noyaux. Au fur et à mesure que le corpus progressait, nous avons pu effectuer des traductions approximatives de CUDA vers CuTeDSL, suivies d'un travail manuel important. En particulier, nous avons constaté que les modèles avaient du mal avec les conversions de pointeur vers tenseur et la vectorisation, restant bloqués en essayant de corriger des erreurs de type pour obtenir la compilation d'un test. Ils ont montré peu de compréhension de la nature basée sur la grille des noyaux, confondant la synchronisation des threads et de la grille. Plus frustrant encore, il y a eu une tentation d'halluciner des traductions inexistantes de bibliothèques spécifiques à CUDA au lieu de compiler avec nos propres remplacements CuTeDSL. Au fil du temps, à mesure que le corpus de noyaux CuTeDSL s'agrandissait, nous avons constaté que les LLM peuvent construire l'essentiel du code squelette et remplir les blancs, ce qui nous permet de nous concentrer sur la conception de haut niveau.

Comparé à CUDA, nous constatons que CuTeDSL offre des performances de compilation bien supérieures, nous permettant de nous spécialiser pour plus de paramètres et d'obtenir des améliorations de performance. Bien que les noyaux soient légèrement plus verbeux en raison de l'absence de certaines abstractions par rapport à Triton, nous estimons que la complexité est justifiée car un contrôle plus fin sur le matériel nous permet de créer des noyaux plus rapides.

Au moment d'écrire ces lignes, l'écart de performance entre CUTLASS C++ et les noyaux CuTeDSL est minimal, ouvrant des perspectives passionnantes pour composer et fusionner des noyaux GEMM haute performance avec la communication inter-appareils, la normalisation ou les activations. En nous engageant dans cette voie avec ce DSL, nous serons en mesure de tirer des performances supplémentaires de nos appareils pour améliorer le débit de nos charges de travail, des plongements aux LLM.

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

L'inférence chez Perplexity AI

Pourquoi CuTeDSL ?

Expérience de développement

Vaste espace de configuration

Pré-remplir et décoder la spécialisation

Synchronisation de la grille

Noyaux CuTeDSL

Norme QK (QK Norm)

Répartition / Combinaison de Mélange d'experts (Mixture-of-Experts Dispatch/Combine)

Local

Mononeud (Singlenode)

Parallèle par données (Data Parallel)

Pair à pair (Peer-to-Peer)

Prise en charge de Blackwell

Génération de noyaux assistée par IA

Conclusion et travaux futurs

Références

Citation