Implement Llama longest prefix cache

[abetlen]

Opening this up to track the development of the new caching behaviour I'm planning to implement. This will leverage 2 significant improvements

• Reduced llama state size which is now a function of evaluated tokens
• Improved efficiency of Llama.generate which now only eval's prompt tokens that are not already in the context window

[abetlen]

This looks like it's good to go now, the cache is in-memory currently so you'll need to set CACHE_SIZE or --cache-size in bytes to the maximum cache size for your available RAM. If anyone wants to test this out let me know if it works as expected, I'll probably push this release later tonight.

[ghost]

Tried this out in oobabooga and it's a game-changer for chats with frequent editing. Before this the model had to spend a lot of time reingesting the entire prompt even on a small edit.

[abetlen]

@eiery glad to hear! I should point out though that text-generation-webui is just using the new generate rewind behaviour and not the full state cache. Basically "rewind" is useful for a single long conversation or generating multiple responses for a single prompt, the full state cache let's you restore completely differrent conversations.

[ibehnam]

@abetlen I'm curious about using this feature in the Python API. I tried this approach but re-running the same llm always returns "... remaining 1 prompt tokens to eval" even though it should not eval anything because the prompt hasn't changed. 🤔

from llama_cpp import Llama, LlamaDiskCache

cache = LlamaDiskCache(cache_dir=".../.cache")


model_path_2 = "...Phi-3-mini-4k-instruct-q4.gguf"
llm = Llama(
    model_path=model_path_2,
    n_gpu_layers=-1,
    flash_attn=True,
    n_threads=10,
    use_mlock=False,
    seed=42,
    verbose=True,
)
llm.reset()
llm.set_cache(cache)

# %%

prompt = "#User:\nFrom now on, your name is Behnam.\n#Assistant:\n"

r = llm(
    prompt=prompt,
    max_tokens=12,
    temperature=0.0,
    stop=[
        "\n\n",
        "#User:",
        "#Assistant:",
    ],
)

print(r["choices"][0]["text"])

Which gives:

Llama._create_completion: cache hit
**Llama.generate: 21 prefix-match hit, remaining 1 prompt tokens to eval**

llama_print_timings:        load time =     125.70 ms
llama_print_timings:      sample time =       1.61 ms [/](https://file+.vscode-resource.vscode-cdn.net/)    12 runs   (    0.13 ms per token,  7453.42 tokens per second)
llama_print_timings: prompt eval time =       0.00 ms [/](https://file+.vscode-resource.vscode-cdn.net/)     0 tokens (     nan ms per token,      nan tokens per second)
llama_print_timings:        eval time =     358.08 ms [/](https://file+.vscode-resource.vscode-cdn.net/)    12 runs   (   29.84 ms per token,    33.51 tokens per second)
llama_print_timings:       total time =     395.63 ms [/](https://file+.vscode-resource.vscode-cdn.net/)    12 tokens
Hello! I'm Behnam, at your service
Llama._create_completion: cache save
Llama.save_state: saving llama state
Llama.save_state: got state size: 13112306
Llama.save_state: allocated state
Llama.save_state: copied llama state: 13112306
Llama.save_state: saving 13112306 bytes of llama state
LlamaDiskCache.__setitem__: called
LlamaDiskCache.__setitem__: set
LlamaDiskCache.__setitem__: trim

I also tried appending the model output to the previous prompt and sending it to the llm again:

...
agg_prompt = prompt + r["choices"][0]["text"] + "#User: What is your name now?\n#Assistant:"

r = llm(
    prompt=agg_prompt,
    max_tokens=32,
    temperature=0.0,
    stop=["\n\n", "#User:", "#Assistant:"]
)

print(r["choices"][0]["text"])

But again got similar results:

**Llama.generate: 33 prefix-match hit, remaining 15 prompt tokens to eval**

llama_print_timings:        load time =     131.47 ms
llama_print_timings:      sample time =       2.94 ms [/](https://file+.vscode-resource.vscode-cdn.net/)    18 runs   (    0.16 ms per token,  6122.45 tokens per second)
llama_print_timings: prompt eval time =     165.82 ms [/](https://file+.vscode-resource.vscode-cdn.net/)    15 tokens (   11.05 ms per token,    90.46 tokens per second)
llama_print_timings:        eval time =     407.14 ms [/](https://file+.vscode-resource.vscode-cdn.net/)    17 runs   (   23.95 ms per token,    41.75 tokens per second)
llama_print_timings:       total time =     589.39 ms [/](https://file+.vscode-resource.vscode-cdn.net/)    32 tokens
 
My name is Behnam. How can I assist you today?
Llama._create_completion: cache save
Llama.save_state: saving llama state
Llama.save_state: got state size: 25695602
Llama.save_state: allocated state
Llama.save_state: copied llama state: 25695602
Llama.save_state: saving 25695602 bytes of llama state
LlamaDiskCache.__setitem__: called
LlamaDiskCache.__setitem__: set
LlamaDiskCache.__setitem__: trim