Comment la hiérarchie mémoire du GPU affecte la vitesse d’inférence des LLM

Traduction automatique

Cet article a été traduit automatiquement depuis la version originale en anglais.

La hiérarchie mémoire du GPU affecte l’inférence des LLM, car transférer un tenseur intermédiaire vers la mémoire principale du GPU coûte plus cher que le réutiliser dans un kernel en cours d’exécution. La HBM contient le grand modèle et le cache d’attention. Les registres, la mémoire partagée et le cache L2 offrent des espaces plus petits pour réutiliser les données pendant le calcul.

Les ingénieurs qui profilent l’inférence doivent distinguer la capacité mémoire du trafic mémoire. Le fait qu’un modèle tienne dans la HBM ne prouve pas que ses kernels transfèrent les données efficacement.

Ce que contient chaque mémoire

MémoirePortéeRôle dans l’inférence
RegistresValeurs utilisées par un threadAccumulateurs et calculs intermédiaires
Mémoire partagéeThreads d’un même bloc de calculRéutilisation coopérative de tuiles matricielles et de réductions
Cache L2Partagé sur l’ensemble du GPURéutilisation entre blocs sans nouvelle lecture en mémoire principale
HBMMémoire principale du GPUPoids, KV cache, entrées, sorties et résultats intermédiaires plus volumineux

Pour le H100 SXM, NVIDIA indique 80 GB de HBM3 à 3.35 TB/s, 50 MB de L2 et jusqu’à 228 KB de mémoire partagée par multiprocesseur de streaming. La quantité disponible pour un seul bloc de threads est plus faible et dépend des paramètres d’allocation. Ces spécifications concernent cette variante du GPU, et ne constituent pas des limites communes à tous les H100. Consultez les spécifications du H100 et le guide d’optimisation de Hopper.

Pourquoi un tenseur intermédiaire compte

Une suite de kernels distincts peut écrire un résultat intermédiaire dans la HBM, puis le relire dans le kernel suivant. La fusion peut conserver ce résultat dans les registres ou la mémoire partagée. FlashAttention applique des principes similaires de réutilisation des données à l’attention et évite les grandes matrices de scores dans la HBM.

Le stockage sur la puce est limité. Un kernel qui utilise trop de registres ou de mémoire partagée peut réduire le nombre de blocs exécutés simultanément. Le débordement des registres vers la mémoire peut ajouter du trafic mémoire. Une fusion plus poussée ne garantit donc pas une exécution plus rapide.

Pour étudier une modification, comparez les octets transférés via la HBM, l’utilisation des registres, l’utilisation de la mémoire partagée, les débordements de registres et le nombre de blocs actifs. Gardez le même modèle, les mêmes formes de tenseurs, la même précision et la même variante de GPU. Mesurez ensuite la requête complète : un kernel d’attention plus rapide peut ne représenter qu’une petite part de la latence totale.

Le guide d’ingénierie : hiérarchie mémoire du GPU montre où se trouvent ces mémoires sur un H100 SXM.