lmstudio-ai / lmstudio-ai/mlx-engine
LRU text cache and image checkpointing for VisionAddOn
- Dominant language
- Python
- Stars
- 1.2k
- Forks
- 133
- Avg merge
- 21h 6m
- Merged PRs (30d)
- 1
Description
# Context
After PR #298 migrated Qwen3.5 to `ModelKit` via `Qwen3_5VisionAddOn`, the model has no cross-turn caching at all. Every conversation turn triggers a full re-prefill from scratch, regardless of shared prefix length. On a long system prompt or a multi-turn image conversation this costs several seconds per message.
This is a follow-up to #287, updated for the ModelKit/VisionAddOn architecture.
# Implementation
A working implementation is available at [AirRunner/mlx-engine · feat/modelkit-image-kv-cache](https://github.com/AirRunner/mlx-engine/tree/feat/modelkit-image-kv-cache).
## Text KV cache: delegate to LRUPromptCache (mlx-lm)
The core change is architectural. `CacheWrapper._find_common_prefix`, `_get_num_tokens_in_cache`, and `_get_unprocessed_tokens` are removed, and prefix matching and trimming are now delegated to `LRUPromptCache.fetch_nearest_cache` (mlx-lm), which correctly handles both trimmable caches (KVCache, RotatingKVCache) and non-trimmable caches (ArraysCache). The previous trim-based implementation silently discarded the cache on any mismatch, causing a full re-prefill every turn on models with non-trimmable layers such as Qwen3.5's GDN/SSM architecture.
`CacheWrapper` also normalizes LRU keys before lookup: complete `…` blocks are stripped so the cache hits on the next turn even when CoT blocks are absent from the next turn's prompt (as is the case with the default Qwen3.5 MLX chat template in LM Studio).
## Image KV cache
On top of the text caching foundation, `ImageCheckpointStore` is added as a standalone class in `cache_wrapper.py`. It records metadata (`image_end_index`, `prefix_hash`, `block_lengths`) after prefilling each image block. On subsequent turns the LRU provides the starting cache, the ViT is skipped for unchanged images, and only new images run through the ViT (partial hit).
### Covered scenarios
| Scenario | Before | After |
|---|---|---|
| Any turn (text or image) | full re-prefill from scratch | LRU hit, only new tokens prefilled |
| Turn with same image(s) | full re-prefill + ViT | ViT skipped, prefill from LRU |
| New image appended | full re-prefill + ViT (all images) | ViT runs for new image only |
| Text turn after image turn | full re-prefill | LRU hit from image checkpoint |
| Long system prompt + first image | system prompt KV discarded | pre-image KV injected from LRU |
### VRAM
`ImageCheckpointStore` is a metadata-only store: it holds `(image_end_index, prefix_hash, block_lengths)` per image block, with no KV tensors. At restore time the text LRU provides the starting cache via `_find_starting_cache(prefer_prev_checkpoint=True)`, which returns `self.cache` directly without a deepcopy. Steady-state VRAM is therefore **1x KV**, and generation does not introduce a second copy.
## Integration note
`CacheWrapper` reads the following duck-typed attributes from the tokenizer object:
- `has_thinking` (bool): enables think-block normalization and generation-prompt exclusion.
- `think_start_id` / `think_end_id` (int): token ids for `` and ``.
- `thinking_prefix_offset` (int, default 3): number of tokens in the role header that precedes ``. Used to determine how many trailing tokens to exclude from the LRU key so it remains a valid prefix of the next turn. The default of 3 is correct for ChatML-based templates (`<|im_start|>`, `assistant`/`user`, `\n`). For thinking models with a different role-header length, LM Studio should set this attribute accordingly.
The first three attributes are already set by LM Studio. Only `thinking_prefix_offset` is new.
## Working implementations
- [AirRunner/mlx-engine · feat/modelkit-image-kv-cache](https://github.com/AirRunner/mlx-engine/tree/feat/modelkit-image-kv-cache)
- [AirRunner/mlx-engine · feat/modelkit-paged-disk-kv-cache](https://github.com/AirRunner/mlx-engine/tree/feat/modelkit-paged-disk-kv-cache)
Contributor guide
Research direction
Start by reading cache_wrapper.py and compare the local ModelKit/VisionAddOn implementation with the referenced feat/modelkit-image-kv-cache branch. Trace how text LRU lookup and image checkpoints cover unchanged images, appended images, and text turns after image turns. Done means these scenarios reuse cached work while preserving the stated steady-state VRAM behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100