abetlen / abetlen/llama-cpp-python

LLamaDiskCache: needs a RO / 'static' disk cache for RAG use cases

オープン
#1,737 コメント 0 件 リアクション 2 件 担当者 0 名 GitHub で見る
主要言語
Python
スター
10.6k
フォーク
1.4k
PR マージ指標
PR 指標を取得中

説明

**Is your feature request related to a problem? Please describe.**
I have a dataset that I'm using for RAG. The user's question determines lookup for the top N most relevant documents, then builds a prompt that looks like:

```

You are an assistant for domain , you summarize information and blah blah blah.

Document 1: Some information that is relevant
Document 2: Other information
Document 3: Final information

User question: "How do I do ?"
```

In order to minimize latency, I've [developed](https://github.com/tc-wolf/llama-cpp-python/blob/bumped_llama_cpp_with_disk_cache/llama_cpp/llama_cache.py#L169-L280) a "static" disk cache that contains the system prompt + the first document as context for every document in my dataset. (An example (though old) [script for doing this](https://github.com/tc-wolf/llama-cpp-python/blob/bumped_llama_cpp_with_disk_cache/examples/high_level_api/create_disk_cache.py) is also in my branch).

This way, I only need to ingest the remaining documents + user question when doing prompt processing, so I save a lot of time in time-to-first-token for this use case.

**Describe the solution you'd like**
I'd like to upstream the [LlamaStaticDiskCache](https://github.com/tc-wolf/llama-cpp-python/blob/bumped_llama_cpp_with_disk_cache/llama_cpp/llama_cache.py#L169-L280) class in my branch. It's very similar to the existing `LlamaDiskCache` but:

* Cache is not mutable once built
* (Does not `pop` in `__getitem__`)
* It uses a trie for finding the longest matching prefix (if any) in the cache
* It has a convenience factory method for building from a list of prompts

So it's well-suited for use cases where you want to build the cache once (for a given model + context size + batch size) and then reload at inference time based on matching the prefix of the prompt.

### Complication / Details with this

I've found that when running locally (Mac OS + Metal GPU) and deploying on different hardware (Linux + CPU), I have had to make a minor change to llama.cpp to avoid serializing / deserializing the RNG state.

I.e., skip loading and set seed for reproducibility: https://github.com/tc-wolf/llama.cpp/commit/ea43d922e4169f4fd7622e4a8e6eca92ef921038

I don't think that this will be a factor anymore because https://github.com/ggerganov/llama.cpp/pull/9294 has removed serializing / deserializing the RNG when saving.

**Describe alternatives you've considered**
* Use lower-level state saving functions (rather than pickling `llama.save_state()` to save less on-disk than full model file
* Use more efficient strategy for saving - right now if each key has the same system prompt (for example), that will be saved independently for every stored prompt. There's a lot of space that could be saved if dedupe and only save each prefix once, but it complicates the saving/loading logic.
* Could also allow for partial matches when checking cache - right now a key has to be a full prefix of the input tokens, but could try and look for partial match to allow for more graceful failure.
* This also complicates the `__getitem__` logic.

コントリビューションガイド

コントリビューションガイドを開く

評価

この issue はまだ評価されていません。

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。