K-Search: CUDA カーネルの専門知識を Apple Silicon MLX に転送

TL;DR

研究者は、K-Search（進化的カーネル最適化フレームワーク）に構造化された変換レイヤーを追加し、Apple Silicon の MLX フレームワーク向けに数十年にわたる CUDA カーネルの専門知識を適応させました。このアプローチにより、AI は高性能な GPU カーネルを自動生成でき、ネイティブな MLX Attention カーネルの 0.97x の速度を達成し、コミュニティ実装と比較して Mamba SSM カーネルのプリフィルで最大 20x の速度向上を実現しました。

クロスプラットフォームカーネル最適化の課題

効率的な GPU カーネルを書くには、何年もの専門知識が必要であり、これらの最適化をあるハードウェアベンダーから別のベンダーに移す場合、通常はゼロから再発見する必要があります。CUDA エコシステムには、注意機構や状態空間モデル（SSM）などの重要な操作に関する広範な手調整済み実装がありますが、Apple Silicon のような新しいエコシステムはこのような最適化カーネルの深さに欠け、MLX フレームワークの能力にもかかわらず、しばしば大きなパフォーマンスを逃しています。

K-Search: 進化的カーネル最適化

K-Search は、反復ループを使用して GPU カーネルを最適化する進化的フレームワークです。このプロセスは、以下の 3 つの主要なフェーズで構成されます：

  1. アクション選択：LLM（この研究では特に Gemini 3.5 Pro Preview）が "GPU カーネル パフォーマンス エンジニア" として機能し、カーネルの分類、データレイアウト、および考えられるボトルネックを分析して、探索木（"ワールドモデル"）から最適化アクションを提案します。
  2. ローカルな改良：コード生成モデルが選択されたアクションに基づいて候補実装を生成し、それらは実際のハードウェアでコンパイルおよびベンチマークされます。
  3. ワールドモデルの更新：LLM が結果を理由付けして、新しいアクションを挿入し、優先度スコアを更新し、または失敗したパスを剪定することで探索木を更新します。

この探索は "Spec" によって根ざしています。Spec は、無効なプリミティブの生成を防ぐためのハードウェアルールと数学的制約を含むドメイン固有のドキュメントです。

CUDA-to-MLX 変換レイヤー

NVIDIA と Apple Silicon のアーキテクチャ間のギャップを埋めるため、研究者は CUDA の概念的知識を MLX/Metal 戦略に変換する変換レイヤーを開発しました。このレイヤーは以下で構成されます：

  • 概念マッピングテーブル：ハードウェア固有の制約を伴う CUDA プリミティブと Metal の等価物を対応付ける用語集（例：Apple Silicon の 32 KB 制約と NVIDIA の 48 KB を考慮し、__shared__ メモリを Metal の threadgroup メモリにマッピング）。
  • MLX 固有のヒント：CUDA と直接対応しないパターンに関するガイダンス。たとえば、simd_shuffle_xor を使用したレジスタベースの行削減や、Apple の高速な fast::exp2() ハードウェア命令を活用する "exp2 トリック"（$e^x$ を $2^{x \log_2 e}$ に置き換える）。
  • 再利用可能なアサーション：エキスパートのカーネル動作を、コードを literally コピーするのではなく、進化的探索が維持しなければならないプロパティとして言い換えること。

パフォーマンスベンチマーク

Attention カーネル結果

変換レイヤーからの完全なコンテキストを進化的探索に提供することで、研究者は専門家レベルに近いパフォーマンスを達成しました。進化されたカーネルは、以下の高度な戦略を独立して発見および実装しました：

  • スレッドグループメモリのタイリング
  • オンライン ソフトマックス
  • メモリアクセスのための K 転置
  • exp2 トリック

これにより、純粋な進化による 0.26x から、Apple の最先端ネイティブ Attention カーネルの 0.97x の速度へとパフォーマンスが跳ね上がりました。

Mamba SSM カーネル結果

K-Search は、一般化をテストするために Mamba 状態空間モデル（SSM）カーネルに適用されました。M1 Max (64GB) で mamba-370m f16 を使用した場合、進化した mlx-mamba カーネルは、コミュニティの mlx-lm 実装と比較してプリフィルスループットが大幅に向上しました：

Metric mlx-mamba (ours) mlx-lm (community) mamba.py
Decode 152 tok/s 116 tok/s 40 tok/s
Prefill L=512 5,751 tok/s 329 tok/s 1,089 tok/s
Prefill L=1024 6,010 tok/s 327 tok/s 1,127 tok/s
Prefill L=2048 6,612 tok/s 1,092 tok/s 1,092 tok/s
Prefill L=4096 6,743 tok/s 339 tok/s 1,042 tok/s

Key Insight: ~20x のプリフィル速度向上は、並列（プレフィックス）スキャンの実装によるものです。コミュニティの mlx-lm 実装はトークンを順次処理しますが、進化したカーネルは結合法則を持つ演算を使用して、$O(\log N)$ の従属ステップでシーケンスを評価し、プリフィルフェーズにおいて Apple Silicon の GPU スループットをフルに活用します。

今後の方向性

研究者は、IBM Spyre AIU などの追加のハードウェアアーキテクチャをサポートするようこの作業を拡張し、 fused MoE ルーティングやページド アテンションなどのより複雑なカーネルを開発しています。主な発見は、AI 駆動のカーネル生成におけるボトルネックは LLM のコーディング能力ではなく、モデルに提供されるアーキテクチャのコンテキストと制約の質であるということです。

Sources