Repository navigation
lean: save and restore the KV cache of a prompt prefix - #10
Merged
Merged
Conversation
KvCache and CpuKvCache gain read_prefix/write_prefix: one layer's K or V for positions [0, n) as [kv_heads, n, head_dim] f32 bytes, the same bytes on both backends. The GPU read copies only the n positions of each head into one staging buffer (Engine::read_ranges); the write uses queue.write_buffer and submits per call so staged copies do not pile up. gguf_header_len gives the offset where tensor data starts.
Lean::prefill(text) makes text resident without generating and returns its token count. save_prefix writes the KV cache of a prefix as a 128-byte header plus each layer's K and V in f32, one tensor per write; restore_prefix reads it back one tensor at a time into a fresh engine, and the next generate prefills only what follows the prefix. prefix_key hashes the format version, dtype, engine build (a build-time hash of the lean and lean-core sources and Cargo.lock), the GGUF header and length, the tokenizer file, the backend and the prefix token ids. A restore checks the same fields and refuses a mismatch with LeanError::PrefixImage, leaving the engine usable. Exported to JS as prefill, prefixKey, savePrefix and restorePrefix.
Native test on the tools prompt of 02_tools_single: the reply after a restore equals a fresh engine's, the KV cache after the reply is bit-identical to the writer's, and images for another engine build, another prefix or cut short are refused. The bench worker's prefix phase runs the same checks in a browser with the image stored in OPFS.
The full header parse checks the remaining file size before every string, three seeks that empty the read buffer each time, and took about 470 ms on a Qwen GGUF. The scan skips every value through one 1 MiB buffer and takes about 4 ms; an ignored test checks it ends where the parse does.
A native build and a wasm build of the same sources run different CPU kernels, so the target architecture and features are hashed too, with rust-toolchain.toml.
…der's header length Drop the scanning header_len and its re-export. GgufReader exposes the data section offset it already computes, both models keep it, and the fingerprint hashes that many bytes after the load, from one more open of the source.
One per-head loop owns the KV layout. The snapshot now reads only the prefix ranges instead of the whole max_ctx buffer.
README: restorePrefix and its read(n) are synchronous, so a browser reads the image from a Web Worker through an OPFS sync access handle.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Save the KV cache of a prompt prefix to bytes and restore it later, so a long fixed prefix (a tools prompt, a system prompt) is computed once and then loaded instead of recomputed.
API (Rust and JS)
prefill(text): run the prompt through the model without generating; returns the token count.prefix_key(text): the key of the image this engine and model would write for that prefix.save_prefix(text, write)/restore_prefix(text, read): stream the image one tensor at a time; F32 only.LeanError::PrefixImage(PrefixImageError); a refused image means recompute and overwrite, never load.Invalidation
Also
Tests