ローカルLLMで長文指示を高速応答させる ～明示的にKVキャッシュ保持して高速化してみた～ (Llama-cpp-python)

こんにちはRcatです。
今回もローカルLLMネタです。

こちらの記事で、様々なローカルLLMを稼働させて実用に足る速度は出るのか実験しました。
その結果として、生成速度は実用的であると結論が出ましたが、ここでは見ていないことがあります。

それが、考える時間(入力プロンプトに対する注意機構の演算)です。
上の記事では、"1000文字の文章書いて"なので考える時間はいらないですが、タスクの自動化を行おうとすると、指示が長くなり応答が始まるまでに時間がかかってきます。

今回はこれを何とかできないか考えていきます。

はじめに

利用規約

情報や作品の活用時は事前に利用規約をご確認ください。

コメントについて

利用規約のガイドラインを確認の上コメントしてください。
則っていないコメントは削除します。

KVキャッシュで"考えずに"回答!?

2回目以降なんか早くない？

LLMを扱っていると、最初は遅いのに同じ指示の2回目からなんか応答が速いなってことありませんか?

なぜかはあまり気にしていなかったのですが、先日購入したinterfaceで語られていました。
Transformer系LLMが文章を生成するときは、前後や指示文との関係性を崩さないように応答するために注意機構が備わっています(G検定で出るよ!)。
そして、プロンプトが全く同じ(最後だけ違うなど)の場合は、この注意機構の演算結果を使いまわしできるそうです。

なので、対応しているプラットフォームでは使いまわしが行われることで次回以降の応答時に考えなくなって高速になるそうです。

これをうまく利用すれば、システムプロンプトに知識を入れて回答させるというような使い方が非常に効率的になります。何せその知識はすでに演算済みで見直さなくていいんですからね。

ただし、Llama-cpp-pythonにおいて実験の結果弱点が見つかりましたので、その内容と対策について紹介していこうと思います。

キャッシュの有無での動作の違い

動作実験

実験にて動作速度を確認します。
上でちょっと何言ってるかわからないという方でも、とりあえずこちらを見てください。

まずは非常に長いシステムプロンプトを入力します。文字数は1万5800文字です。画像生成の記事で使っている文ですね。

初期トークンが24秒程度で出てきています。この時間が指示を読み取り(注意機構の演算)にかかった時間です。

次に内容を一切変えずにもう1度推論を実行します。
すると今度は初期トークンが0.1秒くらいで出ています。
これは入力された内容が完全に一緒なので計算を省略し、ただ出力だけを行った結果です。
これこそが"考えずに回答する"です。

これがキャッシュの実力です。
つまり、どんなに長いシステムプロンプトを入力していたとしても、最初の1回だけで計算を行うので、複雑な処理を短時間で実行させることができます。

しかしここからが弱点です

先ほどとは全く関係ない話をしてみましょう。
以下の例ではシステムプロンプトは入力していません(この場合は自動的にデフォルトのが代入される仕組みです)。

どうでもいい内容なので、超高速で出てますが関係ありません。

問題はここから
関係ない話をした後に、もう1度先ほどの画像生成の指示を投げてみます。
するとどうでしょう、一番最初と同じように24秒の考える時間ができてしまいました。
つまり、これは前回のキャッシュとどうでもいい話が合わなかったので、破棄されてしまったということになります(現在はどうでもいい話がキャッシュされている)。
非常にもったいないです。

キャッシュを明示的に保持して高速化する

どんな方法があるか

というわけで、どうにかキャッシュを保持することができないか調べてみました。
一番簡単なのはモデルを読み込む時に複数のインスタンスを読み込むこと。各タスクごとにモデルを割り当てて使います。
Llama-cppは同じモデルがロードされても、モデルは一回だけ読み込む動作をするらしいので、うまくいくならこれが一番簡単です。
私はGPUだとメモリがうまく共有されず、モデルの数だけメモリを食ってしまったのでこれはやめました。ROCmと相性悪いか?

というわけなので、私はモデルにキャッシュの管理を任せるのではなく、事前に計算し、自分で管理してアサインするという方法を取りました。

キャッシュの計算と割り当ての流れ

今回のプログラムでは次の手順で行います。

• システムプロンプトが長いかどうかを判別
  日本語で500文字程度などグローバル変数で設定します。

• 事前に計算されたキャッシュがあるかどうか確認
  システムプロンプトのハッシュ値を計算して、それが辞書に含まれるかどうかで、前回使ったキャッシュがあるか確認します。

• ない場合はキャッシュを新規作成して保持
  上で計算したハッシュ値を使って辞書で管理します。

• キャッシュを強制的にアタッチ
  使用するキャッシュを明示的にアタッチします

• 通常通り推論する
  すでにシステムプロンプトがある場合は自動的にスキップされるため高速になります。

実際の動作

まずは初回ですね。同じく24秒かかっています。
キャッシュがないので当たり前ですが。

全く同じ内容で2回目投げます。
今度は0.5秒くらいで出てきているので、キャッシュを確認できました。
ここまでは前回と同じです。

ではまた関係ない話を振ってみましょう。
今まではここでキャッシュが破棄されてしまいましたね。

もう一度目的の指示を投げてみます。
今度は0.5秒で初期トークンが返ってきています。
つまり2回目以降と同じスピードです。うまくキャッシュが適用されているようです。

そういえば、今まで完全に入力が同じだったので、ユーザーの入力が少し違う時はどうなの？というところも確認しておきました。
結果は初期トークンが0.8秒なので、ユーザーの部分だけちょっと考えたって感じがしますね。
しかし、キャッシュが適用されていなければ、こんなスピードは出ません。やはりうまくいっているようです。

まとめ

今回はllmのキャッシュについて見ていきました。
ローカルでLLMを使うと、1つのモデルで様々なタスクをこなす必要があります。
色々と自動化しようとすると指示が複雑になり、先ほどのような非常に長いシステムプロンプト持つものは珍しくありません。
さらに、様々なタスクをこなすため、毎回キャッシュが破棄されてしまうため、スループットが大幅に低下します。
しかし、今回の方法を採用したことで、今後は初回のみ2倍の時間がかかるようになりましたが、それ以降は一切の計算が不要で処理が行えることになりました。

追記:少々難点が…

※すでにソースは対策済みです

少し運用して分かったことなのですが、このKVキャッシュ結構でかいです。
128GBのPCでメモリ足りないと言い出して何事だと思ったら、本作がメモリを食い尽くしていました。

現設定だと、1週間はキャッシュを保持にしていたので、いろんなテストでシステムをちょっとだけ変えていたので全部積もった感じですね。

見た感じだと、3000文字くらいで0.4G程度。10000文字で0.8G程度のサイズになっていて線形ではなさそう。
※サイズはメモリ監視して拾った雑な値です。

というわけなので、自プロセスの占有メモリをチェックし、設定値を超えていれば、収まるまで古いキャッシュを破棄するように仕様変更を行いました(Linuxのみ)。
これで無限に増加するのを阻止できます。

ただし、CPUで演算している方はモデルのサイズおよび推論中の増加分も計測されてしまう可能性がありますので、余裕を持った値を設定して下さい。

配布情報

具体的なソースコードは、配布しているものをご確認願います。

こちらの記事で配布しているものをアップデートします。

リンク集

Pythonのインストール方法

【Linux】Systemd追加支援スクリプト

LoggingBOTの導入方法

ミニPCにLinuxをインストール

おまけ