CuTeDSL chez Perplexity

Chez Perplexity, nous utilisons le moteur ROSE (Runtime-Optimized Serving Engine) développé en interne pour servir des modèles allant des plongements lexicaux aux LLM atteignant le billion de paramètres. Situé derrière toutes nos API, y compris Sonar, Search et Embeddings, notre moteur d'exécution 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 de mettre en œuvre, déployer et optimiser l'inférence pour des architectures de pointe à un rythme soutenu. Dans cet article de blog, nous explorons comment les noyaux CuTeDSL intégrés à notre moteur d'inférence nous permettent de construire les noyaux GPU requis pour porter les modèles à des performances optimales sur les GPU NVIDIA.

Des modèles personnalisés, conçus par diverses équipes chez Perplexity, sont hébergés en interne sur des GPU NVIDIA Hopper et Blackwell à l'aide du moteur d'inférence ROSE construit par l'équipe AI Inference. ROSE a été initialement conçu pour mieux servir des modèles Llama personnalisés, tant pour le décodage linguistique que pour la classification, avec une interface compatible avec le serveur NVIDIA Triton Inference Server. Au fil du temps, il est devenu un moteur d'inférence à part entière servant des LLM substantiels, ainsi que des modèles de classement, de classification, de notation et d'autres modèles basés sur les transformeurs pour un large éventail d'équipes afin d'alimenter les produits de Perplexity. L'objectif de conception de ROSE se veut simple : au fil du temps, des efforts sont investis pour atteindre des performances de pointe, mais la priorité principale reste l'adaptation aux architectures et aux cas d'utilisation de nos modèles, en constante évolution.

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

À un haut niveau, 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 la mise en 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 gérés par ces moteurs. Pour les LLM, le moteur est responsable de la mise en lots, du découpage et de l'échantillonnage. Il est également chargé de l'allocation du stockage KV, pour l'attention complète et linéaire, ainsi que de la correspondance de préfixes. Les moteurs d'intégration sont plus simples, car ils n'effectuent qu'une mise en lots en ligne. L'exécution est optimisée de telle sorte que le moteur prépare le cycle de travail suivant 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 des noyaux GPU construits avec CuTeDSL, Triton, CUDA, CUTLASS et cuBLAS. Pour les implémentations de multiplication matricielle et d'attention, nous nous appuyons principalement sur les noyaux NVIDIA de CUTLASS et cuBLAS. Le reste du pipeline, y compris les noyaux d'intégration, de norme, de MoE et d'activation, fait partie de la bibliothèque de noyaux de ROSE. Pour mieux prendre en charge le matériel sur lequel nous exécutons l'inférence, ROSE est équipé de ses propres primitives de communication inter-appareils pour le routage, la répartition et la combinaison MoE. Bien que les autres noyaux soient relativement simples, ils doivent être minutieusement spécialisés pour atteindre des performances de pointe sur toutes les dimensions cachées et tous les appareils sur lesquels ils s'exécutent. De plus, nous comptons sur la spécialisation lors de la compilation pour tenir compte de tous les réglages et ajustements des divers modèles (biais de poids vs aucun 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 en utilisant 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 peu sophistiqués. 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 précise dont le matériel doit le réaliser. Pour construire les meilleurs noyaux, nous avons besoin d'un langage qui offre les bonnes abstractions sur le matériel sous-jacent. Le langage approprié nous permet de spécifier les opérations de haut niveau pour effectuer le calcul, en les spécialisant et en les optimisant agressivement sur tous les points d'un vaste espace de configuration.

CuTeDSL est un DSL basé sur Python pour les GPU NVIDIA, construit sur l'algèbre de disposition CuTe et MLIR. Il se compile juste-à-temps (JIT) en PTX hautement optimisé, tout en conservant le contrôle sur les primitives matérielles de bas niveau. Il vise à offrir les mêmes fonctionnalités et performances que CUTLASS C++, sans les temps de compilation coûteux qui entravent la vitesse de développement. Parmi les langages que nous avons étudiés ou avec lesquels nous avons travaillé (y compris CUDA, cuTile, TVM IR et Triton), nous considérons qu'il est le mieux adapté au développement de nos noyaux.

Dans l'ensemble, bien qu'ils ne soient pas aussi haut niveau et succincts que Triton ou cuTile, nous avons constaté que l'algèbre de disposition CuTe et ses compositions (exposées via cute.Tensor) étaient des substituts raisonnables. De plus, CuTeDSL offre un accès illimité aux primitives matérielles, nous fournissant le niveau de contrôle dont nous avons besoin pour atteindre des performances de pointe dans tous les cas limites. Avec CUDA, les opportunités de spécialisation lors de 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 build. En revanche, avec un DSL compilé juste-à-temps, nous pouvons nous spécialiser plus agressivement 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 descendait vers PTX était plus robuste et plus facile à déboguer qu'un pipeline émettant du CUDA en tant qu'é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.

La majeure partie du coût de compilation des noyaux CUDA est généralement consacrée à la phase d'expansion des modèles des noyaux, car l'optimisation et la compilation à partir de l'étape PTX peuvent être effectuées en quelques secondes. Avec de nombreux paramètres de compilation, cette étape de compilation peut exploser en plusieurs minutes pour un noyau non trivial. Ce goulot d'étranglement est accentué par le fait que les opportunités de compilation parallèle ne surviennent que plus tard, aux niveaux IR et PTX, la gestion des modèles C++ se faisant séquentiellement. 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 distinctes, elles nécessitent des ajustements importants des systèmes de build et l'utilisation de macros lourdes. En revanche, les noyaux CuTeDSL compilés juste-à-temps (JIT) ne nécessitent pas de règles de build complexes et n'entravent pas les temps de compilation, car le cycle de développement peut être concentré sur un seul point de l'espace de configuration, tandis que des tests unitaires minutieux peuvent ultérieurement 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 que quelques points de l'espace de configuration total, entraînant des surcharges de l'ordre de la seconde lors du démarrage/préchauffage du moteur d'inférence. En revanche, la pré-compilation de toutes les configurations d'un noyau pourrait s'étendre sur plusieurs minutes. De plus, les surcharges peuvent être davantage atténuées en mettant en cache les artefacts compilés JIT.

Le coût de compilation élevé de CUDA est accentué par la difficulté de tracer l'origine des erreurs lorsque la répartition et la spécialisation basées sur les 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, facilitant considérablement le débogage. De plus, il est beaucoup plus simple de mapper le code généré à l'instruction DSL qui l'a engendré via des étapes MLIR intermédiaires ou en s'appuyant sur les informations de débogage jointes au PTX généré. Grâce à ces aspects du DSL, nous nous sommes trouvés beaucoup plus productifs lors du développement des noyaux.

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

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

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

Schéma de quantification

Constantes lors de 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 lors de la compilation est utile à la fois pour les performances et le débogage. Particulièrement dans un processeur pipeliné, comme c'est le cas avec les GPU, une branche inexistante vaut mieux qu'une branche toujours empruntée ou jamais empruntée. De plus, la compilation d'un noyau pour un ensemble spécifique de réglages de configuration donne lieu à moins d'instructions PTX/SASS, lesquelles sont plus faciles à analyser et à déboguer lors de l'utilisation de cuda-gdb lors 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 une allocation de grille, une synchronisation et l'utilisation de certaines fonctionnalités matérielles optimales. CuTeDSL nous permet de définir et de réutiliser des composants pour une spécialisation poussée sans encourir une surcharge de temps de compilation inacceptable.

En général, les performances des noyaux d'inférence doivent être optimisées pour deux cas d'utilisation distincts :

Pré-remplissage : petite taille de lot (peu de morceaux, séquences), beaucoup de jetons par séquence. Le débit doit être maximisé pour les plongements et les jetons d'entrée des LLM. Latence de l'ordre de quelques dizaines à quelques centaines de millisecondes.

Décodage : tailles de lot de petite à grande, un (sans décodage spéculatif) à quelques jetons par séquence (spéculation cible-brouillon, prédiction multi-jetons). 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 lot, les temps d'exécution étant 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 tenseurs, contraints par la bande passante mémoire de l'appareil. 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 affectations de grille/bloc avantageraient l'un ou l'autre.

Le paramètre qui est universellement différent entre les noyaux de pré-remplissage et de décodage est le nombre de warps dans un bloc (le nombre de fils, divisé par la taille du warp de 32 fils par warp). Il est souhaitable de faire de ce paramètre une constante de compilation, car il peut aider à l'allocation des registres. Dans CUDA, cela se fait généralement en spécifiant le nombre de fils via __launch_bounds__ et en fournissant le nombre de warps en tant que 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 fils d'un bloc est une contrainte matérielle, 64K sur Hopper et Blackwell. Si le nombre de fils est connu à l'avance, le compilateur sait exactement combien de registres il peut assigner à chaque fil et peut réussir à compiler des noyaux optimisés de manière plus agressive. Cependant, cela se fait au prix d'un coût en temps de compilation, car le noyau doit maintenant être recompilé pour différentes tailles de warp. Avec CuTeDSL, ce n'est plus un problème, car nous n'avons pas besoin de payer le coût de compilation pour tout l'espace de configuration à l'avance. 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 warps constant, il est souhaitable de le faire varier entre le pré-remplissage et le décodage. Pour utiliser pleinement la bande passante mémoire d'un GPU, il ne suffit pas d'assigner systématiquement un nombre élevé de warps 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 donné de jetons et des types de données et dimensions cachées connus, nous choisissons d'abord une dimension vectorielle idéale par fil (8 pour bfloat16, 4 pour float32) pour garantir que le code généré utilise les instructions de chargement/stockage 128 bits (256 bits sur Blackwell). Connaissant la taille vectorielle, nous pouvons calculer le nombre de warps requis par le noyau. Nous essayons d'abord de maximiser le nombre de SM : s'il y a moins de warps que de SM, un seul warp par bloc est choisi. Il s'agit généralement de la variante de décodage du noyau. S'il y a suffisamment de warps, nous augmentons progressivement le nombre de warps à 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 en arrière-plan, nous avons constaté que nous pouvions obtenir des résultats similaires en utilisant soigneusement l'algèbre de disposition CuTe, sans être limités par les défauts d'un compilateur de haut niveau. Particulièrement lorsque des réductions horizontales sous-warp sont impliquées, nous avons constaté que CuTeDSL nous offrait le contrôle granulaire nécessaire pour générer un code optimal.

Dans certaines situations, pour répartir le travail sur autant de blocs que possible, cooperative_groups::this_grid().sync() peut être utilisé pour synchroniser tous les blocs participants. Des restrictions s'appliquent : étant donné qu'un bloc ne peut pas être désordonné d'un SM une fois qu'il commence l'exécution, pour garantir une progression vers l'avant, le nombre de blocs ne peut pas dépasser le nombre de SM disponibles sur l'appareil, afin de garantir 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 3us sur Hopper, plus le coût de vidage de toutes les opérations mémoire précédentes via une barrière mémoire implicite dans certaines situations. Dans l'ensemble, cela 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 par programme (PDL) puisse réduire les surcharges de lancement à partir de Hopper, pour les noyaux de décodage, éviter la 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 noyau. C'est principalement le cas pour le pré-remplissage. Si le nombre de blocs qui sont idéaux pour toutes les charges de travail tient dans la grille, une barrière de grille est préférée, comme c'est le cas avec 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 si nécessaire, comme illustré par 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 être optimisés pour un débit maximal. Contrairement aux solutions basées sur CUDA, cette logique peut être définie plus succinctement et avec moins de surcharge de temps de 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 tout-à-tout, comme détaillé ci-dessous, avoir la flexibilité de construire des noyaux avec ou sans barrières de grille et de les optimiser au moment de 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 à travers la mémoire partagée et les registres, la norme QK nécessite une réduction uniquement à travers les warps, avec la possibilité de compresser plusieurs têtes dans un seul warp. C'est devenu un cas d'utilisation où nous pouvions optimiser agressivement lors de la compilation pour trouver les meilleures affectations jeton-à-warp pour toutes les dimensions de tête et tous les types de données.

La norme RMS est une opération assez directe, 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))

Étant donné que les têtes d'attention sont généralement petites et que les opérations par fil sont fortement vectorisées, un seul warp peut calculer la norme au carré. Afin de maximiser la bande passante mémoire, chaque fil doit traiter 128 bits d'entrée, ce qui se traduit par 8 valeurs bfloat16. Par conséquent, un demi-warp est suffisant pour traiter une tête, et la norme au carré peut être déterminée en utilisant une réduction horizontale effectuée par shuffle_sync_bfly. En commençant par la dimension de la tête et le type de données, nous pouvons déterminer le nombre de têtes par warp et spécifier les arguments et le nombre de cycles 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 automatiquement vectorisés et en utilisant des opérations spécifiques à l'architecture de bas niveau pour la réduction des warps. Le noyau détermine sur quelle tête (Q ou K) travailler et identifie les poids correspondants à appliquer. Il charge un vecteur entier depuis la tête et commence par sommer localement. Ensuite, en fonction du nombre de fils 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 réduits par constante.

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)

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

La répartition et la combinaison MoE mélangent les jetons d'entrée vers les experts vers lesquels ils sont routés, effectuant une accumulation pondérée après que les noyaux GEMM MoE aient effectué les multiplications matricielles pertinentes. Les performances sont d'un intérêt particulier, car divers produits reposent sur des modèles MoE de tailles variées. Sur la base des types de déploiement, nous distinguons différents noyaux à optimiser :

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

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

DP ≠ 1, NVLink : Noyau parallèle aux données avec parallélisme tensoriel 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 H200 ou GB200 parallèles aux données.

DP ≠ 1, InfiniBand : Noyau parallèle aux données et parallèle tensoriel de pair à pair 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 assignés à chaque expert en fonction de la table de routage et assigne 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 optionnellement complétés 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 peu 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, avec des blocs individuels copiant des morceaux de jetons. L'assignation des jetons est déterminée par l'incrémentation atomique de compteurs dans la mémoire partagée. Étant donné que le pré-remplissage doit traiter plus d'informations de routage, le routage est divisé sur plusieurs blocs dans un noyau distinct. Les indices sont d'abord assignés localement à l'aide de compteurs dans la mémoire partagée, puis agrégés sur tous les blocs dans la mémoire globale pour déterminer la position de chaque jeton. La grille de ce noyau est fixée au nombre de SM afin de pouvoir utiliser des barrières par grille. Un noyau ultérieur effectue les copies, lancées sur un nombre beaucoup plus grand de blocs et de warps pour utiliser pleinement la bande passante mémoire disponible.

Le noyau de combinaison est relativement simple, car il effectue une accumulation vectorisée, s'appuyant sur les informations de routage et sur les décalages 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 tensoriel, 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 un léger changement pour router et sélectionner uniquement les jetons assignés aux experts situés sur le rang actuel. Avec quelques paramètres et une spécialisation lors de 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 l'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, suivies d'une réduction globale (all-reduce) pour sommer sur tous les rangs participants. Dans le cas du noyau de pré-remplissage, nous lançons des noyaux distincts, cependant, 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 majeure partie de la logique de réduction et les échanges de tampons de pair à pair 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 aux données s'appuient sur NVLink pour agréger les jetons pour plusieurs appareils. Dans la phase de répartition, avant d'échanger les jetons, les pairs routent d'abord localement pour identifier le décalage cible sur chaque rang pour chaque jeton qu'ils envoient. Ensuite, ils commencent à copier les jetons dans des tampons de réception alloués pour chaque émetteur afin d'éviter les conflits. Des écritures de libération-stockage par jeton sont utilisées pour signaler l'achèvement. Côté récepteur, les informations de routage sont agrégées pour pouvoir condenser tous les jetons dans un seul tampon pour MoE GEMM. Les récepteurs interrogent l'indicateur écrit par le pair avec un chargement-acquisition pour attendre que les écritures soient terminées.

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

Nous avons construit une bibliothèque de wrappers PTX pour exposer des chargements et des stockages de libération-acquisition portés au niveau du système et de l'appareil pour 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 2 fois supérieures par rapport à nos noyaux CUDA précédents.

Nos noyaux de pair à pair atteignaient auparavant 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. Comme les versions CUDA reposaient déjà fortement sur la spécialisation lors de la compilation, elles se sont avérées être des goulots d'étranglement en termes de temps de compilation dans notre dépôt, nous forçant à omettre certains paramètres. En migrant vers le DSL compilé JIT, nous nous sommes spécialisés pour plus de paramètres et avons ajusté les allocations de grille, introduisant des différences supplémentaires entre le pré-remplissage et le décodage pour environ 3 % d'augmentation de la latence.

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

Un GPU Blackwell possède jusqu'à 160 SM, contre 132 pour Hopper. Pour la plupart des noyaux, il s'agit d'une constante lors de la compilation et la logique du code du lanceur côté hôte en tient compte pour déterminer le noyau optimal (pré-remplissage vs décodage) à exécuter. De plus, Blackwell introduit également des instructions de chargement 256 bits, qui peuvent être facilement émises en augmentant la largeur vectorielle (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 significatif 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 une mise à l'échelle par bloc FP8. La prise en charge JIT est ici cruciale, car de nouveaux modes peuvent être pris en charge sans faire exploser l'espace de configuration lors de la compilation. Par le passé, nous avons trouvé Triton maladroit pour exprimer des réductions horizontales sur quelques éléments, sur la moitié ou le quart d'un warp. Un contrôle plus granulaire sur les 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 TMA et les instructions de multiplication matricielle spécifiques à l'architecture. Atteindre un point où les LLM pouvaient aider à l'écriture de noyaux fut 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 appuyer sur des agents pour la conversion de noyaux. À mesure que le corpus progressait, nous avons pu effectuer des traductions approximatives de CUDA vers CuTeDSL, avec un travail manuel important par la suite. En particulier, nous avons constaté que les modèles avaient du mal avec les conversions pointeur-vers-tenseur et la vectorisation, restant bloqués en essayant de corriger les erreurs de type lors de la tentative de 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 fils et de la grille. Plus frustrant encore, il y avait une tentation d'halluciner des traductions inexistantes de bibliothèques spécifiques à CUDA au lieu de construire nos propres remplacements CuTeDSL. Au fil du temps, à mesure que le corpus de noyaux CuTeDSL grandissait, nous avons constaté que les LLM pouvaient construire la majeure partie du code squelette et remplir les vides, nous permettant de nous concentrer sur la conception de haut niveau.

Par rapport à CUDA, nous constatons que CuTeDSL offre de bien meilleures performances en termes de temps de compilation, 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 constatons que la complexité est justifiée car un contrôle plus granulaire sur le matériel nous permet de construire des noyaux plus rapides.

Au moment de la rédaction, l'écart de performance entre les noyaux CUTLASS C++ et CuTeDSL est minime, ouvrant des opportunités passionnantes pour composer et fusionner des noyaux GEMM haute performance avec la communication inter-appareils, la normalisation ou les activations. En nous engageant dans ce DSL, nous serons en mesure d'extraire davantage de performances de nos appareils pour améliorer le débit sur l'ensemble de nos charges de travail, des plongements lexicaux 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 développeur

Vaste espace de configuration

Pré-remplissage et décodage de spécialisation

Synchronisation de grille

Noyaux CuTeDSL

Norme QK

Répartition/Combinaison Mélange-d'Experts (MoE)

Local

Mono-nœud

Parallèle aux données

Pair à pair

Prise en charge de Blackwell

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

Conclusion et travaux futurs

Références

Citation