mirror of
https://github.com/ollama/ollama.git
synced 2026-03-11 20:23:55 -05:00
5c73c4e2ee
The KV cache previously used a tree structure which could store multiple divergent sequences, which is good for cache reuse. However, this is typically used in conjunction with paged attention so each node in the tree can store just a chunk of the KV cache and they can be stitched together later. We don't currently do this, so the cache was storing copies of the full cache for each past sequence. This redundancy plus the lack of resource limits, caused significant memory use as a conversation grew. Instead, this changes to store a single entry for the cache, which can be prefix matched. Although it is less ideal for multiple users, it largely matches Ollama's current behavior. It can be improved as additional pieces are fleshed out.
68 lines
1.6 KiB
Go
68 lines
1.6 KiB
Go
//go:build mlx
|
|
|
|
package mlxrunner
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
|
|
"github.com/ollama/ollama/logutil"
|
|
"github.com/ollama/ollama/x/mlxrunner/cache"
|
|
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
|
)
|
|
|
|
// CacheEntry stores a single sequence
|
|
type CacheEntry struct {
|
|
Tokens []int32
|
|
Caches []cache.Cache
|
|
}
|
|
|
|
// FindNearestCache finds the longest common prefix between tokens and the cached sequence
|
|
func (r *Runner) FindNearestCache(tokens []int32) ([]cache.Cache, []int32) {
|
|
if r.cache == nil {
|
|
slog.Info("Cache miss", "left", len(tokens))
|
|
return nil, tokens
|
|
}
|
|
|
|
// Find longest common prefix
|
|
prefix := 0
|
|
for prefix < len(tokens) && prefix < len(r.cache.Tokens) && tokens[prefix] == r.cache.Tokens[prefix] {
|
|
prefix++
|
|
}
|
|
|
|
switch {
|
|
case prefix == 0:
|
|
for _, c := range r.cache.Caches {
|
|
c.Free()
|
|
}
|
|
r.cache = nil
|
|
slog.Info("Cache miss", "left", len(tokens))
|
|
return nil, tokens
|
|
case prefix < len(r.cache.Tokens):
|
|
trim := len(r.cache.Tokens) - prefix
|
|
for _, c := range r.cache.Caches {
|
|
c.Trim(trim)
|
|
}
|
|
r.cache.Tokens = r.cache.Tokens[:prefix]
|
|
}
|
|
|
|
slog.Info("Cache hit", "total", len(tokens), "cached", prefix, "left", len(tokens[prefix:]))
|
|
return r.cache.Caches, tokens[prefix:]
|
|
}
|
|
|
|
func (r *Runner) InsertCache(tokens []int32, caches []cache.Cache) {
|
|
r.cache = &CacheEntry{
|
|
Tokens: tokens,
|
|
Caches: caches,
|
|
}
|
|
}
|
|
|
|
func (c *CacheEntry) LogCache() {
|
|
var totalBytes int
|
|
for _, kv := range c.Caches {
|
|
k, v := kv.State()
|
|
totalBytes += k.NumBytes() + v.NumBytes()
|
|
}
|
|
logutil.Trace(fmt.Sprintf("kv cache tokens: %d, size: %s", c.Caches[0].Offset(), mlx.PrettyBytes(totalBytes)))
|
|
}
|