Save/restore prompt cache

[Stebalien]

llama.cpp supports saving/restoring the prompt cache by:

1. Passing --slot-save-path /pat/to/cache on the CLI.
2. POSTing to /slots/{slot_id}?action={save,restore}.

It would be great if llama-swap could do this automatically when unloading/loading models (with a separate cache directory per model). This doesn't need to be super reliable, it would just be nice to be able to quickly swap between models without having to reprocess everything every time.

I can implement this with a shell-script wrapper around llama-cpp, but it's messy and unreliable. E.g., restoring the prompt cache races with the first request because, from llama-swap's perspective, llama.cpp is "ready".

[mostlygeek]

This doesn't belong in llama-swap core but I think it can be a nice little tool in cmd/llama-server-cache-intercept/main.go.

It can pass all the cli arguments it is given and start llama-server as a child process. When it gets a SIGINT/SIGTERM it will send the POST request to save the prompt cache before sending its own SIGTERM to llama-server.

To llama-swap this would look like any other process that can be started and stopped.

llama-swap will send a SIGKILL after a timeout so that is something to watch out for. That could live in llama-swap as a configuration option to control how long processes get to shutdown down cleanly.

[Stebalien]

That's the path I started down, it's just a lot of extra work to avoid the race I mentioned. Basically, the shim would need to:

1. Pick a random port and start llama-server on it.
2. Wait for llama-server to pass a health check.
3. Restore the cache.
4. Start proxying requests from llama-swap to the server.

Without this, llama-swap will start proxying requests before we restore the cache.

---

How would you feel about adding a few more hooks into llama-swap? Per-model, synchronous, post-start and pre-stop hooks would make this easy to implement.

• post-start: after the health-check passes, before forwarding any requests.
• pre-stop: after stopping requests, before killing. Technically, cmdStop could be used here instead.

[candrews]

Here's a PR doing just that: #595

[Stebalien]

Whoa, awesome! I should have done more searching.

[candrews]

I'm keeping 🤞 that it'll merged

[WinPooh32]

I've made patch for supporting simple loading/unloading hooks: https://gist.github.com/WinPooh32/3741ca478a32a93965e2204c4cb6bc2a#file-readme-md
It can be helpful for saving and restoring kv caches until mainstream project will support that.