NVIDIA zeigt, wie KI-Agenten Spark-UDFs für GPU-Beschleunigung umschreiben

Warum UDFs GPU-beschleunigtes Spark ausbremsen

Das NVIDIA-cuDF-Plugin für Apache Spark nutzt die Parallelität und den hohen Speicherdurchsatz von GPUs, um Datenanalysen zu beschleunigen. Laut NVIDIA funktioniert das vor allem deshalb, weil die SQL- und DataFrame-Schnittstellen von Spark deklarativ sind: Sie trennen, was eine Abfrage liefern soll, davon, wie das System sie ausführt. So kann das Plugin zeilenweise Operationen auf der GPU parallelisieren, ohne dass der Nutzercode geändert werden muss.

UDFs durchbrechen dieses Modell. Es handelt sich um benutzerdefinierte Routinen mit beliebigem Code, den das System weder allgemein analysieren noch parallelisieren kann. Trifft eine Abfrage auf eine UDF, muss der GPU-Beschleuniger die Daten vom spaltenorientierten in ein zeilenorientiertes Format umwandeln, auf die CPU ausweichen, die Funktion Zeile für Zeile ausführen und die Ergebnisse anschließend zurückkonvertieren.

NVIDIA nennt zwei Nachteile für eine ansonsten beschleunigte Abfrage:

• Der Rückfall auf die CPU verursacht zusätzlichen Overhead durch Datentransfers in beide Richtungen.
• Die UDF verschenkt die Rechenbeschleunigung, die mit GPU-beschleunigten Primitiven möglich wäre.

Drei Wege, eine UDF umzuschreiben

NVIDIA beschreibt drei Ansätze zum Umschreiben, die den Fallback-Overhead beseitigen und Beschleunigungen von mehr als dem 20-Fachen ermöglichen sollen. Sie unterscheiden sich in Komplexität und Flexibilität.

• SQL: geringste Komplexität und höchster Abstraktionsgrad. Am besten geeignet für einfache relationale Logik, etwa kleine oder inlinefähige Funktionen. Die UDF wird durch einen äquivalenten Spark-SQL-Ausdruck ersetzt. Dieser erscheint im Expression Tree von Spark, sodass das Plugin bei der Abfrageplanung GPU-Implementierungen einsetzen kann.
• cuDF: mittlere Komplexität, mehr Kontrolle über die Ausführung. NVIDIA empfiehlt diesen Weg für die meisten Fälle, etwa beim Parsen von Strings oder bei verschachtelten Typen. Die UDF implementiert das RapidsUDF-Interface, einen Vertrag des cuDF-Plugins, der eine einzige Methode verlangt: evaluateColumnar. Das Plugin übergibt ihr einen GPU-residenten ColumnVector, und die Funktion setzt Operationen aus der Java-API für cuDF zusammen.
• CUDA: hohe Komplexität und maximale Flexibilität. Gedacht für besonders komplexe oder performancekritische Logik, die sich nicht effizient aus vorhandenen cuDF-Operationen zusammensetzen lässt. Eine RapidsUDF kann über das Java Native Interface (JNI) einen eigenen CUDA-Kernel aufrufen.

Die ersten beiden Ansätze veranschaulicht der Beitrag an einer einfachen Scala-UDF, die Celsius in Fahrenheit umrechnet. Die SQL-Variante ersetzt den Funktionsaufruf durch den Ausdruck celsius * 9.0 / 5.0 + 32.0, die cuDF-Variante multipliziert und addiert auf GPU-Spalten mit Skalarwerten.

Coding-Agenten übernehmen das Umschreiben

NVIDIA räumt ein, dass das Umschreiben weiterhin erheblichen Engineering-Aufwand erfordert. Nötig sind eine sorgfältige Auswahl der APIs, die Validierung der Semantik und Performance-Messungen. Zudem können spaltenorientierte APIs für Entwickler mit wenig Erfahrung in vektorisierter Programmierung ungewohnt sein. Das Unternehmen argumentiert, dass sich ein Großteil dieser Arbeit für Coding-Agenten eigne: Mit geeigneten Leitplanken könne ein Agent unbekannte APIs durchsuchen, Kandidaten für das Umschreiben implementieren und schnell iterieren, während ausführbare Tests und Performance-Messungen darüber entscheiden, ob eine Änderung akzeptabel ist.

Die cuDF Skills für Apache Spark gliedern das Umschreiben in vier fokussierte Feedback-Schleifen. Jede Skill legt fest, welche Arbeit zu leisten ist und welche Belege vor dem nächsten Schritt vorliegen müssen. Die Skills decken Tests, Konvertierung, Benchmarking und Optimierung ab.

Dazu kommt eine vollständige Projektvorlage mit dem cuDF- und Spark-spezifischen Gerüst für die Entwicklung von Anfang bis Ende, darunter Test- und Benchmark-Harnesses, codierte Routinen zur Überprüfung der Ausgaben und die Build-Konfiguration. NVIDIA zufolge kann sich der Agent so auf die UDF konzentrieren, statt immer wieder spezialisierte Infrastruktur neu aufzubauen. Außerdem enthalten die Skills Best Practices für die GPU-Programmierung mit cuDF, etwa zu Speicherverwaltung und Konfiguration, als beratende Referenz. Der Agent kann weiterhin je nach Kontext verschiedene Implementierungen ausprobieren.

Berichtete Ergebnisse und Bedeutung

NVIDIA zufolge hat ein großes Unternehmen der Konsumtechnologie mit diesem Workflow mehr als 30 produktive UDFs umgestellt. Bei einer der Umschreibungen sank die Operatorzeit bei einer Last von 600 Milliarden Zeilen von 274,55 auf 12,78 Stunden. Diese Zahlen stammen von NVIDIA selbst.

Für Teams, die Spark auf GPUs betreiben, lautet das Versprechen: UDFs, lange eine Quelle für CPU-Fallbacks, lassen sich mit weniger manuellem Aufwand umstellen, wobei die Verifikation in jeden Schritt eingebaut ist. NVIDIA verweist Interessierte auf die Installation der cuDF Skills, eine Session auf dem Data + AI Summit zu Coding-KI-Agenten für die GPU-Beschleunigung von Spark-UDFs sowie auf den Einstiegsleitfaden zum cuDF-Plugin.