lmstudio-ai / lmstudio-ai/mlx-engine

Disk prefix KV caching

Open
#354 1 comment 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.2k
Forks
133
Avg merge
21h 6m
Merged PRs (30d)
1

Description

**The issue:**
When a conversation ends, or the model is unloaded, the KV cache is removed from RAM. Then the user gives a new prompt in the conversation and the whole KV cache has to be built again which can take minutes.

**The solution:**
By storing the KV cache on disk it can be loaded again when needed. This can even happen after the model has been reloaded. I expect great performance improvements on longer conversations with this. Also huge improvements in time-to-first-token since the system prompt will always be cached.

**Implementation:**
Create a unique fingerprint for each model + configuration. Each configuration that impacts the KV cache should result in a different fingerprint. When a conversation ends, write the KV cache as a file to a location in a sub-dir containing the fingerprint. Then when a prompt is received, look for the cache files in the fingerprint dir. For each cache, check if the prompt is a prefix of the cache. For example, a different conversation might be cached, but the system prompt is still the same so the prefix will match. Choose the cache with the longest prefix match. Then load the cache to RAM.

I can supply a PR for this, but I guess it needs to be wired into the actual frontend which isn't open source. I'd suggest making a few things configurable through the UI: The disk location for the cache, the max size/number of cache files (the cache can grow quickly over time). The most recent matches should be kept and the caches not matched recently can be removed when writing a new cache file.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

No files or tests are named in the issue; start by locating the KV-cache lifecycle, model configuration handling, and prompt-processing entry points in the repository. Define the model/configuration fingerprint and disk-cache lookup and eviction behavior, then verify that matching prefixes reload correctly and that the cache can survive model unloads and reloads.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
ai, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.