このフォワードパスを可能な限り速くする

このフォワードパスを可能な限り速くする

Flapping Airplanesは最近、効率的なフォワードパスとバックワードパスを実行できるメガカーネルフレームワークを通じて（なぜか）Doomをプレイしながら、10億ドルを調達しました1。 elsewhereでは、Hyperstitionが「ハイパーカーネル」と呼ぶものを開発しており、これはあらゆる種類のモデルに対して高性能カーネルを生成しています。モデルをより高速なものに変換することは、まさにMLコンパイラが伝統的に行ってきたことであり、このすべての活動は、もはやMLコンパイラが必要なのか、それともエージェントに直接最適化を任せることができるのかという疑問を投げかけます。

TVMに取り組んでいる人々も、同様の疑問を抱いています。

エージェントが低レベルの詳細を自分で処理できるなら、まだそれらの間にコンパイラを置くべきか？
私たちは、その答えは「はい」だが、異なる理由であると信じている。エージェントにカーネル開発の各イテレーションでIRとコンパイラインフラストラクチャを再発明させる代わりに、最小限で安定したMLコンパイラ基盤は、エージェントが独自の最適化戦略を探求することを自由にしながら、コミュニティ全体に共有抽象化と再利用可能なツールを提供できる。

彼らの議論の核心は、カーネル最適化は検索とイテレーションのプロセスであるということです。トークンとツール呼び出しにはコストと時間がかかります。そのため、その検索を最適化し、局所最適解に陥るのを避け、生成されたソリューションが正しいことを検証する方法を提供したいと考えています。彼らはメモリアクセスや同期などのチェックを作成し、LeanでIRを形式化して安価な検証を可能にしました。

この未来では、実際にコンパイルするコンパイラは不要になるかもしれないが、プログラムを理解し検証するためのコンパイラインフラストラクチャは依然として必要になるだろう。共有IRは、エージェントによって構築された分析、検証、および新しいツールの基盤として機能できる。

Togetherの新しいTogether Compiler（TKCC）は、最適化された推論カーネルを構築します。しかし、彼らは中間表現に関する結論にかなり異なる見解を示しています。

従来のIRは表現力が限られている。私たちは、IRがエージェントが複数のコードパスとカーネル起動メカニズムを含む複雑な推論コードベースを理解するのに役立つと仮説を立てた。ソースコードを直接提供すると、エージェントがエンジンファイルを読み取るのに約20%の時間を費やして、関連するコードパスとエンジンコード内のカーネル呼び出しを理解していることが観察される。しかし、無制限のソースコードアクセスがより高いパフォーマンスの天井を解き放つことも発見した。私たちの高度に強化されたInductor IRでさえ、特定の最適化を許可するには表現力が不十分だった。

しかし、彼らは依然として検証のニーズと検索空間の削減に焦点を当てています。彼らが構築したものの多くは、フォワードパスの変動を正規化したカーネルを実行するためのさまざまなティアを備え、高速イテレーションのためにホットスワッピングを提供した経験的測定フレームワークでした。また、並列処理のためのカスタムコストモデルを開発することで検索を削減し、テスト前に分析的にオプションを絞り込んでいます。

私たちは、分析推定とティア測定（理論的なハードウェア機能から測定された計算、メモリ、通信パフォーマンスまで）を使用して、トップkの並列処理構成をランク付けするコストモデルを提案する。これにより、TKCCは高価なエージェント最適化に投資する前に有望な構成を特定できる。

したがって、コンパイラが消えていくというよりも、それらが分解されていくということです。検査と調査がより安価で、はるかに自動化された世界では、多くの従来のコンパイラ技術は有用性が低下します。

今年のTritonカンファレンスで、OpenAIのNeil Dharは（私自身も、以前はMetaのTritonチームに所属していました！）Codexを使用して大規模なコードモッドを行い、Tritonコンパイラからヒューリスティックと暗黙のセマンティクスを削除することを目指していた方法について話しました。そのような技術は、幅広いコードで改善を得るのに優れていますが、コードを予測不可能にする可能性があります――小さな変更が結果に大きな影響を与える可能性があります。なぜなら、それは異なるコンパイラパスをトリガーするからです。

Neilが使用した例の1つは、自動FMA（Fused Multiply-Add）の削除でした。Tritonは、乗算とそれに続く加算のパターンをマッチさせ、それを単一のFMAに置き換えました。しかし！それはFMAが中間丸めをドロップするため、カーネルの数値を変える可能性があります。彼らはこのパスを削除し、それを使用していたすべてのカーネルのソースコードに手動でFMA呼び出しを挿入しました3。これにより、検索ランドスケープがよりスムーズになり、ナビゲートしやすくなります（ただし、大きくなる可能性があります）。

TIRxのTVMは、ゼロから始めることで、より検索・検証フレンドリーなコンパイラスタックを構築できるという賭けです――彼らは意図的にヒューリスティックベースの最適化を適用しませんが、エージェントがその上に新しい機能をイテレーションできるように多くの拡張ポイントを提供します。OpenAIのGluonとMetaのTLXは、どちらも方向性としては同様の賭けです――既存のコンパイラのコアを強化しながら、低レベルの拡張ポイントを追加します。Togetherは、適切なツールを提供することが、検索と検証をガイドするための単一のツールチェーンに賭けることよりも重要であると示唆しています。

コンパイラはユーザーに物事を速くさせるために使用されますが、今では検索を速くさせる必要があります。これはまったく同じことではありません。しかし、勝者を宣言することには躊躇します。この分野はエージェント自体と共に共進化します。Togetherが、彼らのエージェントを助ける上でまったく機能しなかったことのいくつかに言及した特に明確な例では、それは約6ヶ月前のエージェントのベストプラクティスのリストとほとんど同じように読めました4。

1. もちろん、すべてのアイデアがプロンプトで利用可能になったため、誰かがすぐにオープンソース版を公開しました。自由にDoomをGB300で実行してください。
2. そして、完全な開示ですが、以前はMetaのTritonチームに所属していました。現在、私はそこで働いています！
3. つまり、まあ。OpenAIは、この動作に依存していたすべてのカーネルを自動的に修正したかもしれませんが、他のすべての人のカーネルを修正したわけではありません！Neilが変更が既にアップストリームにマージされたことを明らかにしたとき、いくつかの急いでコード検索が行われていました。
4. ナレッジハブを提供すると、同じ最適化時間で25〜31%低いスピードアップにつながりました。専門家が作成した戦略でシードすると、文字通り「このフォワードパスを可能な限り速くする」とプロンプトするよりも劣りました。これは素晴らしいアブレーションです。研究とロギングのために人間が設計したツールを提供すると、スループットが低下しました：エージェントは、最適化を進めるのではなく、システムを理解するためにツールの呼び出しの最大43%を費やしました。