【要約】Multi-Head AttentionとKV Cacheの容量 [Qiita_Trend] | Summary by TechDistill
> Source: Qiita_Trend
Execute Primary Source
// Problem
LLMの推論エンジニアは、コンテキスト長やバッチサイズの増大に伴うGPUメモリの枯渇という課題に直面する。特に、Attention機構における過去のトークン情報の保持（KV Cache）が、以下の問題を引き起こす。
- ・KV Cacheが膨大なGPUメモリを占有し、扱えるコンテキスト長を制限する。
- ・メモリ帯域の消費が激しくなり、推論のレイテンシが悪化する。
- ・バッチサイズを増やした際に、メモリ不足によるシステム停止を招くリスクがある。
// Approach
筆者は、KV Cacheの容量を決定する変数に基づいた計算式を定義し、具体的なモデルを用いた試算を行った。これにより、メモリ消費の構造を可視化している。
- ・容量計算式 $Size = 2 \times B \times L \times T \times H \times D \times S$ を提示した。
- ・Llama 3 8Bを例に、32層で約2GiBの容量が必要であることを算出した。
- ・MQA（Multi-Query Attention）を用い、複数のQuery Headに対してKV Headを1つに共有する手法を解説した。
// Result
MQAの導入により、推論時のメモリ使用量と帯域幅の負荷を劇的に軽減できることが示された。
- ・KV Cacheの容量を、Query Head数に応じて大幅に削減できる。
- ・例として、Queryが32個に対しKVが1個の場合、容量は $1/32$ になる。
- ・実用的な発展形として、GQA（Grouped-Query Attention）が広く採用されている現状に触れている。
Senior Engineer Insight
> 大規模推論基盤の設計において、KV Cacheの管理はVRAM容量とスループットを左右する最重要事項である。コンテキスト長を拡張する際は、メモリ消費が線形に増加することを前提としたリソース設計が不可欠だ。また、MQAやGQAの採用有無は、単なる最適化ではなく、デプロイ可能なバッチサイズや最大コンテキスト長を決定する境界条件となる。モデル選定時には、HiddenSizeだけでなく、KV Headの構成を精査し、インフラコストへの影響を予測すべきである。