Ragged Paged Attention：TPU向け高性能・柔軟なLLM推論カーネル

Abstract

大規模言語モデル（LLM）のデプロイは、性能と総所有コスト（TCO）の両方を重視しつつ、GoogleのTensor Processing Units（TPU）のようなコスト効率の高いアクセラレータへと、ますます移行しています。しかし、既存のLLM推論カーネルおよびサービングシステムは依然として主にGPU中心であり、特に現代のサービングで一般的な動的でラグド（ragged）な実行パターンのもとで、LLMワークロードをTPUアーキテクチャへ効率的にマッピングするための確立された手法は存在しません。本論文では、PallasとMosaicを用いて実装した、高性能かつ柔軟なTPU向けアテンションカーネル「Ragged Paged Attention（RPA）」を提案します。RPAは、次の3つの主要な技術によりこれらの課題に対処します。（1）ラグドなメモリ上で効率的な動的スライシングを可能にするきめ細かなティリング、（2）KVキャッシュ更新とアテンション計算を融合するカスタムのソフトウェアパイプライン、（3）デコード、プリフィル、混在ワークロード向けに特殊化したカーネルを生成する、分布を考慮したコンパイル戦略です。TPU7x上でLlama 3 8Bを評価したところ、RPAはデコードで最大86%のメモリ帯域利用率（MBU）、プリフィルで73%のモデルFLOPs利用率（MFU）を達成します。vLLMおよびSGLangにおける主要なTPUバックエンドとして統合されたRPAは、効率的なTPU推論のためのプロダクション品質の基盤を提供し、カーネル設計に関する実用的な知見をもたらします。