TIP
メニュー
EN
← Research

Research

MLコンパイラ「Cleave」、代数変換とスケジューリングを分離してFlashAttentionを超える最大2.8倍の高速化を実現

2026/10/10

専門家の手書きカーネルを自動生成で超える

FlashAttentionやFlashDecodingをはじめとする高度に最適化されたGPUカーネルは、今日の大規模モデルの推論速度を左右する重要な要素だ。しかしその多くはMLコンパイラが自動生成できないため、専門家が手書きで実装してきた。ニューヨーク大学・コーネル大学の研究チームは、この課題に正面から取り組んだMLコンパイラ「Cleave」を発表した。論文はarXiv（2610.07742）で公開されており、コードもGitHubで入手できる。

問題の核心：探索空間の爆発

高性能なカーネルを自動生成するには、計算グラフの代数変換と、変換後のグラフをGPUスレッドへどう割り当てるかを決めるオペレータスケジューリングという2つの最適化が必要になる。既存の手法はどちらか一方しか扱えないか、両者を同時に探索しようとして空間が爆発的に大きくなるという問題を抱えていた。

先行研究である「Mirage」はこの2つを合同で超最適化（superoptimization）するアプローチを取っているが、論文によれば一部のカーネルの最適化に失敗したり、メモリ不足に陥るケースがあるという。

Cleaveの核心：シンボリック分離

Cleaveが提案する解決策は「シンボリック分離（symbolic decoupling）」と呼ばれる手法だ。

代数変換の有効性はグラフの構造だけに依存し、テンソルの具体的な形状やスケジューリングには依存しないという観察に基づいている。そこでCleaveは、テンソルの次元をすべてシンボル（記号）で表した抽象的なグラフ上で超最適化を実行し、代数的に等価なグラフをすべて列挙する。その後、具体的な形状をシンボルに束縛してから各グラフのスケジューリングを行う。

この分離によって以下の2つの利点が生まれると著者らは説明している。

  • 等価性チェックのコスト削減：確率的テストに使うテンソルサイズを任意の小さな値に設定できるため、計算コストが大幅に下がる
  • 探索空間の縮小：代数探索フェーズはループ構造やタイルサイズを考慮せず、スケジューラは代数的等価性を考慮しない

2つの技術的課題への対応

Split演算子によるSplit-Kの対応

GPUのスループットを最大化するためにはリダクション次元を複数の部分計算に分割する「Split-K」という手法が必要だが、この変換はパーティションサイズというスケジューリング上の決定に依存するように見える。Cleaveはこの問題を、分割数 s を新たなシンボルとして導入したSplitオペレータでシンボリックグラフに組み込むことで解決した。超最適化がSplitの挿入箇所を決定し、具体的な s の値はスケジューラが後から決定する。

複数リダクションのスケジューリング

代数変換後のグラフでは複数のリダクションが含まれる場合があり、中間テンソルのメモリ使用量を抑えつつこれらを融合することが課題となる。Cleaveのスケジューラは以下の2つの戦略を組み合わせて対応すると報告されている。

  1. 反復タイリング（iterative tiling）：リダクション次元を小さなタイルに分割し、中間結果全体をSMEMに展開せずに処理する
  2. 水平融合（horizontal fusion）：共通の上流プロデューサを持つリダクションを同一ループ内で計算し、タイルを再利用する

また、バッチサイズやシーケンス長が動的に変化するワークロードに対しては、代表的な具体形状でスケジューリングを行い、コード生成時には動的な次元をシンボルのままランタイムのループ境界として保持する仕組みを持つ。これにより、オペレータごとに1度のコンパイルでさまざまな形状のリクエストに対応できると説明されている。

評価結果

論文に記載された実験結果は以下のとおりだ。

LLMサブグラフでの性能（GQA、MLA、SwiGLUなど）

  • 最良のベースラインと比較して最大 2.8倍、平均 1.6倍 の高速化
  • コンパイル時間はMirageと比較して平均 5.9倍 短縮

本番サービングトレースからの動的ワークロード（9オペレータ、303形状）

  • FlashInferの手書きFA2バックエンドに対して幾何平均 1.4倍 の高速化
  • FlashInferの手書きFA3バックエンドに対して幾何平均 1.7倍 の高速化

Transformerレイヤー全体（GPT-2、Llama-3 8B、Qwen3-32B）

  • torch.compile と比較してレイヤー実行時間が最大 2.5倍（平均 1.7倍）改善

まとめ

Cleaveは代数変換とスケジューリングを「シンボリック分離」によってデカップルするという明快なアイデアで、手書き専門カーネルの自動生成という長年の課題に対して有望な解法を示した。FlashAttentionやFlashDecodingの開発に費やされてきた専門家の手作業を自動化する可能性を示す研究として注目される。実装はGitHubで公開されており（https://github.com/nyu-systems/cleave）、再現検証も容易な状況だ。

出典: Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling