PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

Im Inneren des modernen Computers fungiert die Grafikverarbeitungseinheit, oder GPU, als ein leistungsstarker Motor, der darauf ausgelegt ist, massive Mengen an Daten gleichzeitig zu verarbeiten. Um das Beste aus diesem Motor herauszuholen, schreiben Softwareentwickler spezialisierte Anweisungen, die der Hardware genau vorgeben, wie sie Daten bewegen und Berechnungen durchführen soll. Diese Anweisungen werden in einer niederschwelligen Sprache namens PTX geschrieben, die dem Programmierer die direkte Kontrolle über die inneren Abläufe der Maschine ermöglicht. Obwohl höherwertige Werkzeuge existieren, um diesen Prozess zu erleichtern, übersehen sie manchmal die einzigartigen Merkmale der neuesten Hardwaregenerationen. Da sich diese Maschinen rasant weiterentwickeln, wird es zu einer schwierigen Herausforderung, die Software perfekt auf deren spezifische Fähigkeiten abzustimmen. Es stellt sich die Frage: Kann künstliche Intelligenz, speziell große Sprachmodelle, lernen, diese präzisen, niederschwelligen Anweisungen eigenständig zu schreiben, oder greift sie weiterhin auf ältere, generische Methoden zurück, die Leistung ungenutzt lassen?

Ein Team von Forschern machte sich daran, dies zu beantworten, indem sie ein neues Testfeld namens PTXBench schufen. Sie wollten sehen, ob diese KI-Modelle nicht nur Code schreiben können, der funktioniert, sondern auch Code, der aktiv die spezifischen, fortschrittlichen Funktionen der neuesten GPUs nutzt. Die Forscher konzentrierten sich auf zwei der leistungsstärksten Grafikkarten, die verfügbar sind, die H100 und die B200, und testeten die Modelle bei Aufgaben, die komplexe mathematische Operationen beinhalten, welche in der künstlichen Intelligenz verwendet werden, wie etwa Matrixmultiplikation und Attention-Mechanismen. Das Ziel war es, drei Dinge zu messen: ob der Code ohne Fehler lief, ob er tatsächlich die spezifischen Hardware-Instruktionen auslöste, die die Forscher angefordert hatten, und ob er schneller lief als die besten existierenden Softwarebibliotheken.

Die Ergebnisse offenbarten eine signifikante Lücke zwischen dem, was die Modelle leisten konnten, und dem, was für Spitzenleistungen erforderlich war. Während die KI-Modelle oft erfolgreich darin waren, Code für einfachere Aufgaben korrekt zu schreiben, hatten sie mit komplexeren Operationen erheblich zu kämpfen, insbesondere bei jenen, die den Backward-Pass von Attention-Mechanismen betreffen, welche entscheidend für das Training von KI-Systemen sind. Selbst wenn die Modelle es schafften, Code zu schreiben, der die angeforderten Hardware-Instruktionen ausführte, erreichte dieser Code selten die Geschwindigkeit der führenden professionellen Bibliotheken. In vielen Fällen war der KI-generierte Code langsamer, was darauf hindeutet, dass das bloße Ausführen der Anweisungen nicht ausreichte, um echte Optimierung zu erreichen. Die Studie fand heraus, dass die Modelle häufig auf generische Codierungsmuster zurückfielen, anstatt die einzigartigen, architekturspezifischen Merkmale der neueren Chips zu nutzen, selbst wenn diese Merkmale explizit beschrieben wurden.

Um zu verstehen, wie man diese Modelle verbessern kann, führten die Forscher ein kontrolliertes Experiment durch, bei dem sie ein spezifisches KI-Modell mittels einer Methode namens Supervised Fine-Tuning unterrichteten. Sie zeigten dem Modell Beispiele seiner eigenen Fehler zusammen mit korrigierten Versionen des Codes und Erklärungen, warum die Korrekturen funktionierten. Dieser Ansatz, den sie Fixit nannten, half dem Modell, sich in mehreren Aufgaben zu verbessern, insbesondere bei der Generierung von funktional korrektem Code. Die Verbesserung war jedoch nicht einheitlich. Das Modell hatte weiterhin Schwierigkeiten, seine neuen Fähigkeiten auf unterschiedliche Problemgrößen oder Variationen der Aufgaben zu generalisieren, die es während des Trainings nicht gesehen hatte. Die Forscher entdeckten, dass die Qualität der Korrekturen und die Balance der Trainingsdaten ebenso wichtig waren wie die schiere Menge der Daten. Ein Modell, das auf einem vielfältigen Satz von Problemen trainiert wurde, schnitt besser ab als eines, das auf einem größeren, aber weniger variierten Datensatz trainiert wurde, was darauf hindeutet, dass die Art der Lernerfahrung wichtiger war als das Volumen.

Die Studie hob auch die Bedeutung der Bereitstellung des richtigen Kontexts für die KI hervor. Wenn die Forscher den Modellen detaillierte Informationen über die spezifischen Fähigkeiten der Hardware gaben, war die Wahrscheinlichkeit deutlich höher, dass die Modelle die angeforderten Instruktionen nutzten. Ohne dieses spezifische Wissen nutzte das Modell das angeforderte PTX selten. Dies deutet darauf hin, dass KI zwar lernen kann, auf dieser niedrigen Ebene zu programmieren, aber präzise Anleitung und qualitativ hochwertige Beispiele benötigt, um dies effektiv zu tun. Die Forscher kamen zu dem Schluss, dass aktuelle KI-Modelle zwar vielversprechend sind, aber noch nicht bereit sind, Experten-Ingenieure zu ersetzen, die Code für die neueste Hardware optimieren. Der Weg nach vorn führt über bessere Trainingsdaten, ausgewogenere Lernerfahrungen und ein tieferes Verständnis dafür, wie man diesen Modellen beibringt, über die spezifischen Einschränkungen sich entwickelnder Computerarchitekturen zu urteilen.