ハクソク

世界を動かす技術を、日本語で。

全てのGPUを購入したので、徹底的にそのGPUを活用します。

239日前原文(hazyresearch.stanford.edu)

概要

  • Llama-70B 向けの高スループット推論用 megakernel をH100 GPUで公開
  • 計算・メモリ・通信処理の重複実行 によるGPUリソースの最大活用
  • Tokasaurusエンジン 統合時、SGLang比で 22%以上のスループット向上
  • 研究用コード として無保証公開、動作環境に敏感
  • 低レイテンシmegakernel の設計思想を高スループット用途に応用

Llama-70B向け低レイテンシ・高スループットMegakernelの公開

  • Llama-70B の推論をH100 GPU上で 高スループット で実行するmegakernelの公開
  • 計算・メモリ・通信 の各オペレーションを積極的に重複させて、GPUの各種リソースを同時活用
  • Tokasaurus inference engine に統合することで、 SGLang より 22%以上高いスループット を実現（ShareGPTベンチマーク65,536プロンプトで測定）
  • コードは 研究目的 での公開であり、 コンパイラバージョンやGPU構成に依存、サポートなし
  • アイデアや手法の参考用としての意図

低レイテンシMegakernelの振り返り

  • 数ヶ月前、 Llama-1B の全forward passを 単一megakernel に融合し、 低レイテンシ推論 を実現
  • 従来の推論エンジン（vLLM, SGLang）は GPU帯域の半分程度しか活用できていなかった
  • 複数カーネル分割 によるオーバーヘッド（セットアップ・ティアダウン）で メモリパイプラインバブル が発生
  • 全forward passを単一megakernelに融合 することで、 オーバーヘッド削減・スループット向上 を達成
  • 命令＋インタプリタモデル によるきめ細かなパイプライン化が鍵

高スループットMegakernelの設計と工夫

  • Llama-70Bの大規模バッチ推論 に特化したmegakernel設計
  • データ並列・テンソル並列 のハイブリッド（sequence parallel）アプローチ
    • 一部処理は データ並列、一部は テンソル並列 で実行
  • 通信コスト隠蔽のための工夫
    • O projectionを データ並列 化し、 分散転置 操作で通信トラフィックを1/8に削減（8GPU時）
    • O projection重複 によるメモリ消費増（GPUごとに9GB増加）、最大バッチサイズ約15%減
  • 命令セットの再設計
    • RMS norm＋all-gather、QKV matmul＋RoPE、Attention＋分散転置、O-proj matmul＋residualなど 複数処理を融合した命令
    • 低レイテンシ用megakernel では主に 行列×ベクトル、高スループットでは 行列×行列 へ最適化
    • 再計算による融合 はスループット重視では非効率、 データ依存性管理 が重要

GPUリソースの重複活用手法

  • SM（Streaming Multiprocessor）内外での重複実行
    • SM内： 命令間パイプライン化 でメモリ転送と計算を重複
    • 複数SM間： 計算集約型・メモリ集約型命令の同時スケジューリング
    • 複数GPU間： 通信コストを専用スレッドで隠蔽 し、他スレッドは次命令へ進行
  • megakernelのインタプリタ型抽象化 は、 低レイテンシ・高スループット両用途で再利用可能

ベンチマークと成果

  • vLLM、SGLang との比較ベンチマークを実施
  • 単一megakernelによる融合 で、 従来手法を大幅に上回るスループット を達成
  • 細粒度なリソース重複 が高スループット化の鍵

注意事項・今後

  • コードは 研究用 であり、 動作保証やサポートなし
  • コンパイラやGPU環境に強く依存、動作の再現性に注意
  • アイデアや設計思想の参考資料 として活用推奨

参考用コード・詳細設計・ベンチマーク結果 はGitHubリポジトリ参照

Hackerたちの意見

図1: ズームmmm 受け入れ！

カーマックがコンソール向けに特化することで得られる効率性について話していたのを思い出すよ。ハードウェアが何かを正確に把握できるのはいいよね。効率性が引き出せるってのも素晴らしい。ただ、本当に難しいのは、異種計算環境でそれを実現するための賢いコンパイラを作ることなんだよね。

└

デモシーンも、どんなハードウェアで動かしているかを完全に把握できれば、どれだけのことができるかの例だよ。問題は、コンソールのようなものでも、通常は「コスト効率」が良いのは、最大限に効果的ではない普通の速く書けるコードを書くことなんだ。コンパイラに魔法をかけさせて、まあそれで十分って感じ。時々、今あるプロセッサに20年間も神秘的に縛られたら、世界がどうなるか夢想することがあるよ。

└

MojoについてのHNのネガティブな意見があるけど、これはModularの最終的な目標なんだ。 https://signalsandthreads.com/why-ml-needs-a-new-programming...

不公平だとは分かってるけど、文体がこのクラシックを思い出させるんだよね: 『境界を越えて: 量子重力の変革的解釈学に向けて』 https://physics.nyu.edu/faculty/sokal/transgress_v2/transgre...

コンパイラのバージョンやGPUの設定、時には見方が違うだけで敏感になってしまうし、サポートするつもりは全くないよ。俺の好きなタイプのコードだね。

└

これも大好きだった、ハハハ。

└

コードを投稿する人がもっと正直になってくれたらいいな。

いい記事だね。インタープリターが気に入った。でも、これらのアイデアはずっと前から主要な研究所で広く実装されていると思うから、2025年にこれが書かれているのはちょっと驚きだよ。これはすべて、理論的な結論に至ることについてであって、難解な魔法じゃない。何十億もGPUに使うなら、CUDAプログラマーの時間にも少しお金を使わない理由がないよね？

└

これらのアイデアはずっと前から主要な研究所で広く実装されていると思うけど、そうじゃないの？ うん、違うよ。そんな主要な研究所に入ったとき、まだ手をつけられていない比較的簡単な課題がたくさんあることに驚いた。でも現実は、やるべき仕事が多すぎて、どれも超重要そうに見えるのに、それをこなす人が足りないってことなんだよね。

ワークロードが実際に（NVIDIA）GPUをフルに活用できない場合は、複数のユーザーで共有できるように分割することが可能だよ：* https://docs.nvidia.com/datacenter/tesla/mig-user-guide/ * https://www.nvidia.com/en-us/technologies/multi-instance-gpu... それとも、1人のユーザーが複数のプロセスで共有することもできる：* https://docs.nvidia.com/deploy/mps/index.html

└

ただし、これはワークステーションやサーバーカードでのみ可能だけど、彼らがラインナップを人工的にセグメント化するために使う手段の一つなんだ。

└

複数のユーザーと共有GPUを使っている場合、情報漏洩のリスクはどのくらい現実的なの？