Zachary、Novita AI アルゴリズムエキスパート
研究重点: 推論高速化
はじめに
昨年、H2O が登場して以来、KV Sparsity に関する研究論文が急増しています。しかし、実際のアプリケーションにおける大きな課題の 1 つは、学術研究と実際の実装との間に大きなギャップがあることです。例えば、vLLM のようなフレームワークは、ページングベースのメモリ技術である PagedAttention を採用しており、これはほとんどのスパーシティアルゴリズムと互換性がないか、PagedAttention よりも性能が劣ります。このような問題により、スパーシティアルゴリズムを本番環境で効果的に活用することができません。
過去 1 年間の KV Sparsity 分野における最近の学術論文をレビューし、KV Sparsity に基づいて vLLM フレームワークに修正を加えました。これらの修正は、Continuous Batching、FlashAttention、PagedAttention など、vLLM フレームワークの固有の最適化機能と整合性を保っています。最終的に、修正されたフレームワークを SOTA 推論フレームワークと比較し、本番環境で最も一般的にデプロイされているモデル、パラメータ、ハードウェアを使用して、推論性能を 1.5 倍高速化することに成功しました。
KV Sparsity について議論する前に、LLM における Massive Activations の特性に触れることが不可欠です。LLM では、少数の活性化値が他の活性化値よりもはるかに活性化されており、その差は 10 万倍以上に及ぶこともあります。言い換えれば、ごく一部のトークンが重要な役割を果たします。この観察結果は、重要なトークンのみを保持することで推論性能を向上させることができる、KV Sparsity 手法を適用するための基礎となります。
Llama2 および Llama3 の具体的な分析データは、図 1 および図 2 に示すように、この観察結果をさらに裏付けています。
図 1. LLaMA2-7B における活性化の大きさ（z軸）。x軸とy軸はそれぞれ系列と特徴量の次元です。この特定のモデルでは、2 つの固定された特徴次元（1415, 2533）と、開始トークン、最初のピリオド（.）または改行トークン（\n）の 2 種類のトークンに、巨大な大きさの活性化が現れていることがわかります。
Llama3 は Llama2 とは異なり、すべての層で Massive Activations の特性を示すわけではありません。下位層では活性化はより均一に分布し、中間層では局所的なアテンションパターンが見られます。Massive Activations の現象が顕著になるのは上位層においてのみです。この観察結果は、Llama3 のようなモデルが階層的スパーシティ手法を採用する理由を説明しています。これは、特定の活性化パターンに応じて異なる層にスパーシティを適用するものです。
図 2. LlaMa (Touvron et al., 2023a;b) における層ごとの検索拡張生成のアテンションパターン。下位層では、モデルは広域スペクトルモードのアテンションを示し、アテンションスコアをすべてのコンテンツに均等に分散させます。中間層では、アテンションは各ドキュメント内でより局所化され、情報の集約が洗練されていることを示しています（層 6 と 10 の点線の赤い三角形の形状）。これは上位層で最高潮に達し、「巨大アテンション」がいくつかの主要トークンに集中し（層 18 以降の集中したアテンションバー）、回答に必要な情報を効率的に抽出します。
システムモデル
LLM における Massive Activations の特性を紹介したところで、KV Sparsity の原理について基本的な理解が得られたはずです。次に、KV Sparsity がどのように実装されるかについて説明します。大規模モデルの推論プロセスでは、一般的に GPU メモリ容量、計算能力、I/O 帯域幅という 3 つの主要なボトルネックに直面します。冗長な計算を避けるために、推論中に KV がキャッシュされることが多く、GPU メモリを大量に消費します。一方、計算負荷と I/O スループットは、それぞれ Prefill 段階と Decode 段階で重要なボトルネックとなります。
vLLM フレームワークの修正は、LLM の Massive Activations 特性に基づいており、層ごとのスパーシティ（つまり、層ごとに異なるスパーシティレベル）を実装しています。これにより、KV キャッシュに関連するオーバーヘッドが大幅に削減され、LLM モデルの推論プロセスが高速化されます。
図 3 に示すように、推論プロセス中に、モデルの異なる層に KV Sparsity を適用します。プルーニング戦略を使用して、スコアの低い KV ペアを削除し、スコアが高く、より近い位置にある KV ペアを保持します。これにより、メモリ使用量、計算オーバーヘッド、I/O オーバーヘッドが削減され、最終的に推論の高速化につながります。
図 3. KV Sparsity がモデルに与える影響
推論高速化の結果
ほとんどの研究者が最も関心を持つのは、vLLM フレームワークにおける KV Sparsity の実際の影響です。詳細に入る前に、まず性能評価の結果を示します。
1. 推論性能
実際の LLM アプリケーションでは、入力/出力長 4000/500 が最も一般的な構成です。また、RTX 4090 GPU は業界で最も広く使用されている GPU の 1 つです。これらの条件に基づいて、バッチ性能比較テストを実施しました。対照群は vLLM 0.6.1.p2、実験群は Novita AI Sparse 0.5.1（vLLM 0.5.1 を KV Sparsity で修正したバージョン）です。両者間で複数回の性能比較テストを実施し、主要な指標として TTFT（最初のトークンまでの時間、ユーザーエクスペリエンスに重要）とスループット（実際の推論速度）を含めました。
最終的なテスト結果を表 1 に示します。KV Sparsity を適用することで、vLLM のスループットを約 1.58 倍向上させることができ、同時に実用的な TTFT（P50 が 1 秒未満）を維持しました。さらに、より大きなバッチサイズのシナリオでは、vLLM 0.6.1.p2 は同時実行数 10 で性能限界に達しましたが、Novita AI Sparse 0.5.1 は同時実行数 20 でも安定した TTFT 性能を維持できました。これは、本番環境における性能安定性を確保する上での KV Sparsity の堅牢性を示しています。
表 1：vLLM 0.6.1.p2 と Novita AI Sparse 0.5.1 の性能比較
2. モデル精度
KV Sparsity は非可逆圧縮アルゴリズムであるため、モデルの性能評価も同様に重要です。Llama3–8B に基づいて性能評価を実施し、MMLU などのベンチマーク評価も含めました。結果は、精度の低下は約 3% 以内に収まっていることを示しています。
表 2 に示すように、MMLU と人文科学ベンチマークの両方を使用してモデル性能をテストしました。
表 2：MMLU および人文科学ベンチマークを使用した Llama3–8B の性能評価結果
次に、実際のアプリケーションでよく見られる長文テキストシナリオにおける QA タスクの評価を、入力長 7k から 30k の範囲で実施しました。圧縮率は 10 倍以上を維持し、表 3 に示す通りです。結果は、全体的な精度の低下が約 10% であることを示しました。
表 3：入力長 7k から 30k の長文テキスト QA タスクの評価結果
主要技術
1. 層固有のスパーシティとテンソル並列処理
vLLM フレームワークは、通常 Continuous Batching と呼ばれる Iteration-Level Scheduling 戦略を採用しています。このスケジューリングアプローチは、バッチ内のすべてのシーケンスが生成を完了するのを待たずに、イテレーションレベルでのスケジューリングを容易にし、静的バッチ処理と比較して GPU 使用率が高くなります。本稿の冒頭で述べたように、Llama3–8B のようなモデルは、本質的に階層的スパーシティなどの特徴を持っています。その結果、Continuous Batching に基づく修正は、図 4 に示すように一連の課題に直面しました。
図 4. 連続バッチ処理を使用した 7 つのシーケンスの完了。左側は 1 回のイテレーション後のバッチ、右側は数回のイテレーション後のバッチを示しています。シーケンスがエンドオブシーケンストークンを生成するとすぐに、その位置に新しいシーケンスを導入します（例：シーケンス S5、S6、S7）。この方法は、すべてのシーケンスが終了してから新しいシーケンスを開始する必要がないため、GPU 使用率を向上させます。
私たちが取り組んだ主な課題は、異なる層にわたるメモリ管理でした。vLLM が統一されたページングメモリ管理モデルを利用しており、キュー戦略が Full Cache または Sliding Window のいずれかのみをサポートするという制限があるため、vLLM のコア構造を変更して両方のオプションをサポートできるようにしました。
この調整により、各層は Full Cache または Sliding Window のいずれかを柔軟に選択でき、図に示すように、層ごとに異なるスパーシティレベルが可能になります。
最初に取り組むべき課題は、異なる層にわたるメモリ管理の問題です。vLLM は全層で統一されたページングメモリ管理モデルを使用し、キューイング戦略では Full Cache または Sliding Window のいずれかのみが許可されるため、vLLM の基盤構造を調整する必要があります。具体的には、Full Cache と Sliding Window の両方をサポートできるようにし、異なる層が好みのメモリ管理戦略を独立して選択できるようにします。この調整により、最終的に異なる層間でさまざまな程度のスパーシティが可能になります。対応する構造図を図 5 に示します。
図 5. vLLM における層ごとのメモリ管理構造
全体として、調整は 3 つの異なるフェーズに分類できます。
サービス初期化フェーズ: サービスの起動プロセス中に、主に残りのスペースの計算とともに、Metadata 構造の階層的な初期化を通じて、3 つの層に割り当てられたメモリが初期化されます。対照的に、元の vLLM は完全な Full Cache として動作し、異なる層のメモリ割り当てロジックを考慮していません。
ステップ準備フェーズ: vLLM のステップの初期化プロセス、特にブロックマネージャーとモデルランナーのメソッドを含むフェーズでは、異なる層に異なるブロックテーブル管理戦略を採用して、KV 更新メカニズムを制御します。前の図に示すように、最初の 3 つの層では Full Cache メカニズムを使用してすべての KV を保存し、3 層目以降では Sliding Window 戦略を使用して最もスコアの低い KV を自動的に置き換えます。このアプローチにより、各層のニーズに合わせてメモリスペースを柔軟に割り当てることができ、スペースを節約しながらスパーシティアルゴリズムの効果を最大化できます。
スパーシフィケーション計算フェーズ: このフェーズは通常、モデル実行後に行われ、計算されたスパーシティスコアに基づいてスパーシティ操作が実行されます。このステップで重要な考慮事項は、Tensor Parallelism、CUDA graphs、GQA などのテクノロジーとの互換性です。実際のエンジニアリング実装では、Tensor Parallelism に特に注意が必要です。以下の図に示すように、シングルマシンマルチ GPU シナリオでは、適切な分割メカニズムを使用して KV スパーシフィケーション操作を各 GPU 内に制限し、GPU 間の通信を回避する必要があります。幸いなことに、Attention ユニットは ColumnLinear ベースの分割アプローチを採用しているため、その実装は比較的簡単です。
図 6. vLLM におけるメモリ管理とスパーシフィケーションの 3 段階プロセス
上記の議論は、主に vLLM で階層的スパーシティをサポートするためにフレームワークレベルで行われた変更に焦点を当てています。次に、CUDA レベルで実装された最適化に注目します。
2. Attention の修正
CUDA レベルの変更は、主に Attention 計算ユニット、特に FlashAttention と PagedAttention に対する拡張に焦点を当てています。PagedAttention は比較的単純ですが、FlashAttention に関連する変更に重点を置きます。
まず、Attention の計算式を示します。
式に示すように、softmax の計算プロセスでは QK の積を計算し、これは 2 パスアルゴリズム（2 回のパスが必要）です。しかし、最終的な計算目標 O は、FlashAttention アルゴリズムを使用して 1 パス（1 回のパス）で達成できます。FlashAttention の実装ロジックは、FlashAttention2 論文の図 7 で説明できます。簡単に説明すると、Q をチャンクごとに走査して KV に対してブロック計算を実行し、ループが終了するまで O、P、rowmax データを徐々に更新します。最後に、O を ℓ で割ることで、FlashAttention の 1 パス計算が可能になります。
図 7. FlashAttention2 の実装ロジック
私たちの KV Sparsity の実装では、FlashAttention の計算プロセス中に、softmax に必要な値、すなわち rowmax/ℓ（グローバル最大値と累積値）がすでに計算されていることに注意することが重要です。ただし、FlashAttention はこれら 2 つの値の結果を返しません。
FlashAttention の基本的な理解を基に、PYRAMIDKV などの論文で提示されているスパーシフィケーションのスコアリング式について説明できます。
ここで、[n − α, n] は指示トークンの範囲です。各層 l と各ヘッド h において、最も高いスコアを持つ上位 k l 個のトークンが選択され、それぞれの KV キャッシュが保持されます。他のすべての KV キャッシュは破棄され、生成プロセス全体を通じていかなる後続の計算にも使用されません。
具体的には、各アテンションヘッド h において、i 番目のトークンを KV キャッシュに保持するためのスコア S を次のように計算します。
上の式で、A は softmax スコアを表し、S は最も近い a 個の softmax 値のスコアの合計を表します。示されているように、softmax 自体は 2 パスアルゴリズムであり、FlashAttention に直接統合することはできません。これは、スパーシフィケーションのスコアリングを 1 パスで実装できないことも意味します。
前述したように、FlashAttention は softmax スコアを出力しませんが、rowmax/ℓ のような中間値を計算します。したがって、Online Softmax: 2 パスアルゴリズムを参照することで、rowmax/ℓ の出力を利用して、KV Sparsity の特定のスコアリング S を実現できます。実際には、図 7 に示す FlashAttention に改良を加えるだけで、この要件を満たすことができます。
図 7. KV Sparsity スコアリングのための FlashAttention の改良
- FlashAttention における rowmax/ℓ の出力: FlashAttention には rowmax/ℓ の出力が組み込まれています。これらの値は最大値と累積合計を表すため、データサイズは最小限であり、FlashAttention の全体的なパフォーマンスに与える影響はごくわずかです。rowmax/ℓ を再利用することで、2 回目の 1 パスを回避し、全体的な計算効率を向上させることができます。
- Online Softmax ステップの追加: Online Softmax ステップを導入します。これは FlashAttention に基づいて、式 p=e((kq)−row_max)/ℓ (1 パスロジック) を使用して softmax スコアを計算します。その後、行方向の合計を実行することで、スパーシフィケーションスコア S を計算できます。
Online Softmax の実装では、2 つの重要な詳細に特に注意する必要があります。
- **ループ順序の調整 **: S は k 次元に基づくスコアの合計であるため、Online Softmax での計算をアテンションと比較して q*k から *k**q に調整する必要があります。この調整により、列の合計を行の合計に変換することで、異なるブロック間の不要な通信を回避し、すべての計算をレジスタ内で完了できます。
- **q の開始位置と終了位置 **: 因果符号化と最も近い a の合計ロジックが存在するため、すべての q 値を k で走査する必要はありません。図 7 で 0/1 とラベル付けされたマスクインジケータを参照することで、q の開始位置と終了位置に基づいて、対応するブロックが計算を必要とするかどうかを判断でき、計算リソースを大幅に節約できます。たとえば、入力プロンプトの長さが 5000、最も近い a の長さが 256 の場合、従来の計算では 5000*5000 の演算が必要になります。しかし、開始位置と終了位置を記録することで、*5000**256 の演算のみが必要となり、計算効率が向上します。
PagedAttention は FlashAttention と同様の処理を実行し、共通の目標は最終的なスコア S を計算することです。ここでは PagedAttention の詳細な説明は省略します。
結論
vLLM 0.5.1 をベースにした修正バージョンである Novita AI Sparse 0.5.1 は、現在主に Llama3–8B や Llama3–70B などのモデルをサポートしています。CUDA graphs を使用した推論モードで動作し、主に RTX 4090 などのコンシューマー向け GPU にデプロイされています。エンジニアリング面とアルゴリズム面の両方において、依然として最適化の余地は大きく残っています。
注意: 以下の機能は vLLM-0.6.2 ではまだサポートされていません。
- Chunked-prefill
- Prefix caching
- FlashInfer およびその他の FlashAttention 以外のバックエンド
Novita AI Sparse 0.5.1 にご興味のある方は、Novita AI にアクセスして、他の方よりも先に体験してください。
参考文献：
[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference
[3] SnapKV: LLM Knows What You are Looking for Before Generation
[4] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
[5] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference
[6] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
[7] Layer-Condensed KV Cache for Efficient Inference of Large Language Models
[8] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
[9] CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion
[10] KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation
