Écrire des kernels GPU à la main demande de la patience, des efforts, et même votre santé mentale. J'ai expérimenté avec des LLM modernes (Claude Opus 5, GPT 5.6 Sol) pour qu'ils m'écrivent littéralement des kernels. Je vais vous parler de ce que j'ai appris en faisant ça.
Honnêtement, écrire des kernels GPU est une tâche embarrassamment vérifiable.
D'abord, vous décidez pour quelle(s) opération(s) vous voulez écrire le kernel. Ensuite, vous écrivez la première version et vous vous assurez qu'elle compile sans erreur évidente. Troisièmement, vous vérifiez la correction à l'aide d'une implémentation de référence lente. Si cela ne correspond pas, vous essayez de résoudre le problème. Une fois cela fait, vous mesurez le temps d'exécution du kernel. Les versions suivantes s'appuient ensuite sur ce travail, et vous continuez à optimiser jusqu'à ce que vous atteigniez les métriques roofline ou que vous soyez satisfait.
Boucle vérifiable : développement de kernels GPU
L'image ci-dessus le présente comme une boucle vérifiable :
- La vérification de compilation est une boucle locale rapide (B <-> C), avant même de penser à la correction.
- La vérification de correction (D <-> E <-> F) est la boucle centrale de récompense vérifiable. C'est ce qui fait de l'écriture de kernels un bon problème pour la vérification automatisée, puisque vous disposez d'une référence de vérité terrain contre laquelle vérifier.
- L'optimisation (G -> H -> retour à D) réutilise la même boucle de correction pour chaque nouvelle version : un kernel rapide mais incorrect ne vaut rien, la correction doit donc être vérifiée à chaque version.
- La vérification roofline/satisfaction (I) est la boucle externe qui décide de continuer à optimiser ou de s'arrêter.
Première version
L'image ci-dessus n'est encore qu'une vue de haut niveau, et le diable se cache dans les détails. Nous devons nous assurer que notre agent LLM dispose de tout le contexte nécessaire pour commencer à écrire une bonne première version.
C'est ici que les DSL CUDA entrent en jeu. Triton, CuTeDSL et Tilelang sont ceux avec lesquels il est très facile de démarrer, en Python. La courbe d'apprentissage est moins raide qu'avec CUDA C++ ; en revanche, les abstractions de ces DSL peuvent dérouter notre agent encore davantage. Nous avons besoin d'un moyen de transmettre le contexte de ces abstractions à l'agent.
Les LLM modernes savent déjà écrire du « bon » code Triton. Ils peuvent très bien travailler avec les abstractions de Triton, même sans aucun contexte. En revanche, pour d'autres DSL comme CuTeDSL (qui offre beaucoup plus de contrôle que Triton), j'ai constaté qu'avoir un répertoire de contexte dans lequel l'agent peut chercher à comprendre les abstractions du DSL aide énormément.
Par exemple, cloner le dépôt
NVIDIA cutlass
dans le répertoire de contexte est un bon moyen de laisser l'agent chercher des abstractions liées à
l'algèbre des layouts, aux atomes Copy/GEMM, à la hiérarchie mémoire, aux kernels d'exemple
, et ainsi de suite lors de l'écriture de kernels en CuTeDSL.
D'après mon expérience, une bonne première version du kernel compile sans erreur évidente et passe le test de correction dont je parlerai plus bas.
Test, benchmark et profiling
Une fois que l'agent a reçu suffisamment de contexte, le véritable goulot d'étranglement se déplace vers la validation. L'implémentation de référence elle-même, et la validation par rapport à celle-ci, deviennent de plus en plus importantes. J'appelle cette phase le test de correction, ou simplement le test. La vitesse de l'implémentation de référence importe moins que son intention. Ce que vous avez l'intention de mesurer et de vérifier est ce que votre agent cherchera à optimiser.
En général, lorsque le calcul n'est pas censé se faire en basse précision, je mesure l'erreur absolue/relative maximale (MAE), l'erreur quadratique moyenne (MSE/RMSE) et le PSNR (rapport signal sur bruit de crête). Lorsque la basse précision est en jeu, j'ai tendance à mesurer le PSNR et la similarité cosinus (cossim).
La façon de faire exécuter les versions du kernel sur un GPU dépend du fait que le GPU soit disponible en local ou dans le cloud. Quoi qu'il en soit, notre agent doit pouvoir accéder à ses sorties d'une manière ou d'une autre.
Je trouve que la méthode des rungs (échelons) ci-dessous est un bon moyen d'avoir N fonctions de test :
1def rung(name):2 def deco(fn):3 try:4 out = fn()5 results[name] = {"ok": True, **(out or {})}6 print(f"[{name}] ok " + " ".join(f"{k}={v}" for k, v in (out or {}).items()))7 except Exception as e:8 results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}9 print(f"[{name}] FAILED {type(e).__name__}: {e}")10 traceback.print_exc()11 return fn12 return deco
que vous pouvez appeler ainsi :
1out = {}23@rung("pre-checks")4def _():5 run_pure_checks()6 run_dsl_checks()78@rung("run")9def _():10 out["o"] = custom_kernel(*inputs)11 torch.cuda.synchronize()12 return {"shape": tuple(out["o"].shape),13 "finite": bool(torch.isfinite(out["o"]).all())}
Pour le rung de benchmark, vous pouvez faire plusieurs choses :
- Benchmarker le temps d'exécution de bout en bout du kernel afin d'obtenir le temps global passé
- Utiliser le traçage intra-kernel pour benchmarker des sections d'un kernel et les écrire dans la sortie (à l'aide d'un traceur personnalisé ou de CUPTI)
- Extraire l'IR, le PTX, le SASS et le CUBIN générés vers un répertoire dumps et laisser l'agent les parcourir
Ce dernier point mérite d'être développé. Il peut arriver que le DSL abaisse le code pour générer un PTX sous-optimal (puis un SASS sous-optimal), et que vous trouviez de meilleures instructions ou de meilleures formes à utiliser à la place. Notre agent peut parcourir les fichiers texte PTX/SASS et intégrer directement du code de plus bas niveau au lieu de laisser le DSL gérer la partie sous-optimale. Là encore, fournir la documentation PTX « consultable » comme contexte est très utile.
La dernière étape pour tout rassembler est le profiling. Si votre agent peut accéder à la CLI NCU (Nsight Compute Systems), vous pouvez lui demander de profiler et de générer un rapport pour votre kernel dans le cadre de la boucle de rétroaction vérifiable ci-dessus.
Réflexions finales
« Le développement de kernels GPU est-il donc mort ? »
« Eh bien oui, mais en fait non. »
Oui, parce que la partie difficile des layouts, de l'indexation, des abstractions et de la structure globale peut être en grande partie résolue par des agents disposant de suffisamment de contexte. Vous pouvez facilement réduire un travail de 2 à 3 semaines à 1 ou 2 jours. Non, parce que le véritable goulot d'étranglement s'est déplacé des kernels vers les validations. Il n'y a plus une seule façon de faire : plus votre contexte et votre harnais de test sont bons, plus le processus est rapide. Les cas spécialisés en bénéficieront encore plus, et un bon harnais de test est tout ce dont vous aurez besoin.
Enfin, plutôt que de traiter les agents comme autonomes, il reste nécessaire de les considérer comme des assistants très intelligents que vous pouvez guider. C'est là que votre compréhension fondamentale des GPU et des kernels s'avère utile. La part humaine (vous) est toujours nécessaire.
C'est un sentiment doux-amer, je sais :)
