Von CUDA zu MLX: Die Expertise von K-Search im Dienste von Apple Silicon

Eine neue Ära im Rechnen

Die Welt des Rechnens entwickelt sich rasant. Die Hardware-Innovationen beschränken sich nicht auf schnellere GPUs, sondern führen auch zu einer Vielzahl von Chips verschiedener Hersteller, von denen jeder seine eigene Architektur hat, die oft für spezifische Arbeitslasten im Zusammenhang mit künstlicher Intelligenz (KI) konzipiert ist. Parallel dazu schreitet die Softwareentwicklung ebenso schnell voran, mit KI-Codierungstools, die in der Lage sind, das, was zuvor Monate dauerte, in wenigen Minuten zu generieren.

In einer Umgebung, in der KI eine zentrale Rolle spielt, werden GPU-Kernels unerlässlich. Diese Programme auf niedriger Ebene, die innerhalb von GPUs ausgeführt werden, erfordern erhebliche Expertise, um effektiv optimiert zu werden. Die Migration eines Kernels von einer Hardware zu einer anderen stellt eine noch grössere Herausforderung dar, die oft eine Neuentdeckung der spezifischen Optimierungen für jede Architektur erfordert.

Die Expertise der CUDA-Kernels

Das CUDA-Ökosystem, unterstützt durch Jahrzehnte an Forschung und Optimierungen, hat die Entwicklung hochgradig effizienter Implementierungen kritischer Funktionen wie Aufmerksamkeit oder Zustandsraum-Modelle ermöglicht. Diese Errungenschaften repräsentieren Tausende von Ingenieurstunden. Allerdings entwickeln sich neuere Hardware-Ökosysteme, wie Apple Silicon und andere massgeschneiderte KI-Beschleuniger, schnell, verfügen jedoch nicht über diese Tiefe an Expertise.

Es wird daher entscheidend, die Frage zu stellen, ob diese Expertise automatisch übertragen werden kann.

K-Search: Ein Fortschritt in der Optimierung von Kernels

K-Search, ein skalierbares Optimierungsframework für Kernels, das an der UC Berkeley entwickelt wurde, stellt sich dieser Herausforderung. Diese Methode nutzt künstliche Intelligenz, um GPU-Kernels zu optimieren. Durch die Erweiterung von K-Search mit einem Backend für MLX, dem maschinellen Lern-Framework von Apple, haben wir eine Übersetzungsschicht von CUDA zu MLX geschaffen. Dieser Ansatz ermöglicht es K-Search, auf vorhandene CUDA-Kernels zurückzugreifen, um neue für Apple Silicon zu generieren, wodurch der Bedarf verringert wird, Kernels von Grund auf neu zu erstellen.

Die Ergebnisse sprechen für sich: Unsere Methode erreicht eine Leistung, die der eines Experten auf Apple Silicon nahekommt, und bietet einen Geschwindigkeitsgewinn von 0,97x im Vergleich zum nativen MLX-Aufmerksamkeitskernel und bis zu 20x in der Vorabfüllgeschwindigkeit im Vergleich zur Community-Implementierung mlx-lm auf dem Mamba SSM-Kernel.

Warum MLX wählen?

Seit Ende 2023 hat das MLX-Framework von Apple eine bemerkenswerte Akzeptanz erfahren. Mit Millionen von MacBooks und Mac Studios, die mit Apple Silicon ausgestattet sind, ermöglicht MLX die lokale KI-Inferenz ohne Cloud-Kosten. Seine einheitliche Speicherarchitektur macht es besonders geeignet für Modelle mittlerer Grösse (von 7 bis 70 Milliarden Parametern auf den M-Chips).

Allerdings bleibt eine bedeutende Lücke bestehen: Viele leistungsrelevante Kernels, die häufig im NVIDIA-Ökosystem zu finden sind, sind entweder nicht vorhanden oder auf MLX ohne hardware-spezifisches Tuning unteroptimiert. Obwohl MLX Modelle korrekt ausführt, lässt es oft erhebliches Leistungspotenzial ungenutzt.

Die Funktionsweise von K-Search

K-Search zeichnet sich durch seine Fähigkeit aus, auf einem naiven Kernel basierend auf einer Hardware-Spezifikation zu iterieren. Ein KI-Modell überlegt sich, welche Optimierungen anzuwenden sind, generiert Kandidaten-Kernels, die dann auf echter Hardware kompiliert und getestet werden. Die gewonnenen Messwerte fliessen in den Suchprozess ein und ermöglichen eine kontinuierliche Verbesserung, bis die Leistung konvergiert.

Dieser Prozess wird durch ein domänenspezifisches Dokument geleitet, das die Hardware-Regeln und mathematischen Einschränkungen kodiert, um die Generierung ungültigen Codes zu verhindern und sicherzustellen, dass die Kandidaten effizient kompiliert und ausgeführt werden. In der Praxis übernimmt ein einzelnes Modell (wie Gemini 3.5 Pro Preview) sowohl die Rolle des Denkens als auch der Code-Generierung, indem es den Kernel analysiert, bevor es Optimierungen vorschlägt.

Fazit

Zusammenfassend eröffnet die Kombination aus der angesammelten Expertise im CUDA-Ökosystem und der Innovation, die K-Search mit sich bringt, neue Perspektiven für die Leistungsoptimierung auf Apple Silicon. Die Anwendungen dieser Methode beschränken sich nicht auf MLX, sondern erstrecken sich auf jedes Ökosystem, in dem CUDA-Expertise nützlich übertragen werden kann.

Wenn Sie mehr über die Optimierung der Leistung Ihrer KI-Infrastruktur oder die Umsetzung ähnlicher Strategien erfahren möchten, zögern Sie nicht, Kontaktieren Sie mich.