あなたはOpenAIのAIエンジニアの面接を受けています。 面接官はこう尋ねます: 「私たちのGPTモデルは42秒で100個のトークンを生成します。 どうやって5倍速くするんだ?」 あなた:「もっと多くのGPUを割り当てて、より高速な発電をします。」 面接終了。 見落とした内容は以下の通りです:

本当のボトルネックは計算ではなく冗長計算です。 KVキャッシュがなければ、モデルは各トークンのキーと値を再計算し、作業を繰り返します。 - KVキャッシュ→9秒 - KVキャッシュなし→42秒(~5倍遅い) それでは、その仕組みを理解していきましょう!

KV キャッシュを理解するには、LLM がトークンをどのように出力するかを知る必要があります。 - Transformer は、すべてのトークンに対して隠れ状態を生成します。 - 非表示の状態は語彙スペースに投影されます。 - 最後のトークンのロジットは、次のトークンを生成するために使用されます。 - 後続のトークンについても繰り返します。 これを👇チェックしてください

したがって、新しいトークンを生成するには、最新のトークンの隠し状態のみが必要です。 その他の非表示状態はどれも必要ありません。 次に、アテンションメカニズムからトランスフォーマー層内で最後の隠れ状態がどのように計算されるかを見てみましょう。

注意中: query-key-product の最後の行には、次のものが含まれます。 - 最後のクエリベクトル。 - すべてのキーベクトル。 また、最終的なアテンション結果の最後の行には、次のものが含まれます。 - 最後のクエリベクトル。 - すべてのキーと値のベクトル。 理解を深めるために、次のビジュアルを確認してください。

上記の洞察は、新しいトークンを生成するには、ネットワーク内のすべてのアテンション操作に必要なのは次のことだけであることを示唆しています。 - 最後のトークンのクエリベクトル。 - すべてのキーと値のベクトル。 しかし、ここにはもう1つの重要な洞察があります。

新しいトークンを生成すると、次のようになります。 - 以前のすべてのトークンに使用される KV ベクトルは変更されません。 したがって、1ステップ前に生成されたトークンのKVベクトルを生成するだけで済みます。 残りのKVベクトルはキャッシュから取得して、計算と時間を節約できます。

これをKVキャッシュと呼びます! 繰り返しになりますが、すべてのコンテキストトークンのKVベクトルを冗長に計算する代わりに、それらをキャッシュします。 トークンを生成するには: - 1ステップ前に生成されたトークンのQKVベクトルを生成します。 - キャッシュから他のすべてのKVベクトルを取得します。 - アテンションを計算します。 これを👇チェックしてください

KVキャッシュは、トークンを生成する前にプロンプトのKVキャッシュを計算することで推論を高速化します。 これがまさに、ChatGPTが最初のトークンを生成するのに他のトークンよりも時間がかかる理由です。 この遅延は「初回トークンまでの時間(TTFT)」として知られています。 TTFTの改善はまた別の機会に話しましょう!