LLMのトークンは同じコストではない：KV Cacheの電力をSIMPLEがモデル化

生成AIの料金は通常トークン単位で表示される。しかし、すべてのトークンが同じコストで処理されるわけではない。生成中の各トークンはKV Cacheに追加され、後続のデコードでは増え続ける履歴を参照する。そのため、トークン自体が同じように見えても、シーケンスの後半ほど計算とメモリアクセスの負担が大きくなり得る。

アーキテクチャから推論エネルギーを推定

論文「The SIMPLE Law for Energy Cost of LLM Inference」は、実機で測定する前に、モデル構成と入出力長から推論エネルギーを見積もれるかを検討した。研究者は計算量とメモリアクセス量を解析し、マルチヘッド注意機構（MHA）とGrouped-Query Attention（GQA）の双方を対象に式を導出した。そこから、LLMのエネルギーを予測する4パラメータモデル、SIMPLEを構築している。

SIMPLEは推論のコストを、入力を処理するprefill、デコード時のKV Cacheアクセス、次のトークンを計算する処理の3つに分解する。入力には隠れ次元、層数、クエリヘッド数、KVヘッド数、さらに入力・出力トークン数を使う。個々のLLMごとに係数を合わせる経験式とは異なり、アーキテクチャの違いがコストにどう反映されるかを説明できる点が特徴だ。

研究が示したこと

• NVIDIA H100上で、6系統13個の密モデルの推論トレースを評価した結果、R2は0.969、MAPEは11.97%だった。報告された既存の最先端モデルより誤差は約3.1～3.5倍小さく、LLMごとの校正も必要としない。
• Prefillのエネルギーは隠れ次元に対して二次的に増加する。したがって、長いプロンプトでは大規模モデルほど入力処理の負担が重くなる。
• GQAはKVヘッド数を減らし、キャッシュ関連の処理を抑える。より積極的なグループ化では、標準MHAに対してトークン当たりエネルギーを最大約3分の1まで削減できた。
• 推論トークン数を増やせば常に効率が上がるわけではない。出力が長くなるほどKV Cacheも拡大し、解析的に求めた最小エネルギー点を超えるとエネルギーは超線形に増える。精度が伸び続けても、1ジュール当たりの知能は中間的な予算で最大になる。

推論基盤への意味

この研究は、推論エネルギーをデプロイ後に測るだけの指標から、事前に比較できる設計指標へ近づける。サービス提供者はモデルやコンテキスト長を比較し、開発者は最大出力長や推論予算を設定する際に、KV Cacheの容量、帯域幅、アクセス回数を考慮できる。

ただし、評価対象は論文で示された13個の密モデルとH100である。量子化、バッチ処理、カーネル実装、別のアクセラレータでは結果が変わる可能性がある。SIMPLEは実測を置き換える万能な電力計ではなく、アーキテクチャとワークロードを比較するための推定ツールとして捉えるのが適切だ。

出典：Hugging Face Daily Papers