AIシステム性能エンジニアリング (Japanese Edition)

第18章 高度な プリフィル・デコードとキーバリューキャッシュのチューニング

この作品はAIを使って翻訳されている。ご意見、ご感想をお待ちしている：[e-mail]

本章では第17章の内容を基に、推論のプリフィルおよびデコード段階における高度な最適化についてさらに深く掘り下げる。高レベルのスケール戦略を基盤とし、低レベルのテクニックについて解説する。具体的には、単一デコードの「メガカーネル」、知能なKVキャッシュの調整とGPU間での共有、プロンプト状態の高速GPU間転送、適応型リソーススケジューリング、プリフィルワーカーとデコードワーカー間の動的ルーティングなどを扱う。

また、新たなレベルの性能と効率性を実現するハードウェア・ソフトウェアの革新技術にも焦点を当てる。これらのテクニックを適用することで、デコード遅延を大幅に削減し、GPUあたりのスループットを向上させ、大規模環境でも厳しい遅延SLOを達成できる。

最適化されたデコードカーネル

これまで我々は、 といった高レベルなシステムおよびクラスタ最適化戦略に焦点を当ててきた。超大規模推論を拡大する際に考慮すべき別の手法セットは、低レベルなカーネルとメモリ管理のチューニングである。特にデコードフェーズにおいて重要だ。

デコードフェーズは分散処理され、しばしばメモリがボトルネックとなる。このため研究者や実務者は、デコードフェーズを可能な限り高速化し、特定のハードウェア向けに最適化しようとしてきた。このスペースにおける顕著な革新として、FlashMLA（DeepSeek）、ThunderMLA（スタンフォード大学）、FlexDecoding（PyTorch）が挙げられる。これらは特に、LLMワークロードで一般的な可変シーケンスシナリオにおけるデコード時のトランスフォーマーのマルチヘッドアテンション効率をターゲットとしている。次にこれらをそれぞれ見ていこう。

FlashMLA (DeepSeek)

Flash Multi-Latent Attention（FlashMLA） は、DeepSeekが導入した最適化されたデコードカーネルである。これは特に単一トークンデコードステップに焦点を当てており、これは本質的に次のトークンを生成するために使用されるトランスフォーマー層のフォワードパスである。FlashMLAは操作を融合し、GPUメモリ階層をより効果的に活用することでデコードを高速化する。

FlashMLA（デコード）は推論において、FlashAttention（プリフィル）がトレーニングにおいて果たす役割に相当する。メモリアクセスオーバーヘッドと遅延を低減する。FlashMLAを用いることで、標準カーネルと比較してデコードフェーズの大幅な遅延削減を実現できる。

FlashMLAは複数のアテンション演算を1つに融合することで算術演算を高める。これにより、複数のヘッドと複数の時間ステップを1回の融合カーネル起動でプロセスできる。これにより、小さなバッチサイズでも演算ユニットを稼働させ続け、デコード中のGPU利用率を向上させる。図18-1は、Hopper H100 GPU上で、グループクエリ注意（GQA）やマルチクエリ注意（MQA）といった他の注意実装と比較したMLAの演算集約度の向上を示す。（注：Blackwellはより高いTFLOPsとHBM帯域幅で両方のルーフラインを押し上げる。）

図18-1. MLAが 演算制約領域に近づく様子（NVIDIA ...