KV Cache Persistence
Relevant source files
The following files were used as context for generating this wiki page:
KV Cache Persistence saves the compressed Multi-Head Latent Attention (MLA) Key-Value state to disk. This system allows conversations to reopen in a "warm" state across engine restarts. By persisting this data, the engine avoids the need for a full re-prefill of the prompt, ensuring the session is byte-identical to an uninterrupted state.
Sources: c/kv_persist.h:1-5, README.md:143-147, CHANGELOG.md:180-182
Disk Format and Header Structure
The persistence system uses a specific binary format identified by a magic string. The header contains model configuration parameters to ensure that saved cache files are only loaded by compatible model versions.
File Header Mapping
| Offset (int32_t index) | Field | Description |
|---|---|---|
| Magic (Bytes 0-7) | KV_MAGIC | String identifier "COLIKV1\0" |
| 0 | n_layers | Number of transformer layers in the model |
| 1 | kv_lora | Dimension of the KV LoRA projection |
| 2 | qk_rope | Dimension of the QK RoPE projection |
| 3 | index_hd | DSA sparse attention index dimension (if applicable) |
| 4 | nic | Number of layers with specific Ic components |
| 5 | vocab | Vocabulary size of the model |
| 6 | nrec | Number of tokens/records currently saved in the file |
| 7 | Padding | Reserved for future use (set to 0) |
Sources: c/kv_persist.h:8-16
Persistence Flow
The following diagram illustrates the lifecycle of a KV cache record from memory to disk.
The engine appends compressed MLA KV-cache data incrementally after every turn.
Sources: c/kv_persist.h:2-5, c/kv_persist.h:47-75
Record Management
The system manages records based on the number of tokens processed. Each record contains the token ID and the projected key-value components for all layers.
Data Components per Token Record
A single record in the .coli_kv file consists of:
- Token ID: A 4-byte integer representing the vocabulary index.
- Lc Tensors: KV LoRA data for each layer.
- Rc Tensors: QK RoPE data for each layer.
- Ic Tensors: Optional DSA (Dense-Sparse Attention) index data.
Sources: c/kv_persist.h:18-23, c/kv_persist.h:60-70
Key Functions
| Function | Action |
|---|---|
kv_disk_open | Opens the .coli_kv file in read/write binary mode or creates it with a valid header. |
kv_disk_append | Writes new token KV data to the end of the file and updates the record count. |
kv_disk_load | Validates the header, reads records into memory, and resumes the conversation state. |
kv_disk_truncate | Reduces the file record count and updates the header, effectively resetting or shortening the cache. |
kv_disk_reset | Resets the disk cache by truncating the record count to zero. |
Sources: c/kv_persist.h:25-104
Memory and State Integration
The persistence module integrates directly with the Model and KVState structures. It uses a temporary disk buffer (disk_buf) to stage writes, reducing I/O overhead by batching tensor copies before performing fwrite operations.
Integration Constraints
- Compatibility Check: During
kv_disk_load, the engine compares the file header against the currentCfgparameters. If a mismatch inn_layers,vocab, or projection dimensions occurs, the engine ignores the file to prevent state corruption. - Context Limit: The engine will not load a saved conversation if the token count (
nrec) exceeds the available context space minus a draft buffer. - MTP Head Support: When Multi-Token Prediction (MTP) is used, loading a KV cache resets
kv_startto ensure proper draft verification.
Sources: c/kv_persist.h:77-104, c/inkling.c:643-648
Summary
KV Cache Persistence provides crash-safe, incremental storage for the model's internal attention state. By persisting the compressed MLA KV-cache, the colibri engine enables instant session resumption, significantly improving user experience by eliminating the latency of re-processing long conversation histories.
Sources: c/kv_persist.h:1-5, README.md:143-147