昨日、私たちの『推論エンジニアリング・マスタークラス』ポッドキャストで、Megakernel をめぐる激しい議論がありました。
Megakernel は死んだ
Megakernel はなぜ有用なのか？ 2か月かけて 1 つの Kernel を書くのは、起動オーバーヘッドを削減し、Kernel 間のオーバーラップ実行がうまくいかない問題を解決するためにすぎない。
以前は PDL があったが、後になって、それも完璧ではないと言われるようになった。末尾に残る CTA（straggler CTA）の影響により、依然としてわずかながら改善の余地があるからだ。だから——あ、すみません、忘れていた。Rubin がこの問題を解決したのだった。（Kernel 2 には 10 個の CTA が必要だが、Kernel 1 はすでに 7 個を完了しており、3 個が末尾で足止めされている。このとき Kernel 2 は、まずその 7 個の CTA を起動する。）
タイムラインが十分に長ければ、最終的にはすべてが均衡する。真剣に推論サービスを提供している企業で、6.7 万行のコードによる手作業で融合された前向き伝播 Kernel を本番環境で使うところなどない。そうしたことをしているチームは、純粋に研究をしているだけだ。
死んだ。
2週間前、私は @swyx のポッドキャストに出演し、いくつか……言うべきではなかったことを話しました。
それ以来、いろいろなことが起こりました。皆さんに謝らなければなりません。
申し訳ありません。私が言ったことは、すべて正しいと証明されてしまいました。
a）「Megakernel は死んだ」について
Megakernel はなぜ有用なのか？ 2か月かけて……
最後まで聞いてくださる方のために、以下に議論の全容を掲載します。
Ali： 融合 Kernel では問題を解決できません。たとえばテンソル並列（tensor parallelism）では、行列の半分が一方の GPU にあり、残りの半分が別の GPU にあります。次のステップで行列全体を必要とする非線形演算を実行する場合——たとえばアテンション計算で softmax を実行するときや、指数演算を行うとき——データの 1 行全体を持っていなければなりません。したがって、次の段階で softmax を実行するには、GPU 2 側の部分的な結果と GPU 1 側の部分的な結果を把握する必要があります。
つまり、融合 Kernel を使ったとしても、それらの間で通信を行わなければなりません。各部分の内部に非線形演算が存在するからです。Megakernel については、正直なところ、私はあまり期待していません。研究の方向性としては面白いですし、直感的にも理論的にも、とても美しく見えます。起動する Kernel を 1 つだけにすることで、大幅な起動オーバーヘッドを削減できます——データを絶えず融合・移動させ、すべてを 1 つにまとめるわけです。
しかし、Kernel 自体が複雑になるため、高度に最適化された Megakernel を書くのは非常に困難です。本当に、とても難しい。特定の企業を指しているわけではありませんが、融合 Megakernel を開発した企業であっても、あるいは私が接してきた、そうした企業で働く人たちであっても、最終的には本番環境でそれらの Kernel を実行しないことが多いのです。TensorRT-LLM とモジュール型 Kernel のほうが起動が速いからです。各コンポーネントを個別に最適化できますし、それらを並列に実行することもできます。
NVIDIA の技術責任者の 1 人が、「Rubin のベールを脱がせ、その仕様を紹介する」という趣旨の投稿をしていました。3 番目の投稿に関連する内容が示されています。具体的な技術の詳細については、まだあまり話したくありませんし、もう一度注意深く読み直す必要もあります。しかし、この GPU の設計によって、Megakernel は意味を失うことになります。したがって、この研究分野全体が今後発展しなくなるように見えます。
彼が引用していたのは、（この番組の友人でもある！）Kyle Kranen による依存関係トリガー（dependency trigger）の紹介です。これは、これまでパイプラインをブロックし、Kernel の融合を必要にしていた要因の 1 つでした。
Kernel のオーバーラップ実行を改善： Rubin は、タイル単位の依存関係トリガー（tile-level dependency trigger）を含む、より粒度の細かい Kernel の協調をサポートします。つまり、ある処理に必要なデータの一部が利用可能になった時点で、その部分を処理する Kernel を直ちに起動できるということです！
番組内でも話したとおり、現時点では、まだ解決されていない物理的な制約もいくつかあります。しかし、Kernel 分野で進められているさまざまな極限的最適化により適応できるよう、NVIDIA が Rubin の設計を更新したことは、まったく理にかなっています。
Ben Spector の Megakernel の共同研究者である Stuart Sul は、現在 Mixture of Kittens のチームを率いています。Mixture of Kittens は、Cursor が本日公開したオープンソースの Megakernel で、混合エキスパート（Mixture of Experts、MoE）の学習に使用されます。
この名前は、Ben の個性が色濃く表れたプロジェクト ThunderKittensにちなんだものです。このプロジェクトは Dan Fu の研究チームによるものです。
Mixture-of-Kittens（MoK）をオープンソースとして公開します。これは、NVL72 向けの MoE 学習 Megakernel です。
混合エキスパートにおける通信と計算のすべてを、完全に決定論的な 1 つの Kernel に融合しており、速度は最も強力な公開ベースラインの最大 2.37 倍に達します。