Stanford CS336 Language Modeling from Scratch: Inference Engines and Full-Stack Innovation

Stanford CS336 Language Modeling from Scratch: Inference Engines and Full-Stack Innovation

Inference: The Engine Turning Electricity into Intelligence

Inference は、訓練された言語モデルを静的な数学的オブジェクト（操作の有向非循環グラフ）から実用的なツールへと変換する重要な段階です。基盤となる推論エンジンと GPU カーネルを理解することは「フルスタックイノベーション」に不可欠であり、これらのコンポーネントが実運用時の効率、レイテンシ、機能を決定します。

The Lifetime of a Token

リクエストが推論システムに送られると、以下のパイプラインをたどります:

  1. Scheduling: リクエストは特定の GPU にルーティングされ、場合によっては prefill と decode の処理を別々のマシンに分散させます。
  2. KV Cache Lookup: システムは同一リクエストまたはそのバージョンが過去に処理されたかを確認し、計算を節約します。
  3. Execution: コアとなる機械学習コードが実行され、モデルサイズに応じてノードや GPU 間で並列化されます。
  4. Token Generation: システムはトークンを出力し、ストップシーケンスや安全性チェックのためにさらに処理されます。

Prefill vs. Decode Workloads

推論は、計算特性が根本的に異なる 2 つのフェーズから構成されます:

  • Prefill: 初期入力プロンプト（例: 10,000 トークン）を処理し、初期活性化を計算します。これは 計算量に依存 し、逆伝搬がない点で学習に似ており、高い FLOPs を使用します。
  • Decode: トークンを 1 つずつ生成します。これは メモリ帯域幅に依存 しており、単一トークン生成のためにモデル全体をメモリからロードする必要がありますが、必要な FLOPs は比較的少ないです。

これらの違いにより、最新のスタックはしばしば prefill と decode を別々のワーカーや専用ハードウェア（例: prefill に NVIDIA GPU、decode に LPU/Grok チップ）に分散させます。

Production Challenges and System Optimizations

1 兆トークン以上を日々提供する規模になると、小規模では顕在化しない微細なバグやボトルネックが露呈します。たとえば、カーネルエラーによりモデルがトークンを無限に繰り返す「doom loop」や、未初期化 GPU メモリを読み込むカーネルのオフバイワンエラーで言語が英語から中国語に急変するケースがあります。

KV Cache Management

スループット最大化のため、システムはキー・バリュー（KV）キャッシュを効率的に管理しなければなりません。GPU メモリは限られているため、階層型ストレージが採用されます:

  • GPU Memory: 最速だが容量最小。
  • CPU DRAM: やや遅いが、頻度の低い活性化をオフロードするのに使用。
  • SSD/Disk: 最も遅く、セッションデータの長期保存に利用。

これは従来の OS のメモリ管理に似ており、しばしば Least Recently Used（LRU）ヒューリスティックで、どの活性化をスローストレージへ追い出すかを決定します。

Cache-Aware Disaggregation

シンプルながら効果的な最適化として、キャッシュヒット率に基づいてリクエストを振り分けます。"fresh"（キャッシュヒット率低く prefill コスト高）リクエストをある GPU 群へ、"warm"（キャッシュヒット率高く prefill コスト低）リクエストを別の GPU 群へ送ることで、サービス速度が最大 40% 向上します。

Megakernels: Achieving "Speed of Light" Inference

従来の推論エンジンは操作を逐次実行します（例: Norm カーネル → MatMul カーネル）。この方式はカーネル起動オーバーヘッドや "tail effect"（バッチ内最長シーケンスが終わるまで GPU が待機）により大きなダウンタイムを招きます。

The Megakernel Approach

Megakernel は複数の操作を単一カーネルに融合し、GPU をシーケンシャル演算子ではなく大規模分散システムとして扱います。これにより以下のようなオーバーラップが可能になります:

  • KV キャッシュのロードと同時に QKV 投影や RoPE スケーリングを実行。
  • アテンション演算が完全に終わる前に O 投影用の重みをロード。

低レベル CUDA 制御を提供する ThunderKittens ライブラリを用いることで、Megakernel は H100 GPU 上で帯域利用率 72% に達し、ほぼ "光速" の推論性能を実現します。

Parcae: Stabilized Recurrent Architectures

現在の多くの LLM はパラメータとデータを増やすことでスケールしますが、Parcae の研究は 再帰（トランスフォーマーのブロックをループさせる）によるスケーリングを探ります。

The Stability Problem

単純なループトランスフォーマーは非常に不安定で、学習率などハイパーパラメータを僅かに変えるだけで損失が急上昇し NaN が発生します。この不安定性は変換行列のスペクトル半径に起因し、行列ノルムが 1 を超えるとループを通じて活性化が指数的に爆発します。

The Parcae Solution

Parcae は再帰システムの A 行列と B 行列を再パラメータ化して安定化します:

  • Matrix A: 負の対角行列として制約し、システムが安定し活性化が爆発せずに徐々に減衰するようにします。
  • Matrix B: シンプルな線形ノルムで制御。

この数学的制約により損失曲線が安定し、従来のトランスフォーマーよりもパラメータあたりの品質が向上します。

Scaling Laws for Recurrence

初期の調査結果は、再帰はデータ量とパラメータ数に合わせてスケールすべきことを示唆しています。具体的には、学習データが増えるほど再帰回数も増やすことで計算効率が最適化されます。したがって、固定パラメータ予算の下で、ループ付き事前学習を行うことでより高品質なモデルが得られる可能性があります。

Summary

ゲスト講師の Dan Fu は、LLM を数学的オブジェクトから実用的なインテリジェンスへと変換する上で、推論エンジンと GPU カーネルの重要な役割を論じ、Megakernel や Parcae 再帰アーキテクチャといった最適化手法を紹介しました。

Sources