LLMにおけるKVキャッシュ：ゼロからの理解と実装

KVキャッシュは、以前のアテンション計算から得たキー（K）とバリュー（V）ベクトルを保存し、それ以降のトークン生成時に再利用することで、再計算を避けて推論を高速化します。KVキャッシュがない場合、各生成ステップで過去のすべてのトークンに対するキーとバリューを再計算しますが、キャッシュがある場合は、新しいトークンのベクトルのみを計算してキャッシュに追加します。この記事では、著者の著書に基づくGPTライクなモデルのコード実装を提供しています。主な変更点は、MultiHeadAttentionクラスにキャッシュバッファ（cache_kとcache_v）を追加し、forwardメソッドを変更してキャッシュを条件付きで使用するようにし、resetメソッドを追加し、use_cacheフラグをモデル全体に伝播させることです。生成時、use_cacheが真の場合、モデルは最初のプロンプト後は新しいトークンのみを処理しますが、キャッシュがない場合は毎ステップでシーケンス全体を処理します。簡単なパフォーマンス比較では、KVキャッシュにより、テストした例の生成速度がおよそ2倍になることが示されています。