AutoKernel : système de recherche automatique de kernels GPU

• Application de l’Autoresearch de Karpathy à l’optimisation de kernels GPU
• Outil de recherche automatique de kernels GPU qui optimise automatiquement des kernels Triton ou CUDA C++ à partir d’un modèle PyTorch en entrée
• Fournit un pipeline entièrement automatisé pour profiler, extraire, optimiser et valider les kernels goulots d’étranglement du modèle
• Ajuste les priorités via une orchestration fondée sur la loi d’Amdahl ; chaque expérience prend environ 90 secondes, ce qui permet des centaines d’itérations pendant la nuit
• Intégré à KernelBench, avec 50 à 300 itérations d’expérimentation sur plus de 250 problèmes, pour une exploration systématique au-delà de la génération ponctuelle
• Prend en charge un double backend Triton et CUDA C++, une validation priorisant la correction et une structure de modification sur fichier unique, pour une optimisation de kernels efficace et reproductible

---

Principe de fonctionnement

• AutoKernel prend un modèle PyTorch en entrée et exécute les étapes suivantes
  • Profilage pour identifier les kernels GPU qui constituent les goulots d’étranglement
  • Extraction pour les isoler en kernels Triton ou CUDA C++ indépendants
  • Boucle d’optimisation automatique qui répète modification, benchmark, conservation ou annulation
  • Validation de la correction puis rapport du gain global de vitesse
• En s’appuyant sur les instructions de program.md, il modifie kernel.py et effectue, via bench.py, une vérification de correction en 5 étapes ainsi qu’une analyse roofline
• Chaque expérience prend environ 90 secondes, soit environ 40 essais par heure et environ 320 essais sur une nuit

Exécution et composants

• Environnement requis : GPU NVIDIA (H100/A100/RTX 4090), Python 3.10+, paquet uv
• Composition des scripts principaux
  • profile.py : classement des kernels selon le temps GPU
  • extract.py : extraction des principaux kernels goulots d’étranglement
  • bench.py : benchmark de correction et de performance en 5 étapes
  • orchestrate.py : planification multi-kernel fondée sur la loi d’Amdahl
  • verify.py : validation du modèle complet et rapport de gain de vitesse
• program.md inclut une stratégie d’optimisation en 6 étapes, la gestion des conflits et un cadre de prise de décision permettant une exécution automatique sur de longues durées

Kernels pris en charge et modèles d’exemple

• Prise en charge de 9 types de kernels : matmul, softmax, layernorm, rmsnorm, flash_attention, fused_mlp, cross_entropy, rotary_embedding, reduce
• Chaque kernel inclut une implémentation de référence PyTorch (reference.py) et des versions initiales Triton/CUDA (kernels/, kernels/cuda/)
• Modèles d’exemple : GPT-2, LLaMA (y compris 7B), BERT-base, modèles personnalisés
• Les modèles HuggingFace peuvent aussi être intégrés avec uv sync --extra models

Intégration à KernelBench

• Intégré à KernelBench (Stanford Scaling Intelligence Lab) pour exécuter le benchmark standard des kernels GPU générés par IA
• AutoKernel explore systématiquement l’espace d’optimisation avec 50 à 300 itérations d’expérimentation par problème
• Outils principaux
  • bridge.py : chargement des problèmes et génération du kernel initial
  • bench_kb.py : évaluation de la correction et des performances
  • scorer.py : calcul du score global par niveau
  • program_kb.md : instructions d’agent dédiées à KernelBench

Export vers HuggingFace Hub

• Les kernels optimisés peuvent être exportés vers HuggingFace Hub et chargés simplement avec get_kernel()
• export_hf.py permet l’export et l’envoi de kernels CUDA

Principes de conception

• Double backend Triton + CUDA C++ : Triton pour itérer rapidement, CUDA pour les performances maximales
• Priorité à la correction : si le résultat ne correspond pas à celui de PyTorch, retour immédiat en arrière
• Orchestration fondée sur la loi d’Amdahl pour définir les priorités selon la contribution aux performances globales
• Structure de modification sur fichier unique (kernel.py) facilitant le suivi des changements et la restauration
• Journal TSV (results.tsv) pour enregistrer les résultats d’expérience de manière simple et lisible

Format d’enregistrement des résultats

• results.tsv enregistre pour chaque expérience le numéro, le type de kernel, le débit (TFLOPS), la latence, le ratio par rapport au pic GPU, l’accélération par rapport à PyTorch, la correction, l’utilisation de VRAM et la description

Contexte du projet

• Inspiré par le concept d’autoresearch d’Andrej Karpathy, qui applique une architecture d’agent IA autonome destinée à la recherche sur les LLM à l’optimisation de kernels GPU
• L’intégration de KernelBench s’appuie sur les travaux du Stanford Scaling Intelligence Lab, et AutoKernel réalise une optimisation itérative plutôt qu’une génération ponctuelle
• Le projet est développé par l’équipe Forge de RightNow AI et publié sous licence MIT