KernelAgent: Ottimizzazione dei Kernel GPU con Guida Hardware

Il team di PyTorch ha presentato KernelAgent, un sistema open source basato su agenti per l'ottimizzazione dei kernel GPU. Questo sistema integra i segnali di performance dell'hardware in un flusso di lavoro multi-agente a ciclo chiuso per guidare l'ottimizzazione dei kernel Triton.

KernelAgent ha dimostrato un miglioramento di 2.02x rispetto ai kernel generati dalle versioni precedenti e un miglioramento medio di 1.56x rispetto a torch.compile. In particolare, ha superato le prestazioni di torch.compile nel 65% dei task L1 di KernelBench e ha raggiunto l'89% dell'efficienza teorica dell'hardware su NVIDIA H100.

L'ottimizzazione dei kernel GPU è fondamentale per i carichi di lavoro AI moderni, dove le performance sono spesso limitate dall'efficienza dei kernel. KernelAgent automatizza il processo di ottimizzazione, tradizionalmente intensivo in termini di competenze, profilando i kernel, diagnosticando i colli di bottiglia e proponendo ottimizzazioni.

Pipeline di Lavoro di Ottimizzazione

KernelAgent automatizza il flusso di lavoro seguito dagli ingegneri esperti, scomponendolo in una serie di agenti cooperanti. Ogni agente è responsabile di una fase ben definita del ciclo di ottimizzazione, formando un sistema di feedback guidato dall'hardware.

Il flusso di lavoro comprende le seguenti fasi:

Profile → Diagnose → Prescribe → Orchestrate → Explore → Measure

Ogni fase produce output strutturati che alimentano direttamente la successiva, consentendo un'iterazione rapida e guidata dai dati.

Come i Dati Fluiscono Attraverso il Sistema

  • Profiling: L'agente di Profiling utilizza NVIDIA Nsight Compute (NCU) per acquisire metriche di performance a livello hardware.
  • Diagnosis: L'agente di Diagnosi interpreta le metriche di profiling per classificare il collo di bottiglia dominante del kernel.
  • Prescribing Fixes: L'agente di Analisi genera raccomandazioni di ottimizzazione concrete, tenendo conto delle specifiche della GPU.
  • Orchestration: L'agente Orchestrator sintetizza le diagnosi correnti con i dati di ottimizzazione storici per formulare una strategia di ricerca.
  • Exploration: L'Optimization Manager esegue la fase di exploration, mantenendo i kernel con le migliori performance ed esplorando diverse correzioni in parallelo.
  • Measure: Il Benchmarking Agent convalida la correttezza e misura le performance reali di ogni variante del kernel.

Conclusioni

KernelAgent dimostra che i principi degli agenti profondi si estendono naturalmente all'ottimizzazione delle performance. Aggiungendo il profiling hardware e la working memory al ciclo, consentendo ai multi-agenti di sviluppare ed esplorare diversi percorsi di ottimizzazione, possiamo spingere i kernel verificati da "corretti" a "corretti e veloci".

Per chi valuta deployment on-premise, esistono trade-off da considerare. AI-RADAR offre framework analitici su /llm-onpremise per valutare questi aspetti.