LLM推論性能モデリング - fltech - 富士通研究所の技術ブログ

こんにちは、富士通研究所のコンピューティング研究所の福本です。最近、LLM推論サービスが活用が進む中で、LLM推論サービスを高速化する研究開発が活発に行われています。一方で推論処理をサービス化するにあたって、LLM推論のスループットやレイテンシ（応答速度）がどのパラメータによって決まり変動するのか分かりづらいです。この課題は想定するLLM推論サービスに適切なGPUを選定するときにも同様に発生します。そこで今回は、簡単な性能モデルを作って調べてみました。その結果、簡単な近似でも、ボトルネックの切り替わりや性能の変化を説明できて、意外と役に立ちそうなことがわかりました。本記事ではLLM推論サービスの性能モデルの作り方と使用例について簡単に説明します。

LLM推論の動作

LLM推論では以下の2つのフェーズがあります。

Prefill: 入力プロンプト（ユーザがLLM推論サービスに入力した文字列）を読み込んで文字列生成の準備をして1トークン生成します
Decode: 準備済みの状態から1トークン(*1)ずつ文字列を生成します

(*1) LLM推論におけるトークンとは、テキストを処理する最小単位（単語や文字の一部）のことで、どういった区切りになるかは、使用するLLMによって変わります。

Prefill: 入力プロンプトを一括処理して、最初の1トークンを出力

Prefill は、ユーザが渡した入力プロンプト（例：数百〜数千トークン）をまとめて処理するものです。ここでモデルはプロンプト全体を読み取り、各レイヤで必要な内部状態（後述のKV Cache）を作りながら、プロンプト末尾の次に来るトークンを予想します。 Prefillは行列積処理の割合が大きいため、演算性能（OPS）が重要です。また、プロンプト長が長いほど、まとめて処理するトークン数が増えるため、Prefillの計算量や時間が増えます。

Decode：生成トークンと履歴（KV Cache）から次の1トークンを出す

Decode は、すでに生成したトークン列に対して次の1トークンを出す処理を繰り返すものです。このとき重要になるのが KV Cache（Key Value Cache）です。TransformerのAttention Head部では過去に生成したトークンのKeyとValueという2つの行列データ (KV Cache)を使用します。毎回プロンプト先頭から計算し直すこともできるのですが、計算するかわりに過去に生成したKV Cacheを再利用することでより短時間で処理できます。

Decodeは行列ベクトル積（もしくは片側の行列の1辺が小さい行列積）の処理の割合が大きいです。行列ベクトル積はデータの読み書きがボトルネックとなる有名な処理で、メモリバンド幅が重要です。LLMの重みデータとKV Cacheを読み出し、1トークン生成するという動作になるため、1トークンあたりの計算量は小さく、大量の重みデータを読み込み、さらに KV Cache が生成の進行とともに増えていきます。複数のリクエストをまとめて処理することで、オンチップメモリに読み込んだ重みの再利用ができ、効率よく処理することができます。

LLM推論サービス

LLM推論はPrefill（入力を一括処理して最初のトークンを生成）とDecode（KV Cacheを使って1トークンずつ生成）にわかれます。LLM推論サービスでは、単に高いスループットを達成すればよいわけではなく、ユーザ向けに守るべきレイテンシの制約を満たしつつ、GPUを効率的に使い大量のリクエストを処理することが重要です。

LLM推論サービスの満たすべき制約

Time To First Token (TTFT)

TTFTはリクエストを受けてから最初のトークンを返すまでの時間です。チャットAIサービスの場合、最初の1トークンを返すまでの時間はユーザ体験に大きく影響するため非常に重要です。MLPerf Inferenceという著名なベンチマークのLLM（Llama2 70B）のServerシナリオという条件では、TTFTは2秒以下に抑える必要があります。他のクラウドAPIのベンチマークではTTFT1秒以下という例もあります。

Time per Output Token（TPOT）

TPOTは、生成中に次のトークンをどの程度の間隔で出力できるかを表す指標です。ユーザ視点では、出力が出始めたあとに、スムーズに文字が出力されるかというものになり、こちらも非常に重要です。MLPerf Inference の Llama 2 70B（serverシナリオ）では、TTFTに加えて TPOT 200ms以下の制約があります。他のMLPerfのモデルやシナリオでは制約が変わっています。たとえば MLPerf Inference 5.1 の Small LLM（Llama 3.1 8B）では serverシナリオのしきい値として TPOT ≤ 100ms、より厳しい interactive シナリオではTPOT ≤ 30ms となっています。

レイテンシ制約を満たしながらGPUシステムで効率よくリクエストを処理する

上記で説明したTTFT, TPOTの制約を満たしながら、リクエストをまとめて処理して、推論スループット（Tokens/sec）を高めることが目標です。Decodeはリクエストをまとめて処理するほどスループットを高められます。Decode処理はメモリバンド幅ネックのため、一度オンチップメモリに読み込んだLLMの重みを使って、複数のリクエストをまとめて処理することで、効率よく処理することができます。しかし、リクエストをまとめると以下の理由でTPOTが増えます。

リクエストをまとめるために、処理の完了やリクエストの到達を待つ時間が増える
重みのロード以外の再利用できないデータの読み込みや、演算回数の増加により1回の処理時間が増える

したがって、TPOT制約を満たせるようにリクエストをまとめて処理する数を決める必要があります。

　TTFTはリクエスト到達後にPrefill処理を始めるまでの時間とその後Prefill時間そのもので決まります。よってなるべく早く新規リクエストの処理を始めたいですが、すでに処理中のリクエスト(=Decode中のリクエスト)にとっては、Prefillが割り込んで長時間GPUを占有すると、TPOTが悪化します。そこで、走行中のリクエストのDecodeと新規リクエストのPrefillをうまくスケジューリングして、TPOTとTTFTを守る必要があります。

　Prefillの処理時間は長いため、一括で処理するとDecode中のTPOT制約を満たすことが難しくなります。そこでPrefillの分割処理（Chunked prefill）が行われます。大きいPrefillを小さく分割し、Decode処理でまぜて処理することで、スループットとレイテンシを両立します。 vLLMでは、Chunked prefillがデフォルトで有効となり、まとめて処理可能なトークン数（max_num_batched_tokens）という指標を設定し、これを超えないようにDecodeのリクエストを優先してまとめて、余裕があればPrefillを分割し実行するというスケジューリング方法が取られているようです(https://docs.vllm.ai/en/stable/configuration/optimization/#chunked-prefill)。

性能モデリング

LLM推論性能を見積もるために、実行時間を単純な数式で表します。細かく考えると実行時間は、HW実装やSW実装によって変わってしまうのですが、このような詳細な動作まで見込んだ性能モデルを今回は対象としません。ここではGPUのどういったパラメータで性能が大きく変わるのかについて考えるために、なるべく簡単な数式で性能を表します。

　LLMの各レイヤの演算回数 $OPs$ とデータの入出力容量 $Bytes$ を対象モデルの定義（hidden size, FFN次元、head数、など）から計算します。あるレイヤの実行時間 $T_{\mathrm{layer}}$ は以下の形で近似できます。

$\displaystyle T_{\mathrm{layer}} = \mathrm {max} \left( T_{\mathrm{com}} , T_{\mathrm{mem}} \right)$

$T_{\mathrm{com}}$ : 演算時間
$T_{\mathrm{mem}}$ : データ転送時間

これは演算がボトルネックであれば演算時間によって実行時間が決まり、データ転送がボトルネックであればデータ転送時間によって実行時間が決まるというものです。演算時間は以下のように近似します。

$\displaystyle T_{\mathrm{com}} = \dfrac {OPs} {P_{\mathrm{peak}} \times Eff_{\mathrm{com}} }$

$OPs$ : 演算回数
$P_{\mathrm{peak}}$ : GPUの理論演算性能
$Eff_{\mathrm{com}}$ : 演算効率（理論演算性能に対する実行効率）

データ入出力時間は以下のように近似します。

$\displaystyle T_{mem} = \dfrac {Bytes} { BW_{\mathrm{peak}} \times Eff_{\mathrm{mem}}}$

$BW_{\mathrm{peak}}$ : GPUの理論メモリバンド幅
$Eff_{\mathrm{mem}}$ : メモリ転送効率（理論メモリバンド幅に対する実行効率)
$Bytes$ : 当該レイヤで実際に読み書きするデータ量（重み、KV Cache, レイヤの入出力の合計値）

全体の実行時間Tは以下のように表します。

$\displaystyle T=\displaystyle \sum^{\mathrm{num\_layers}} \left(T_{\mathrm{layer}} \right)$

この性能モデルはLLM推論の特性を見るためには便利ですが、現実のLLM推論はより複雑です。たとえば、オンチップメモリの効果でHBMへのアクセス量が変わりますし、行列形状やリクエストをまとめる個数によって行列積の実行効率が変わります。実行時間の推定精度を上げたい場合、専用のシミュレータを開発するか、実機を使ってLLM推論サービスを動作させて計測するなどしたほうがよいです。また、ArXivなどにLLM推論の性能モデリングの論文は多数投稿されているので、こういったものを参考にするのもよいと思います。

　LLM推論サービスの性能推定で次に重要になるのはサービス制約（TTFT/TPOT）を満たせるバッチサイズを決めることです。ここでバッチサイズとは、LLM推論サービスにおいてリクエストをいくつまとめて処理するかを示す指標のことです。バッチサイズを大きくすると、スループットが改善しますが、サービス制約を満たしにくくなります。先程作った性能モデルにおいて、バッチサイズを増やすと各レイヤの演算回数やデータ入出容量が変わり、実行時間が変わります。Input Token数やOutput Token数、PrefillとDecodeのスケジューリング方針を決めることで特定のバッチサイズでTTFT, TPOT制約を満たせるかを計算することができます。

演算性能とメモリバンド幅がバッチサイズやTTFT/TPOTに与える影響

Prefillは演算がボトルネックとなるため、GPUの演算性能があがるとPrefillの実行時間が短くなります。したがって演算性能の向上により、同じTTFT制約でも、より大きいバッチサイズを設定できます。また、Chunked prefillのようにPrefillをDecodeと同じバッチに混在させる運用では、Prefillが短いほどDecodeを待たせにくくなり、TPOT制約も満たしやすくなります。さらに、バッチサイズがある程度大きくなると、Decodeでも大きな行列積を扱うレイヤがメモリバンド幅ネックではなく、演算ネックとなることがあります。この場合、演算性能向上はDecodeのTPOTの削減にも効果があります。一方で、Decodeフェーズは重みとKV Cacheの読み込みが支配的となりやすいため、メモリバンド幅向上によりDecodeの処理時間（TPOT）が短くなります。つまり、LLM推論サービスにおいては、演算性能もメモリバンド幅もそれぞれ重要です。求めるレイテンシ制約（TTFT/TPOT）、対象モデル、ワークロード（Input Tokens数、Output トークン数、リクエストの到達頻度）がわかれば、必要なGPUの演算性能とメモリバンド幅も概算できます。

ここで、演算性能の違いがLLM推論性能に与える影響を簡単な性能推定モデルを使って計算します。AMDが公開しているMLPerf v4.1　Llama2-70BのServerシナリオの結果ではTTFT 2秒以下、TPOT 200ms以下のレイテンシ制約を満たす必要があります。MI300X 1GPUの結果（https://www.amd.com/en/blogs/2024/engineering-insights-unveiling-mlperf-results-on.html ）において、最大バッチサイズ(max_num_seqs)が768で2520 Tokens/secという結果が報告されています。この結果をベースに、GPUがMI300Aとなった場合の推論スループットを推定します。MI300XとMI300Aは、メモリバンド幅は5.3TB/sと同等ですが、FP8の演算性能はMI300Xが2614.9 [TFLOPS], MI300Aが1961.2 [TFLOPS]であり、演算性能比は0.75となります。次に、MI300Xのバッチサイズが768のとき、Llama2 70BのDecodeフェーズの実行時間のうち演算ネックとなる時間の割合を70%、Prefillをバッチサイズ1で処理したときの実行時間のうち演算ネックとなる時間の割合の85%とします（これはLlama2 70Bのモデル定義とMI300Xの演算性能とメモリバンド幅から、各レイヤの演算時間とデータ入出力時間を計算した結果をもとに設定した値です）。Decodeの時間は、全体の実行時間の70%の性能が75%となるため、アムダールの法則のように近似すると、1.233倍（=0.7/0.75 + 0.3）になります。同様にPrefillの実行時間は1.283倍となります。この実行時間の増加により、MI300AでもTPOT/TTFT制約を満たすためには、バッチサイズを小さくする必要があります。

ここで簡単に計算するために、MI300XがDecode処理のみ768個同時に処理する場合に、TPOT制約をちょうど満たせているという仮定をおいて、MI300Aで許容できるバッチサイズを計算します。Decode処理時間は、バッチサイズに比例しない重みの読込時間、および、バッチサイズに比例するKV Cacheの読み込み時間、演算時間、その他の読み書きの時間に分けられます。バッチサイズが大きい場合、KV Cacheの読み書きの時間および演算ネックの処理の割合が大きくなります。ここでは簡単に計算するために、Decode時間はバッチサイズに比例するとします。このときTPOTを同程度に維持するための許容バッチサイズは768/1.233 = 623となります。同様にスループットは2520 * 623/786 = 1997 Tokens/secと推定できます。

おわりに

本記事で見てきたように、LLM推論の性能はレイテンシに関する2種類の制約が存在するため、バッチサイズやハードウェア構成によって性能の支配要因が切り替わり、挙動は単純ではありません。しかし一方で、今回紹介したような簡単な性能モデルを用いるだけでも、LLM推論のスループットやボトルネックを大まかに見積もることができます。推論サービスを行う計算機システムを構築するような場合、リクエスト到達頻度、TTFT/TPOTを決めれば、必要なGPU個数などを概算できます。また、最近では、Prefill/Decodeを異なるGPUで処理するような事が行われています。こういった場合にどのような特徴を持つGPUを採用するかを考えることができるようになります。現在、ArXivなどではLLM推論性能のモデリングに関する論文が数多く公開されています。より厳密な分析や最適化に興味のある方は、ぜひこうした論文も読んでみてください。また、LLMのパラメータから演算回数やデータ入出力容量をどのように見積もるのかは、関心を持つ方が多いようであれば、別途追記したいと考えています。