Jalapeño MLA カーネル と AI によるコンパイラ最適化の新時代
HotChips での Jalapeño MLA カーネル の発表と、それに続く SemiAnalysis の解説を受けて、多くの議論が交わされています。OpenAI のハードウェアチームとして、私たちはこの小さな金塊にほんの少しだけ触れました。それは、AI が私たちのカーネルを記述しているという事実、そして AI がそうする場合、私たちはカーネルが何を行っているかを 1 行 1 行理解する必要が 本当は ないということです。私たちは明らかに次の点を説明しませんでした：どうしてそんなことが可能なのか？従来のコード生成方法と比較して、これをどのように考えるべきなのか？最適化されたカーネルは、最適化されていないカーネルと同様に健全なのか？
私の経歴について言えば、私は 10 年以上にわたってアクセラレータ向けコンパイラに取り組んできました。私は XLA を立ち上げました。これは優れたコンパイラ基盤であり、素晴らしいクロスカンパニーチームと取り組みが進められています。過去 2 年以上 OpenAI で、私は AI の時代においてコンパイラがどのように機能すべきかを再概念化しようとしてきました。新しいコンパイラの定式化は既存の強みを活かすでしょうが、ツールキットに活用すべき強力な新しいツールがあることを否定することはできません。
これは少し長い旅になりますが、AI が コンピュータプログラムの改善、すなわち最適化コンパイルの自動化にどのように使用されているかを明らかにしたいと思います。私は、AI によって、私たちが「コンパイラ 2.0」と呼ぶものを経験するかもしれないと考えています。AI は提案できる内容において根本的に制約が少なく、その提案はモデルのトレーニングとコンテキストの結果です。このことから、私はこれを「確率的オプティマイザ」と分類しています。これは課題を提起する可能性がありますが、後述するように、大きな強みの源でもあります…
多くの学術研究と産業応用はすでにこの方向に向かっており、最適化コンパイラの領域における AI の関与の可能性を急速に明らかにしていますが、私たちは広範な説明が必要な時点にいます。
背景
コンパイラはプログラムを受け取り、それらのプログラムの翻訳版または改善版を出力します。
プログラムは、入力側と出力側の両方で、プログラムが何を意味するのか、何ができるのか、そしてそれらの可能な動作についてどのように推論するのかを教えてくれる セマンティクス を持っています。
コンパイラに取り組む私たちは、コンパイラを純粋関数のように考えます。つまり、データ構造を受け取り、対応するセマンティクスを持つべきデータ構造を出力します。
時には、私たちのコンパイラは「低レベル化」または「翻訳」に焦点を当てます。例えば、C を受け取り x86-64 アセンブリを出力することがあり、これはしばしば「より低レベル」と見なされます。しかし、多くの場合、そのプロセスの一部として単なる翻訳以上のことを行っています…
私たちのコンパイラは、実際には、「最適化」に焦点を当てています。プログラムを表すデータ構造（私たちの用語では「中間表現」（IR））を受け取り、そのプログラムのより良いバージョンを生成しようとします。「より良い」とは、実行にかかるサイクルが少ないこと、不要なコードが少ないこと、プログラムについて「真でなければならない」と証明できることに特化すること（部分評価）を意味する場合があります。
さて、簡単に考えてみてください。LLM はもともと、人間のテキストをある言語から別の言語に翻訳するために作られました。明らかに翻訳は彼らの得意分野です。そして、日常業務での LLM の使用を通じて、LLM が新しいソリューションを記述したり、既存のソリューションを改善したりできることもわかります。私たちプログラマーの多くは、LLM に「このコードスニペットを最適化して」と依頼した経験があり、驚くべきことにそれができるのです。（ただし、コードが正しく最適化されたことを確認する必要がありますが、これについては後述します！）これは単に、LLM が最適化コンパイラに求められる能力を持っていることを強調するためです。
最適化と最適性
最適化コンパイラは、言うまでもなく、何らかの目的（通常は実行時間）によって、作業中のプログラムの 最適性を高める ことを試みます。任意のプログラムに対して一般的なケースでこれを行うことは 非常に 困難であり、コンパイラエンジニアのための完全雇用定理と呼ばれる定理があります。（私はコンパイラエンジニアになることを選んだ後にこれを知りましたが、それでも安心しました！）
「スーパーオプティマイザ」は、最適化コンパイラの素晴らしい小さなサブフィールドです。与えられたプログラムがあり、セマンティクスを通じてそれが何を行うかを言えると想像してください。同じセマンティクスを持つ 最も 最適なプログラムは何でしょうか？それがスーパーオプティマイザが取り組もうとしていることであり、実質的に探索問題です…
同じセマンティクスを持つ最短のプログラムを見つけようとしていて、候補プログラムが同じセマンティクスを持つかどうかを尋ねる方法があると想像してください。仮説的には、すべてのプログラムを目的順に列挙し、同じセマンティクスを持つ最小のものを選ぶことができます。
しかし、すべてのプログラムを目的順に列挙することは、非常に扱いにくいように思えます。2013 年に作成された私のお気に入りの学術論文の 1 つである「STOKE」（確率的スーパー最適化）は、次のように問いかけました：「では、プログラムを何度もランダムに 微調整 すれば、最終的に最良のプログラムを観測できるだろうか？」彼らは、ランダムウォーク（そして私たちの OG 機械学習の友であるマルコフ連鎖モンテカルロ / メトロポリス・ヘイスティングス）を介して、最終的にその最適なプログラムを見ることができると提案しました。
モンテカルロ微調整は通常は愚かですが（ランダムに微調整を選びます）、高速でもあります。LLM は非常に賢いですが（多くの推論トークンを使用します）、比較的遅いです。
愚かで高速なモンテカルロ微調整の代わりに、LLM にプログラムをどの方向に進めるべきかを考えさせたらどうでしょうか？プログラムを最適化されたプログラム空間を通じて導く、非常に インテリジェントな確率的オプティマイザが得られるでしょう。
最適化の直感
一歩下がって考えてみましょう。「コードスニペットを徹底的に最適化する」という言葉を最も体現している人を思い浮かべてください。簡単に「最適化するオリ」と呼びましょう。オリはおそらく、どのようなコード微調整が実を結ぶかについての直感を持っています。オリはおそらくいくつかのことを試して、うまくいかなければロールバックして別のことを試します。しかし、どのような種類のことが可能で、どのようにコンパイラを打ち負かすことができるかについての直感を持っています。
オリが持つこれらの直感は、多くの場合、コンパイラが行うことを超えています。現代の最適化コンパイラはその結果において非常に印象的ですが、かなり単純なルールとヒューリスティックに基づいています。技術用語で言えば、固定点まで実行されるローカルデータフロー変換のアイデアに基づいています。また、考慮事項をフェーズ順序付けします。つまり、A を考慮し 次に B を考慮するようにコンパイラパイプラインを構築しますが、複合的な AB 問題は考慮しません。スケジューラとレジスタアロケータはその悪名高い例であり、複合スケジューラ・レジスタアロケータ（フェーズ順序付けの利点を得るため）に関する多くの博士論文が試みられてきましたが、実際に機能させることは困難でした。
これが、オリの専門知識が価値がある理由です。多くの場合、オリはいくつかの NP 完全問題を、状況に合わせたヒューリスティックでバランスを取る方法を知っています。そのため、よりカスタマイズされたコンテキスト認識と感度があります。オリはまた、最適化コンパイラが有益に適用できない可能性のあるテクニック、特に組み合わせて使用するテクニック（アウトライン化、カスタム ABI の作成、ベクトル化を可能にする変換、その他「コンパイラにこれを行う方法があればいいのに」と私たちがつぶやくようなこと）を採用することができます。
ここで、AI が、それが持つ何らかの推論機能を通じて、ミニオリとして機能できる可能性を考えてみましょう。何が実を結ぶかについての一致した直感は持っていないかもしれませんが、何が有益であるかの兆候を持っており、非常に多くの試行を行うことができます。
このアプローチでは、STOKE の論文とは異なり、時間が無限に経過するにつれて 最適な プログラムを見ることができることを 保証 できませんが、AI が「より人間らしい」推論機能を持っているため、単位時間あたりにかなりの人間らしい牽引力を得ることができます。
関連付け：MLA カーネル
まず最初に言っておきます：私は AI が Jalapeño MLA カーネルに対して出力した低レベルコードを 知りません が、MLA の numpy を入力する方法は 知っています。
私が以前取り組んでいた XLA コンパイラでは、これらの numpy 操作を塊に融合し、次に「エミッタ」と呼ばれるメタプログラムを使用して、ループ、命令、および低レベルのプリミティブに低レベル化していました。
XLA コンパイラ / エミッタプログラムがそれを行ったとき、私はどのようなアセンブリが出力されるかを気にする必要はありませんでした。私たちの確率的オプティマイザでは、AI は概念的にエミッタメタプログラムの役割を果たします。つまり、低レベル化と最適化の両方を行い、ルーフラインに向けてさらに最適化するように依頼することができます。
これにより、AI がどこに位置づけられ、既存の最適化コンパイラシステムのコンポーネントとどのように類似しているかが明確になることを願っています。また、「アセンブリコード」と見なすレイヤーが現在 上方に移動している と考えることも役立ちます。通常の C++ を入力し、-O3（通常の最高最適化レベル）でコンパイルする場合、入力した C++ を理解していても、出力されるアセンブリを理解するとは期待しません。私たちはここで類似のことを行っていますが、より高レベルでより数学的な入力仕様を使用しています。
さて、重要な疑問は、AI から得られるプログラムが高レベルの説明 / numpy と実際に同等であることを どのように 確認するかです。その確認メカニズムが、AI の確率的最適化プロセスの健全性を確立します。将来のブログ投稿でこれについてより詳細に説明されることを期待していますが、今のところ、セマンティック等価性のテストが可能であり、私たちはそれを実行していると言えば十分です。アクセラレータプログラムは、検証可能な強力で完全な契約に特に適しており、「AI 最適化プログラムが正確に行うこと」を確認できます。なぜなら、それらは広いコンテキストにおいて非常に数学的でデータフロー指向だからです。
この分野の多くの関連技術は、プログラム合成のサブフィールドでの取り組みによって開拓されたことに注意してください。最適化コンパイラが「ここにセマンティクスを持つプログラムがあります。同等のセマンティクスでより良くしてください！」と言うのに対し、プログラム合成は「これらのセマンティクスを持つプログラムが存在すると信じています。見つけられる最良のものを見つけてください」と言います。プログラム合成は最適化コンパイルよりも難しい問題ですが、根本的な制約も少ないです。これは実質的に、オリのような人間が最適化コンパイラよりも優れた結果を出すときに行っていることであり、AI が現在私たちの自動化を支援できることです。AI は元のプログラムから「インスピレーション」を得ることができますが、それにマイナーなローカル変換を実行する必要はありません。古典的な最適化コンパイラは「ああ、バブルソートを書いたね」とは見なさず、契約を理解してクイックソートに切り替えることはありませんが、オリと AI の両方はそれができます。これにより、私たちは確率的最適化を用いたプログラム合成の領域に、古典的な最適化コンパイラの領域よりも近づいています。
これらすべてが、「numpy からそれほど遠くない」ものから始めて、48 時間待つと、同じセマンティクスを持つ最適化されたカーネルが得られるという事実に集約されます。これは HotChips のトークで示した通りです：
スライドにも記載されているように、私たちのマシンでは、かなりうまくチューニングされていると思っていたカーネルでも、AI が人間の専門家を超えてパフォーマンスを向上させるのを観察できることがよくあります。探索する必要がある組み合わせ / 順列の多様性のために、達成可能なかなりのパーセンテージがまだ残っていることがよくあります。これらは、人間のパフォーマンスエンジニアにとってはしばしば扱いにくいほど退屈です。
まとめと結論
コンパイラは、結局のところ、単なる関数です。プログラムをその関数に渡すと、プログラムのより良いバージョンが返ってきます。出力されるプログラムと入力するプログラムは、同じセマンティクスを持つことが期待されます。
従来の最適化コンパイラは、データフロールールとヒューリスティックを介してプログラムを改善します。これらはその起源において完全に理解可能ですが、実行できる操作においてより制限される可能性があります。
対照的に、AI は 確率的 オプティマイザとして、「深く考え」て何かを出力するだけです。その操作は根本的に制限されておらず、人間の専門家オプティマイザにより類似しています。AI が出力するプログラムが健全で、入力したのと同じセマンティクスを実装していることを確認する方法が必要であり、私たちはそれを整えています。そして、この種の AI 最適化は、非常に強力な契約を持つ数学的操作に特に適しています。契約により、カーネルが何を行っているかを 1 行 1 行理解する必要がなくなります。
それが、AI 生成の MLA カーネルを手に入れた方法です！
